================================================================================
PYTAE COMPREHENSIVE KNOWLEDGE BASE & AGENT TRAINING REFERENCE MANUAL
================================================================================

TABLE OF CONTENTS
--------------------------------------------------------------------------------
1.  OVERVIEW & CORE PHILOSOPHY
2.  PACKAGE ARCHITECTURE & IMPORT PATTERNS
3.  SAMPLE DATASETS
4.  PYTHON LIBRARY REFERENCE (API VERBS)
    4.1.  pt.qry() / df.pt.qry() — Row Filtering
    4.2.  pt.mutate() / df.pt.mutate() — Column Creation & Feature Engineering
    4.3.  pt.select() / df.pt.select() — Column Selection & Reordering
    4.4.  pt.agg_df() / df.pt.agg_df() — Grouped Aggregation
    4.5.  pt.long() / pt.wide() / df.pt.long() / df.pt.wide() — Reshaping (Pure 1-to-1)
    4.6.  pt.pivot() / df.pt.pivot() — 2D Pivot Tables (Excel-Style)
    4.7.  pt.sql() / df.pt.sql() — DuckDB SQL Engine
    4.8.  Other Utilities:
          - pt.cols() / df.pt.cols()
          - pt.handle_missing() / df.pt.handle_missing()
          - pt.clean_columns() / df.pt.clean_columns()
          - pt.replace_values() / df.pt.replace_values()
          - df.to_clip() (DataFrame & Series method)
    4.9.  Plotter — Visualization & Dashboarding (Plotting API)
5.  CLI REFERENCE (COMMAND LINE INTERFACE)
    5.1.  Execution Model & Pipeline Architecture
    5.2.  File Formats & Zero-Cost Metadata Inspection
    5.3.  All CLI Flags & Syntax Rules
    5.4.  Quoting & Delimiter Standards (: vs =)
    5.5.  Multi-File Operations (-file, -merge, -concat, -sql)
    5.6.  Batch Processing & Conversions
6.  COMMON PITFALLS, SYNTAX RULES, AND EDGE CASES
7.  END-TO-END RECIPES & EXAMPLES (PYTHON & CLI)
8.  REPOSITORY DIRECTORY TREE
9.  BUILD & CONFIGURATION (pyproject.toml)
10. UNDERLYING PYTHON IMPLEMENTATION SOURCE CODE (src/pytae/)
11. TEST SUITE IMPLEMENTATION SOURCE CODE (tests/)
================================================================================


1. OVERVIEW & CORE PHILOSOPHY
--------------------------------------------------------------------------------
`pytae` is an ergonomic, high-performance tabular data manipulation library and 
command-line tool built on top of Pandas, PyArrow, DuckDB, and Matplotlib.

Key principles:
1. Dual Calling Convention: Every library verb exists both as a functional 
   function `pt.verb(df, ...)` and as a Pandas DataFrame accessor `df.pt.verb(...)`.
   The accessor enables seamless method chaining with standard Pandas methods:
   `df.rename(...).pt.qry(...).pt.select(...).pt.agg_df(...)`.
2. Unix Pipeline CLI: The command-line tool `pytae` processes operations in 
   the exact order flags are passed on the terminal:
   `pytae data.parquet -qry "..." -mutate "..." -select "..." -head 10`.
3. Fast & Zero-Cost Metadata Inspection: For formats containing file-level metadata 
   (Parquet, SAS7BDAT), commands like `-shape`, `-cols`, `-dtype`, `-head` inspect 
   headers without parsing full table rows.
4. Expressive & Human-Friendly: Minimizes boilerplate. Filtering supports 
   intervals `[3000, 4000]`, string operations, and keyword arguments. Grouped 
   aggregations express grouping via `by=` and concise aggregations (e.g. `tip=mean`, 
   `n=n`) without tedious `.groupby().agg()` boilerplate.
5. Non-Destructive: All library verbs return copies of DataFrames and never 
   mutate inputs in place. Original indices are preserved during filtering.
6. Honest Error Handling ("Be generous in what you accept, but strictly honest 
   when ambiguous"): Flexible inputs are accepted where unambiguous, but collisions, 
   conflicts, and invalid arguments raise clear exceptions rather than silently 
   inventing artificial column suffixes or guessing.
7. One and Only One Right Way: Distinct responsibilities per verb—pure 1:1 unmelting 
   is `wide()`, 2D aggregation is `pivot()`, column selection is `select()`.


2. PACKAGE ARCHITECTURE & IMPORT PATTERNS
--------------------------------------------------------------------------------
Installation (PyPI: https://pypi.org/project/pytae/):
  pip install pytae            # latest release: 3.7.0
  pip install 'pytae[plot]'    # includes matplotlib, scipy
  pip install 'pytae[sql]'     # includes duckdb

Importing:
  import pytae as pt
  import pandas as pd

Top-Level Functions exported by `pytae`:
  - pt.select(df, *args, **kwargs)
  - pt.everything() (sentinel for select)
  - pt.qry(df, *args, **kwargs)
  - pt.mutate(df, *args, **kwargs)
  - pt.agg_df(df, *args, **kwargs)
  - pt.agg(df, *args, **kwargs) (alias for agg_df)
  - pt.long(df, cols=None, id_vars=None, c="variable", v="value")
  - pt.wide(df, c="variable", v="value")
  - pt.pivot(df, r=None, c=None, v=None, a="sum", dropna=False, fill_value=None)
  - pt.sql(df, query, **frames)
  - pt.handle_missing(df, fillna=".")
  - pt.cols(df, ascending=True)
  - pt.clean_columns(df, **kwargs)
  - pt.replace_values(df, v, c=None, exact=True)
  - pt.sample(name)
  - pt.sample_data (Mapping of bundled datasets)
  - pt.Plotter (lazy loaded when matplotlib is installed)
  - pt.plot(df, *args, **kwargs)
  - pt.finalize(plotter=None, **kwargs)

DataFrame Accessor:
  Importing pytae automatically registers the `.pt` accessor on `pd.DataFrame`.
  Methods on `df.pt`:
    df.pt.select(...)
    df.pt.qry(...)
    df.pt.mutate(...)
    df.pt.agg_df(...)
    df.pt.agg(...)
    df.pt.long(...)
    df.pt.wide(...)
    df.pt.pivot(...)
    df.pt.sql(...)
    df.pt.handle_missing(...)
    df.pt.cols(...)
    df.pt.clean_columns(...)
    df.pt.replace_values(...)
    df.pt.plot(...)
    df.pt.finalize(...)


3. SAMPLE DATASETS
--------------------------------------------------------------------------------
`pytae` includes pre-packaged datasets stored in parquet format inside the library.
Access methods:
  df = pt.sample("penguins")
  df = pt.sample_data["penguins"]

List available dataset names:
  pt.sample_data.keys()
  Available datasets include:
  - 'penguins': Palmer penguins dataset (species, island, bill_length_mm, 
    bill_depth_mm, flipper_length_mm, body_mass_g, sex, year).
  - 'tips': Restaurant tips dataset (total_bill, tip, sex, smoker, day, time, size).
  - 'titanic': Titanic passenger survival data (survived, pclass, sex, age, ...).
  - 'healthexp': Healthcare spending and life expectancy.
  - 'fmri': Functional MRI time points and BOLD signal.
  - 'flights', 'diamonds', 'iris', 'planets', etc.


4. PYTHON LIBRARY REFERENCE (API VERBS)
--------------------------------------------------------------------------------

4.1. pt.qry() / df.pt.qry() — Row Filtering
--------------------------------------------------------------------------------
Filters DataFrame rows based on keyword arguments, string expressions, 
plain dictionaries, or any combination of the three.

Signature:
  pt.qry(df: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame
  df.pt.qry(*args: Any, **kwargs: Any) -> pd.DataFrame

Calling Styles:
  1. Keyword Arguments:
     df.pt.qry(species="Adelie", body_mass_g="> 5000")
  2. String Expressions (matches CLI -qry):
     df.pt.qry("body_mass_g > 5000, species = 'Adelie'")
     df.pt.qry("bill length mm > 40")   # Handles columns with spaces natively!
  3. Plain Dictionaries:
     df.pt.qry({"bill length mm": "> 40", "species": "Adelie"})
  4. Mixed Calling:
     df.pt.qry("bill length mm > 40", {"island": "Biscoe"}, species="Gentoo")
  5. Library Function:
     pt.qry(penguins, species="Adelie")

Supported Condition Types & Operators:
  - Equality:
    species="Adelie"
    "species = Adelie" or "species = 'Adelie'"
    {"species": "Adelie"}
  - Membership (in list):
    species=["Adelie", "Gentoo"]
    species=("in", ["Adelie", "Gentoo"])
  - Exclusion (not in list):
    sex=("not in", ["Male"])
  - Comparison Operators:
    Supported operators: '>', '>=', '<', '<=', '==', '!='
    As string condition: body_mass_g="> 5000"
    As tuple condition:  body_mass_g=(">", 5000)
    Note: For numeric columns, values are converted to float automatically.
  - Interval Notation (range queries on a single column):
    Open interval:       body_mass_g="(3500,4500)"    -> 3500 < x < 4500
    Half-open interval:  bill_length_mm="[40,50)"     -> 40 <= x < 50
    Closed interval:     body_mass_g="[3500,4000]"    -> 3500 <= x <= 4000
    String interval:     island="(Biscoe,Dream]"      -> Alphabetical range
  - String Accessor Operations (requires string-dtype column; na=False):
    Prefix match:        species=("startswith", "Ad")
    Multi-prefix match:  species=("startswith", ["Ad", "Ge"])  # accepts list/tuple
    Suffix match:        code=("endswith", "3")
    Substring match:     species=("contains", "in")    # regex=False, search-anywhere
    Regex match:         species=("regex", r"^[AB]")   # regex=True, search-anywhere
  - Null Checks (unary 1-element tuples):
    Missing values:      sex=("isna",)
    Non-missing values:  sex=("notna",)

Multiple Conditions on the Same Column:
  Python kwargs cannot repeat parameter names. Solve this using:
  1. Comma-separated string:  df.pt.qry("bill_length_mm > 40, bill_length_mm < 50")
  2. Multiple string args:    df.pt.qry("bill_length_mm > 40", "bill_length_mm < 50")
  3. Interval syntax:         df.pt.qry(bill_length_mm="[40,50)")
  4. Chained calls:           df.pt.qry(bill_length_mm="> 40").pt.qry(bill_length_mm="< 50")

Error Handling:
  If a column name is misspelled, qry() raises a KeyError with fuzzy match suggestions:
  KeyError: "unknown column 'speceis' (did you mean 'species'?)"


4.2. pt.mutate() / df.pt.mutate() — Column Creation & Feature Engineering
--------------------------------------------------------------------------------
Creates or overwrites columns using expressions evaluated in order via pandas eval() 
with automatic fallback to plain Python eval on Series objects.

Signature:
  pt.mutate(df: pd.DataFrame, *args: Any, by: str | Sequence[str] | None = None, dropna: bool = False, observed: bool = True, params: dict | None = None, **kwargs: Any) -> pd.DataFrame
  df.pt.mutate(*args: Any, by: str | Sequence[str] | None = None, dropna: bool = False, observed: bool = True, params: dict | None = None, **kwargs: Any) -> pd.DataFrame

Calling Styles:
  1. Keyword Arguments:
     df.pt.mutate(bmi="body_mass_g / bill_length_mm ** 2", mass_kg="body_mass_g / 1000")
  2. Dictionary Unpacking (for column names containing spaces):
     df.pt.mutate(**{"body mass kg": "body_mass_g / 1000"})
  3. External Spec File:
     df.pt.mutate("@features.txt")
  4. Callables / Lambdas:
     df.pt.mutate(is_heavy=lambda d: d["body_mass_g"] > 4000)

Key Behaviors & Features:
  - Sequential Chaining: Later expressions can reference columns created earlier 
    in the SAME mutate() call:
    df.pt.mutate(mass_kg="body_mass_g / 1000", mass_lb="mass_kg * 2.20462")
  - Columns With Spaces:
    When referencing existing columns that contain spaces in expressions, 
    surround them in SQL brackets `[col]` or backticks `` `col` ``:
    df.pt.mutate(ratio="[body mass g] / [bill length mm]")
    df.pt.mutate(ratio="`body mass g` / `bill length mm`")
  - Local Scope Variables (@local_var):
    Reference local variables from caller scope using `@name`:
    threshold = 4000
    df.pt.mutate(heavy="body_mass_g >= @threshold")
    Or pass explicitly via params:
    df.pt.mutate(heavy="body_mass_g >= @threshold", params={"threshold": 4000})

Injected Helper Functions in Expressions:
  - if_else(condition, true_value, false_value):
    df.pt.mutate(category="if_else(body_mass_g > 4000, 'Heavy', 'Light')")
  - case_when((cond1, val1), (cond2, val2), ..., default=val):
    Evaluated in order, first match wins. Accepts tuples or flat alternating pairs:
    df.pt.mutate(size="case_when((body_mass_g >= 4500, 'L'), (body_mass_g >= 3500, 'M'), default='S')")
    df.pt.mutate(size="case_when(body_mass_g >= 4500, 'L', body_mass_g >= 3500, 'M', default='S')")
  - coalesce(col1, col2, ..., default):
    Returns the first non-null value per row:
    df.pt.mutate(phone="coalesce(mobile, home, work, 'None')")
  - map(column, {key: value, ...}[, default]):
    Recodes values via dictionary lookup:
    df.pt.mutate(code="map(species, {'Adelie': 'A', 'Gentoo': 'G'}, default='Other')")
  - Series Methods Fallback:
    Method chains on Series work seamlessly:
    df.pt.mutate(initial="species.str[0]", rounded="bill_length_mm.round()")

Grouped Window Operations (by=..., dropna=False, observed=True):
  Evaluates window calculations and group summaries broadcast back to every row (N -> N):
  - Injected Aggregations: mean(x), sum(x), median(x), min(x), max(x), std(x), var(x), n (or n())
  - Sequential Derivation:
    df.pt.mutate("avg_tip = mean(tip), diff = tip - avg_tip, sz = n", by="day")
  - Multi-Column Grouping:
    df.pt.mutate("avg = mean(val)", by=["species", "island"])
  - Handling Missing Group Keys (dropna):
    - dropna=False (default across pytae): Rows with NA in grouping columns form their own group
      and compute group metrics across all NA rows together.
    - dropna=True: Rows with NA in grouping columns receive NaN in new columns;
      row count and alignment are strictly preserved.
      df.pt.mutate("avg = mean(val)", by="grp", dropna=True)
      (In CLI: pass -dropna true to exclude NA groups)

External Spec File Format (@file.txt):
  Lines in `@specs.txt` are evaluated sequentially. Supports `#` comments and empty lines:
  # Feature engineering specs
  mass_kg = body_mass_g / 1000
  mass_lb = mass_kg * 2.20462
  bill_ratio = [bill length mm] / [bill depth mm]


4.3. pt.select() / df.pt.select() — Column Selection & Reordering
--------------------------------------------------------------------------------
Selects and reorders columns using names, ranges/slices, data types, patterns, or predicates.

Signature:
  pt.select(
      df: pd.DataFrame,
      *args: Any,
      dtype: str | type | Sequence[str | type] | None = None,
      exclude_dtype: str | type | Sequence[str | type] | None = None,
      contains: str | Sequence[str] | None = None,
      startswith: str | Sequence[str] | None = None,
      endswith: str | Sequence[str] | None = None,
      regex: str | Sequence[str] | None = None,
  ) -> pd.DataFrame

Argument Types for *args:
  - Exact Column Names:
    df.pt.select("species", "island")
  - List of Names:
    df.pt.select(["species", "island"])
  - Slice Range ('start:end'):
    Contiguous slice from start to end (inclusive of both bounds):
    df.pt.select("bill_length_mm:body_mass_g")
    Open slices: ":body_mass_g" (from start), "body_mass_g:" (to end)
  - pt.everything():
    Pulls in all remaining unselected columns in their original order:
    df.pt.select("species", pt.everything())  # species first, then the rest
  - Predicate Callable:
    df.pt.select(lambda c: c.startswith("b"))

Keyword Arguments:
  - regex: Regex pattern matching column names.
    df.pt.select(regex=r"^bill")
    df.pt.select("species", regex=r"mm$")
    IMPORTANT: Regex must be specified via `regex=...`. A positional string containing 
    regex characters like "^bill" is treated as an exact name and raises a KeyError!
  - contains: Substring(s) in column names (OR if list).
    df.pt.select(contains="mm")
    df.pt.select(contains=["bill", "mass"])
  - startswith: Prefix(es) in column names.
    df.pt.select(startswith="bill")
  - endswith: Suffix(es) in column names.
    df.pt.select(endswith="mm")
  - dtype: Include columns by dtype shorthand or type.
    Shorthands: 'numeric', 'datetime', 'category', 'bool'
    df.pt.select(dtype="numeric")
  - exclude_dtype: Exclude columns by dtype. Cannot be combined with other criteria.
    df.pt.select(exclude_dtype="numeric")      # keep non-numeric
    df.pt.select(exclude_dtype="non_numeric")  # keep numeric

Column Order Rules:
  Explicitly listed column names preserve YOUR specified order.
  Columns pulled in by contains, regex, or everything() follow original DataFrame order.


4.4. pt.agg_df() / df.pt.agg_df() / df.pt.agg() — Grouped Aggregation
--------------------------------------------------------------------------------
Aggregates a DataFrame grouped by explicit `by=` column(s) (or None for whole-table
summary) and aggregates numeric columns. Supports string mapping specifications
with bracketed columns `[col]` (matching CLI `-agg`), keyword arguments, or
whole-frame aggregations.

Signature:
  pt.agg_df(
      df: pd.DataFrame,
      by: str | Sequence[str] | None,
      *args: Any,
      a: str | list[str] | None = None,
      dropna: bool = False,
      observed: bool = True,
      **kwargs: Any,
  ) -> pd.DataFrame

Calling Styles:
  1. String Mapping Specification (Zero Dict Unpacking! Perfect for Spaced Columns):
     df.pt.agg("smoker", "tip = mean, [total bill] = mean, n = n")
     df.pt.agg("smoker", "tip = mean", "[total bill] = mean", "n = n")
     df.pt.agg("[smoker status]", "[avg bill] = [total bill]:mean, [total count] = n")
  2. Whole-Frame Aggregation:
     df.pt.agg("species", "mean")                  # applies mean to all numeric cols
     df.pt.agg(["species", "island"], ["mean", "sum", "n"]) # multi-column grouping
     df.pt.agg(None, a=["mean", "n"])              # whole-table summary (1 row)
  3. Specific Column Aggregations (kwargs):
     df.pt.agg("species", body_mass_g="mean", flipper_length_mm="max", n="n")
     df.pt.agg("species", body_mass_g=["mean", "sum"], row_count="n")
  - Row Count ('n'):
     Assigning a key with value 'n' produces the group row count.
     If the key is `n="n"`, the count column is named `n`.
     If the key is `count="n"`, the count column is named `count`.
  - Column Names:
     If a single aggregation is applied to a column, its original name is preserved:
     `body_mass_g='mean'` -> output column: `body_mass_g`
     If multiple aggregations are applied to a column, names become `{col}_{agg}`:
     `body_mass_g=['mean', 'max']` -> output columns: `body_mass_g_mean`, `body_mass_g_max`
  - Group Counts Only:
     If the DataFrame has no numeric columns and only 'n' is requested, it returns 
     the frequency counts of the groups.
  - dropna Default:
     dropna=False by default retains missing category groups across aggregations.


4.5. pt.long() / pt.wide() / df.pt.long() / df.pt.wide() — Reshaping (Pure 1-to-1)
--------------------------------------------------------------------------------
Pure structural reshaping between long and wide formats without aggregation.

pt.long(df, cols=None, id_vars=None, c="variable", v="value"):
  Melts columns to rows. If cols is None, melts all numeric columns, keeping 
  non-numeric columns as ID variables.
  Parameters:
    cols: Column(s) to melt to rows. Can be string or list/tuple of strings.
    id_vars: Identifier column(s) to keep as rows.
    c: column name for melted metric names (default: "variable")
    v: column name for melted metric values (default: "value")
  Example:
    tall = df.pt.long(c="metric", v="measurement")

pt.wide(df, c="variable", v="value"):
  Pivots long-format records back into column headers (exact reverse of pt.long).
  Strictly 1-to-1 without aggregation.
  Parameters:
    c: column whose unique values become headers (default: "variable")
    v: column containing values (default: "value")
  Rules:
    - Requires unique (index, c) pairs. If duplicate keys exist, raises ValueError
      guiding the user to use pt.pivot() instead.
    - c and v cannot be the same column.
  Example:
    wide = tall.pt.wide(c="metric", v="measurement")


4.6. pt.pivot() / df.pt.pivot() — 2D Pivot Tables (Excel-Style)
--------------------------------------------------------------------------------
Full multi-dimensional 2D pivot table and frequency cross-tabulation engine.
Automatically flattens multi-level column hierarchies to clean 1D string names 
and resets the index to standard RangeIndex(0, 1, 2, ...).

Signature:
  pt.pivot(
      df: pd.DataFrame,
      r: str | Sequence[str] | None = None,
      c: str | Sequence[str] | None = None,
      v: str | Sequence[str] | None = None,
      a: str | Callable | Sequence[str | Callable] = "sum",
      dropna: bool = False,
      fill_value: Any = None,
  ) -> pd.DataFrame
  df.pt.pivot(r=None, c=None, v=None, a="sum", dropna=False, fill_value=None) -> pd.DataFrame

Parameters:
  r: Row grouping column(s). String (e.g. "island", "species,island") or list of strings.
     Alias: index=, rows=
  c: Column dimension(s). String (e.g. "sex") or list of strings.
     Alias: columns=, cols=
  v: Value column(s) to aggregate. Optional when counting (a="n").
     Alias: values=
  a: Aggregation function (default: "sum"). Supports standard aggregations
     ("mean", "sum", "median", "min", "max", "std", "var") and "n" for row counts.
     Alias: aggfunc=
  dropna: Whether to drop NA categories in row/column keys (default: False).
  fill_value: Scalar value for missing grid intersections. For a="n", defaults to 0.

Key Features & Conventions:
  - 1D Flat Headers: Multi-column pivots format headers cleanly (e.g. "MALE_2007", "FEMALE_2007").
  - Pure Integer Counts: Row counts (a="n") return standard int64 (or Int64) integers, not floats.
  - Zero Count Fill: Missing grid combinations in frequency tables default to 0.
  - Omit v for Counts: For frequency cross-tabulations (a="n"), v is optional.
  - RangeIndex: Index is always reset to 0, 1, 2, ... with no lingering MultiIndex.

Examples:
  # 2D summary matrix
  df.pt.pivot(r="island", c="species", v="body_mass_g", a="mean")

  # Frequency cross-tabulation (count rows across dimensions)
  df.pt.pivot(r="island", c="species", a="n")

  # Multi-dimensional pivot with custom fill
  df.pt.pivot(r=["species", "island"], c="sex", v="body_mass_g", a="median", fill_value=0)


4.7. pt.sql() / df.pt.sql() — DuckDB SQL Engine
--------------------------------------------------------------------------------
Executes SQL queries directly over in-memory DataFrames using DuckDB.
Requires: pip install 'pytae[sql]'

Signature:
  pt.sql(df: pd.DataFrame, query: str, **frames: pd.DataFrame) -> pd.DataFrame
  df.pt.sql(query: str, **frames: pd.DataFrame) -> pd.DataFrame

Rules & Features:
  - Table 'data': The current/calling DataFrame is automatically registered as table `data`.
  - Joining Extra DataFrames: Additional DataFrames can be passed as keyword arguments:
    pt.sql(orders, "select * from data join cust using (cust_id)", cust=customers)
  - Columns with Spaces: Supported using double quotes `"col a"`, brackets `[col a]`, 
    or backticks `` `col a` `` in the SQL query:
    df.pt.sql('select [bill length mm], avg(body_mass_g) from data group by 1')
  - Query From File: Pass an external SQL file path prefixed with `@`:
    df.pt.sql("@analysis.sql")


4.8. Other Utilities
--------------------------------------------------------------------------------
- pt.cols(df, ascending=True):
  Returns list of column names.
  ascending=True  -> Alphabetical A-Z
  ascending=False -> Alphabetical Z-A
  ascending=None  -> Original DataFrame file order

- pt.handle_missing(df, fillna="."):
  Sanitizes missing values:
  - Category columns are converted to object.
  - String columns: fills NaN with `fillna` (default '.') and strips whitespace.
  - Numeric columns: fills NaN with `0`.

- pt.clean_columns(df, strip=False, strip_special=False, squeeze=False, fill=None, case=None, dedupe=False):
  Systematically sanitizes DataFrame column header names in fixed order:
  strip:          trim leading/trailing whitespace
  strip_special:  remove characters other than alphanumeric, underscore, whitespace, or fill
  squeeze:        collapse multi-space runs into single spaces
  fill:           replace whitespace characters with this string (e.g. "_")
  case:           'lower', 'upper', or 'proper' (Title Case)
  dedupe:         auto-number duplicates (e.g. 'rev', 'rev_1', 'rev_2')

- pt.replace_values(df, v: dict, c: str | list | None = None, exact=True):
  Replaces cell values in the DataFrame (wraps pandas replace).
  v:      dictionary of {old_value: new_value}
  c:      scoped column name or list of columns (default None = entire DataFrame)
  exact:  True = exact match of cell value; False = substring replacement (regex)

- df.to_clip() / s.to_clip():
  Copies DataFrame or Series to system clipboard as tab-separated values without index. Attached directly to pd.DataFrame and pd.Series on import pytae (no pt. or df.pt. needed).


4.9. Plotter — Visualization & Dashboarding (Plotting API)
--------------------------------------------------------------------------------
Method-chainable plotting system wrapping `pandas.plot` and `matplotlib`.
Requires: pip install 'pytae[plot]'

Key Concepts:
  - Fluent Accessor & Functional Shortcuts:
    tips.pt.agg("day,sex", total_bill="mean").pt.plot(kind="bar", x="day", y="total_bill", by="sex", palette="tab10").pt.finalize()
    tips.pt.agg("day", total_bill="mean").pt.plot(kind="bar", x="day", y="total_bill").pt.finalize()
  - Chain Pattern:
    k = pt.Plotter(figsize=(6, 4))
    (k
     .data(df)
     .plot(kind="scatter", x="bill_length_mm", y="bill_depth_mm", by="species", palette="Set1")
     .finalize()
    )
    k.fig  # access matplotlib figure
  - Clean Separation of Concerns (No In-Plot Aggregation):
    In accordance with Unix philosophy ("one and only one right way"), Plotter plots data directly without performing aggregation.
    Summarize first with `df.pt.agg()` or `df.pt.pivot()`:
    tips_avg = tips.pt.agg("day,sex", total_bill="mean")
    k.data(tips_avg).plot(kind="bar", x="day", y="total_bill", by="sex").finalize()
  - Box Plots (True Unaggregated Quartiles):
    pt.Plotter(penguins).plot(kind="box", x="species", y="body_mass_g", palette="Set1").finalize()
  - Heatmap Matrix:
    corr = penguins.select_dtypes("number").corr()
    pt.Plotter(corr).plot(kind="heatmap", annot=True, fmt=".2f", cmap="coolwarm", title="Correlation").finalize()
  - Multi-Panel Dashboards (Mosaic Layouts):
    mosaic = '''
    AB
    CD
    '''
    k = pt.Plotter(mosaic, figsize=(10, 8))
    (k
     .data(penguins).plot(on="A", kind="scatter", x="bill_length_mm", y="bill_depth_mm", by="species")
     .data(tips.pt.agg("day", total_bill="mean")).plot(on="B", kind="bar", x="day", y="total_bill")
     .finalize(consolidate_legends=True)
    )
  - Secondary Y-Axis ('^'):
    Append `^` to panel name (e.g. `on="A^"` or `secondary_y=True`) to create a secondary y-axis:
    (k
     .plot(kind="bar", x="day", y="total_bill", on="A", color="skyblue")
     .plot(kind="line", x="day", y="tip", secondary_y=True, color="crimson")
     .finalize()
    )
  - Small Multiples / Auto-Faceting:
    pt.plot(penguins, by="species", ncols=3, kind="scatter", x="bill_length_mm", y="bill_depth_mm").finalize()
    pt.Plotter.facet(df, by="species", ncols=2, kind="scatter", x="bill_length_mm", y="bill_depth_mm")
  - Saving:
    k.save("chart.png", dpi=300)


5. CLI REFERENCE (COMMAND LINE INTERFACE)
--------------------------------------------------------------------------------

5.1. Execution Model & Pipeline Architecture
--------------------------------------------------------------------------------
The CLI command `pytae` processes operations sequentially in the order given on 
the command line:
  pytae input.parquet -qry "body_mass_g > 3000" -select "species,body_mass_g" -head 5

Pipeline Rules:
1. Positional Path: The first unflagged argument is the input file path.
   (Except in `-file` multi-file mode, where `-file` replaces the positional path).
2. Sequential Mutation: Each operation transforms the intermediate dataset in memory.
   For example, `-qry` before `-select -col` allows filtering on a column that is subsequently excluded.
3. Final Operation Determines Output: Only the last operation in the chain prints or exports, 
   unless an intermediate operation explicitly writes to disk.
4. Shell Quoting: Always quote multi-word arguments and expressions in your shell.


5.2. File Formats & Zero-Cost Metadata Inspection
--------------------------------------------------------------------------------
Supported input & output file types:
  - .parquet, .pq (columnar storage)
  - .csv (comma-delimited text)
  - .txt (tab-delimited text or custom delimiter)
  - .dat (pipe-delimited text, latin-1 encoding)
  - .jsonl, .ndjson (line-delimited JSON records)
  - .csv.gz, .txt.gz, .dat.gz, .jsonl.gz (transparent gzip compression & decompression)
  - .sas7bdat (SAS binary dataset, read-only)

Fast Zero-Cost Inspection (Metadata-Only):
  For Parquet and SAS7BDAT, the following flags inspect metadata without reading 
  the full file rows into RAM:
  - `-shape`: Returns (rows, columns)
  - `-cols`: Prints column names in order
  - `-dtype`: Shows data types
  - `-head [N]`: Reads only the first N rows from the file stream
  - `-meta`: Displays complete zero-scan Parquet metadata (row groups count, 
    column compression codecs, compression ratio, Arrow & Pandas schema) 
    without loading table records.


5.3. All CLI Flags & Syntax Rules
--------------------------------------------------------------------------------
Inspection & Comparison Flags:
  -head [N]             Print first N rows (default 5)
  -tail [N]             Print last N rows (default 5)
  -shape                Print row and column counts
  -cols [ORDER]         Print column names (ORDER: asc, desc, none)
  -dtype [ORDER]        Print column data types
  -nulls [ORDER]        Print null value counts per column
  -describe             Descriptive statistics for numeric columns
  -info                 DataFrame summary (dtypes, non-null counts, memory)
  -meta                 Display zero-scan Parquet metadata (row groups, schema, compression)
  -diff PATH            Compare schema, shape, null counts, and cell values against another file
  -value_counts         Value counts for categorical/string columns
  -unique               Count of unique values per column
  -sample [N]           Random sample of N rows
  -seed N               Random seed for sampling
  -frac P               Sample fraction P (0.0 to 1.0)
  -nrows N              Limit reading to first N rows of file

Row Filtering Flags:
  -qry CONDITIONS       pytae filter expressions (e.g. "species='Adelie', body_mass_g > 3500")
  -query EXPR           pandas query() expression passed to numexpr

Column Transformation Flags:
  -select SPEC          Select columns (names, slices 'a:b', contains, dtypes, negative '-col', exclude=)
  -mutate SPEC          Create/overwrite columns ("new_col = expression")
  -rename OLD:NEW,...   Rename columns using ':' delimiter (e.g. "a:alpha, b:beta")
  -replace_values SPEC  Replace values: "v='old:new', c='col', exact=True"
  -clean_columns SPECS  Clean headers: "strip=True, fill='_', case='lower'"

Aggregation & Reshaping Flags:
  -by, -group_by COLS   Grouping columns for -agg and -mutate (e.g. -by species -agg mean)
  -agg, -agg_df SPECS   Aggregate columns (e.g. -by species -agg "body_mass_g=mean, count=n" or -agg mean)
  -dropna BOOL          Control NA group handling for -agg, -mutate, -pivot, -wide (true/false, default: false)
  -long [SPECS]         Melt numeric columns to rows: "c=variable, v=value"
  -wide [SPECS]         Pure 1-to-1 unmelting into headers: "c=variable, v=value"
  -pivot SPECS          2D pivot table: "r=species, c=island, v=body_mass_g, a=mean" (or counts: "r=species, c=island, a=n")
  -crosstab SPECS       Frequency crosstabulation (delegates to -pivot): "index=species, columns=sex"

SQL & Multi-File Flags:
  -sql QUERY            Run DuckDB SQL on table 'data' (or @query.sql)
  -file SPECS           Multi-file mode: "file1.parquet=df1; file2.csv=df2"
  -merge SPECS          Join files: "left=df1, right=df2, on=id, how=inner"
  -concat SPECS         Stack files: "frames='df1,df2'"

Output & Formatting Flags:
  -o TARGET             Output destination: file path (.csv, .parquet, .jsonl, .csv.gz, etc.),
                        format for in-place or batch conversion ('csv', 'parquet', 'txt', 'dat',
                        'jsonl', 'csv.gz', 'jsonl.gz'), or 'clip'/'clipboard'
  -out_dir, -od DIR     Target directory for exported files (created if missing; requires -o)
  -dlim CHAR            Delimiter character for input/output text files
  -encoding ENC         File encoding (e.g. utf-8, latin-1)
  -pretty               Pretty-print tables with bordered markdown formatting
  -pager                Pipe table or inspect outputs through system pager ($PAGER or less)
  -round N              Round numeric output columns to N decimal places
  -progress [N]         Show row-reading/writing progress (default: 200,000 rows per chunk)


5.4. Quoting & Delimiter Standards (: vs =)
--------------------------------------------------------------------------------
pytae enforces strict delimiters depending on the semantic meaning:

1. Mapping / Translation uses COLON (`:`):
   - `-rename`: "old_col:new_col, col_b:col_beta"
   - `-replace_values`: "v='old_val:new_val', c='col_name'"
   Rule: Old is mapped to New via `:`.

2. Assignment / Specification uses EQUALS (`=`):
   - `-mutate`: "new_col = expression, mass_kg = body_mass_g / 1000"
   - `-qry`: "species = 'Adelie', body_mass_g > 3500"
   - `-clean_columns`: "strip=True, fill='_', case='lower'"
   - `-merge`: "left=df1, right=df2, on=id, how=left"

Quoting Rules:
- String literals in expressions require quotes: `species = 'Adelie'`.
- Column names inside expressions must remain unquoted so they resolve as variables:
  `mass_kg = body_mass_g / 1000` (NOT `'body_mass_g' / 1000`).
- Column names containing spaces must use brackets `[col]` or backticks `` `col` ``.


5.5. Multi-File Operations (-file, -merge, -concat, -sql)
--------------------------------------------------------------------------------
When operating across multiple files, `-file` replaces the positional file path:

Syntax:
  pytae -file "PATH=ALIAS; PATH=ALIAS" -merge ...

1. Joining Files (-merge):
   pytae -file "orders.parquet=ord; customers.csv=cust"          -merge "left=ord, right=cust, on=cust_id, how=inner"          -select "order_id,cust_name,amount" -head 10

   Joining on differing column names:
   -merge "left=df1, right=df2, on='left_id:right_id', how=left"

   Joining multiple files in sequence ('data' represents previous merge result):
   pytae -file "a.csv=a; b.csv=b; c.csv=c"          -merge "left=a, right=b, on=id"          -merge "left=data, right=c, on=id"

2. Stacking Files (-concat):
   pytae -file "jan.csv=m1; feb.csv=m2; mar.csv=m3"          -concat "frames='m1,m2,m3'" -shape

3. SQL Across Multiple Files:
   pytae -file "orders.parquet=ord; customers.csv=cust"          -sql "select ord.id, cust.name, ord.total from ord join cust using (cust_id)"


5.6. Batch Processing & Conversions (-o)
--------------------------------------------------------------------------------
Convert files between formats with zero python code via -o:
  # Convert parquet to CSV
  pytae penguins.parquet -o penguins.csv

  # In-place conversion (saves penguins.csv next to penguins.parquet)
  pytae penguins.parquet -o csv

  # Batch convert all parquet files in a directory to CSV
  pytae 'data/*.parquet' -o csv

  # Filter and save to parquet
  pytae penguins.parquet -qry "body_mass_g > 4000" -o heavy_penguins.parquet

  # Copy to clipboard directly without printing to terminal
  pytae penguins.parquet -head -o clip


6. COMMON PITFALLS, SYNTAX RULES, AND EDGE CASES
--------------------------------------------------------------------------------
1. Spaced Column Names:
   - In `df.pt.qry()`: String expressions handle spaces directly:
     `df.pt.qry("bill length mm > 40")`.
   - In `df.pt.mutate()`: Python kwargs syntax cannot have spaces (`col name = ...` is a SyntaxError).
     Must unpack a dictionary:
     `df.pt.mutate(**{"col name": "a + b"})`.
   - Inside expressions: Always enclose spaced names in brackets `[col name]` or backticks `` `col name` ``.
2. Select Regex vs Positional String:
   - Calling `pt.select(df, "^bill")` will search for a literal column named `^bill` and fail.
   - You MUST use `regex=`: `pt.select(df, regex=r"^bill")`.
3. Aggregation Grouping Logic:
   - `agg_df` automatically groups by EVERY non-numeric column when by= is omitted.
   - If your DataFrame has 10 string columns, it groups by all 10.
   - Always filter columns first if you only want to group by 1 or 2 columns, or specify `by=`:
     `df.pt.select("species", "body_mass_g").pt.agg_df("mean")`.
4. The Count Column 'n':
   - 'n' is a reserved aggregation token in `agg_df`, `pivot`, and grouped `mutate(..., by=...)`.
   - If your input table already has a real numeric column named 'n', rename it first to avoid collision.
5. In-Place Modification:
   - pytae NEVER modifies DataFrames in place. Always assign the returned DataFrame:
     `df = df.pt.mutate(...)` or chain operations.
6. Null Checks in qry():
   - Null checks are 1-element tuples: `sex=('isna',)` or `sex=('notna',)`.
   - Do NOT pass a second element.
7. Reshaping: wide() vs pivot():
   - `wide()` is strictly for pure 1-to-1 structural unmelting (reversing long()) without aggregation.
   - For 2D cross-tabulations, aggregations across dimensions, or when duplicate keys exist, use `pivot()`:
     `df.pt.pivot(r="island", c="species", v="body_mass_g", a="mean")`.


7. END-TO-END RECIPES & EXAMPLES
--------------------------------------------------------------------------------

Recipe 1: Full Python Data Engineering Pipeline
--------------------------------------------------------------------------------
```python
import pytae as pt

# Load bundled dataset
penguins = pt.sample("penguins")

# Clean pipeline: filter -> derive features -> select -> aggregate
summary = (
    penguins
    # 1. Filter: Adelie & Gentoo with known sex and body mass >= 3500g
    .pt.qry(
        species=["Adelie", "Gentoo"],
        sex=("notna",),
        body_mass_g=">= 3500",
    )
    # 2. Derive features: convert mass, calculate ratio, categorize size
    .pt.mutate(
        mass_kg="body_mass_g / 1000",
        bill_ratio="bill_length_mm / bill_depth_mm",
        size_class="if_else(mass_kg > 4.5, 'Large', 'Regular')",
    )
    # 3. Create a column with spaces using dict unpacking
    .pt.mutate(**{
        "normalized bill": "bill_length_mm / bill_length_mm.max()",
    })
    # 4. Narrow to columns of interest
    .pt.select("species", "size_class", "mass_kg", "bill_ratio")
    # 5. Grouped aggregation: mean mass, max ratio, group counts
    .pt.agg_df(
        mass_kg="mean",
        bill_ratio="max",
        n="n",
    )
)
print(summary)
```

Recipe 2: DuckDB SQL Query on DataFrame
--------------------------------------------------------------------------------
```python
import pytae as pt

tips = pt.sample("tips")

query = """
select 
    day, 
    time, 
    count(*) as total_orders,
    round(avg(total_bill), 2) as avg_bill,
    round(avg(tip), 2) as avg_tip
from data
where total_bill > 15
group by day, time
order by avg_bill desc
"""

result = tips.pt.sql(query)
print(result)
```

Recipe 3: Advanced CLI Pipeline One-Liners
--------------------------------------------------------------------------------
```bash
# 1. Inspect schema and row count without loading full data into memory
pytae raw_data.parquet -shape -cols

# 2. Filter, engineer columns, clean headers, and inspect head
pytae penguins.parquet \
  -qry "species = 'Adelie', body_mass_g > 3500" \
  -mutate "mass_kg = body_mass_g / 1000, is_heavy = mass_kg > 4.0" \
  -clean_columns "strip=True, fill='_', case='lower'" \
  -select "species,island,mass_kg,is_heavy" \
  -head 5

# 3. Fast Join Across Two Files and Export to Parquet
pytae -file "transactions.csv=tx; customers.parquet=cust" \
      -merge "left=tx, right=cust, on='customer_id', how='inner'" \
      -qry "amount > 100" \
      -o high_value_tx.parquet

# 4. Grouped Window Calculation with Missing Group Key Retention (-dropna false default)
pytae penguins.parquet \
  -by sex \
  -mutate "avg_mass = mean(body_mass_g), diff = body_mass_g - avg_mass" \
  -select "species,sex,body_mass_g,avg_mass,diff" \
  -head 5
```

Recipe 4: 2D Pivot Tables & Frequency Matrices (Python & CLI)
--------------------------------------------------------------------------------
```python
import pytae as pt

penguins = pt.sample("penguins")

# 1. 2D Summary Grid (Mean Body Mass by Island and Species)
grid = penguins.pt.pivot(r="island", c="species", v="body_mass_g", a="mean")
print(grid)

# 2. 2D Frequency Matrix (Counts with a="n")
counts = penguins.pt.pivot(r="island", c="species", a="n")
print(counts)
```

```bash
# CLI 2D Pivot Summary
pytae penguins.parquet -pivot "r=island, c=species, v=body_mass_g, a=mean"

# CLI 2D Frequency Counts
pytae penguins.parquet -pivot "r=island, c=species, a=n"
```


================================================================================
8. REPOSITORY DIRECTORY TREE
================================================================================

.
├── .githooks
│   └── pre-commit
├── .github
│   └── workflows
│       └── ci.yml
├── .gitignore
├── LICENSE
├── changelog.md
├── contributing.md
├── docs
│   ├── cli
│   │   ├── aggregate.md
│   │   ├── clean_replace.md
│   │   ├── diff.md
│   │   ├── export_io.md
│   │   ├── filter.md
│   │   ├── inspect.md
│   │   ├── multi_file.md
│   │   ├── mutate.md
│   │   ├── pivot.md
│   │   ├── plotting.md
│   │   ├── reshape.md
│   │   ├── select.md
│   │   └── sql.md
│   ├── cli.md
│   ├── library
│   │   ├── agg.ipynb
│   │   ├── mutate.ipynb
│   │   ├── other_utilities.ipynb
│   │   ├── pivot.ipynb
│   │   ├── plotting.ipynb
│   │   ├── qry.ipynb
│   │   ├── reshape.ipynb
│   │   ├── select.ipynb
│   │   └── sql.ipynb
│   └── library.md
├── pyproject.toml
├── readme.md
├── reference
│   └── pytae_reference.txt
├── scripts
│   ├── generate_reference.py
│   └── run_notebooks.py
├── src
│   └── pytae
│       ├── __init__.py
│       ├── __main__.py
│       ├── _text.py
│       ├── accessor.py
│       ├── agg_df.py
│       ├── cli.py
│       ├── cli_parsing.py
│       ├── cli_pipeline.py
│       ├── cli_run.py
│       ├── datasets
│       │   ├── anagrams.parquet
│       │   ├── anscombe.parquet
│       │   ├── attention.parquet
│       │   ├── brain_networks.parquet
│       │   ├── car_crashes.parquet
│       │   ├── diamonds.parquet
│       │   ├── dots.parquet
│       │   ├── dowjones.parquet
│       │   ├── exercise.parquet
│       │   ├── flights.parquet
│       │   ├── fmri.parquet
│       │   ├── geyser.parquet
│       │   ├── glue.parquet
│       │   ├── healthexp.parquet
│       │   ├── iris.parquet
│       │   ├── mpg.parquet
│       │   ├── penguins.parquet
│       │   ├── planets.parquet
│       │   ├── seaice.parquet
│       │   ├── taxis.parquet
│       │   ├── tips.parquet
│       │   └── titanic.parquet
│       ├── mutate.py
│       ├── other_utilities.py
│       ├── plotting.py
│       ├── py.typed
│       ├── qry.py
│       ├── readers.py
│       ├── select.py
│       ├── shape.py
│       └── sql.py
├── tests
│   ├── __init__.py
│   ├── cli_helpers.py
│   ├── conftest.py
│   ├── test_accessor.py
│   ├── test_agg_df.py
│   ├── test_cli_agg.py
│   ├── test_cli_clean.py
│   ├── test_cli_enhancements.py
│   ├── test_cli_filter.py
│   ├── test_cli_inspect.py
│   ├── test_cli_merge.py
│   ├── test_cli_mutate.py
│   ├── test_cli_parse.py
│   ├── test_cli_reshape.py
│   ├── test_cli_sample_datasets.py
│   ├── test_cli_select.py
│   ├── test_cli_sql.py
│   ├── test_import.py
│   ├── test_mutate.py
│   ├── test_pivot.py
│   ├── test_plotter.py
│   ├── test_qry.py
│   ├── test_readers.py
│   ├── test_select.py
│   ├── test_shape.py
│   ├── test_sql.py
│   └── test_utilities.py
└── uv.lock

================================================================================
9. BUILD & CONFIGURATION (pyproject.toml)
================================================================================

File: pyproject.toml
----------------------------------------
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "pytae"
version = "3.7.0"
description = "Fast, ergonomic pandas tools and zero-code CLI for tabular data manipulation and plotting"
authors = [{ name = "Madhukar Kumar Jha", email = "madhukar.k.jha@gmail.com" }]
readme = "readme.md"
requires-python = ">=3.10"
dependencies = [
    "pandas>=2.0",
    "pyarrow>=15.0.0",
]
classifiers = [
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.10",
    "Programming Language :: Python :: 3.11",
    "Programming Language :: Python :: 3.12",
    "Programming Language :: Python :: 3.13",
    "Programming Language :: Python :: 3.14",
    "License :: OSI Approved :: MIT License",
    "Operating System :: OS Independent",
]

[project.optional-dependencies]
plot = ["matplotlib", "scipy"]
sql = ["duckdb>=0.9"]
dev = ["pytest>=8.0.0", "matplotlib", "duckdb>=0.9", "scipy", "ruff>=0.8", "mypy>=1.13"]
notebooks = ["nbclient", "nbformat", "ipykernel", "scipy"]

[project.scripts]
pytae = "pytae.cli:main"

[project.urls]
"Homepage" = "https://github.com/maddytae/pytae"
"Bug Tracker" = "https://github.com/maddytae/pytae/issues"
"CLI Docs" = "https://github.com/maddytae/pytae/blob/master/docs/cli.md"
"Library Docs" = "https://github.com/maddytae/pytae/blob/master/docs/library.md"

[tool.hatch.build.targets.wheel.force-include]
"reference" = "pytae/reference"

[tool.ruff]
target-version = "py310"
line-length = 120
src = ["src", "tests"]

[tool.ruff.lint]
select = ["E", "F", "I"]
ignore = [
    "E501",  # line length — CLI specs and docs strings are long on purpose
    "E731",  # clip_action = lambda ... in the CLI pipeline
]

[tool.ruff.lint.per-file-ignores]
"src/pytae/__init__.py" = ["F401"]  # re-exports
"tests/*" = ["E402"]  # sys.path insert before pytae import
"tests/test_mutate.py" = ["E402", "F841"]  # @locals are used via pandas eval
"src/pytae/plotting.py" = ["E741", "F841"]
"docs/library/*" = ["E402", "F401", "E401", "E701", "I001"]
"docs/plotting/*" = ["E402", "F401", "E401", "E701", "I001"]

[tool.mypy]
python_version = "3.12"
files = ["src/pytae"]
ignore_missing_imports = true
check_untyped_defs = true
warn_unused_ignores = true
pretty = true
# plotting is matplotlib-heavy; typed incrementally with the rest of the library
exclude = ["src/pytae/plotting.py"]

[[tool.mypy.overrides]]
module = ["pytae.plotting", "pytae.plotting.*"]
ignore_errors = true

================================================================================
10. UNDERLYING PYTHON IMPLEMENTATION SOURCE CODE (src/pytae/)
================================================================================

################################################################################
# FILE: src/pytae/__init__.py
################################################################################
"""pytae: package functions (`pt.select(df, ...)`) and DataFrame accessor `df.pt`, plus a CLI."""

from __future__ import annotations

from collections.abc import Mapping
from pathlib import Path

import pandas as pd

from .accessor import PtAccessor  # noqa: F401  — registers df.pt
from .agg_df import agg, agg_df
from .mutate import mutate
from .other_utilities import (
    clean_columns,
    cols,
    handle_missing,
    replace_values,
)
from .qry import qry
from .select import everything, select
from .sql import sql

DATA_PATH = Path(__file__).resolve().parent / "datasets"
_DATASET_NAMES = tuple(sorted(p.stem for p in DATA_PATH.glob("*.parquet")))
_cache: dict[str, pd.DataFrame] = {}


def sample(name: str) -> pd.DataFrame:
    """Load a bundled sample dataset by name (cached after the first call).
    Returns a copy each time — mutating the result does not affect later calls."""
    if name not in _DATASET_NAMES:
        available = ", ".join(_DATASET_NAMES) or "(none)"
        raise KeyError(f"unknown dataset {name!r}; available: {available}")
    if name not in _cache:
        _cache[name] = pd.read_parquet(DATA_PATH / f"{name}.parquet")
    return _cache[name].copy()


class _SampleData(Mapping):
    """Bundled datasets, loaded from parquet on first access."""

    def __getitem__(self, name):
        return sample(name)

    def __contains__(self, name):
        return name in _DATASET_NAMES

    def __iter__(self):
        return iter(_DATASET_NAMES)

    def __len__(self):
        return len(_DATASET_NAMES)

    def keys(self):
        """Dataset names as a plain tuple (not a KeysView) so it's printable/readable
        directly, e.g. print(pytae.sample_data.keys())."""
        return _DATASET_NAMES


sample_data = _SampleData()


def _bind_shape():
    from .shape import long, pivot, wide
    globals()["long"] = long
    globals()["wide"] = wide
    globals()["pivot"] = pivot
    return long, wide, pivot


def __getattr__(name):
    if name in ("Plotter", "plot", "finalize"):
        try:
            from .plotting import Plotter, finalize, plot
        except ImportError as exc:
            raise ImportError(
                f"{name} requires matplotlib. Install with: pip install 'pytae[plot]'"
            ) from exc
        if name == "Plotter":
            return Plotter
        elif name == "plot":
            return plot
        else:
            return finalize
    if name in ("long", "wide", "pivot"):
        _bind_shape()
        return globals()[name]
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


__all__ = [
    "sample_data",
    "sample",
    "Plotter",
    "plot",
    "finalize",
    "select",
    "qry",
    "mutate",
    "agg",
    "agg_df",
    "long",
    "wide",
    "pivot",
    "handle_missing",
    "cols",
    "clean_columns",
    "replace_values",
    "sql",
    "everything",
    "PtAccessor",
]

################################################################################
# FILE: src/pytae/__main__.py
################################################################################
"""Executable module entrypoint for pytae: python -m pytae ..."""

from __future__ import annotations

import sys

from pytae.cli import main

if __name__ == "__main__":
    sys.exit(main())

################################################################################
# FILE: src/pytae/_text.py
################################################################################
"""Quote-aware string splitting shared by mutate() and the CLI parser."""

from __future__ import annotations


def tokenize(raw: str, seps: str, *, keep_quotes: bool = False, track_brackets: bool = False) -> list[str]:
    """Split raw into segments at top-level occurrences of any character in `seps`,
    respecting quotes (a matched quote pair is never split inside) and, if
    track_brackets, (), [], {} nesting depth. keep_quotes controls whether the quote
    characters themselves are kept in each segment's text (needed by callers that
    re-scan a segment for a second, nested split) or dropped as they're consumed.
    Every segment is returned (including empty ones) with surrounding whitespace
    stripped -- callers filter/validate as needed.
    """
    segments: list[str] = []
    buf: list[str] = []
    quote: str | None = None
    depth = 0
    for ch in raw:
        if quote:
            if ch == quote:
                quote = None
                if keep_quotes:
                    buf.append(ch)
            else:
                buf.append(ch)
            continue
        if ch in "'\"":
            quote = ch
            if keep_quotes:
                buf.append(ch)
        elif track_brackets and ch in "([{":
            depth += 1
            buf.append(ch)
        elif track_brackets and ch in ")]}":
            depth -= 1
            buf.append(ch)
        elif depth == 0 and ch in seps:
            segments.append("".join(buf).strip())
            buf = []
        else:
            buf.append(ch)
    segments.append("".join(buf).strip())
    return segments


def _is_enclosed_pair(raw: str, open_ch: str, close_ch: str) -> bool:
    if len(raw) < 2 or raw[0] != open_ch or raw[-1] != close_ch:
        return False
    if open_ch == close_ch:
        for idx in range(1, len(raw)):
            if raw[idx] == open_ch:
                return idx == len(raw) - 1
        return False
    depth = 0
    for idx, ch in enumerate(raw):
        if ch == open_ch:
            depth += 1
        elif ch == close_ch:
            depth -= 1
            if depth == 0:
                return idx == len(raw) - 1
    return False


def unquote_name(raw: str) -> str:
    raw = raw.strip()
    changed = True
    while changed and len(raw) >= 2:
        changed = False
        if raw[0] in "'\"" and _is_enclosed_pair(raw, raw[0], raw[0]):
            raw = raw[1:-1].strip()
            changed = True
        elif raw[0] == "[" and _is_enclosed_pair(raw, "[", "]"):
            raw = raw[1:-1].strip()
            changed = True
    return raw

################################################################################
# FILE: src/pytae/accessor.py
################################################################################
"""DataFrame accessor `df.pt` — mix with pandas methods in one chain.

    df.rename(columns={...}).pt.agg_df(a="mean")
    df.pt.select("species", contains="bill").head()
    df.pt.select("species", "body_mass_g").pt.agg_df(a=["mean", "n"])
    df.pt.sql("select species, avg(body_mass_g) from data group by species")
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import pandas as pd

from .agg_df import _UNSET
from .agg_df import agg_df as _agg_df
from .mutate import mutate as _mutate
from .other_utilities import (
    clean_columns as _clean_columns,
)
from .other_utilities import (
    cols as _cols,
)
from .other_utilities import (
    handle_missing as _handle_missing,
)
from .other_utilities import (
    replace_values as _replace_values,
)
from .qry import qry as _qry
from .select import select as _select
from .sql import sql as _sql


@pd.api.extensions.register_dataframe_accessor("pt")
class PtAccessor:
    """Pandas DataFrame accessor (`df.pt`) providing ergonomic tabular verbs:
    - select: Column selection, slicing, negative exclusion, and type filtering.
    - qry: Row filtering via string expressions, dicts, intervals, and kwargs.
    - mutate: Feature engineering, formulas, and grouped transforms (`by=`).
    - agg_df / agg: Grouped aggregations (`by=`) and whole-table summaries (`by=None`).
    - long / wide: Reshaping (melting to rows and pivoting to columns).
    - sql: DuckDB SQL queries over the DataFrame (as table `data`).
    - clean_columns: Standardizing and normalizing column header names.
    - replace_values: Cell value replacement (exact or substring).
    - handle_missing: Missing value imputation with type-safe defaults.
    - cols: Column names sorted alphabetically or preserved in order.
    - plot: Method-chainable visualization engine with faceting and mosaics.
    """

    def __init__(self, pandas_obj: pd.DataFrame):
        self._obj = pandas_obj

    def select(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
        """Select, reorder, or exclude columns by name, pattern, slice, or dtype.

        Supports exact names, negative prefixes ('-col', '~col'), negative slices
        ('-start:end'), regex, dtype filters ('dtype=numeric'), and 'exclude=' kwargs.
        """
        return _select(self._obj, *args, **kwargs)

    def qry(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
        """Filter rows using human-readable conditions, dicts, or kwargs.

        Supports comparisons ('>', '<=', '=='), intervals ('[min, max]'), set
        membership (['a', 'b']), string operations ('startswith', 'contains'),
        and safe handling of column names with spaces.
        """
        return _qry(self._obj, *args, **kwargs)

    def mutate(
        self,
        *args: Any,
        by: str | Sequence[str] | None = None,
        dropna: bool = False,
        observed: bool = True,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Create or overwrite columns using formulas, callables, or grouped window transforms.

        Expressions evaluate sequentially via pandas eval() with automatic fallback.
        When `by=` is given, aggregations (mean, sum, n, etc.) evaluate per group and broadcast
        back to each row without collapsing the dataset.
        """
        return _mutate(self._obj, *args, by=by, dropna=dropna, observed=observed, **kwargs)

    def sql(self, query: str, /, **frames: pd.DataFrame) -> pd.DataFrame:
        """Run DuckDB SQL queries over this DataFrame (registered as table `data`).

        Additional DataFrames can be registered as keyword arguments. Spaced column
        names can use brackets `[col]`, double quotes `"col"`, or backticks.
        """
        return _sql(self._obj, query, **frames)

    def agg_df(
        self,
        by: str | Sequence[str] | None = _UNSET,  # type: ignore[assignment]
        *args: Any,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Aggregate numeric columns grouped by explicit `by` column(s).

        Pass `by=None` for a whole-table summary (1 row). Supports whole-frame
        functions ('mean', ['mean', 'n']), column mapping strings with brackets
        ('tip = mean, [total bill] = mean, n = n'), and keyword arguments.
        """
        return _agg_df(self._obj, by, *args, **kwargs)

    def agg(
        self,
        by: str | Sequence[str] | None = _UNSET,  # type: ignore[assignment]
        *args: Any,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Alias for `df.pt.agg_df()`: aggregate numeric columns grouped by `by` column(s)."""
        return _agg_df(self._obj, by, *args, **kwargs)

    def long(
        self,
        cols: str | Sequence[str] | None = None,
        id_vars: str | Sequence[str] | None = None,
        c: str = "variable",
        v: str = "value",
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Melt columns to rows (unpivot wide format into long format).

        Default melts all numeric columns, storing variable names in `c` (default 'variable')
        and values in `v` (default 'value').
        """
        from .shape import long as _long
        return _long(self._obj, cols=cols, id_vars=id_vars, c=c, v=v, **kwargs)

    def wide(
        self,
        c: str = "variable",
        v: str = "value",
        index: str | Sequence[str] | None = None,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Pivot long-form records back to wide format (columns).

        Headers are created from `c` (default 'variable'), cells from `v` (default 'value').
        Strictly 1-to-1 structural reshaping without aggregation.
        For multi-dimensional aggregation, use `df.pt.pivot()`.
        """
        from .shape import wide as _wide
        return _wide(self._obj, c=c, v=v, index=index, **kwargs)

    def pivot(
        self,
        r: str | Sequence[str] | None = None,
        c: str | Sequence[str] | None = None,
        v: str | Sequence[str] | None = None,
        a: str = "sum",
        dropna: bool = False,
        fill_value: Any = None,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """Summarize and aggregate DataFrame across dimensions (Excel-style pivot table).

        Parameters:
        -----------
        r : str or sequence of str, optional
            Row dimension(s) to group by. Aliases: `rows`, `row`, `index`, `by`.
        c : str or sequence of str, optional
            Column dimension(s) to spread as headers. Aliases: `cols`, `col`, `columns`.
        v : str or sequence of str
            Value column(s) to aggregate. REQUIRED. Aliases: `values`, `value`, `val`, `vals`.
        a : str, default 'sum'
            Aggregation function ('sum', 'mean', 'median', 'min', 'max', 'count', 'std', etc.).
            'n' is accepted as an alias for 'size' (group row count). Aliases: `agg`, `aggfunc`.
        dropna : bool, default False
            Whether to drop NA categories from grouping keys.
        fill_value : any, optional
            Value to replace missing grid intersections with (e.g. 0).

        Returns:
        --------
        pd.DataFrame
            Clean, flattened DataFrame with standard RangeIndex and 1D column names.
        """
        from .shape import pivot as _pivot
        return _pivot(self._obj, r=r, c=c, v=v, a=a, dropna=dropna, fill_value=fill_value, **kwargs)

    def handle_missing(
        self,
        fillna: str = ".",
        numeric_fill: Any = 0,
        cols: Sequence[str] | None = None,
        preserve_categories: bool = True,
    ) -> pd.DataFrame:
        """Fill missing values with type-appropriate defaults across columns."""
        return _handle_missing(
            self._obj,
            fillna=fillna,
            numeric_fill=numeric_fill,
            cols=cols,
            preserve_categories=preserve_categories,
        )

    def cols(self, ascending: bool | None = True) -> list:
        """Return the column names sorted alphabetically (True/False) or in original order (None)."""
        return _cols(self._obj, ascending=ascending)

    def clean_columns(self, **kwargs: Any) -> pd.DataFrame:
        """Clean and standardize column header names (strip, squeeze, fill, case, dedupe)."""
        return _clean_columns(self._obj, **kwargs)

    def replace_values(
        self, v: dict[Any, Any], c: str | Sequence[str] | None = None, exact: bool = True
    ) -> pd.DataFrame:
        """Replace cell values via `{old: new}` mapping, optionally scoped to columns `c`."""
        return _replace_values(self._obj, v, c=c, exact=exact)

    def plot(self, *args: Any, **kwargs: Any) -> Any:
        """Create a pytae Plotter pre-loaded with this DataFrame.

        Can be called with layout parameters to start a pipeline:
            df.pt.plot(mosaic="AB", figsize=(10, 5)).plot(on="A", ...)

        Or with direct plot arguments to plot immediately:
            df.pt.agg("day", total_bill="mean").pt.plot(kind="bar", x="day", y="total_bill").pt.finalize()

        Or with faceting parameters (e.g. by="species", ncols=3) to automatically
        produce a small-multiples grid across groups:
            df.pt.plot(by="species", ncols=3, kind="scatter", x="a", y="b").pt.finalize()
        """
        try:
            from .plotting import Plotter
        except ImportError as exc:
            raise ImportError(
                "df.pt.plot() requires matplotlib. Install with: pip install 'pytae[plot]'"
            ) from exc

        by_col = kwargs.get("by") if kwargs.get("by") is not None else kwargs.get("col")
        has_facet_trigger = (
            kwargs.get("ncols") is not None
            or kwargs.get("facet") is True
            or "col" in kwargs
        )
        is_plot_call = bool(
            {"kind", "x", "y", "column"}.intersection(kwargs.keys())
        )

        if (
            isinstance(self._obj, pd.DataFrame)
            and by_col is not None
            and by_col in self._obj
            and has_facet_trigger
            and is_plot_call
        ):
            if "mosaic" in kwargs or (len(args) > 0 and args[0] is not None):
                raise ValueError("Cannot combine 'mosaic' with faceting ('by' with 'ncols'/'facet'/'col').")
            if "nrows" in kwargs:
                raise ValueError("Cannot combine 'nrows' with faceting ('by' with 'ncols'/'facet'/'col'). Use 'ncols' to control grid columns.")

            facet_kwargs = dict(kwargs)
            facet_kwargs.pop("by", None)
            facet_kwargs.pop("col", None)
            facet_kwargs.pop("facet", None)
            ncols = facet_kwargs.pop("ncols", None)
            return Plotter.facet(self._obj, by=by_col, ncols=ncols, **facet_kwargs)

        save_path = kwargs.pop("save", None)
        plotter_keys = {"mosaic", "figsize", "aggregate", "sharex", "sharey", "nrows", "ncols"}
        init_kwargs = {k: v for k, v in kwargs.items() if k in plotter_keys}
        plot_kwargs = {k: v for k, v in kwargs.items() if k not in plotter_keys}

        plotter = Plotter(*args, **init_kwargs)
        plotter.data(self._obj)
        if plot_kwargs:
            plotter.plot(**plot_kwargs)
        if save_path:
            plotter.finalize().save(save_path)
        return plotter

################################################################################
# FILE: src/pytae/agg_df.py
################################################################################
from __future__ import annotations

import difflib
from collections.abc import Sequence
from typing import Any

import pandas as pd

from pytae._text import tokenize as _tokenize
from pytae._text import unquote_name as _unquote_name
from pytae.cli_parsing import parse_agg as _parse_agg

_UNSET = object()


def _safe_parse_agg(raw: str) -> Any:
    try:
        return _parse_agg(raw)
    except SystemExit as exc:
        msg = str(exc)
        if msg.startswith("-agg: "):
            msg = f"agg(): {msg[6:]}"
        raise ValueError(msg) from exc
_KNOWN_AGGS = {
    "mean",
    "sum",
    "std",
    "var",
    "min",
    "max",
    "count",
    "size",
    "median",
    "first",
    "last",
    "n",
    "prod",
    "sem",
    "skew",
    "kurt",
}


def _agg_df_list(
    df: pd.DataFrame,
    group_cols: list[str],
    agg_types: str | list[str],
    dropna: bool,
    observed: bool,
) -> pd.DataFrame:
    """Helper function to handle string/list aggfunc for agg_df.

    Applies aggregations to all numeric columns (excluding group_cols).
    """
    agg_types = [agg_types] if isinstance(agg_types, str) else list(agg_types)
    unique_agg_types = list(dict.fromkeys(agg_types))  # Remove duplicates, preserve order
    remaining_agg_types = [agg for agg in unique_agg_types if agg != "n"]
    has_n = "n" in unique_agg_types

    numeric_cols = [col for col in df.select_dtypes(include=["number"]).columns if col not in group_cols]

    if has_n and "n" in numeric_cols:
        raise ValueError(
            "agg_df: cannot compute 'n' (row count) because the input already has a "
            "numeric column named 'n'; rename that column first, or use the kwargs form "
            "(e.g. df.pt.agg(by='...', n_col='n')) to pick a different count-column name."
        )
    if has_n and "n" in group_cols:
        raise ValueError(
            "agg_df: cannot compute 'n' (row count) because the input already has a "
            "group column named 'n'; rename that column first, or use the kwargs form "
            "(e.g. df.pt.agg(by='...', n_col='n')) to pick a different count-column name."
        )

    if not numeric_cols and not has_n:
        raise ValueError("No numeric columns to aggregate and 'n' not specified")

    if group_cols:
        if has_n and not numeric_cols:
            grouped_df = df.groupby(group_cols, dropna=dropna, observed=observed).size().reset_index(name="n")
            return grouped_df

        agg_operations = {col: remaining_agg_types for col in numeric_cols}
        try:
            grouped_df = df.groupby(group_cols, as_index=False, dropna=dropna, observed=observed).agg(agg_operations)
        except (AttributeError, ValueError) as exc:
            raise ValueError(f"agg(): invalid aggregation function ({exc})") from exc

        # Flatten MultiIndex in columns
        if len(remaining_agg_types) > 1:
            flattened = [
                col[0] if (len(col) > 1 and not col[1]) else (f"{col[0]}_{col[1]}" if len(col) > 1 else str(col[0]))
                for col in grouped_df.columns.values
            ]
        else:
            flattened = [col[0] for col in grouped_df.columns.values]

        non_group_cols = [c for i, c in enumerate(flattened) if i >= len(group_cols)]
        for col in non_group_cols:
            if col in group_cols:
                raise ValueError(f"agg_df: output column name '{col}' collides with group column '{col}'")
        if len(non_group_cols) != len(set(non_group_cols)):
            counts: dict[str, int] = {}
            for col in non_group_cols:
                counts[col] = counts.get(col, 0) + 1
            dups = [col for col, count in counts.items() if count > 1]
            raise ValueError(f"agg_df: output column name '{dups[0]}' is duplicated in aggregation output")
        if has_n and "n" in non_group_cols:
            raise ValueError("agg_df: output column name 'n' collides with count column 'n'")

        grouped_df.columns = flattened

        if has_n:
            grouped_df["n"] = df.groupby(group_cols, dropna=dropna, observed=observed).size().values

        g_cols = group_cols + (["n"] if has_n else [])
        remaining_cols = [col for col in grouped_df.columns if col not in g_cols]
        return grouped_df.reindex(columns=g_cols + remaining_cols)
    else:
        # Whole-table summary (1 row, no groups)
        if has_n and not numeric_cols:
            return pd.DataFrame({"n": [len(df)]})

        row_dict: dict[str, list[Any]] = {}
        if has_n:
            row_dict["n"] = [len(df)]

        for col in numeric_cols:
            for agg in remaining_agg_types:
                col_name = f"{col}_{agg}" if len(remaining_agg_types) > 1 else col
                if col_name in row_dict:
                    raise ValueError(f"agg_df: output column name '{col_name}' is duplicated in aggregation specification")
                s = df[col]
                try:
                    val = getattr(s, agg)() if hasattr(s, agg) and callable(getattr(s, agg)) else s.agg(agg)
                except Exception as exc:
                    raise ValueError(f"agg(): invalid aggregation function {agg!r} ({exc})") from exc
                row_dict[col_name] = [val]

        return pd.DataFrame(row_dict)


def _agg_df_dict(
    df: pd.DataFrame,
    group_cols: list[str],
    agg_types: dict[str, Any],
    dropna: bool,
    observed: bool,
) -> pd.DataFrame:
    """Helper function to handle keyword/dictionary aggfunc for agg_df.

    Applies aggregations to specified columns, with 'n' keys used for count column names.
    Supports named aggregations like total='v1:sum' or total=('v1', 'sum').
    Output columns follow dictionary order after group columns.
    """
    output_cols = []
    named_aggs: dict[str, tuple[str, str]] = {}
    count_cols = []

    for out_name, spec in agg_types.items():
        clean_out_name = _unquote_name(out_name)
        if spec == "n" or spec == ["n"] or spec == ("n",):
            if clean_out_name in group_cols:
                raise ValueError(
                    f"agg_df: count output name '{clean_out_name}' collides with group column '{clean_out_name}'"
                )
            if clean_out_name in output_cols:
                raise ValueError(f"agg_df: output column name '{clean_out_name}' is duplicated in aggregation specification")
            count_cols.append(clean_out_name)
            output_cols.append(clean_out_name)
            continue

        if isinstance(spec, str) and ":" in spec:
            src_col, aggfunc = spec.split(":", 1)
            src_col, aggfunc = _unquote_name(src_col.strip()), aggfunc.strip()
            aggfunc_clean = _unquote_name(aggfunc)
            if "," in aggfunc_clean:
                aggs_list = [s.strip() for s in aggfunc_clean.split(",") if s.strip()]
            else:
                aggs_list = [aggfunc_clean]
        elif isinstance(spec, tuple) and len(spec) == 2 and _unquote_name(spec[0]) in df.columns:
            src_col = _unquote_name(spec[0])
            aggfunc = spec[1]
            aggs_list = [aggfunc] if isinstance(aggfunc, str) else list(aggfunc)
        else:
            src_col = clean_out_name
            if isinstance(spec, str):
                spec_clean = _unquote_name(spec)
                if "," in spec_clean:
                    aggs_list = [s.strip() for s in spec_clean.split(",") if s.strip()]
                else:
                    aggs_list = [spec]
            elif isinstance(spec, (list, tuple, set)):
                aggs_list = list(spec)
            else:
                aggs_list = [spec]

        aggs_list = list(dict.fromkeys(aggs_list))  # preserve order
        if "n" in aggs_list:
            raise ValueError(f"'n' cannot be used as an aggregation function for column '{clean_out_name}' in a dictionary aggfunc")
        if src_col not in df.columns:
            raise KeyError(f"Column '{src_col}' does not exist in DataFrame")
        if src_col in group_cols:
            raise ValueError(f"Column '{src_col}' is already a grouping column in 'by'")
        if not pd.api.types.is_numeric_dtype(df[src_col]):
            raise ValueError(f"Column '{src_col}' is not numeric")

        if not aggs_list:
            raise ValueError(f"No valid aggregation functions specified for column '{src_col}'")

        def _record_output_col(name: str, src: str, fn: str) -> None:
            if name in group_cols:
                raise ValueError(f"agg_df: output column name '{name}' collides with group column '{name}'")
            if name in output_cols:
                raise ValueError(f"agg_df: output column name '{name}' is duplicated in aggregation specification")
            output_cols.append(name)
            named_aggs[name] = (src, fn)

        if src_col == clean_out_name:
            if len(aggs_list) > 1:
                for agg_fn in aggs_list:
                    _record_output_col(f"{src_col}_{agg_fn}", src_col, agg_fn)
            else:
                _record_output_col(src_col, src_col, aggs_list[0])
        else:
            if len(aggs_list) > 1:
                for agg_fn in aggs_list:
                    _record_output_col(f"{clean_out_name}_{agg_fn}", src_col, agg_fn)
            else:
                _record_output_col(clean_out_name, src_col, aggs_list[0])


    if group_cols:
        if not named_aggs and count_cols:
            grouped_df = df.groupby(group_cols, dropna=dropna, observed=observed).size().reset_index(name="n")
            result = grouped_df[group_cols].copy()
            for count_col in count_cols:
                result[count_col] = grouped_df["n"]
            return result.reindex(columns=group_cols + count_cols)

        if not named_aggs:
            raise ValueError("No valid numeric aggregations specified")

        grouped_df = df.groupby(group_cols, as_index=False, dropna=dropna, observed=observed).agg(**named_aggs)

        if count_cols:
            counts = df.groupby(group_cols, dropna=dropna, observed=observed).size().values
            for count_col in count_cols:
                grouped_df[count_col] = counts

        final_cols = group_cols + output_cols
        return grouped_df.reindex(columns=final_cols)
    else:
        # Whole-table summary (1 row, no groups)
        if not named_aggs and not count_cols:
            raise ValueError("No aggregations specified")

        row_dict: dict[str, list[Any]] = {}
        for count_col in count_cols:
            row_dict[count_col] = [len(df)]

        for out_name, (src_col, agg_fn) in named_aggs.items():
            s = df[src_col]
            val = getattr(s, agg_fn)() if hasattr(s, agg_fn) and callable(getattr(s, agg_fn)) else s.agg(agg_fn)
            row_dict[out_name] = [val]

        df_out = pd.DataFrame(row_dict)
        return df_out.reindex(columns=output_cols)


def agg_df(
    df: pd.DataFrame,
    by: str | Sequence[str] | None = _UNSET,  # type: ignore[assignment]
    *args: Any,
    **kwargs: Any,
) -> pd.DataFrame:
    """Aggregate a DataFrame grouped by explicit `by` column(s).

    Parameters:
    -----------
    df : pd.DataFrame
        The pandas DataFrame to aggregate.
    by : str, Sequence[str], or None
        Grouping column name(s). Pass None to perform a whole-table summary (grand total, 1 row).
        Can be passed positionally as the first argument or via keyword `by=` / `group_by=`.
    *args : str or list of str
        Aggregation specification(s):
        - Column mapping string(s) (e.g., 'tip = mean, [total bill] = mean, n = n' or
          'tip = mean', '[total bill] = mean', 'n = n'). Brackets [...] enclose columns with spaces.
        - Whole-frame aggregation function(s) applied to all numeric columns:
          - If str (e.g., 'mean', or 'mean, n'): Apply to all numeric columns.
          - If list (e.g., ['mean', 'n']): Apply listed aggregations. 'n' computes group row counts.
    **kwargs :
        - Column aggregations passed as keyword arguments (e.g. df.pt.agg('species', body_mass_g='mean', count='n')).
          Keys with value 'n' specify the output column name for group row counts.
          Supports named aggregations like total='v1:sum'.
        - a (str or list, optional): Whole-frame or mapping aggregation function(s) when passed as keyword `a=`.
        - dropna (bool): Whether to drop NA values in groupby. Defaults to False.
        - observed (bool): Whether to show only observed values for categorical groupby columns. Defaults to True.

    Returns:
    --------
    pd.DataFrame
        Aggregated DataFrame with group columns, count 'n' (if requested), and aggregated metrics.
    """
    if "group_by" in kwargs:
        if "by" in kwargs or by is not _UNSET:
            raise TypeError("agg() got multiple values for argument 'by'/'group_by'")
        by = kwargs.pop("group_by")
    elif "by" in kwargs:
        if by is not _UNSET:
            raise TypeError("agg() got multiple values for argument 'by'")
        by = kwargs.pop("by")

    if by is _UNSET:
        raise TypeError(
            "agg() missing required argument: 'by'. "
            "Specify grouping column(s) as first argument (e.g. df.pt.agg('species', 'mean')), "
            "or None for whole-table summary (e.g. df.pt.agg(None, 'mean'))."
        )

    if isinstance(by, dict):
        raise TypeError(
            "agg() no longer accepts dictionaries. Pass column aggregations as string mapping "
            "(e.g. df.pt.agg('by_col', 'col = mean, n = n')) or keyword arguments: "
            "df.pt.agg('by_col', body_mass_g='mean', count='n')."
        )

    # Check if user mistakenly passed an aggregation name as 'by'
    if isinstance(by, str):
        clean_by = _unquote_name(by)
        if clean_by in df.columns:
            by = clean_by
        else:
            tokens = [_unquote_name(t.strip()) for t in _tokenize(by, ",", keep_quotes=True, track_brackets=True) if t.strip()]
            if len(tokens) > 1:
                by = tokens
            else:
                if by.lower() in _KNOWN_AGGS:
                    raise ValueError(
                        f"agg(): '{by}' is not a column in DataFrame, but is a known aggregation function. "
                        "The first argument to agg() must be 'by' (group column(s), or None for whole-table summary). "
                        f"Did you mean: df.pt.agg(None, {by!r}) or df.pt.agg(by='col', a={by!r})?"
                    )
                close = difflib.get_close_matches(clean_by, df.columns, n=1)
                hint = f" (did you mean '{close[0]}'?)" if close else ""
                raise KeyError(f"agg(): group column '{by}' not found in DataFrame{hint}")

    if isinstance(by, (list, tuple)):
        by = [_unquote_name(x) if isinstance(x, str) else x for x in by]
        if by and all(isinstance(x, str) and x.lower() in _KNOWN_AGGS for x in by) and not any(x in df.columns for x in by):
            raise ValueError(
                f"agg(): {list(by)!r} looks like an aggregation list, but the first argument to agg() must be 'by' "
                "(group column(s), or None for whole-table summary). "
                f"Did you mean: df.pt.agg(None, {list(by)!r}) or df.pt.agg(by='col', a={list(by)!r})?"
            )
        for col in by:
            if col not in df.columns:
                close = difflib.get_close_matches(str(col), df.columns, n=1)
                hint = f" (did you mean '{close[0]}'?)" if close else ""
                raise KeyError(f"agg(): group column '{col}' not found in DataFrame{hint}")


    if by is None:
        group_cols: list[str] = []
    elif isinstance(by, str):
        group_cols = [by]
    elif isinstance(by, (list, tuple)):
        group_cols = list(by)
    else:
        raise TypeError(f"'by' must be a column name (str), list of column names, or None, got {type(by).__name__}")

    # Extract aggregation spec
    col_kwargs = {k: v for k, v in kwargs.items() if k not in ("dropna", "observed", "a")}
    has_a = "a" in kwargs
    a_arg = kwargs.pop("a", None)

    if args and has_a:
        raise TypeError("agg() got multiple values for aggregation spec ('a' and positional args)")

    agg_types: Any
    if args:
        if isinstance(args[0], dict):
            raise TypeError(
                "agg() no longer accepts dictionaries. Pass column aggregations as string mapping "
                "(e.g. df.pt.agg('by_col', 'col = mean, n = n')) or keyword arguments: "
                "df.pt.agg('by_col', body_mass_g='mean', count='n')."
            )
        if all(isinstance(x, str) for x in args):
            raw_spec = ", ".join(args)
            parsed = _safe_parse_agg(raw_spec)
            if isinstance(parsed, dict):
                if col_kwargs:
                    parsed.update(col_kwargs)
                agg_types = parsed
            else:
                if col_kwargs:
                    raise ValueError(
                        "agg(): cannot mix whole-frame aggregation string with column keyword arguments; "
                        "use one or the other"
                    )
                agg_types = parsed
        elif len(args) == 1 and isinstance(args[0], (list, tuple)):
            if col_kwargs:
                raise ValueError(
                    "agg(): cannot mix whole-frame aggregation list with column keyword arguments; "
                    "use one or the other"
                )
            agg_types = list(args[0])
        elif len(args) == 1:
            agg_types = args[0]
        else:
            raise ValueError(f"agg(): unsupported positional arguments {args!r}")
    elif has_a:
        if isinstance(a_arg, dict):
            raise TypeError(
                "agg() no longer accepts dictionaries. Pass column aggregations as string mapping "
                "(e.g. df.pt.agg('by_col', 'col = mean, n = n')) or keyword arguments: "
                "df.pt.agg('by_col', body_mass_g='mean', count='n')."
            )
        if isinstance(a_arg, str):
            parsed = _safe_parse_agg(a_arg)
            if isinstance(parsed, dict):
                if col_kwargs:
                    parsed.update(col_kwargs)
                agg_types = parsed
            else:
                if col_kwargs:
                    raise ValueError(
                        "agg(): cannot mix whole-frame aggregation string with column keyword arguments; "
                        "use one or the other"
                    )
                agg_types = parsed
        elif isinstance(a_arg, (list, tuple)):
            if col_kwargs:
                raise ValueError(
                    "agg(): cannot mix whole-frame aggregation list with column keyword arguments; "
                    "use one or the other"
                )
            agg_types = list(a_arg)
        else:
            agg_types = a_arg
    else:
        agg_types = col_kwargs if col_kwargs else ["sum"]

    dropna = kwargs.get("dropna", False)
    observed = kwargs.get("observed", True)

    if isinstance(agg_types, dict):
        return _agg_df_dict(df, group_cols, agg_types, dropna, observed)
    else:
        return _agg_df_list(df, group_cols, agg_types, dropna, observed)


# Alias: agg is identical to agg_df
agg = agg_df

################################################################################
# FILE: src/pytae/cli.py
################################################################################
"""Command line tool for inspecting tabular files (parquet/csv/txt/sas7bdat) and converting between formats."""

from __future__ import annotations

import argparse
import sys
from collections.abc import Sequence
from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as _pkg_version
from pathlib import Path

import pandas as pd

from pytae.cli_parsing import (
    expand_paths,
    parse_bool_text,
    parse_concat_arg,
    parse_file_arg,
    parse_fraction,
    parse_list_order,
    parse_merge_arg,
    parse_positive_int,
    parse_qry,
    parse_rename,
    parse_replace_values_arg,
    parse_select_spec,
)
from pytae.cli_pipeline import (
    _OrderedAppend,
    _OrderedFlag,
    _OrderedStore,
    _OrderedValue,
)
from pytae.cli_run import (
    NON_DF_TERMINAL_OPS,
    _encoding_error_message,
    _fail,
    _process_path,
)
from pytae.readers import get_reader

try:
    __version__ = _pkg_version("pytae")
except PackageNotFoundError:
    __version__ = "0.0.0-dev"

_DATA_SUFFIXES = frozenset((".parquet", ".pq", ".csv", ".txt", ".dat", ".sas7bdat", ".jsonl", ".ndjson", ".gz", ".yaml", ".yml"))
_FORMAT_SHORTHANDS = frozenset((
    "csv", "parquet", "pq", "txt", "dat", "jsonl", "ndjson",
    "csv.gz", "txt.gz", "dat.gz", "jsonl.gz", "ndjson.gz",
))


def _is_path_like(token: str) -> bool:
    if token == "-":
        return True
    if token.startswith("-"):
        return False
    if any(ch in token for ch in "*?[]"):
        return True
    if "/" in token or "\\" in token:
        return True
    p = Path(token)
    if p.exists():
        return True
    if p.suffix.lower() in _DATA_SUFFIXES:
        return True
    return False


_CLI_FLAGS = {
    "-version", "--version", "-head", "--head", "-tail", "--tail", "-shape", "--shape",
    "-cols", "--cols", "-dtype", "--dtype", "-nulls", "--nulls", "-describe", "--describe",
    "-info", "--info", "-meta", "--meta", "-diff", "--diff", "-value_counts", "--value_counts",
    "-unique", "--unique", "-sample", "--sample", "-seed", "--seed", "-frac", "--frac",
    "-sort_by", "--sort_by", "-by", "--by", "-group_by", "--group_by", "-nrows", "--nrows",
    "-limit", "--limit", "-select", "--select", "-agg", "--agg",
    "-agg_df", "--agg_df", "-handle_missing", "--handle_missing",
    "-clean_columns", "--clean_columns", "-long", "--long", "-wide", "--wide", "-pivot", "--pivot",
    "-dropna", "--dropna", "-o", "--output", "-out_dir", "--out_dir",
    "-od", "--out-dir", "-dlim", "--dlim", "-encoding", "--encoding", "-rename", "--rename",
    "-file", "--file", "-query", "--query", "-qry", "--qry", "-mutate", "--mutate",
    "-sql", "--sql", "-replace_values", "--replace_values", "-merge", "--merge",
    "-concat", "--concat", "-progress", "--progress", "-pretty", "--pretty",
    "-pager", "--pager", "-round", "--round", "-h", "--help",
    "-fmt", "--fmt", "-freq", "--freq", "-hist", "--hist",
    "-plot", "--plot", "-finalize", "--finalize",
}


def _normalize_cli_args(argv: list[str]) -> list[str]:
    out: list[str] = []
    i = 0
    while i < len(argv):
        arg = argv[i]
        if arg in ("-progress", "--progress"):
            if i + 1 < len(argv):
                nxt = argv[i + 1]
                try:
                    int(nxt)
                    out.append(arg)
                    out.append(nxt)
                    i += 2
                    continue
                except ValueError:
                    out.append("--progress=200000")
                    i += 1
                    continue
            else:
                out.append("--progress=200000")
                i += 1
                continue
        elif arg in ("-select", "--select") and "=" not in arg:
            if i + 1 < len(argv):
                nxt = argv[i + 1]
                if nxt.startswith("-") and nxt not in _CLI_FLAGS:
                    out.append(f"--select={nxt}")
                    i += 2
                    continue
        out.append(arg)
        i += 1
    return out


class PytaeParser(argparse.ArgumentParser):
    """Custom parser that normalizes optional-value flags such as -progress and negative -select specs."""

    def parse_known_args(  # type: ignore[override]
        self,
        args: Sequence[str] | None = None,
        namespace: argparse.Namespace | None = None,
    ) -> tuple[argparse.Namespace, list[str]]:
        if args is None:
            args = sys.argv[1:]
        normalized = _normalize_cli_args(list(args))
        return super().parse_known_args(normalized, namespace)


def build_parser() -> argparse.ArgumentParser:
    parser = PytaeParser(
        prog="pytae",
        description="Inspect and convert parquet/csv/txt/dat/sas7bdat files (glob patterns convert multiple files at once), or read a Databricks table / remote SSH file via a .yaml connection config.",
        allow_abbrev=False,
    )
    parser.add_argument("path", nargs="*", default=None,
                         help="path to a .parquet, .csv, .txt, .dat, or .sas7bdat file, "
                              "a .yaml/.yml connection config (Databricks table or remote SSH file), "
                              "or glob patterns/multiple files for batch operations; "
                              "omit when using -file with -merge/-concat/-sql")
    parser.add_argument("-version", "--version", action="version", version=f"%(prog)s {__version__}")
    parser.add_argument("-head", "--head", nargs="?", const=5, type=parse_positive_int, default=None, metavar="N",
                         action=_OrderedValue, help="print the first N rows (default 5)")
    parser.add_argument("-tail", "--tail", nargs="?", const=5, type=parse_positive_int, default=None, metavar="N",
                         action=_OrderedValue, help="print the last N rows (default 5)")
    parser.add_argument("-shape", "--shape", action=_OrderedFlag, help="print (rows, cols)")
    parser.add_argument("-cols", "--cols", nargs="?", const="file", default=None, metavar="ORDER",
                         type=parse_list_order, action=_OrderedValue,
                         help="print column names; optional asc|desc sorts names (default: file order)")
    parser.add_argument("-dtype", "--dtype", nargs="?", const="file", default=None, metavar="ORDER",
                         type=parse_list_order, action=_OrderedValue,
                         help="print column dtypes; optional asc|desc sorts by name (default: file order)")
    parser.add_argument("-nulls", "--nulls", nargs="?", const="file", default=None, metavar="ORDER",
                         type=parse_list_order, action=_OrderedValue,
                         help="print null/NaN count per column; optional asc|desc sorts by name (default: file order)")
    parser.add_argument("-describe", "--describe", dest="describe", action=_OrderedFlag,
                         help="print pandas describe() summary (count/mean/std/min/quartiles/max)")
    parser.add_argument("-info", "--info", dest="info", action=_OrderedFlag,
                         help="print pandas info() (columns, non-null counts, dtypes, memory)")
    parser.add_argument("-meta", "--meta", dest="meta", action=_OrderedFlag,
                         help="display Parquet metadata (row groups, column statistics, compression, schema) without loading data")
    parser.add_argument("-diff", "--diff", dest="diff", default=None, metavar="PATH", action=_OrderedStore,
                         help="compare schema and contents against another dataset file")
    parser.add_argument("-value_counts", "--value_counts", dest="value_counts", action=_OrderedFlag,
                         help="show pandas value_counts across current working columns (use -select to choose columns)")
    parser.add_argument("-freq", "--freq", dest="freq", action=_OrderedStore, default=None, metavar="COLUMN",
                         help="render horizontal frequency bars with counts and percentages for a categorical column")
    parser.add_argument("-hist", "--hist", dest="hist", action=_OrderedStore, default=None, metavar="COLUMN[:BINS]",
                         help="render an in-terminal distribution histogram for a numeric column (e.g. -hist mass or -hist mass:10)")
    parser.add_argument("-unique", "--unique", dest="unique", action=_OrderedFlag,
                         help="drop duplicate rows and print unique rows")
    parser.add_argument("-sample", "--sample", nargs="?", const=5, type=parse_positive_int, default=None, metavar="N",
                         action=_OrderedValue, help="print N randomly sampled rows (default 5)")
    parser.add_argument("-seed", "--seed", dest="seed", type=int, default=None, metavar="N",
                         help="random seed for -sample, for reproducible rows (default: random each run)")
    parser.add_argument("-frac", "--frac", dest="frac", type=parse_fraction, default=None, metavar="P",
                         help="sample a fraction of rows instead of a count (0 < P <= 1, e.g. 0.1 for 10 percent); "
                              "requires -sample and overrides its N")
    parser.add_argument("-sort_by", "--sort_by", dest="sort_by", default=None, metavar="SPEC",
                         action=_OrderedStore,
                         help="sort rows by a comma-separated column list, optionally ending with "
                              "asc or desc (default: asc), e.g. species,body_mass_g desc")
    parser.add_argument("-by", "--by", "-group_by", "--group_by", dest="by", default=None, metavar="COLUMNS",
                         help="grouping columns (comma-separated), e.g. -by species or -by 'species,island'. "
                              "Used by -agg and -mutate.")
    parser.add_argument("-nrows", "--nrows", "-limit", "--limit", dest="nrows", type=parse_positive_int, default=None, metavar="N",
                         help="cap the number of rows loaded (default: no cap)")
    parser.add_argument("--select", "-select", dest="select", action=_OrderedAppend, default=None, metavar="SPEC",
                         help="restrict columns at this point in the pipeline (union of tokens in one SPEC): "
                              "names, start:end slices, -negated names/slices (-col, ~col), and key=value (dtype, "
                              "exclude, contains, startswith, endswith, regex, exclude_dtype); repeat to filter remaining columns, "
                              'including after -agg/-long/-wide, e.g. -select "-species" or -select "contains=bill"')
    parser.add_argument("-agg", "--agg", "-agg_df", "--agg_df", dest="agg", nargs="?", const="sum", default=None,
                         metavar="AGGFUNC", action=_OrderedValue,
                         help="aggregate numeric columns using pytae agg; if -by is given, groups by those columns; "
                              "without -by, computes a whole-table summary (grand total). "
                              "Accepts a name (mean), a comma list (mean,sum,n), or column mappings "
                              "(v1 = sum, total = v1:sum, count = n). Defaults to sum.")
    parser.add_argument("-handle_missing", "--handle_missing", dest="handle_missing", nargs="?", const=".", default=None,
                         metavar="FILL", action=_OrderedValue,
                         help="fill NaN using pytae handle_missing(): FILL (default '.') for object/category "
                              "columns, 0 for numeric columns")
    parser.add_argument("-clean_columns", "--clean_columns", dest="clean_columns", action=_OrderedStore,
                         metavar="KEY=VALUE,...",
                         help="clean column header names, in order strip -> strip_special -> squeeze -> "
                              "fill -> case -> dedupe: strip/squeeze/strip_special/dedupe are bools (bare "
                              "key means true, e.g. \"strip\"), fill[=STR] replaces whitespace in a header "
                              "(bare \"fill\" defaults to '_', omit entirely for no fill), "
                              "case=lower|upper|proper (always needs a value); e.g. "
                              "\"strip,fill,case=lower\" or \"fill='$',case=proper,dedupe=true\"")
    parser.add_argument("-long", "--long", dest="long", nargs="?", const="", default=None,
                         metavar="KEY=VALUE,...", action=_OrderedValue,
                         help="melt numeric columns to long form (pytae long()); defaults c=variable, "
                              "v=value; e.g. c=metric,v=reading")
    parser.add_argument("-wide", "--wide", dest="wide", nargs="?", const="", default=None,
                         metavar="KEY=VALUE,...", action=_OrderedValue,
                         help="pivot long form to wide for 1-to-1 reshaping (pytae wide()); defaults c=variable, v=value; "
                              "for aggregation or 2D contingency tables, use -pivot")
    parser.add_argument("-pivot", "--pivot", dest="pivot", nargs="?", const="", default=None,
                         metavar="KEY=VALUE,...", action=_OrderedValue,
                         help="2-D pivot table (pytae pivot()); r=row_cols (optional), c=col_cols "
                              "(optional), v=value_cols (required), a=aggfunc (default sum; n=size), "
                              "dropna=false (default), fill_value=val; e.g. -pivot \"r=Region,c=Year,v=Sales,a=sum\"")
    parser.add_argument("-dropna", "--dropna", dest="dropna", type=parse_bool_text, default=None,
                         metavar="BOOL",
                         help="for -agg, -mutate, -value_counts, -freq, -pivot, -plot (default: false): "
                              "include NA keys when false; set true to drop NA keys; accepts true or false")
    parser.add_argument("-o", "--output", dest="output", default=None, metavar="TARGET",
                         help="output destination: a file path (e.g. 'out.csv', 'out.parquet'), "
                              "a format for in-place or batch conversion ('csv', 'parquet', 'txt', 'dat', 'jsonl', 'csv.gz', 'jsonl.gz'), "
                              "or 'clip'/'clipboard' to copy to system clipboard")
    parser.add_argument("-out_dir", "--out_dir", "-od", "--out-dir", dest="out_dir", type=Path, default=None, metavar="DIR",
                         help="target directory for exported files (created if it does not exist); requires -o/--output")
    parser.add_argument("-fmt", "--fmt", dest="fmt", default=None, metavar="FORMAT",
                         help="input format when reading from STDIN or extensionless files: csv, parquet, jsonl, txt, tsv")
    parser.add_argument("-dlim", "--dlim", dest="dlim", default=None, metavar="CHAR",
                         help="field delimiter for reading/writing .csv/.txt/.dat (default: ',' for .csv, "
                              "tab for .txt, '|' for .dat); not used for .parquet or .sas7bdat")
    parser.add_argument("-encoding", "--encoding", dest="encoding", default=None, metavar="ENC",
                         help="text encoding for .csv/.txt/.dat/.sas7bdat, e.g. latin-1 "
                              "(default: utf-8 for .sas7bdat, latin-1 for .dat, pandas infer for .csv/.txt); "
                              "not used for .parquet")
    parser.add_argument("-rename", "--rename", dest="rename", action=_OrderedAppend, default=None, metavar="OLD:NEW,...",
                         help="rename columns at this point in the pipeline (or during conversion), e.g. \"old_a:new_a,old_b:new_b\"")
    parser.add_argument("-file", "--file", dest="file", default=None, metavar="PATH=ALIAS;...",
                         help="load multiple named files for -merge/-concat/-sql, instead of the positional path; "
                              "';'-separated entries, each PATH=ALIAS optionally followed by "
                              ",dlim=/,encoding= overrides for that file, e.g. "
                              "\"data1.parquet=df1; data2.parquet=df2,encoding='latin-1'\"; "
                              "requires -merge, -concat, or -sql as the first operation, and can't be "
                              "combined with the positional path")
    parser.add_argument("-query", "--query", dest="query", action=_OrderedAppend, default=None, metavar="EXPR",
                         help="filter rows at this point in the pipeline using pandas query(), e.g. \"col > 5\"")
    parser.add_argument("-qry", "--qry", dest="qry", action=_OrderedAppend, default=None, metavar="CONDITIONS",
                         help="filter rows at this point in the pipeline using pytae qry(); conditions like "
                              "\"col = 'val', col > 5\"")
    parser.add_argument("-mutate", "--mutate", dest="mutate", action=_OrderedAppend, default=None, metavar="SPEC",
                         help="create/overwrite columns at this point in the pipeline using pytae mutate(); "
                              "\"new_col = expression\" entries, comma-separated, quoting the key optional; "
                              "the expression is pandas eval() syntax and column names in it must stay unquoted, e.g. "
                              "\"bmi = body_mass_g / bill_length_mm ** 2\"")
    parser.add_argument("-sql", "--sql", dest="sql", action=_OrderedAppend, default=None, metavar="QUERY",
                         help="run a SQL query (via duckdb) against the current view at this point in "
                              "the pipeline; the view is queryable as table `data`; standard SQL identifier "
                              "quoting applies (double quotes for names with spaces, e.g. \"col a\"; single "
                              "quotes are string literals, not identifiers), e.g. "
                              "\"select \\\"col a\\\" from data where \\\"col a\\\" > 10\"; "
                              "in -file/-merge mode, may be used instead of -merge as the first op, with "
                              "every -file alias queryable by its own name (e.g. \"select * from df1 "
                              "inner join df2 on df1.\\\"col a\\\" = df2.cola\") — `data` becomes queryable "
                              "too once something later in the pipeline has produced a current view")
    parser.add_argument("-replace_values", "--replace_values", dest="replace_values", action=_OrderedAppend, default=None, metavar="SPEC",
                         help="replace values at this point in the pipeline; key=value tokens: v= (required) "
                              "an old:new mapping, e.g. \"v='old:new,alpha:bravo'\"; c= (optional) restrict "
                              "to specific columns, e.g. \"c='col a,col b',v='old:new'\"; exact= (optional bool, "
                              "default true) — true matches whole cell values, false matches a substring "
                              "anywhere in the cell, e.g. \"v='old:new',exact=false\"")
    parser.add_argument("-merge", "--merge", dest="merge", action=_OrderedAppend, default=None, metavar="KEY=VALUE,...",
                         help="merge two frames into the pipeline (pandas merge()); repeatable, to fold in "
                              "one more file at a time; must be the first op when using -file (unless -sql/"
                              "-concat starts it instead); key=value tokens: left=/right= (required — -file "
                              "aliases, or 'data' for the pipeline's current result so far), "
                              "on= (required; shared column name(s), or 'left:right' pairs if they "
                              "differ between sides — quote on= if it has more than one column/pair, e.g. "
                              "\"on='col a:cola,colb:colb'\"), how= (optional, default 'inner': "
                              "inner/left/right/outer/cross), validate= (optional, e.g. one_to_one)")
    parser.add_argument("-concat", "--concat", dest="concat", action=_OrderedAppend, default=None, metavar="KEY=VALUE,...",
                         help="stack frames row-wise into the pipeline (pandas concat(), always with "
                              "ignore_index=True); repeatable; must be the first op when using -file "
                              "(unless -sql/-merge starts it instead); key=value tokens: frames= (required) "
                              "an ordered comma-separated list of -file aliases (or 'data' for the pipeline's "
                              "current result so far), quoted since it has internal commas, e.g. "
                              "\"frames='df1,df2,df3'\"")
    parser.add_argument("-plot", "--plot", dest="plot", action=_OrderedStore, default=None, metavar="SPEC",
                         help="render visualization using pytae Plotter; e.g. -plot \"kind=scatter, x=col1, y=col2\"")
    parser.add_argument("-finalize", "--finalize", dest="finalize", default=None, metavar="SPEC",
                         help="layout/legend options for -plot, e.g. -finalize \"consolidate_legends=True, style=True\"")
    parser.add_argument("-progress", "--progress", nargs="?", const=200_000, type=parse_positive_int,
                         default=None, metavar="N",
                         help="show row-count progress while converting large files (default: 200000 rows per chunk; optional N sets chunk size)")
    parser.add_argument("-pretty", "--pretty", action="store_true",
                         help="render tables as a bordered markdown table instead of plain pandas text")
    parser.add_argument("-pager", "--pager", action="store_true",
                         help="pipe long table or inspect outputs through system pager ($PAGER or less)")
    parser.add_argument("-round", "--round", dest="round_ndigits", type=int, default=None, metavar="N",
                         help="round numeric columns to N decimal places before printing/copying; "
                              "non-numeric columns are left unchanged")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args, extras = parser.parse_known_args(argv)
    if extras:
        if any(e in ("-to_clip", "--to_clip") for e in extras):
            parser.error("'-to_clip' has been removed; use '-o clip' or '-o clipboard' instead")
        if any(e in ("-convert", "--convert") for e in extras):
            parser.error("'-convert' has been removed; use '-o <filename>' or '-o <format>' instead")
        if args.file is None:
            path_extras = [e for e in extras if _is_path_like(e)]
            if path_extras:
                if args.path is None:
                    args.path = []
                elif not isinstance(args.path, list):
                    args.path = [args.path]
                args.path.extend(path_extras)
                extras = [e for e in extras if not _is_path_like(e)]

        if extras:
            msg = f"unrecognized arguments: {' '.join(extras)}"
            if extras and all(not e.startswith("-") for e in extras):
                # likely an unquoted value with a space (e.g. a column name) split by the
                # shell into separate argv tokens -- the whole spec needs one pair of quotes
                msg += (
                    "\nIf this is part of a value with a space (e.g. a column name), wrap the "
                    'whole spec in quotes, e.g. -select "col a,col b" -- see docs/cli.md#quoting.'
                )
            parser.error(msg)

    has_progress = args.progress is not None
    chunk_size = args.progress or 200_000
    args.progress = has_progress
    args.chunk_size = chunk_size

    op_order = getattr(args, "op_order", [])
    last_idx = len(op_order) - 1
    for idx, op in enumerate(op_order):
        if op in NON_DF_TERMINAL_OPS and idx != last_idx:
            parser.error(
                f"-{op} does not return a DataFrame/Series, so no flag may follow it "
                f"(except '-o clip'; matches pandas: you can't chain another call off "
                f"df.shape/df.columns/df.dtypes/df.info())"
            )

    out_target = str(args.output).strip() if args.output is not None else None
    args.output = out_target
    is_clip = out_target is not None and out_target.lower() in ("clip", "clipboard")
    is_file = out_target is not None and not is_clip

    if args.out_dir is not None:
        if args.output is None:
            parser.error("-out_dir requires -o/--output")
        if is_clip:
            parser.error("-out_dir cannot be used with '-o clip'")
        args.out_dir.mkdir(parents=True, exist_ok=True)

    if is_file and any(op in NON_DF_TERMINAL_OPS for op in op_order):
        bad_op = next(op for op in op_order if op in NON_DF_TERMINAL_OPS)
        parser.error(
            f"-{bad_op} does not produce a tabular DataFrame, so it cannot be exported to a file; "
            f"use '-o clip' or view in terminal"
        )

    raw_paths = args.path if isinstance(args.path, list) else ([args.path] if args.path else [])

    if args.meta and args.file is not None:
        parser.error("-meta cannot be used with -file")
    if args.merge and args.file is None:
        parser.error("-merge requires -file")
    if args.concat and args.file is None:
        parser.error("-concat requires -file")
    if args.file is not None:
        if raw_paths:
            parser.error("-file/-merge can't be combined with a positional path; list every input via -file instead")
        if not op_order or op_order[0] not in ("merge", "sql", "concat"):
            parser.error("-file requires -merge, -concat, or -sql as its first operation")
        if is_file and out_target is not None and out_target.lower() in _FORMAT_SHORTHANDS:
            parser.error(f"-o {out_target}: in -file/-merge mode, an explicit output file path is required")
    elif not raw_paths:
        parser.error("the following arguments are required: path")

    show_all = not any([args.shape, args.cols, args.dtype, args.nulls, args.describe, args.info,
                         args.value_counts, args.unique, args.freq is not None, args.hist is not None,
                         args.head is not None, args.tail is not None, args.sample is not None, args.sort_by is not None,
                         args.agg is not None, args.handle_missing is not None,
                         args.long is not None, args.wide is not None, args.pivot is not None,
                         args.select, args.qry, args.query, args.sql, args.replace_values,
                         args.rename, args.clean_columns is not None, args.merge, args.concat,
                         args.meta, args.diff is not None, args.plot is not None])

    wants_df = any([args.cols, args.dtype, args.nulls, args.describe, show_all,
                     args.value_counts, args.unique, args.freq is not None, args.hist is not None,
                     args.head is not None, args.tail is not None, args.sample is not None, args.sort_by is not None,
                     args.agg is not None, args.handle_missing is not None,
                     args.long is not None, args.wide is not None, args.pivot is not None,
                     args.clean_columns is not None, args.merge, args.concat, args.plot is not None])
    if is_clip and args.shape and wants_df:
        parser.error("-o clip can't combine -shape (not a DataFrame/Series) with a DataFrame-producing flag "
                     "like -head/-tail/-cols/-dtype/-nulls/-describe/-value_counts/-unique/-sample/-sort_by/"
                     "-agg/-handle_missing/-long/-wide/-pivot/-clean_columns/-merge/-concat/-freq/-hist/-plot; "
                     "run -shape separately")
    if args.by is not None and args.agg is None and not args.mutate:
        parser.error("-by requires -agg or -mutate")
    if args.frac is not None and args.sample is None:
        parser.error("-frac requires -sample")

    rename_specs = [parse_rename(raw) for raw in (args.rename or [])]
    select_specs = [parse_select_spec(raw) for raw in (args.select or [])]
    qry_specs = [parse_qry(raw) for raw in (args.qry or [])]
    mutate_specs = list(args.mutate or [])
    query_specs = list(args.query or [])
    sql_specs = list(args.sql or [])
    replace_specs = [parse_replace_values_arg(raw) for raw in (args.replace_values or [])]
    merge_specs = [parse_merge_arg(raw) for raw in (args.merge or [])]
    concat_specs = [parse_concat_arg(raw) for raw in (args.concat or [])]

    if args.file is not None:
        frames: dict[str, pd.DataFrame] = {}
        for entry in parse_file_arg(args.file):
            entry_path = Path(entry["path"])
            if not entry_path.exists():
                parser.error(f"-file: file not found: {entry_path}")
            try:
                reader = get_reader(entry_path, sep=entry["dlim"], encoding=entry["encoding"], chunk_size=args.chunk_size or 200_000)
            except ValueError as exc:
                parser.error(f"-file: {exc}")
            try:
                frames[entry["alias"]] = reader.to_dataframe(nrows=args.nrows, progress=args.progress, chunk_size=args.chunk_size or 200_000)
            except UnicodeError as exc:
                parser.error(_encoding_error_message(entry_path, entry["encoding"], exc))
        failed = _process_path(None, args, parser, False, show_all=show_all, select_specs=select_specs,
                                qry_specs=qry_specs, mutate_specs=mutate_specs, query_specs=query_specs, sql_specs=sql_specs,
                                replace_specs=replace_specs, rename_specs=rename_specs, frames=frames,
                                merge_specs=merge_specs, concat_specs=concat_specs)
        return 1 if failed else 0

    if not raw_paths:
        parser.error("the following arguments are required: path")
    paths = expand_paths(raw_paths)
    batch = len(paths) > 1

    if batch and args.output is not None:
        if is_clip:
            parser.error("-o clip cannot be used with multiple matched files")
        if out_target is not None and out_target.lower() not in _FORMAT_SHORTHANDS:
            parser.error(
                "-o/--output with multiple matched files requires a format (e.g. '-o csv' or '-o parquet'), "
                "not a single file path"
            )
        assert out_target is not None
        fmt = out_target.lower()
        ext = ".parquet" if fmt == "pq" else f".{fmt}"
        dest_map: dict[Path, Path] = {}
        for p in paths:
            clean_name = p.name[:-3] if p.name.lower().endswith(".gz") else p.name
            dest_name = f"{Path(clean_name).stem}{ext}"
            dest = (args.out_dir / dest_name).resolve() if args.out_dir is not None else p.with_name(dest_name).resolve()
            if dest in dest_map:
                parser.error(
                    f"-o {out_target}: multiple input files resolve to the same destination path '{dest}': "
                    f"'{dest_map[dest]}' and '{p}'"
                )
            dest_map[dest] = p

    exit_code = 0

    for path in paths:
        if batch:
            print(f"== {path} ==")
        try:
            failed = _process_path(path, args, parser, batch, show_all=show_all, select_specs=select_specs,
                                    qry_specs=qry_specs, mutate_specs=mutate_specs, query_specs=query_specs, sql_specs=sql_specs,
                                    replace_specs=replace_specs, rename_specs=rename_specs)
        except UnicodeError as exc:
            failed = _fail(parser, batch, _encoding_error_message(path, args.encoding, exc))
        if failed:
            exit_code = 1

    return exit_code


if __name__ == "__main__":
    sys.exit(main())

################################################################################
# FILE: src/pytae/cli_parsing.py
################################################################################
"""Parsing helpers for pytae CLI flag values (pure string parsing; no pandas/pipeline state)."""

from __future__ import annotations

import argparse
import ast
import difflib
import glob
import re
from pathlib import Path
from typing import Any

from pytae._text import _is_enclosed_pair
from pytae._text import tokenize as _tokenize
from pytae._text import unquote_name as _unquote_name

SELECT_KEYS = ("dtype", "exclude_dtype", "contains", "startswith", "endswith", "regex", "exclude")


def parse_columns(raw: str) -> list[str]:
    """Parse a column list like "'col a','col b'" or "col_a,col_b" or "[col a], [col b]" into a list of names."""
    raw = raw.strip()
    if _is_enclosed_pair(raw, "[", "]") and "," in raw:
        raw = raw[1:-1].strip()
    try:
        parsed = ast.literal_eval(f"[{raw}]")
        cols = list(parsed) if isinstance(parsed, (list, tuple)) else [parsed]
        return [_unquote_name(str(c).strip()) for c in cols]
    except (ValueError, SyntaxError):
        tokens = _tokenize(raw, ",", track_brackets=True)
        return [_unquote_name(c.strip()) for c in tokens if c.strip()]



def parse_sort_by(raw: str) -> tuple[list[str], str]:
    """Parse -sort_by SPEC: a comma-separated column list, optionally ending with
    asc or desc as a trailing word (default asc). e.g. "species,body_mass_g desc"."""
    raw = (raw or "").strip()
    if not raw:
        raise SystemExit("-sort_by: expected a column list, optionally followed by asc or desc")
    order = "asc"
    head, sep, tail = raw.rpartition(" ")
    if sep and tail.lower() in ("asc", "desc"):
        raw, order = head.rstrip(",").strip(), tail.lower()
    cols = parse_columns(raw)
    if not cols:
        raise SystemExit("-sort_by: expected a column list, optionally followed by asc or desc")
    return cols, order


def _split_tokens(raw: str, sep: str = ",") -> list[str]:
    """Split on sep, respecting quotes and brackets."""
    return [t for t in _tokenize(raw, sep, track_brackets=True) if t]


def _split_groups(raw: str, sep: str = ";") -> list[str]:
    """Split on sep but keep quote characters so inner comma-split still sees them."""
    return [t for t in _tokenize(raw, sep, keep_quotes=True) if t]


def parse_select_spec(raw: str) -> tuple[list[str], dict]:
    """Parse -select into positional names/slices and select() kwargs.

    Tokens without '=' are column names or start:end slices. Tokens like
    dtype=numeric / contains=bill / regex=^flip become kwargs. Repeated keys
    become a list. Union of all tokens, matching pt.select().
    """
    raw = raw.strip()
    if _is_enclosed_pair(raw, "[", "]") and "," in raw:
        raw = raw[1:-1].strip()
    tokens = _split_tokens(raw)

    names: list[str] = []
    kwargs: dict = {}
    for token in tokens:
        if "=" in token:
            key, _, value = token.partition("=")
            key = key.strip()
            value = value.strip()
            if key in SELECT_KEYS:
                if not value:
                    raise SystemExit(f"-select: {key}= needs a value")
                if key in kwargs:
                    prev = kwargs[key]
                    kwargs[key] = (prev if isinstance(prev, list) else [prev]) + [value]
                else:
                    kwargs[key] = value
                continue
        names.append(token)
    if not names and not kwargs:
        raise SystemExit("-select: expected column names and/or key=value tokens")
    if "exclude_dtype" in kwargs and (names or len(kwargs) > 1):
        raise SystemExit("-select: exclude_dtype cannot be combined with other selection criteria")
    return names, kwargs


def _select_unknown_names(tokens: list[str], available: list[str]) -> list[str]:
    """Exact-name tokens (and slice endpoints) that are not in the file."""
    unknown: list[str] = []
    for token in tokens:
        clean = _unquote_name(token)
        if clean in available:
            continue
        if token.startswith(("-", "~")):
            raw = _unquote_name(token[1:].strip())
            if raw in available:
                continue
            if ":" in raw:
                start, end = raw.split(":", 1)
                start, end = _unquote_name(start.strip()), _unquote_name(end.strip())
                if start and start not in available:
                    unknown.append(start)
                if end and end not in available:
                    unknown.append(end)
            else:
                unknown.append(raw)
        elif ":" in token:
            start, end = token.split(":", 1)
            start, end = _unquote_name(start.strip()), _unquote_name(end.strip())
            if start and start not in available:
                unknown.append(start)
            if end and end not in available:
                unknown.append(end)
        else:
            unknown.append(clean)
    return unknown


def unknown_columns_message(flag: str, requested: list[str], available: list[str]) -> str:
    """Build an 'unknown column(s)' error message, suggesting a close match for likely typos."""
    unknown = [c for c in requested if c not in available]
    parts = []
    for c in unknown:
        close = difflib.get_close_matches(c, available, n=1)
        parts.append(f"'{c}'" + (f" (did you mean '{close[0]}'?)" if close else ""))
    return f"{flag}: unknown column(s): {', '.join(parts)}"


def parse_rename(raw: str) -> dict[str, str]:
    """Parse a rename mapping like "old_a:new_a,old_b:new_b" into a dict.
    Quoting either side (e.g. "'old a':'new a'") is optional and stripped if present—
    plain old_a:new_a already handles spaces, quoting just needs to not break things."""
    mapping: dict[str, str] = {}
    for pair in raw.split(","):
        pair = pair.strip()
        if not pair:
            continue
        if ":" in pair:
            old, new = pair.split(":", 1)
        elif "=" in pair:
            raise SystemExit(f"invalid -rename mapping '{pair}'; use ':' (e.g. -rename 'old:new'). '=' is not supported.")
        else:
            raise SystemExit(f"invalid -rename mapping '{pair}'; expected old:new")
        old_clean = _unquote_name(old)
        new_clean = _unquote_name(new)
        if old_clean in mapping and mapping[old_clean] != new_clean:
            raise SystemExit(f"invalid -rename: conflicting mappings for '{old_clean}': '{mapping[old_clean]}' and '{new_clean}'")
        mapping[old_clean] = new_clean
    return mapping


def expand_paths(patterns: str | list[str]) -> list[Path]:
    """Expand glob pattern(s) into matching paths, or wrap plain paths as-is."""
    items = [patterns] if isinstance(patterns, str) else list(patterns)
    results: list[Path] = []
    seen: set[Path] = set()
    for item in items:
        p_literal = Path(item)
        if p_literal.exists():
            if p_literal not in seen:
                results.append(p_literal)
                seen.add(p_literal)
        elif any(ch in item for ch in "*?["):
            matches = sorted(Path(p) for p in glob.glob(item))
            if not matches:
                raise SystemExit(f"no files matched pattern: {item}")
            for m in matches:
                if m not in seen:
                    results.append(m)
                    seen.add(m)
        else:
            if p_literal not in seen:
                results.append(p_literal)
                seen.add(p_literal)
    return results


def _split_qry_entries(raw: str) -> list[tuple[str, str]]:
    """Split a -qry body into raw (key, value) text pairs on top-level commas and equals,
    respecting quotes and nested (), [], {} so tuples/lists/intervals inside a value
    aren't mistaken for entry or key/value separators. Direct comparisons like
    'col > 5' or assignments like 'col = > 3500' or 'col = 3500' are accepted."""
    entries: list[tuple[str, str]] = []
    for raw_entry in _tokenize(raw, ",", keep_quotes=True, track_brackets=True):
        if not raw_entry:
            continue
        # Direct comparison first: col > 5, col >= 5, col <= 5, etc.
        m_cmp = re.match(r"^([^>=<!:]+?)\s*(>=|<=|!=|==|>|<)\s*(.+)$", raw_entry)
        if m_cmp:
            col = m_cmp.group(1).strip()
            op = m_cmp.group(2).strip()
            val = m_cmp.group(3).strip()
            entries.append((col, f"{op} {val}"))
            continue
        # Assignment with '=': col = val, col = > 5, col = ['a', 'b']
        m_eq = re.match(r"^([^>=<!:]+?)\s*=\s*(.+)$", raw_entry)
        if m_eq:
            col = m_eq.group(1).strip()
            val = m_eq.group(2).strip()
            entries.append((col, val))
            continue
        # If colon was used, give a helpful error
        pieces = _tokenize(raw_entry, ":", keep_quotes=True, track_brackets=True)
        if len(pieces) >= 2:
            raise SystemExit(f"invalid -qry condition '{raw_entry}': use '=' (e.g. -qry 'species = Adelie'). Colon ':' is not supported.")
        raise SystemExit(f"invalid -qry conditions: missing '=' in entry '{raw_entry}'")
    return entries


def parse_qry(raw: str) -> list[tuple[str, Any]]:
    """Parse -qry conditions like "col = ('>', 5), other = ['a','b']"; wrapping {} and
    quotes around column names are both optional (matching -select), e.g.
    "sex='Male'" == "'sex'='Male'". Prefix operators (e.g. "col = > 5" or "col > 5") and bare string
    values (e.g. "species = Adelie") are also accepted."""
    stripped = raw.strip()
    if stripped.startswith("{") and stripped.endswith("}"):
        stripped = stripped[1:-1]
    conditions: list[tuple[str, Any]] = []
    for key_raw, value_raw in _split_qry_entries(stripped):
        key = _unquote_name(key_raw)
        if not key:
            raise SystemExit("invalid --qry conditions: empty column name")
        val_strip = value_raw.strip()
        if not val_strip:
            raise SystemExit(f"invalid --qry conditions: '{key}' has no value")

        # Check interval syntax: e.g. [3000, 4500], (3000, 4500), [a, c), (a, c]
        m_int = re.match(r"^([\[(])\s*([^,()\[\]]+)\s*,\s*([^,()\[\]]+)\s*([\])])$", val_strip)
        if m_int:
            left, b1, b2, right = m_int.groups()
            b1_is_quoted = (b1.startswith("'") and b1.endswith("'")) or (b1.startswith('"') and b1.endswith('"'))
            b2_is_quoted = (b2.startswith("'") and b2.endswith("'")) or (b2.startswith('"') and b2.endswith('"'))
            if not (b1_is_quoted and b2_is_quoted and left == "[" and right == "]"):
                conditions.append((key, val_strip))
                continue

        op_match = re.match(r"^(>=|<=|!=|==|>|<)\s*(.+)$", val_strip)
        if op_match:
            op = op_match.group(1)
            sub_raw = op_match.group(2).strip()
            try:
                sub_val = ast.literal_eval(sub_raw)
            except (ValueError, SyntaxError):
                sub_val = _unquote_name(sub_raw)
            conditions.append((key, (op, sub_val)))
            continue

        try:
            value = ast.literal_eval(val_strip)
        except (ValueError, SyntaxError) as exc:
            raise SystemExit(
                f"invalid --qry conditions: value for '{key}' ('{val_strip}') must be quoted "
                "(e.g. 'Male') or a valid literal (number/tuple/list)"
            ) from exc
        conditions.append((key, value))
    if not conditions:
        raise SystemExit("-qry expects keyword entries, e.g. \"col = ('>', 5)\" or \"col > 5\" (quotes around column name optional)")
    return conditions


def parse_agg(raw: str):
    """Parse -agg: a name (mean), a comma list (mean,sum,n), or a col=aggfunc
    mapping (val = sum, n = n, total = v1:sum). Quote a mapping key only to protect a comma."""
    raw = (raw or "").strip()
    if not raw:
        return "sum"
    if raw.startswith("{") or (raw.startswith("[") and "=" not in raw):
        raise SystemExit(
            "-agg: use a name (mean), a comma list (mean,sum), or a mapping "
            "(col = mean, n = n)"
        )
    entries = [e for e in _tokenize(raw, ",", keep_quotes=True, track_brackets=True) if e]
    if not entries:
        raise SystemExit("-agg: expected a name, a comma list, or a mapping")

    def _is_mapping(entry: str) -> bool:
        return (
            len(_tokenize(entry, "=", keep_quotes=True, track_brackets=True)) >= 2
            or len(_tokenize(entry, ":", keep_quotes=True, track_brackets=True)) >= 2
        )

    mapped = [_is_mapping(e) for e in entries]
    if all(mapped):
        out: dict[str, str] = {}
        for entry in entries:
            if "=" in entry:
                parts = _tokenize(entry, "=", keep_quotes=True, track_brackets=True)
                key = _unquote_name(parts[0])
                value = _unquote_name("=".join(parts[1:]))
                if key in ("column", "aggfunc"):
                    raise SystemExit(
                        "-agg: the old 'column=...,aggfunc=...' syntax is retired. "
                        "Use '-by <cols> -agg \"col = aggfunc\"' or '-agg \"total = col:aggfunc\"' instead."
                    )
            elif ":" in entry:
                raise SystemExit(f"-agg: invalid mapping entry '{entry}'; use '=' (e.g. -agg 'col = mean'). Colon ':' is not supported.")
            else:
                raise SystemExit(f"-agg: invalid mapping entry {entry!r}")
            if not key or not value:
                raise SystemExit(f"-agg: invalid mapping entry {entry!r}")
            out[key] = value
        return out
    if any(mapped):
        raise SystemExit(
            "-agg: mix of names and col=aggfunc mappings; use one or the other"
        )
    names = [_unquote_name(e) for e in entries]
    return names[0] if len(names) == 1 else names


_AGG_KEYS = ("column", "aggfunc", "as")


def parse_group_agg(raw: str) -> list[tuple[str, str, str]]:
    """Parse -agg as key=value specs: column=, aggfunc=, optional as=.

    Several specs are separated by ';'. Several source columns in one spec share
    the same aggfunc: column='value,val_growth',aggfunc=sum. as= needs a single column.
    Returns (source_col, output_name, aggfunc) rows.
    """
    raw = (raw or "").strip()
    if not raw:
        raise SystemExit("-agg: expected key=value specs, e.g. column=value,aggfunc=sum,as=v")
    if raw.startswith("{"):
        raise SystemExit(
            "-agg: use key=value specs, e.g. column=value,aggfunc=sum,as=v "
            "(not a dict literal)"
        )
    rows: list[tuple[str, str, str]] = []
    for group in _split_groups(raw, ";"):
        canon = parse_reshape_kwargs(group, keys=_AGG_KEYS, flag="-agg")
        if "column" not in canon or "aggfunc" not in canon:
            raise SystemExit("-agg: each spec needs column= and aggfunc=")
        cols = parse_columns(canon["column"])
        if not cols:
            raise SystemExit("-agg: column= needs at least one column")
        out = canon.get("as")
        if out and len(cols) != 1:
            raise SystemExit("-agg: as= requires a single column")
        aggfunc = canon["aggfunc"]
        for col in cols:
            rows.append((col, out or col, aggfunc))
    names = [name for _, name, _ in rows]
    if len(names) != len(set(names)):
        raise SystemExit("-agg: duplicate output names")
    return rows


_LONG_KEYS = ("c", "v", "cols", "values", "id_vars", "by")
_WIDE_KEYS = ("c", "v", "index", "by", "r", "rows", "cols", "values")


def parse_reshape_kwargs(raw: str | None, *, keys: tuple[str, ...], flag: str) -> dict:
    """Parse -long/-wide/-pivot as key=value tokens, e.g. c=metric,v=reading,a=mean.
    Comma-separated lists like frames=a,b,c or on=id:id,code:code work with or without quotes."""
    raw = (raw or "").strip()
    if not raw:
        return {}
    kwargs: dict = {}
    current_key: str | None = None
    for token in _split_tokens(raw):
        if "=" in token:
            key, _, value = token.partition("=")
            key = key.strip()
            value = _unquote_name(value)
            if flag == "-wide" and key in ("a", "agg", "aggfunc"):
                raise SystemExit(
                    "-wide: -wide is strictly for 1-to-1 reshaping without aggregation. "
                    "Use -pivot for aggregations (e.g. -pivot 'r=...,c=...,v=...,a=...')."
                )
            if key not in keys:
                raise SystemExit(f"{flag}: unknown key {key!r}; expected {', '.join(keys)}")
            if not value:
                raise SystemExit(f"{flag}: {key}= needs a value")
            if key in kwargs:
                raise SystemExit(f"{flag}: {key}= given more than once")
            current_key = key
            kwargs[key] = value
        else:
            if current_key is None:
                raise SystemExit(f"{flag}: expected key=value tokens ({', '.join(keys)})")
            kwargs[current_key] = f"{kwargs[current_key]},{_unquote_name(token)}"
    for key, value in list(kwargs.items()):
        if key in ("margins", "exact", "dropna"):
            kwargs[key] = parse_bool_text(value)
        elif key in ("fill", "fill_value"):
            if isinstance(value, str):
                v_lower = value.strip().lower()
                if v_lower in ("none", "null"):
                    kwargs[key] = None
                else:
                    try:
                        kwargs[key] = int(value)
                    except ValueError:
                        try:
                            kwargs[key] = float(value)
                        except ValueError:
                            raise SystemExit(f"{flag}: {key}= must be a numeric value or 'none'")
        else:
            kwargs[key] = value
    return kwargs


def parse_long_arg(raw: str | None) -> dict:
    return parse_reshape_kwargs(raw, keys=_LONG_KEYS, flag="-long")


def parse_wide_arg(raw: str | None) -> dict:
    parsed = parse_reshape_kwargs(raw, keys=_WIDE_KEYS, flag="-wide")
    if "cols" in parsed and "c" not in parsed:
        parsed["c"] = parsed.pop("cols")
    if "values" in parsed and "v" not in parsed:
        parsed["v"] = parsed.pop("values")
    return parsed


_PIVOT_KEYS = (
    "r", "c", "v", "a", "dropna", "fill_value", "fill",
    "rows", "row", "index", "by",
    "cols", "col", "columns",
    "values", "value", "val", "vals",
    "agg", "aggfunc",
)


def parse_pivot_arg(raw: str | None) -> dict:
    return parse_reshape_kwargs(raw, keys=_PIVOT_KEYS, flag="-pivot")


_REPLACE_KEYS = ("c", "v", "exact")


def parse_value_map(raw: str) -> dict[str, str]:
    """Parse a -replace_values v= mapping like "old_a:new_a,old_b:new_b" into a dict."""
    mapping: dict[str, str] = {}
    for pair in raw.split(","):
        pair = pair.strip()
        if not pair:
            continue
        if ":" in pair:
            old, new = pair.split(":", 1)
        elif "=" in pair:
            raise SystemExit(f"-replace_values: invalid v= mapping '{pair}'; use ':' (expected old:new). '=' is not supported.")
        else:
            raise SystemExit(f"-replace_values: invalid v= mapping '{pair}'; expected old:new")
        mapping[_unquote_name(old)] = _unquote_name(new)
    if not mapping:
        raise SystemExit("-replace_values: v= needs at least one old:new pair")
    return mapping


def parse_replace_values_arg(raw: str) -> tuple[list[str] | None, dict[str, str], bool]:
    """Parse -replace_values as key=value tokens: c= (optional column scope), v= (required
    old:new mapping), exact= (optional bool, default true — whole-cell match vs
    substring match anywhere in the cell). Returns (cols_or_None, mapping, exact)."""
    kwargs = parse_reshape_kwargs(raw, keys=_REPLACE_KEYS, flag="-replace_values")
    if "v" not in kwargs:
        raise SystemExit("-replace_values: expected v='old:new,...'")
    cols = parse_columns(kwargs["c"]) if "c" in kwargs else None
    mapping = parse_value_map(kwargs["v"])
    exact = kwargs.get("exact", True)
    return cols, mapping, exact


def parse_bool_text(raw: str) -> bool:
    """Parse a required true/false text flag value into bool."""
    lowered = raw.strip().lower()
    if lowered == "true":
        return True
    if lowered == "false":
        return False
    raise argparse.ArgumentTypeError("expected 'true' or 'false'")


def parse_positive_int(raw) -> int:
    """Parse a row-count flag; reject 0 and negatives (e.g. -head 0, -head -5)."""
    try:
        n = int(raw)
    except (TypeError, ValueError):
        raise argparse.ArgumentTypeError(f"invalid integer: {raw!r}") from None
    if n <= 0:
        raise argparse.ArgumentTypeError(f"must be > 0, got {n}")
    return n


def parse_fraction(raw) -> float:
    """Parse -frac: a fraction of rows in (0, 1] (e.g. -frac 0.1 for 10%)."""
    try:
        frac = float(raw)
    except (TypeError, ValueError):
        raise argparse.ArgumentTypeError(f"invalid number: {raw!r}") from None
    if not (0 < frac <= 1):
        raise argparse.ArgumentTypeError(f"must be > 0 and <= 1, got {frac}")
    return frac


def parse_list_order(raw: str) -> str:
    """Parse optional listing order for -cols/-dtype/-nulls: asc or desc.

    'file' is the internal sentinel used when the flag is present with no value
    (argparse 3.14 type-converts const).
    """
    lowered = str(raw).strip().lower()
    if lowered in ("asc", "desc", "file"):
        return lowered
    raise argparse.ArgumentTypeError("expected 'asc' or 'desc'")


_CLEAN_COLUMNS_BOOL_KEYS = ("strip", "squeeze", "strip_special", "dedupe")
_CLEAN_COLUMNS_KEYS = _CLEAN_COLUMNS_BOOL_KEYS + ("fill", "case")
_CLEAN_COLUMNS_CASE_VALUES = ("lower", "upper", "proper")


def parse_clean_columns_arg(raw: str) -> dict:
    """Parse -clean_columns as key[=value] tokens, comma-separated:

    - strip=, squeeze=, strip_special=, dedupe= — bools (default false); a bare
      key with no '=' means true, e.g. "strip" is short for "strip=true".
    - fill[=STR] — replace whitespace in a header with STR; bare "fill" defaults
      to '_'; omit the key entirely for no fill.
    - case=lower|upper|proper — no bare/default form, always needs a value.
    """
    raw = (raw or "").strip()
    if not raw:
        raise SystemExit(
            "-clean_columns: expected at least one key, e.g. \"strip\" or \"case=lower\""
        )
    kwargs: dict = {}
    for token in _split_tokens(raw):
        key, has_value, value = token.partition("=")
        key = key.strip()
        if key not in _CLEAN_COLUMNS_KEYS:
            raise SystemExit(f"-clean_columns: unknown key {key!r}; expected {', '.join(_CLEAN_COLUMNS_KEYS)}")
        if key in kwargs:
            raise SystemExit(f"-clean_columns: {key}= given more than once")
        if key in _CLEAN_COLUMNS_BOOL_KEYS:
            kwargs[key] = parse_bool_text(_unquote_name(value)) if has_value else True
        elif key == "fill":
            kwargs[key] = _unquote_name(value) if has_value else "_"
        else:  # case
            case_value = _unquote_name(value).strip().lower()
            if not has_value or not case_value:
                raise SystemExit("-clean_columns: case= needs a value (lower, upper, or proper)")
            if case_value not in _CLEAN_COLUMNS_CASE_VALUES:
                raise SystemExit(f"-clean_columns: case= must be one of {', '.join(_CLEAN_COLUMNS_CASE_VALUES)}")
            kwargs[key] = case_value
    return kwargs


_FILE_ENTRY_KEYS = ("dlim", "encoding")


def parse_file_arg(raw: str) -> list[dict]:
    """Parse -file as ';'-separated entries (at least two), each PATH=ALIAS optionally
    followed by ,dlim=/,encoding= overrides for reading that file, e.g.
    "data1.parquet=df1,dlim='|'; data2.parquet=df2,encoding='latin-1'".
    Returns a list of {'path', 'alias', 'dlim', 'encoding'} dicts (dlim/encoding are
    None when not given).
    """
    raw = (raw or "").strip()
    if not raw:
        raise SystemExit("-file: expected at least one PATH=ALIAS entry")
    entries: list[dict] = []
    aliases_seen: set[str] = set()
    for group in _split_groups(raw, ";"):
        tokens = _split_tokens(group, ",")
        if not tokens:
            raise SystemExit("-file: expected PATH=ALIAS")
        head, has_value, alias = tokens[0].partition("=")
        if not has_value or not alias.strip():
            raise SystemExit(f"-file: expected PATH=ALIAS, got {tokens[0]!r}")
        alias = alias.strip()
        if alias in aliases_seen:
            raise SystemExit(f"-file: duplicate alias '{alias}'")
        aliases_seen.add(alias)
        entry = {"path": head.strip(), "alias": alias, "dlim": None, "encoding": None}
        for token in tokens[1:]:
            key, has_kv, value = token.partition("=")
            key = key.strip()
            if not has_kv or key not in _FILE_ENTRY_KEYS:
                raise SystemExit(f"-file: unknown option {token!r}; expected dlim= or encoding=")
            entry[key] = _unquote_name(value)
        entries.append(entry)
    if len(entries) < 2:
        raise SystemExit("-file: need at least two PATH=ALIAS entries")
    return entries


_MERGE_KEYS = ("left", "right", "on", "how", "validate")


def parse_merge_on(raw: str) -> tuple[list[str] | None, list[str] | None, list[str] | None]:
    """Parse -merge's on= value: shared column names (comma list) when names match on
    both sides, or 'left:right' pairs when they differ between sides. Quote the whole
    on= value if it has more than one column/pair (protects the inner commas), e.g.
    on='col a:cola,colb:colb'. Returns (on_cols, left_cols, right_cols) — on_cols is
    set XOR (left_cols, right_cols) are set.
    """
    pairs = [p.strip() for p in raw.split(",") if p.strip()]
    if not pairs:
        raise SystemExit("-merge: on= needs at least one column (or left:right pair)")
    has_colon = [":" in p for p in pairs]
    if any(has_colon) and not all(has_colon):
        raise SystemExit("-merge: on= mixes plain columns and left:right pairs; use one style")
    if all(has_colon):
        left_cols, right_cols = [], []
        for pair in pairs:
            left, _, right = pair.partition(":")
            left_cols.append(left.strip())
            right_cols.append(right.strip())
        return None, left_cols, right_cols
    return pairs, None, None


def parse_merge_arg(raw: str) -> dict:
    """Parse -merge as key=value tokens: left=/right= (required -file aliases, or the
    literal 'df' to reference the pipeline's current result so far — lets repeated
    -merge calls fold in one more file at a time), on= (required; shared column
    name(s), or left:right pairs if they differ between sides), how= (optional,
    default 'inner', passed straight to pandas merge()), validate= (optional, passed
    straight to pandas merge()).
    """
    kwargs = parse_reshape_kwargs(raw, keys=_MERGE_KEYS, flag="-merge")
    how = kwargs.get("how", "inner")
    required = ["left", "right"] if how == "cross" else ["left", "right", "on"]
    missing = [k for k in required if k not in kwargs]
    if missing:
        raise SystemExit(f"-merge: missing required key(s): {', '.join(missing)}")
    if how == "cross":
        if "on" in kwargs:
            raise SystemExit("-merge: on= cannot be used with how=cross (pandas cross joins don't take on=)")
        on_cols, left_cols, right_cols = None, None, None
    else:
        on_cols, left_cols, right_cols = parse_merge_on(kwargs["on"])
    return {
        "left": kwargs["left"],
        "right": kwargs["right"],
        "on": on_cols,
        "left_on": left_cols,
        "right_on": right_cols,
        "how": how,
        "validate": kwargs.get("validate"),
    }


_CONCAT_KEYS = ("frames",)


def parse_concat_arg(raw: str) -> dict:
    """Parse -concat as key=value tokens: frames= (required) — an ordered,
    comma-separated list of -file aliases (or 'df' for the pipeline's current
    result, to stack more files onto something already produced), e.g.
    frames='df1,df2,df3'. Quote the whole value — it has internal commas.
    """
    kwargs = parse_reshape_kwargs(raw, keys=_CONCAT_KEYS, flag="-concat")
    if "frames" not in kwargs:
        raise SystemExit("-concat: expected frames='alias1,alias2,...'")
    names = [n.strip() for n in kwargs["frames"].split(",") if n.strip()]
    if len(names) < 2:
        raise SystemExit("-concat: frames= needs at least two names")
    return {"frames": names}


def parse_kv_spec(raw: str | None, *, flag: str) -> dict[str, Any]:
    """Parse comma-separated key=value tokens into a kwargs dict, type-casting numbers and booleans."""
    if not raw or not raw.strip():
        return {}
    kwargs: dict[str, Any] = {}
    for token in _tokenize(raw, ",", keep_quotes=True, track_brackets=True):
        token = token.strip()
        if not token:
            continue
        if "=" not in token:
            kwargs[token.strip()] = True
            continue
        key, _, raw_val = token.partition("=")
        key = key.strip()
        raw_val = raw_val.strip()
        val: Any
        if (raw_val.startswith("'") and raw_val.endswith("'")) or (raw_val.startswith('"') and raw_val.endswith('"')):
            val = raw_val[1:-1]
        elif raw_val.lower() == "true":
            val = True
        elif raw_val.lower() == "false":
            val = False
        elif raw_val.lower() in ("none", "null"):
            val = None
        else:
            try:
                val = int(raw_val)
            except ValueError:
                try:
                    val = float(raw_val)
                except ValueError:
                    val = raw_val
        kwargs[key] = val
    return kwargs

################################################################################
# FILE: src/pytae/cli_pipeline.py
################################################################################
"""The _Pipeline view and the argparse Actions that record CLI flag order (op_order)."""

from __future__ import annotations

import argparse
import re
from typing import Any

import pandas as pd

from pytae._text import unquote_name as _unquote_name
from pytae.cli_parsing import _select_unknown_names, unknown_columns_message
from pytae.mutate import mutate
from pytae.other_utilities import replace_values, safe_reset_index
from pytae.qry import qry
from pytae.select import select


def _sql_string_literal(value: str) -> str:
    """Quote a plain Python string as a SQL string literal (escaping embedded quotes) —
    duckdb table functions like read_parquet()/read_csv() can't be parameterized via
    prepared-statement placeholders, so the path/delimiter must be inlined as a literal."""
    return "'" + value.replace("'", "''") + "'"


def _mask_quoted(spec: str) -> str:
    """Blank out the contents of quoted substrings in spec (keeping overall length/
    positions), so a regex check on the result only ever sees text outside string
    literals -- e.g. an '@' inside a quoted value like 'a@b' is masked out."""
    out = []
    quote: str | None = None
    for ch in spec:
        if quote:
            out.append(ch if ch == quote else " ")
            if ch == quote:
                quote = None
            continue
        if ch in "'\"":
            quote = ch
        out.append(ch)
    return "".join(out)


class _Pipeline:
    """Flag order is the method chain. Each -select/-qry/-query/-head/… call
    runs on the current view, same as pt.select(df, …) then pt.qry(df, …) then head.
    Only the last flag prints. Schema-only ops (-cols/-dtype/-shape) avoid loading
    row data until something actually requires it. -shape/-cols/-dtype/-nulls/-info
    don't return a DataFrame/Series in pandas either, so (like main()'s validation)
    nothing may follow them except -o clip.
    """

    def __init__(self, reader=None, *, nrows=None, progress=False, frames=None, chunk_size: int = 200_000) -> None:
        self._reader = reader
        self._nrows = nrows
        self._progress = progress
        self._chunk_size = chunk_size
        self._df: pd.DataFrame | None = None
        self._pending_exact: list[str] | None = None
        self._frames: dict[str, pd.DataFrame] = frames or {}

    def _available_columns(self) -> list[str]:
        if self._df is not None:
            return list(self._df.columns)
        if self._pending_exact is not None:
            return list(self._pending_exact)
        return list(self._reader.columns())

    def apply_select(self, names: list[str], kwargs: dict) -> str | None:
        """Apply one -select spec to the current view. Returns an error message or None."""
        available = self._available_columns()
        unknown = _select_unknown_names(names, available)
        if unknown:
            return unknown_columns_message("-select", unknown, available)

        needs_frame = (
            self._df is not None
            or bool(kwargs)
            or any(":" in token for token in names)
            or any(token.startswith(("-", "~")) for token in names)
        )
        if not needs_frame:
            self._pending_exact = [_unquote_name(n) for n in names]
            return None

        df = self._df if self._df is not None else self.dataframe()
        try:
            self._df = select(df, *names, **kwargs)
        except (ValueError, KeyError) as exc:
            return f"-select: {exc}"
        return None

    def apply_rename(self, mapping: dict[str, str]) -> str | None:
        """Apply one -rename mapping to the current view. Returns an error message or None."""
        available = self._available_columns()
        unknown = [c for c in mapping if c not in available]
        if unknown:
            return unknown_columns_message("-rename", unknown, available)
        df = self.dataframe()
        self._df = df.rename(columns=mapping)
        return None

    def apply_qry(self, conditions: list[tuple[str, Any]] | dict) -> str | None:
        """Apply one -qry spec to the current view. Returns an error message or None."""
        df = self.dataframe()
        try:
            if isinstance(conditions, list):
                self._df = qry(df, conditions)
            else:
                self._df = qry(df, **conditions)
        except Exception as exc:
            return f"-qry: {exc}"
        return None

    def apply_mutate(
        self, spec: str, by: list[str] | None = None, dropna: bool = False
    ) -> str | None:
        """Apply one -mutate spec to the current view. Returns an error message or None."""
        raw_spec = spec.strip()
        if (raw_spec.startswith("'") and raw_spec.endswith("'")) or (raw_spec.startswith('"') and raw_spec.endswith('"')):
            raw_spec = raw_spec[1:-1].strip()
        if raw_spec.startswith("@"):
            path = raw_spec[1:].strip()
            if (path.startswith("'") and path.endswith("'")) or (path.startswith('"') and path.endswith('"')):
                path = path[1:-1].strip()
            from pathlib import Path
            file_path = Path(path)
            if not file_path.is_file():
                return f"-mutate: spec file not found: '{path}'"
            try:
                raw_spec = file_path.read_text(encoding="utf-8").strip()
            except Exception as exc:
                return f"-mutate: failed to read spec file '{path}': {exc}"
        if re.search(r"@\w", _mask_quoted(raw_spec)):
            return "-mutate: '@name' local-variable references are library-only (pt.mutate() from Python), not available on the CLI"
        df = self.dataframe()
        try:
            from pytae.mutate import parse_mutate_spec
            parsed = parse_mutate_spec(raw_spec)
            self._df = mutate(df, parsed, by=by, dropna=dropna)
        except Exception as exc:
            return f"-mutate: {exc}"
        return None

    def apply_query(self, expr: str) -> str | None:
        """Apply one -query expression to the current view. Returns an error message or None."""
        df = self.dataframe()
        try:
            self._df = df.query(expr)
        except Exception as exc:
            return f"-query: {exc}"
        return None

    def apply_replace_values(self, cols: list[str] | None, mapping: dict[str, str], exact: bool) -> str | None:
        """Apply one -replace_values spec to the current view. Returns an error message or None."""
        df = self.dataframe()
        if cols is not None:
            available = list(df.columns)
            unknown = [c for c in cols if c not in available]
            if unknown:
                return unknown_columns_message("-replace_values", unknown, available)
        try:
            self._df = replace_values(df, v=mapping, c=cols, exact=exact)
        except Exception as exc:
            return f"-replace_values: {exc}"
        return None

    def apply_sql(self, query: str) -> str | None:
        """Apply one -sql query via duckdb. Single-file mode: the current view is registered
        as table `data` (the file itself is already named on the command line, so there's no
        separate file-derived alias) — if nothing has touched the view yet, duckdb scans the
        source parquet/csv/txt/dat file directly instead of first materializing it through
        pandas (much faster; pandas is only used as a fallback, see _register_source_view).
        -file/-merge mode: every -file alias is registered under its own name instead, so
        -sql can do the join itself; `data` is also registered once something (e.g. -merge)
        has produced a current view. Returns an error message or None."""
        try:
            import duckdb
        except ImportError as exc:
            return f"-sql requires duckdb. Install with: pip install 'pytae[sql]' ({exc})"
        q = query.strip()
        if (q.startswith("'") and q.endswith("'")) or (q.startswith('"') and q.endswith('"')):
            q = q[1:-1].strip()
        if q.startswith("@"):
            path = q[1:].strip()
            if (path.startswith("'") and path.endswith("'")) or (path.startswith('"') and path.endswith('"')):
                path = path[1:-1].strip()
            from pathlib import Path
            file_path = Path(path)
            if not file_path.is_file():
                return f"-sql: query file not found: '{path}'"
            try:
                q = file_path.read_text(encoding="utf-8").strip()
            except Exception as exc:
                return f"-sql: failed to read query file '{path}': {exc}"
        from pytae.sql import _normalize_sql
        normalized_query = _normalize_sql(q)
        con = duckdb.connect()
        try:
            for alias, frame in self._frames.items():
                con.register(alias, frame)
            if self._df is not None:
                con.register("data", self._df)
            elif self._reader is not None:
                if not self._register_source_view(con):
                    con.register("data", self.dataframe())
            try:
                self._df = con.sql(normalized_query).df()
            except Exception as exc:
                return f"-sql: {exc}"
        finally:
            con.close()
        return None

    def _register_source_view(self, con) -> bool:
        """Best-effort: have duckdb scan the source file directly and register the
        result as view `df`, instead of first materializing it through pandas
        (self.dataframe()). Returns False when the fast path doesn't apply (a prior
        op already narrowed columns, -progress was requested, or the source format
        has no fast native duckdb reader e.g. .sas7bdat) — the caller then falls
        back to the pandas-backed view."""
        from pytae.readers import CsvReader, ParquetReader, TxtReader

        if self._pending_exact is not None or self._progress:
            return False
        reader = self._reader
        path_sql = _sql_string_literal(str(reader.path))
        if isinstance(reader, ParquetReader):
            scan = f"read_parquet({path_sql})"
        elif isinstance(reader, (CsvReader, TxtReader)):
            if reader.encoding not in (None, "utf-8", "utf8"):
                return False  # duckdb's CSV reader doesn't support arbitrary encodings
            scan = f"read_csv({path_sql}, delim={_sql_string_literal(reader.sep)}, header=true)"
        else:
            return False  # e.g. .sas7bdat -- duckdb has no native reader for it
        limit_sql = f" LIMIT {int(self._nrows)}" if self._nrows is not None else ""
        con.execute(f"CREATE VIEW data AS SELECT * FROM {scan}{limit_sql}")
        return True

    def dataframe(self) -> pd.DataFrame:
        if self._df is None:
            df = self._reader.to_dataframe(
                columns=self._pending_exact, nrows=self._nrows, progress=self._progress,
                chunk_size=self._chunk_size,
            )
            self._df = df
        return self._df

    def columns(self) -> list[str]:
        if self._df is not None:
            return list(self._df.columns)
        return self._available_columns()

    def dtypes(self) -> pd.Series:
        if self._df is not None:
            return self._df.dtypes
        dtypes = self._reader.dtypes()
        return dtypes[self._pending_exact] if self._pending_exact is not None else dtypes

    def shape(self) -> tuple[int, int]:
        if self._df is not None:
            return self._df.shape
        rows = self._reader.shape()[0]
        cols = len(self._pending_exact) if self._pending_exact is not None else self._reader.shape()[1]
        return (rows, cols)

    def head(self, n: int) -> pd.DataFrame:
        if self._df is None:
            if self._nrows is not None:
                df = self.dataframe().head(n)
            else:
                df = self._reader.head(n)
                if self._pending_exact is not None:
                    df = select(df, *self._pending_exact)
        else:
            df = self.dataframe().head(n)
        self._df = df
        return df

    def tail(self, n: int) -> pd.DataFrame:
        if self._df is None:
            if self._nrows is not None:
                df = self.dataframe().tail(n)
            else:
                df = self._reader.tail(n)
                if self._pending_exact is not None:
                    df = select(df, *self._pending_exact)
        else:
            df = self.dataframe().tail(n)
        self._df = df
        return df

    def sample(self, n: int, *, seed: int | None = None, frac: float | None = None) -> pd.DataFrame:
        df = self.dataframe()
        if not (isinstance(df.index, pd.RangeIndex) and df.index.name is None):
            df = safe_reset_index(df)
        if frac is not None:
            sampled = df.sample(frac=frac, random_state=seed)
        else:
            n = min(n, len(df))
            sampled = df.sample(n=n, random_state=seed) if n else df.iloc[0:0]
        sampled = sampled.reset_index(drop=True)
        self._df = sampled
        return sampled


class _OrderedFlag(argparse.Action):
    """Boolean flag (like store_true) that also records its dest in namespace.op_order, in CLI order."""

    def __init__(self, option_strings, dest, **kwargs) -> None:
        super().__init__(option_strings, dest, nargs=0, default=False, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None) -> None:
        setattr(namespace, self.dest, True)
        namespace.op_order = getattr(namespace, "op_order", []) + [self.dest]


class _OrderedValue(argparse.Action):
    """Optional-value flag (nargs='?') that also records its dest in namespace.op_order, in CLI order."""

    def __call__(self, parser, namespace, values, option_string=None) -> None:
        val = self.const if values is None else values
        setattr(namespace, self.dest, val)
        namespace.op_order = getattr(namespace, "op_order", []) + [self.dest]
        op_values = getattr(namespace, "_op_values", None)
        if op_values is None:
            op_values = {}
            setattr(namespace, "_op_values", op_values)
        op_values.setdefault(self.dest, []).append(val)


class _OrderedStore(argparse.Action):
    """Required-value flag (nargs=default) that also records its dest in namespace.op_order, in CLI order."""

    def __call__(self, parser, namespace, values, option_string=None) -> None:
        setattr(namespace, self.dest, values)
        namespace.op_order = getattr(namespace, "op_order", []) + [self.dest]
        op_values = getattr(namespace, "_op_values", None)
        if op_values is None:
            op_values = {}
            setattr(namespace, "_op_values", op_values)
        op_values.setdefault(self.dest, []).append(values)


class _OrderedAppend(argparse.Action):
    """Collect repeated flags into a list and record each occurrence in op_order."""

    def __call__(self, parser, namespace, values, option_string=None) -> None:
        items = getattr(namespace, self.dest, None)
        items = [] if items is None else list(items)
        items.append(values)
        setattr(namespace, self.dest, items)
        namespace.op_order = getattr(namespace, "op_order", []) + [self.dest]

################################################################################
# FILE: src/pytae/cli_run.py
################################################################################
"""Execute one CLI pipeline: display/convert/agg ops against a file or named frames."""

from __future__ import annotations

import argparse
import io
import subprocess
import sys
from pathlib import Path
from typing import Any

import pandas as pd

from pytae.agg_df import agg_df
from pytae.cli_parsing import (
    parse_agg,
    parse_clean_columns_arg,
    parse_columns,
    parse_kv_spec,
    parse_long_arg,
    parse_pivot_arg,
    parse_sort_by,
    parse_wide_arg,
    unknown_columns_message,
)
from pytae.cli_pipeline import _Pipeline
from pytae.other_utilities import clean_columns, handle_missing, safe_reset_index
from pytae.readers import StdinReader, _split_path_suffixes, get_reader, write_dataframe


def _apply_round(df: pd.DataFrame, ndigits: int | None) -> pd.DataFrame:
    """Round numeric columns to ndigits; non-numeric columns pass through unchanged."""
    return df.round(ndigits) if ndigits is not None else df


def _dataframe_info(df: pd.DataFrame) -> str:
    buf = io.StringIO()
    df.info(buf=buf)
    return buf.getvalue().rstrip()


def _format_table(df: pd.DataFrame, *, index: bool | None = None, pretty: bool = False) -> str:
    """Render a DataFrame as standard pandas text, or a markdown/bordered table when pretty=True."""
    if index is None:
        index = not (isinstance(df.index, pd.RangeIndex) and df.index.name is None)
    if not pretty:
        return df.to_string(index=index)
    try:
        return df.to_markdown(index=index)
    except ImportError:
        return df.to_string(index=index)


def _compute_freq(source_df: pd.DataFrame, col: str | None, dropna: bool) -> pd.DataFrame:
    if col is None:
        if len(source_df.columns) == 1:
            col = source_df.columns[0]
        else:
            raise SystemExit("-freq: requires a column name (or a single-column DataFrame)")
    col = col.strip()
    if col not in source_df.columns:
        raise SystemExit(unknown_columns_message("-freq", [col], list(source_df.columns)))

    counts = source_df[col].value_counts(dropna=dropna)
    total = counts.sum()
    max_cnt = counts.max() if not counts.empty else 1
    bar_max_width = 20
    data = []
    for category, count in counts.items():
        bar_len = int(round((count / max_cnt) * bar_max_width)) if max_cnt > 0 else 0
        bar_str = "█" * bar_len
        pct_str = f"({(count / total * 100):.1f}%)" if total > 0 else "(0.0%)"
        dist_str = f"{bar_str} {pct_str}".strip()
        data.append({col: str(category), "Count": count, "Distribution": dist_str})
    return pd.DataFrame(data)


def _compute_hist(source_df: pd.DataFrame, hist_arg: str | None) -> pd.DataFrame:
    raw = (hist_arg or "").strip()
    bins = 10
    col = None
    if ":" in raw:
        col_part, _, b_str = raw.partition(":")
        col = col_part.strip()
        try:
            bins = int(b_str.strip())
        except ValueError:
            pass
    elif "," in raw and "bins=" in raw:
        col_part, _, b_part = raw.partition(",")
        col = col_part.strip()
        b_str = b_part.replace("bins=", "").strip()
        try:
            bins = int(b_str)
        except ValueError:
            pass
    elif raw:
        col = raw
    else:
        if len(source_df.columns) == 1:
            col = source_df.columns[0]
        else:
            raise SystemExit("-hist: requires a column name (e.g. -hist mass or -hist mass:10)")

    if col not in source_df.columns:
        raise SystemExit(unknown_columns_message("-hist", [col], list(source_df.columns)))

    s = pd.to_numeric(source_df[col], errors="coerce").dropna()
    if s.empty:
        return pd.DataFrame(columns=["Range", "Count", "Distribution"])

    cut_series = pd.cut(s, bins=bins, include_lowest=True)
    counts = cut_series.value_counts(sort=False)
    total = len(s)
    max_cnt = counts.max() if not counts.empty else 1
    bar_max_width = 20
    data = []
    for interval, count in counts.items():
        bar_len = int(round((count / max_cnt) * bar_max_width)) if max_cnt > 0 else 0
        bar_str = "█" * bar_len
        pct_str = f"({(count / total * 100):.1f}%)" if total > 0 else "(0.0%)"
        dist_str = f"{bar_str} {pct_str}".strip()
        left_fmt = f"{interval.left:.2f}".rstrip("0").rstrip(".")
        right_fmt = f"{interval.right:.2f}".rstrip("0").rstrip(".")
        range_str = f"[{left_fmt}, {right_fmt}]"
        data.append({"Range": range_str, "Count": count, "Distribution": dist_str})
    return pd.DataFrame(data)



def _copy_to_clipboard(text: str) -> None:
    """Copy text to the system clipboard via pbcopy/clip/xclip; warns on stderr if none is available."""
    if sys.platform == "darwin":
        cmd = ["pbcopy"]
    elif sys.platform == "win32":
        cmd = ["clip"]
    else:
        cmd = ["xclip", "-selection", "clipboard"]
    try:
        subprocess.run(cmd, input=text.encode(), check=True)
    except (OSError, subprocess.CalledProcessError):
        print("pytae: unable to copy to clipboard (no clipboard utility found)", file=sys.stderr)


def _output_text(text: str, args: argparse.Namespace) -> None:
    """Output text to stdout, piping through pydoc.pager if -pager is requested."""
    if getattr(args, "pager", False):
        import pydoc

        pydoc.pager(text)
    else:
        print(text)


def _format_bytes(num: int | float) -> str:
    """Format byte counts into human-readable strings (e.g. 1.2 KB, 3.4 MB)."""
    val = float(num)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if abs(val) < 1024.0:
            return f"{val:3.1f} {unit}" if unit != "B" else f"{int(val)} B"
        val /= 1024.0
    return f"{val:.1f} PB"


def _extract_metadata(path: Path, *, sep: str | None = None, encoding: str | None = None) -> str:
    """Extract file metadata without scanning full row data (fastest on Parquet)."""
    if not path.exists():
        return f"File not found: {path}"
    suffix, compression = _split_path_suffixes(path)
    file_bytes = path.stat().st_size
    lines: list[str] = [f"File: {path.name}", f"File size: {_format_bytes(file_bytes)}"]
    if suffix in (".parquet", ".pq"):
        try:
            import pyarrow.parquet as pa_parquet

            pf = pa_parquet.ParquetFile(path)
            md = pf.metadata
            lines.append(f"Format: Parquet (version {md.format_version})")
            lines.append(f"Rows: {md.num_rows:,}")
            lines.append(f"Columns: {md.num_columns}")
            lines.append(f"Row groups: {md.num_row_groups}")
            codecs: set[str] = set()
            total_uncompressed = 0
            for j in range(md.num_row_groups):
                rg = md.row_group(j)
                total_uncompressed += rg.total_byte_size
                for i in range(md.num_columns):
                    codecs.add(str(rg.column(i).compression))
            codec_str = ", ".join(sorted(codecs)) if codecs else "None"
            lines.append(f"Compression: {codec_str}")
            lines.append(f"Uncompressed data size: {_format_bytes(total_uncompressed)}")
            ratio = (file_bytes / total_uncompressed * 100) if total_uncompressed else 100
            lines.append(f"Space saving: {100 - ratio:.1f}%")
            lines.append("")
            lines.append("Schema:")
            schema = pf.schema_arrow
            max_name_len = max(len(name) for name in schema.names) if schema.names else 10
            lines.append(f"  {'#':<4} {'Column':<{max_name_len}}  {'Type'}")
            lines.append(f"  {'-'*4} {'-'*max_name_len}  {'-'*20}")
            for idx, field in enumerate(schema):
                lines.append(f"  {idx:<4} {field.name:<{max_name_len}}  {field.type}")
        except Exception as exc:
            lines.append(f"Error reading Parquet metadata: {exc}")
    else:
        fmt_name = "CSV" if suffix == ".csv" else ("JSON Lines" if suffix in (".jsonl", ".ndjson") else suffix.lstrip(".").upper())
        comp_str = f" ({compression})" if compression else ""
        lines.append(f"Format: {fmt_name}{comp_str}")
        try:
            reader = get_reader(path, sep=sep, encoding=encoding)
            cols = reader.columns()
            lines.append(f"Columns: {len(cols)}")
            lines.append(f"Column names: {', '.join(cols[:15])}{'...' if len(cols) > 15 else ''}")
        except UnicodeError as exc:
            lines.append(f"Error reading encoding: {_encoding_error_message(path, encoding, exc)}")
        except Exception as exc:
            lines.append(f"Error reading column info: {exc}")
    return "\n".join(lines)


def _compute_diff(
    left_df: pd.DataFrame,
    right_path: Path,
    left_name: str,
    *,
    sep: str | None = None,
    encoding: str | None = None,
) -> str:
    """Compare the current pipeline result against another tabular file."""
    if not right_path.exists():
        return f"Diff target not found: {right_path}"
    try:
        reader = get_reader(right_path, sep=sep, encoding=encoding)
        right_df = reader.to_dataframe()
    except UnicodeError as exc:
        return _encoding_error_message(right_path, encoding, exc)
    except Exception as exc:
        return f"Cannot read diff target '{right_path}': {exc}"

    r1, c1 = left_df.shape
    r2, c2 = right_df.shape
    cols1 = list(left_df.columns)
    cols2 = list(right_df.columns)

    lines: list[str] = [
        "Comparing:",
        f"  Left (source):  {left_name} ({r1:,} rows, {c1} cols)",
        f"  Right (target): {right_path.name} ({r2:,} rows, {c2} cols)",
        "",
        "Shape:",
        f"  Rows: {r1:,} vs {r2:,} ({'+' if r1 >= r2 else ''}{r1 - r2:,})",
        f"  Cols: {c1} vs {c2} ({'+' if c1 >= c2 else ''}{c1 - c2})",
    ]

    added_cols = [c for c in cols1 if c not in cols2]
    removed_cols = [c for c in cols2 if c not in cols1]
    common_cols = [c for c in cols1 if c in cols2]

    lines.append("")
    lines.append("Columns:")
    if added_cols:
        lines.append(f"  + Added in left ({len(added_cols)}):   {', '.join(added_cols)}")
    if removed_cols:
        lines.append(f"  - Removed in left ({len(removed_cols)}): {', '.join(removed_cols)}")
    lines.append(f"  Common ({len(common_cols)}):        {', '.join(common_cols)}")

    dtype_diffs: list[str] = []
    for c in common_cols:
        d1 = str(left_df[c].dtype)
        d2 = str(right_df[c].dtype)
        if d1 != d2:
            dtype_diffs.append(f"  * {c}: {d2} (right) -> {d1} (left)")

    lines.append("")
    lines.append("Schema Drift:")
    if dtype_diffs:
        lines.extend(dtype_diffs)
    else:
        lines.append("  None (all common columns have matching dtypes)")

    null_diffs: list[str] = []
    for c in common_cols:
        n1 = int(left_df[c].isna().sum())
        n2 = int(right_df[c].isna().sum())
        if n1 != n2:
            null_diffs.append(f"  * {c}: {n1:,} vs {n2:,} nulls ({'+' if n1 >= n2 else ''}{n1 - n2:,})")

    if null_diffs:
        lines.append("")
        lines.append("Null Counts:")
        lines.extend(null_diffs)

    if r1 == r2 and common_cols:
        try:
            mismatches = 0
            for c in common_cols:
                s1 = left_df[c]
                s2 = right_df[c]
                diff = ~((s1 == s2) | (s1.isna() & s2.isna()))
                mismatches += int(diff.sum())
            lines.append("")
            lines.append("Values:")
            if mismatches == 0 and cols1 == cols2:
                lines.append("  Identical: all cell values match exactly.")
            elif mismatches == 0:
                lines.append("  Matching: all cell values in common columns match.")
            else:
                lines.append(f"  Mismatches found in {mismatches:,} cells across common columns.")
        except Exception:
            pass

    return "\n".join(lines)


# Ops whose pandas equivalent does not return a DataFrame (shape -> tuple, cols -> Index,
# dtype -> Series, nulls -> Series, info() -> None, meta -> str, diff -> str). Like real method
# chaining, nothing can follow them except -o clip.
NON_DF_TERMINAL_OPS = frozenset({"shape", "cols", "dtype", "nulls", "info", "meta", "diff"})


def _list_order_names(names, order):
    if order == "asc":
        return sorted(names)
    if order == "desc":
        return sorted(names, reverse=True)
    return list(names)



def _list_order_index(series: pd.Series, order) -> pd.Series:
    if order == "asc":
        return series.sort_index(ascending=True)
    if order == "desc":
        return series.sort_index(ascending=False)
    return series





def cmd_convert(
    df: pd.DataFrame,
    source: Path,
    output: Path | None,
    *,
    rename: dict[str, str] | None = None,
    sep: str | None = None,
    encoding: str | None = None,
    progress: bool = False,
    announce: bool = True,
    chunk_size: int = 200_000,
    index: bool | None = None,
) -> None:
    if rename:
        df = df.rename(columns=rename)
    dest = output if output is not None else source.with_suffix(".csv")
    if dest.resolve() == source.resolve():
        raise SystemExit(
            f"refusing to overwrite the source file '{source}'; "
            f"specify a different format or an explicit path with -o/--output"
        )
    try:
        write_dataframe(df, dest, sep=sep, encoding=encoding, progress=progress, chunk_size=chunk_size, index=index)
    except ValueError as exc:
        raise SystemExit(str(exc)) from exc
    if announce:
        print(f"Wrote {len(df)} rows to {dest}")


def _fail(parser: argparse.ArgumentParser, batch: bool, msg: str) -> bool:
    """Abort immediately outside batch mode; in batch mode, warn on stderr and signal failure instead."""
    if not batch:
        parser.error(msg)
    print(f"pytae: {msg}", file=sys.stderr)
    return True


_COMMON_ENCODINGS = ("utf-8", "utf-8-sig", "latin-1", "cp1252")


def _encoding_error_message(path: Path, encoding: str | None, exc: UnicodeError) -> str:
    used = encoding or "utf-8"
    suggestions = ", ".join(enc for enc in _COMMON_ENCODINGS if enc != used)
    return f"can't decode '{path}' with encoding '{used}'; try -encoding {suggestions}"


def _resolve_alias(alias: str, frames: dict[str, pd.DataFrame], pipeline: _Pipeline, flag: str) -> tuple[pd.DataFrame | None, str | None]:
    """Look up a -merge/-concat alias: a -file alias, or 'df' for the pipeline's current
    result so far (lets repeated -merge/-concat calls fold in one more file at a time).
    Returns (frame, error_message)."""
    if alias == "data":
        if pipeline._df is None:
            return None, f"{flag}: 'data' isn't available yet — nothing has produced a pipeline result yet"
        return pipeline._df, None
    if alias not in frames:
        return None, f"{flag}: unknown -file alias '{alias}'"
    return frames[alias], None


def _process_path(
    path: Path | None,
    args: argparse.Namespace,
    parser: argparse.ArgumentParser,
    batch: bool,
    *,
    show_all: bool,
    select_specs: list[tuple[list[str], dict]],
    qry_specs: list[list[tuple[str, Any]]],
    mutate_specs: list[str],
    query_specs: list[str],
    sql_specs: list[str],
    replace_specs: list[tuple[list[str] | None, dict[str, str], bool]],
    rename_specs: list[dict[str, str]],
    frames: dict[str, pd.DataFrame] | None = None,
    merge_specs: list[dict] | None = None,
    concat_specs: list[dict] | None = None,
) -> bool:
    """Run every requested display/export/-agg action against one file, or (when frames
    is given, i.e. -file/-merge mode) against named in-memory frames instead. Returns True
    if an error occurred."""
    merge_specs = merge_specs or []
    concat_specs = concat_specs or []
    if frames is None:
        chunk_size = getattr(args, "chunk_size", None) or 200_000
        if str(path) == "-":
            data_bytes = sys.stdin.buffer.read()
            try:
                reader = StdinReader(
                    data_bytes,
                    fmt=getattr(args, "fmt", None),
                    sep=args.dlim,
                    encoding=args.encoding,
                    chunk_size=chunk_size,
                )
            except (ValueError, ImportError) as exc:
                return _fail(parser, batch, str(exc))

            pipeline = _Pipeline(
                reader, nrows=args.nrows, progress=args.progress, chunk_size=chunk_size,
            )
        elif path is None or not path.exists():
            return _fail(parser, batch, f"file not found: {path}")
        else:
            try:
                reader = get_reader(path, sep=args.dlim, encoding=args.encoding, chunk_size=chunk_size)
            except (ValueError, ImportError) as exc:
                return _fail(parser, batch, str(exc))

            pipeline = _Pipeline(
                reader, nrows=args.nrows, progress=args.progress, chunk_size=chunk_size,
            )
    else:
        pipeline = _Pipeline(frames=frames)

    out_target = str(args.output).strip() if args.output is not None else None
    image_exts = (".png", ".svg", ".pdf", ".jpg", ".jpeg", ".webp", ".tif", ".tiff")
    is_image_out = out_target is not None and any(out_target.lower().endswith(ext) for ext in image_exts)
    is_clip = out_target is not None and out_target.lower() in ("clip", "clipboard")
    is_file = out_target is not None and not is_clip and not is_image_out
    clip_action = None
    emit_stdout = not is_clip and not is_file and not is_image_out

    if show_all:
        df = _apply_round(pipeline.dataframe(), args.round_ndigits)
        if emit_stdout:
            print(_format_table(df, pretty=args.pretty))
        if is_clip:
            clip_action = lambda d=df: d.to_clipboard(index=False)

    op_order = getattr(args, "op_order", [])
    last_idx = len(op_order) - 1
    op_values = getattr(args, "_op_values", {})
    op_iters = {k: iter(v) for k, v in op_values.items()}

    def _next_op_val(op_name: str, fallback: Any = None) -> Any:
        it = op_iters.get(op_name)
        if it is not None:
            try:
                return next(it)
            except StopIteration:
                pass
        return fallback

    select_iter = iter(select_specs)
    qry_iter = iter(qry_specs)
    mutate_iter = iter(mutate_specs)
    query_iter = iter(query_specs)
    sql_iter = iter(sql_specs)
    replace_iter = iter(replace_specs)
    rename_iter = iter(rename_specs)
    merge_iter = iter(merge_specs)
    concat_iter = iter(concat_specs)

    def should_print(idx: int) -> bool:
        return emit_stdout and idx == last_idx

    def emit_frame(idx: int) -> None:
        nonlocal clip_action
        if not (should_print(idx) or is_clip):
            return
        df = _apply_round(pipeline.dataframe(), args.round_ndigits)
        if should_print(idx):
            print(_format_table(df, pretty=args.pretty))
        if is_clip:
            clip_action = lambda d=df: d.to_clipboard(index=False)

    for idx, op in enumerate(op_order):
        if op == "select":
            names, kwargs = next(select_iter)
            err = pipeline.apply_select(names, kwargs)
            if err:
                return _fail(parser, batch, err)
            emit_frame(idx)
        elif op == "rename":
            err = pipeline.apply_rename(next(rename_iter))
            if err:
                return _fail(parser, batch, err)
            emit_frame(idx)
        elif op == "qry":
            err = pipeline.apply_qry(next(qry_iter))
            if err:
                return _fail(parser, batch, err)
            emit_frame(idx)
        elif op == "mutate":
            by_cols = parse_columns(args.by) if args.by else None
            if by_cols:
                df_cur = pipeline.dataframe()
                if any(c not in df_cur.columns for c in by_cols):
                    return _fail(parser, batch, unknown_columns_message("-by", by_cols, list(df_cur.columns)))
            err = pipeline.apply_mutate(next(mutate_iter), by=by_cols, dropna=False if args.dropna is None else args.dropna)
            if err:
                return _fail(parser, batch, err)
            emit_frame(idx)
        elif op == "query":
            err = pipeline.apply_query(next(query_iter))
            if err:
                return _fail(parser, batch, err)
            emit_frame(idx)
        elif op == "sql":
            err = pipeline.apply_sql(next(sql_iter))
            if err:
                return _fail(parser, batch, err)
            emit_frame(idx)
        elif op == "replace_values":
            cols, mapping, exact = next(replace_iter)
            err = pipeline.apply_replace_values(cols, mapping, exact)
            if err:
                return _fail(parser, batch, err)
            emit_frame(idx)
        elif op == "merge":
            spec = next(merge_iter)
            if frames is None:
                return _fail(parser, batch, "-merge requires -file")
            left_df, err = _resolve_alias(spec["left"], frames, pipeline, "-merge")
            if err or left_df is None:
                return _fail(parser, batch, err or "-merge: missing left frame")
            right_df, err = _resolve_alias(spec["right"], frames, pipeline, "-merge")
            if err or right_df is None:
                return _fail(parser, batch, err or "-merge: missing right frame")
            merge_kwargs = {"how": spec["how"]}
            if spec["validate"]:
                merge_kwargs["validate"] = spec["validate"]
            if spec["how"] == "cross":
                pass  # cross joins don't take on=/left_on=/right_on=
            elif spec["on"] is not None:
                missing = [c for c in spec["on"] if c not in left_df.columns or c not in right_df.columns]
                if missing:
                    return _fail(parser, batch, f"-merge: on= column(s) not in both frames: {', '.join(missing)}")
                merge_kwargs["on"] = spec["on"]
            else:
                missing_left = [c for c in spec["left_on"] if c not in left_df.columns]
                missing_right = [c for c in spec["right_on"] if c not in right_df.columns]
                if missing_left:
                    return _fail(parser, batch, unknown_columns_message("-merge (left)", missing_left, list(left_df.columns)))
                if missing_right:
                    return _fail(parser, batch, unknown_columns_message("-merge (right)", missing_right, list(right_df.columns)))
                merge_kwargs["left_on"] = spec["left_on"]
                merge_kwargs["right_on"] = spec["right_on"]
            try:
                result = pd.merge(left_df, right_df, **merge_kwargs)
            except Exception as exc:
                return _fail(parser, batch, f"-merge: {exc}")
            pipeline._df = result
            if should_print(idx):
                print(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty))
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "concat":
            spec = next(concat_iter)
            if frames is None:
                return _fail(parser, batch, "-concat requires -file")
            dfs = []
            for alias in spec["frames"]:
                df, err = _resolve_alias(alias, frames, pipeline, "-concat")
                if err or df is None:
                    return _fail(parser, batch, err or f"-concat: missing frame '{alias}'")
                dfs.append(df)
            try:
                result = pd.concat(dfs, ignore_index=True)
            except Exception as exc:
                return _fail(parser, batch, f"-concat: {exc}")
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "shape":
            shape_str = str(pipeline.shape())
            if should_print(idx):
                _output_text(shape_str, args)
            if is_clip:
                clip_action = lambda s=shape_str: _copy_to_clipboard(s)
        elif op == "cols":
            names = _list_order_names(pipeline.columns(), args.cols)
            if should_print(idx):
                _output_text("\n".join(names), args)
            if is_clip:
                clip_action = lambda n=names: pd.Series(n).to_clipboard(index=False, header=False)
        elif op == "dtype":
            dtypes = _list_order_index(pipeline.dtypes(), args.dtype)
            if should_print(idx):
                _output_text(dtypes.to_string(), args)
            if is_clip:
                clip_action = lambda s=dtypes: s.to_clipboard()
        elif op == "nulls":
            nulls = _list_order_index(pipeline.dataframe().isna().sum(), args.nulls)
            if should_print(idx):
                _output_text(nulls.to_string(), args)
            if is_clip:
                clip_action = lambda s=nulls: s.to_clipboard()
        elif op == "describe":
            described = pipeline.dataframe().describe()
            pipeline._df = described
            if should_print(idx):
                _output_text(_format_table(_apply_round(described, args.round_ndigits), index=True, pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=described: _apply_round(d, args.round_ndigits).to_clipboard(index=True)
        elif op == "info":
            info_str = _dataframe_info(pipeline.dataframe())
            if should_print(idx):
                _output_text(info_str, args)
            if is_clip:
                clip_action = lambda s=info_str: _copy_to_clipboard(s)
        elif op == "meta":
            meta_str = _extract_metadata(path if path is not None else Path("<merged>"), sep=args.dlim, encoding=args.encoding)
            if should_print(idx):
                _output_text(meta_str, args)
            if is_clip:
                clip_action = lambda s=meta_str: _copy_to_clipboard(s)
        elif op == "diff":
            diff_path = Path(args.diff)
            diff_str = _compute_diff(pipeline.dataframe(), diff_path, path.name if path is not None else "<pipeline>", sep=args.dlim, encoding=args.encoding)
            if should_print(idx):
                _output_text(diff_str, args)
            if is_clip:
                clip_action = lambda s=diff_str: _copy_to_clipboard(s)
        elif op == "value_counts":
            source_df = pipeline.dataframe()
            value_count_cols = list(source_df.columns)

            if "count" in value_count_cols and "n" in value_count_cols:
                return _fail(parser, batch, "-value_counts: columns contain both 'count' and 'n'; rename one to avoid ambiguous count headers")

            try:
                if len(value_count_cols) == 1:
                    col = value_count_cols[0]
                    count_name = "count" if col != "count" else "n"
                    result = source_df[col].value_counts(dropna=False if args.dropna is None else args.dropna).rename(count_name).reset_index()
                    result.columns = [col, count_name]
                else:
                    count_name = "count" if "count" not in value_count_cols else "n"
                    result = source_df.value_counts(subset=value_count_cols, dropna=False if args.dropna is None else args.dropna).rename(count_name).reset_index()
            except Exception as exc:
                return _fail(parser, batch, str(exc))
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "freq":
            freq_col = _next_op_val("freq", getattr(args, "freq", None))
            source_df = pipeline.dataframe()
            result = _compute_freq(source_df, freq_col, False if args.dropna is None else args.dropna)
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "hist":
            hist_arg = _next_op_val("hist", getattr(args, "hist", None))
            source_df = pipeline.dataframe()
            result = _compute_hist(source_df, hist_arg)
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "unique":
            source_df = safe_reset_index(pipeline.dataframe())
            unique_df = source_df.drop_duplicates().reset_index(drop=True)
            pipeline._df = unique_df
            if should_print(idx):
                _output_text(_format_table(_apply_round(unique_df, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=unique_df: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "head":
            head_val = _next_op_val("head", args.head)
            df = pipeline.head(head_val)
            if should_print(idx):
                _output_text(_format_table(_apply_round(df, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_idx = not (isinstance(df.index, pd.RangeIndex) and df.index.name is None)
                clip_action = lambda d=df, ci=clip_idx: _apply_round(d, args.round_ndigits).to_clipboard(index=ci)
        elif op == "tail":
            tail_val = _next_op_val("tail", args.tail)
            df = pipeline.tail(tail_val)
            if should_print(idx):
                _output_text(_format_table(_apply_round(df, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_idx = not (isinstance(df.index, pd.RangeIndex) and df.index.name is None)
                clip_action = lambda d=df, ci=clip_idx: _apply_round(d, args.round_ndigits).to_clipboard(index=ci)
        elif op == "sample":
            sample_val = _next_op_val("sample", args.sample)
            sampled = pipeline.sample(sample_val, seed=args.seed, frac=args.frac)
            n = len(sampled)
            if should_print(idx):
                _output_text(_format_table(_apply_round(sampled, args.round_ndigits), pretty=args.pretty) if n else "(no rows)", args)
            if is_clip and n:
                clip_action = lambda d=sampled: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "sort_by":
            source_df = safe_reset_index(pipeline.dataframe())
            pipeline._df = source_df
            sort_by_arg = _next_op_val("sort_by", args.sort_by)
            sort_cols, order = parse_sort_by(sort_by_arg)
            if any(c not in source_df.columns for c in sort_cols):
                return _fail(parser, batch, unknown_columns_message("-sort_by", sort_cols, list(source_df.columns)))
            ascending = order != "desc"
            sorted_df = source_df.sort_values(by=sort_cols, ascending=ascending).reset_index(drop=True)
            pipeline._df = sorted_df
            if should_print(idx):
                _output_text(_format_table(_apply_round(sorted_df, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=sorted_df: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op in ("agg", "agg_df"):
            agg_val = _next_op_val("agg", args.agg)
            if isinstance(agg_val, str):
                parse_agg(agg_val)
            df_cur = pipeline.dataframe()
            by_cols = parse_columns(args.by) if args.by else None
            if by_cols and any(c not in df_cur.columns for c in by_cols):
                return _fail(parser, batch, unknown_columns_message("-by", by_cols, list(df_cur.columns)))
            try:
                result = agg_df(df_cur, by_cols, agg_val, dropna=False if args.dropna is None else args.dropna)
            except Exception as e:
                return _fail(parser, batch, str(e))
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "handle_missing":
            hm_val = _next_op_val("handle_missing", args.handle_missing)
            result = handle_missing(pipeline.dataframe(), fillna=hm_val)
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "clean_columns":
            cc_val = _next_op_val("clean_columns", args.clean_columns)
            opts = parse_clean_columns_arg(cc_val)
            result = clean_columns(pipeline.dataframe(), **opts)
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "long":
            from pytae.shape import long as long_fn
            source_df = pipeline.dataframe()
            long_val = _next_op_val("long", args.long)
            long_kwargs = parse_long_arg(long_val)
            cols_arg = long_kwargs.get("cols")
            if cols_arg:
                cols_list = parse_columns(cols_arg) if isinstance(cols_arg, str) else list(cols_arg)
                missing = [c for c in cols_list if c not in source_df.columns]
                if missing:
                    return _fail(parser, batch, unknown_columns_message("-long", missing, list(source_df.columns)))
            try:
                result = long_fn(source_df, **long_kwargs)
            except (KeyError, ValueError) as exc:
                return _fail(parser, batch, str(exc))
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "wide":
            from pytae.shape import wide as wide_fn
            source_df = pipeline.dataframe()
            wide_val = _next_op_val("wide", args.wide)
            wide_kwargs = parse_wide_arg(wide_val)
            to_check = [wide_kwargs.get("c", "variable"), wide_kwargs.get("v", "value")]
            by_cols = wide_kwargs.get("by") or wide_kwargs.get("index") or wide_kwargs.get("r") or wide_kwargs.get("rows")
            if by_cols:
                to_check.extend(parse_columns(by_cols) if isinstance(by_cols, str) else list(by_cols))
            missing = [c for c in to_check if c not in source_df.columns]
            if missing:
                return _fail(parser, batch, unknown_columns_message("-wide", missing, list(source_df.columns)))
            try:
                result = wide_fn(source_df, **wide_kwargs)
            except (KeyError, ValueError) as exc:
                return _fail(parser, batch, str(exc))
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "pivot":
            from pytae.shape import pivot as pivot_fn
            source_df = pipeline.dataframe()
            pivot_val = _next_op_val("pivot", getattr(args, "pivot", None))
            pivot_kwargs = parse_pivot_arg(pivot_val)
            if "dropna" not in pivot_kwargs:
                pivot_kwargs["dropna"] = False if args.dropna is None else args.dropna
            to_check = []
            for k in ("r", "rows", "index", "by", "c", "cols", "columns", "v", "values", "val", "vals"):
                val = pivot_kwargs.get(k)
                if val:
                    to_check.extend(parse_columns(val) if isinstance(val, str) else list(val))
            missing = [col for col in to_check if col not in source_df.columns]
            if missing:
                return _fail(parser, batch, unknown_columns_message("-pivot", missing, list(source_df.columns)))
            try:
                result = pivot_fn(source_df, **pivot_kwargs)
            except Exception as exc:
                return _fail(parser, batch, str(exc))
            pipeline._df = result
            if should_print(idx):
                _output_text(_format_table(_apply_round(result, args.round_ndigits), pretty=args.pretty), args)
            if is_clip:
                clip_action = lambda d=result: _apply_round(d, args.round_ndigits).to_clipboard(index=False)
        elif op == "plot":
            plot_spec = _next_op_val("plot", getattr(args, "plot", None))
            plot_kwargs = parse_kv_spec(plot_spec, flag="-plot")
            finalize_kwargs = parse_kv_spec(getattr(args, "finalize", None), flag="-finalize")

            source_df = pipeline.dataframe()
            import matplotlib

            from pytae.plotting import Plotter
            if is_image_out or not should_print(idx):
                matplotlib.use("Agg")
            import matplotlib.pyplot as plt

            if isinstance(finalize_kwargs.get("style"), str):
                try:
                    plt.style.use(finalize_kwargs["style"])
                except Exception as exc:
                    return _fail(parser, batch, f"-finalize style: {exc}")

            plotter_keys = {"mosaic", "figsize", "sharex", "sharey", "nrows", "ncols"}
            init_kwargs = {k: v for k, v in plot_kwargs.items() if k in plotter_keys}
            chart_kwargs = {k: v for k, v in plot_kwargs.items() if k not in plotter_keys}
            if "dropna" not in chart_kwargs and getattr(args, "dropna", None) is not None:
                chart_kwargs["dropna"] = args.dropna

            try:
                p = Plotter(source_df, **init_kwargs)
                p.plot(**chart_kwargs)
                p.finalize(**finalize_kwargs)
            except Exception as exc:
                return _fail(parser, batch, str(exc))

            if is_image_out:
                assert out_target is not None
                dest_path = Path(out_target)
                if args.out_dir is not None:
                    args.out_dir.mkdir(parents=True, exist_ok=True)
                    dest_path = args.out_dir / dest_path.name
                p.save(dest_path)
                print(f"Saved plot to {dest_path}")
            elif should_print(idx):
                p.show()

    if is_clip and clip_action is not None:
        clip_action()
    elif is_file:
        assert out_target is not None
        fmt = out_target.lower()
        out_dir: Path | None = getattr(args, "out_dir", None)
        valid_formats = (
            "csv", "parquet", "pq", "txt", "dat", "jsonl", "ndjson",
            "csv.gz", "txt.gz", "dat.gz", "jsonl.gz", "ndjson.gz",
        )
        if fmt in valid_formats:
            if path is None or str(path) == "-":
                return _fail(parser, batch, f"-o {out_target}: when reading from STDIN or -file mode, an explicit output file path is required")
            ext = ".parquet" if fmt == "pq" else f".{fmt}"
            clean_name = path.name[:-3] if path.name.lower().endswith(".gz") else path.name
            dest_name = f"{Path(clean_name).stem}{ext}"
            dest = (out_dir / dest_name) if out_dir is not None else path.with_name(dest_name)
        else:
            dest = Path(out_target)
            if out_dir is not None and not dest.is_absolute():
                dest = out_dir / dest

        df_to_write = _apply_round(pipeline.dataframe(), args.round_ndigits)
        try:
            cmd_convert(
                df_to_write,
                path if path is not None else Path("<merged>"),
                dest,
                sep=args.dlim,
                encoding=args.encoding,
                progress=args.progress,
                announce=True,
                chunk_size=getattr(args, "chunk_size", None) or 200_000,
            )
        except SystemExit as exc:
            return _fail(parser, batch, str(exc))

    return False

################################################################################
# FILE: src/pytae/mutate.py
################################################################################
from __future__ import annotations

import ast
import difflib
import inspect
import re
from collections.abc import Sequence
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from pandas.errors import UndefinedVariableError

from pytae._text import tokenize as _tokenize
from pytae._text import unquote_name as _unquote_name


def _split_mutate_entries(raw: str) -> list[tuple[str, str]]:
    """Split a mutate spec into raw (key, expression) text pairs on top-level
    commas/newlines/equals/colons, respecting quotes and nested (), [], {} — same tokenizing
    convention as qry()'s conditions dict. Lines starting with # are stripped."""
    lines = [line for line in raw.splitlines() if not line.strip().startswith("#")]
    cleaned = "\n".join(lines)
    entries: list[tuple[str, str]] = []
    for raw_entry in _tokenize(cleaned, ",\n", keep_quotes=True, track_brackets=True):
        if not raw_entry:
            continue
        # Support '=' as primary assignment operator (ignoring ==, <=, >=, !=)
        m_eq = re.match(
            r"^((?:[^\x22\x27\[\]=]|\x22[^\x22]*\x22|\x27[^\x27]*\x27|\[[^\]]*\])+?)(?<![<>!=])=(?!=)(.+)$",
            raw_entry,
        )
        if m_eq:
            entries.append((m_eq.group(1).strip(), m_eq.group(2).strip()))
            continue
        pieces = _tokenize(raw_entry, ":", keep_quotes=True, track_brackets=True)
        if len(pieces) >= 2:
            raise ValueError(f"invalid mutate spec: use '=' for assignment (e.g. 'col = expr'). Colon ':' is not supported in entry '{raw_entry}'")
        raise ValueError(f"invalid mutate spec: missing '=' in entry '{raw_entry}'")
    return entries


def parse_mutate_spec(raw: str) -> dict[str, str]:
    """Parse a mutate spec string like "bmi = body_mass_g / bill_length_mm ** 2,
    mass_kg = body_mass_g / 1000" into an ordered {new_col: expression} dict.
    If raw starts with '@', it reads the spec from the specified file path.
    Quoting the key is optional (matches qry()); the expression is kept as raw
    text — column names inside it must stay unquoted, since eval() treats a
    quoted name as a string literal, not a column reference."""
    stripped = raw.strip()
    if (stripped.startswith("'") and stripped.endswith("'")) or (stripped.startswith('"') and stripped.endswith('"')):
        stripped = stripped[1:-1].strip()
    if stripped.startswith("@"):
        path = stripped[1:].strip()
        if (path.startswith("'") and path.endswith("'")) or (path.startswith('"') and path.endswith('"')):
            path = path[1:-1].strip()
        file_path = Path(path)
        if not file_path.is_file():
            raise FileNotFoundError(f"mutate spec file not found: '{path}'")
        stripped = file_path.read_text(encoding="utf-8").strip()
    expressions: dict[str, str] = {}
    for key_raw, expr_raw in _split_mutate_entries(stripped):
        key = _unquote_name(key_raw)
        if not key:
            raise ValueError("invalid mutate spec: empty column name")
        expr = expr_raw.strip()
        if not expr:
            raise ValueError(f"invalid mutate spec: '{key}' has no expression")
        expressions[key] = expr
    if not expressions:
        raise ValueError('mutate expects entries like "col = expr"')
    return expressions


_NO_MAP_DEFAULT = object()


def _map_values(series, mapping, default=_NO_MAP_DEFAULT):
    """map(column, mapping[, default]) helper injected into the fallback eval
    namespace: recode a column's values through a lookup (dict or Series), like
    pandas Series.map(). Keys with no match become `default` if given, else
    NaN."""
    s = pd.Series(series)
    result = s.map(mapping)
    if default is not _NO_MAP_DEFAULT:
        keys = set(mapping.keys()) if hasattr(mapping, "keys") else (set(mapping.index) if hasattr(mapping, "index") else set(mapping))
        is_missing_key = ~s.isin(keys)
        result = result.mask(is_missing_key, default)
    return result


def _result_index(*candidates):
    """Index for a helper's Series result: the first Series among its inputs
    (a condition/branch/choice), so the result aligns on assignment and supports
    .str/.dt/etc. chaining. None (RangeIndex) if no input is a Series."""
    for candidate in candidates:
        if isinstance(candidate, pd.Series):
            return candidate.index
    return None


def _clean_condition(cond):
    if isinstance(cond, pd.Series):
        return cond.fillna(False).to_numpy(dtype=bool)
    if isinstance(cond, (list, tuple)):
        return pd.Series(cond).fillna(False).to_numpy(dtype=bool)
    return cond


def _if_else(condition, true_value, false_value):
    """if_else(condition, true_value, false_value) helper injected into the
    fallback eval namespace -> np.where(...), returned as a Series so it chains
    and composes like any pandas Series."""
    clean_cond = _clean_condition(condition)
    try:
        result = np.where(clean_cond, true_value, false_value)
    except TypeError:
        # incompatible dtypes across branches have no common numpy dtype --
        # object arrays accept anything
        result = np.where(
            clean_cond,
            np.asarray(true_value, dtype=object),
            np.asarray(false_value, dtype=object),
        )
    if getattr(result, "ndim", None) == 0:
        return result.item()
    return pd.Series(result, index=_result_index(condition, true_value, false_value))


def _case_when(*entries, default=None):
    """case_when((cond, value), ..., [default=...]) helper injected into the fallback
    eval namespace -> np.select(...), returned as a Series so it chains and
    composes. Each positional arg can be a (condition, value) pair checked in order
    (first match wins); a trailing non-tuple arg or `default=` keyword argument
    is the catch-all default (unmatched rows are NaN if omitted).
    Also supports flat alternating arguments: case_when(c1, v1, c2, v2, default=d)."""
    if not entries:
        raise ValueError("case_when expects at least one (condition, value) pair")
    conditions, choices = [], []
    has_tuple = any(isinstance(e, tuple) for e in entries)
    if not has_tuple:
        flat_entries = list(entries)
        if len(flat_entries) % 2 == 1:
            if default is not None:
                raise ValueError("case_when received an odd number of flat arguments and an explicit default=")
            default = flat_entries.pop()
        if len(flat_entries) < 2:
            raise ValueError("case_when needs at least one (condition, value) pair")
        for i in range(0, len(flat_entries), 2):
            conditions.append(flat_entries[i])
            choices.append(flat_entries[i + 1])
    else:
        for i, entry in enumerate(entries):
            if isinstance(entry, tuple):
                if len(entry) != 2:
                    raise ValueError(f"case_when pair must be (condition, value), got a {len(entry)}-tuple")
                condition, value = entry
                conditions.append(condition)
                choices.append(value)
            elif i == len(entries) - 1:
                if default is not None:
                    raise ValueError("case_when default specified both positionally and via keyword")
                default = entry
            else:
                raise ValueError("case_when default (a non-tuple) must be the last argument")
    if not conditions:
        raise ValueError("case_when needs at least one (condition, value) pair")
    clean_conditions = [_clean_condition(c) for c in conditions]
    try:
        result = np.select(clean_conditions, choices, default=default)
    except TypeError:
        # incompatible dtypes across choices/default -- object arrays accept anything
        object_choices = [np.asarray(choice, dtype=object) for choice in choices]
        object_default = default if default is None else np.asarray(default, dtype=object)
        result = np.select(clean_conditions, object_choices, default=object_default)  # type: ignore[arg-type]
    if getattr(result, "ndim", None) == 0:
        return result.item()
    return pd.Series(result, index=_result_index(*conditions, *choices))


def _coalesce(*candidates):
    """coalesce(col1, col2, ..., default) helper injected into the fallback eval namespace:
    returns the first non-null value for each row, like SQL COALESCE() or dplyr::coalesce()."""
    if not candidates:
        raise ValueError("coalesce expects at least one argument")
    series_candidates = [c for c in candidates if isinstance(c, pd.Series)]
    if not series_candidates:
        for c in candidates:
            if c is not None and not (isinstance(c, float) and np.isnan(c)) and not pd.isna(c):
                return c
        return None

    target_index = series_candidates[0].index
    res = None
    for cand in candidates:
        if isinstance(cand, pd.Series):
            s = cand if cand.index.equals(target_index) else cand.reindex(target_index)
            res = s.copy() if res is None else res.combine_first(s)
        else:
            s = pd.Series(cand, index=target_index)
            res = s if res is None else res.fillna(cand)
    return res


def _normalize_col_brackets(expr: str, columns: set[str]) -> str:
    """Rewrite [col] to `col` when col is a known DataFrame column, outside string literals."""
    out_chars: list[str] = []
    quote = None
    i, n = 0, len(expr)
    while i < n:
        ch = expr[i]
        if quote is not None:
            out_chars.append(ch)
            if ch == "\\" and i + 1 < n:
                out_chars.append(expr[i + 1])
                i += 2
                continue
            if ch == quote:
                quote = None
            i += 1
            continue
        if ch in "'\"":
            quote = ch
            out_chars.append(ch)
            i += 1
            continue
        if ch == "[":
            j = expr.find("]", i + 1)
            if j != -1:
                inner = expr[i + 1 : j]
                if inner in columns:
                    out_chars.append(f"`{inner}`")
                    i = j + 1
                    continue
        out_chars.append(ch)
        i += 1
    return "".join(out_chars)


def _rewrite_fallback_col_refs(expr: str, columns: set[str]) -> tuple[str, dict[str, str]]:
    """Rewrite `col` and [col] for known columns to synthetic identifiers like __col_safe_0__
    for the plain Python eval() fallback, outside string literals."""
    out_chars: list[str] = []
    quote = None
    i, n = 0, len(expr)
    mapping: dict[str, str] = {}
    col_idx = 0
    while i < n:
        ch = expr[i]
        if quote is not None:
            out_chars.append(ch)
            if ch == "\\" and i + 1 < n:
                out_chars.append(expr[i + 1])
                i += 2
                continue
            if ch == quote:
                quote = None
            i += 1
            continue
        if ch in "'\"":
            quote = ch
            out_chars.append(ch)
            i += 1
            continue
        if ch == "`":
            j = expr.find("`", i + 1)
            if j != -1:
                inner = expr[i + 1 : j]
                if inner in columns:
                    alias = f"__col_safe_{col_idx}__"
                    col_idx += 1
                    mapping[alias] = inner
                    out_chars.append(alias)
                    i = j + 1
                    continue
        if ch == "[":
            j = expr.find("]", i + 1)
            if j != -1:
                inner = expr[i + 1 : j]
                if inner in columns:
                    alias = f"__col_safe_{col_idx}__"
                    col_idx += 1
                    mapping[alias] = inner
                    out_chars.append(alias)
                    i = j + 1
                    continue
        out_chars.append(ch)
        i += 1
    return "".join(out_chars), mapping


def _strip_at_refs(expr: str) -> tuple[str, set[str]]:
    """Rewrite pandas-eval `@name` local-variable refs to bare `name` for the
    plain Python eval() fallback (which has no `@` syntax), and report which
    names were referenced. `@` inside string literals is left untouched."""
    out_chars: list[str] = []
    names: set[str] = set()
    quote = None
    i, n = 0, len(expr)
    while i < n:
        ch = expr[i]
        if quote is not None:
            out_chars.append(ch)
            if ch == "\\" and i + 1 < n:
                out_chars.append(expr[i + 1])
                i += 2
                continue
            if ch == quote:
                quote = None
            i += 1
            continue
        if ch in "'\"":
            quote = ch
            out_chars.append(ch)
            i += 1
            continue
        if ch == "@" and i + 1 < n and (expr[i + 1].isalpha() or expr[i + 1] == "_"):
            j = i + 1
            while j < n and (expr[j].isalnum() or expr[j] == "_"):
                j += 1
            names.add(expr[i + 1 : j])
            out_chars.append(expr[i + 1 : j])
            i = j
            continue
        out_chars.append(ch)
        i += 1
    return "".join(out_chars), names


class _BoolOpRewriter(ast.NodeTransformer):
    """Rewrite `and`/`or`/`not` to `&`/`|`/`~` for the Python-eval fallback, so
    boolean conditions on a Series work (plain `and`/`or` on a Series raises
    "ambiguous truth value"). Done on the AST, so operand grouping/precedence is
    preserved automatically. Matches pandas eval()'s own and/or translation."""

    def visit_BoolOp(self, node: ast.BoolOp) -> ast.AST:
        self.generic_visit(node)
        op: ast.operator = ast.BitAnd() if isinstance(node.op, ast.And) else ast.BitOr()
        folded = node.values[0]
        for right in node.values[1:]:
            folded = ast.BinOp(left=folded, op=op, right=right)
        return folded

    def visit_UnaryOp(self, node: ast.UnaryOp) -> ast.AST:
        self.generic_visit(node)
        if isinstance(node.op, ast.Not):
            return ast.UnaryOp(op=ast.Invert(), operand=node.operand)
        return node


def _compile_fallback(expr: str):
    """Parse a fallback expression and rewrite its boolean operators, returning a
    code object ready for eval()."""
    tree = ast.parse(expr, mode="eval")
    tree = _BoolOpRewriter().visit(tree)
    ast.fix_missing_locations(tree)
    return compile(tree, "<pytae-mutate>", "eval")


class _RowCount(int):
    """Integer that can also be called as a function n() to match dplyr/polars conventions."""

    def __call__(self) -> int:
        return int(self)


def _eval(out: pd.DataFrame, expr: str, local_dict: dict, global_dict: dict):
    """out.eval(expr), resolving @local_var references against the scope that
    called mutate() rather than mutate()'s own frame — mutate() sits between
    the user's call and this eval(), so pandas' default frame-walking would
    otherwise look in the wrong place.

    pandas' eval() can't parse some Python constructs (notably dict literals,
    e.g. `species.map({'a': 1})`, string slicing like `col.str[:8]`, and custom
    helpers like `if_else`, `case_when`, `coalesce`); when it rejects one
    (NotImplementedError, ValueError, SyntaxError, or UndefinedVariableError) we fall back
    to a plain Python eval() with each column exposed as a Series plus the caller's scope,
    so natural pandas method chains work. Bracketed `[col]` and backtick `col`
    names are safely mapped to valid identifiers in fallback mode. `@name` refs
    are rewritten to the caller-scope value, and `and`/`or`/`not` are rewritten
    to `&`/`|`/`~` so boolean conditions on a Series work."""
    known_cols = set(out.columns)
    normalized = _normalize_col_brackets(expr, known_cols)
    try:
        return out.eval(normalized, local_dict=local_dict, global_dict=global_dict)
    except (NotImplementedError, ValueError, SyntaxError, UndefinedVariableError):
        safe_expr, col_map = _rewrite_fallback_col_refs(expr, known_cols)
        rewritten, at_names = _strip_at_refs(safe_expr)
        namespace = {**global_dict, **local_dict}
        namespace["map"] = _map_values  # prefix map(col, {...}[, default]) form
        namespace["if_else"] = _if_else
        namespace["case_when"] = _case_when
        namespace["coalesce"] = _coalesce
        namespace["mean"] = lambda s: s.mean() if hasattr(s, "mean") else np.mean(s)
        namespace["sum"] = lambda s: s.sum() if hasattr(s, "sum") else np.sum(s)
        namespace["median"] = lambda s: s.median() if hasattr(s, "median") else np.median(s)
        namespace["min"] = lambda s: s.min() if hasattr(s, "min") else np.min(s)
        namespace["max"] = lambda s: s.max() if hasattr(s, "max") else np.max(s)
        namespace["std"] = lambda s: s.std() if hasattr(s, "std") else np.std(s)
        namespace["var"] = lambda s: s.var() if hasattr(s, "var") else np.var(s)
        if "n" not in out.columns:
            namespace["n"] = _RowCount(len(out))
        namespace.update({col: out[col] for col in out.columns if col.isidentifier()})
        for alias, orig_col in col_map.items():
            namespace[alias] = out[orig_col]
        for name in at_names:  # @-refs mean caller-scope vars, so they win over columns
            if name in local_dict:
                namespace[name] = local_dict[name]
            elif name in global_dict:
                namespace[name] = global_dict[name]
            else:
                raise UndefinedVariableError(f"local variable '{name}' is not defined")
        return eval(_compile_fallback(rewritten), namespace)  # noqa: S307 - spec is developer-authored, same trust as pandas eval


def mutate(
    df: pd.DataFrame,
    *args: Any,
    by: str | Sequence[str] | None = None,
    dropna: bool = False,
    observed: bool = True,
    **kwargs: Any,
) -> pd.DataFrame:
    """
    Create or overwrite columns using string expressions (`"col = expr"` or `col="expr"`),
    callables (`col=callable`), or dictionaries. Each evaluated in order via pandas eval() —
    no lambda needed for plain arithmetic/boolean column assignments.

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame to mutate columns on.
    *args : str or dict
        Positional expressions:
        - String expression(s), e.g. "bmi = body_mass_g / bill_length_mm ** 2".
        - File path prefixed with '@' (e.g. "@transforms.txt").
        - Dictionary of column expressions, e.g. {"bmi": "body_mass_g / 1000"}.
    by : str or sequence of str, optional
        Grouping column(s) across which to evaluate the mutations (like dplyr's `.by` or
        pandas transform/window operations). When specified, aggregations (e.g. `mean(x)`,
        `sum(x)`, `n`) and window expressions (e.g. `x - mean(x)`) are evaluated per group
        and broadcast back to each row without collapsing the DataFrame.
    dropna : bool, default False
        Whether to drop NA groups when grouping with `by=`. Defaults to False.
    observed : bool, default True
        Whether to observe categorical levels when grouping with `by=`.
    params : dict, optional
        Explicit dictionary of parameters/variables to make available for `@name` references.
        Can also be passed via `_params=`.
    **kwargs : Any
        Column expressions or callables passed as keyword arguments,
        e.g. `df.pt.mutate(bmi="body_mass_g / bill_length_mm ** 2", rank=1)`.

    Returns:
    --------
    pd.DataFrame
        A copy of df with each key assigned the result of its expression,
        applied in order.

    Examples:
    ---------
    >>> import pandas as pd
    >>> import pytae as pt
    >>> df = pd.DataFrame({'body_mass_g': [3000.0, 4000.0], 'bill_length_mm': [30.0, 40.0]})
    >>> df.pt.mutate(bmi="body_mass_g / bill_length_mm ** 2")
       body_mass_g  bill_length_mm       bmi
    0       3000.0            30.0  3.333333
    1       4000.0            40.0  2.500000

    >>> # Grouped mutation: compute group mean and deviation without collapsing rows
    >>> tips = pt.sample("tips")
    >>> tips.pt.mutate("avg_tip = mean(tip), diff = tip - avg_tip, group_size = n", by="day").head(3)

    >>> # Handling missing groups: dropna=False (default) computes on NA group; dropna=True sets NaN
    >>> sample = pd.DataFrame({'grp': ['A', 'A', None, None], 'val': [10.0, 20.0, 30.0, 50.0]})
    >>> sample.pt.mutate("avg = mean(val)", by="grp")
       grp   val   avg
    0    A  10.0  15.0
    1    A  20.0  15.0
    2  NaN  30.0  40.0
    3  NaN  50.0  40.0
    """
    _here = inspect.currentframe()
    caller_frame = None if _here is None else _here.f_back
    del _here
    while caller_frame is not None:
        # skip pytae internals (df.pt.mutate() sits in the accessor)
        name = caller_frame.f_globals.get("__name__") or ""
        if not name.startswith("pytae"):
            break
        caller_frame = caller_frame.f_back
    local_dict = caller_frame.f_locals.copy() if caller_frame is not None else {}
    global_dict = caller_frame.f_globals if caller_frame is not None else {}
    del caller_frame  # avoid holding a reference cycle via the frame object

    by = kwargs.pop("by", by)
    dropna = kwargs.pop("dropna", dropna)
    observed = kwargs.pop("observed", observed)

    by_cols: list[str] | None = None
    if by is not None:
        if isinstance(by, str):
            by_cols = [_unquote_name(c.strip()) for c in by.split(",") if c.strip()]
        elif isinstance(by, (list, tuple, set)):
            by_cols = [_unquote_name(c) for c in by]
        else:
            raise TypeError(f"mutate(): by must be a column name or sequence of names, got {type(by).__name__}")
        all_cols = list(df.columns)
        for col in by_cols:
            if col not in df.columns:
                close = difflib.get_close_matches(str(col), [str(c) for c in all_cols], n=1)
                hint = f" (did you mean '{close[0]}'?)" if close else ""
                raise KeyError(f"mutate: grouping column '{col}' not found in DataFrame{hint}")

    params = kwargs.pop("_params", None)
    if params is None and "params" in kwargs and isinstance(kwargs["params"], dict):
        params = kwargs.pop("params")

    if params is not None:
        local_dict.update(params)

    expressions: dict[str, Any] = {}
    for arg in args:
        if isinstance(arg, str):
            expressions.update(parse_mutate_spec(arg))
        elif isinstance(arg, dict):
            expressions.update(arg)
        elif callable(arg):
            raise TypeError(
                "mutate() expects callables as keyword arguments with the column name, "
                "e.g. df.pt.mutate(new_col=lambda d: d['x'] * 2)"
            )
        else:
            raise TypeError(
                f"mutate() received unexpected positional argument of type {type(arg).__name__}. "
                "Pass string expressions (e.g. 'col = expr'), a dict, or keyword arguments."
            )

    if kwargs:
        expressions.update(kwargs)

    if not expressions:
        raise ValueError(
            "mutate() expects at least one expression, e.g. df.pt.mutate('col = expr') "
            "or df.pt.mutate(col='expr')"
        )

    out = df.copy()
    for col, expr in expressions.items():
        try:
            if by_cols is None:
                if callable(expr):
                    out[col] = expr(out)
                elif isinstance(expr, str):
                    out[col] = _eval(out, expr, local_dict, global_dict)
                else:
                    out[col] = expr
            else:
                work_df = out.copy(deep=False)
                work_df["__pt_pos__"] = np.arange(len(out))
                pieces = []
                for _, group_df in work_df.groupby(by_cols, dropna=dropna, observed=observed):
                    pos = group_df["__pt_pos__"].to_numpy()
                    clean_group = group_df.drop(columns=["__pt_pos__"])
                    if callable(expr):
                        val = expr(clean_group)
                    elif isinstance(expr, str):
                        val = _eval(clean_group, expr, local_dict, global_dict)
                    else:
                        val = expr
                    if isinstance(val, pd.Series):
                        pieces.append(pd.Series(val.to_numpy(), index=pos))
                    else:
                        pieces.append(pd.Series(val, index=pos))
                if pieces:
                    combined = pd.concat(pieces).sort_index().reindex(np.arange(len(out)))
                    out[col] = combined.array
                else:
                    out[col] = pd.Series(index=out.index, dtype=float)
        except NameError as exc:  # UndefinedVariableError (pandas) and plain NameError (fallback)
            parts = str(exc).split("'")
            missing = parts[1] if len(parts) > 1 else ""
            close = difflib.get_close_matches(missing, list(out.columns), n=1) if missing else []
            hint = (
                f" (did you mean '{close[0]}'?)"
                if close
                else f" (if you intended a string literal, quote it like '{missing}')"
            )
            raise KeyError(f"mutate: '{col}': {exc}{hint}") from exc
    return out

################################################################################
# FILE: src/pytae/other_utilities.py
################################################################################
"""General tabular utilities: clipboard export, missing-value handling, column sorting, header cleaning, and value replacement."""

from __future__ import annotations

import re
from collections.abc import Hashable, Sequence
from typing import Any

import pandas as pd


def to_clip(df: pd.DataFrame | pd.Series) -> None:
    """Copy the DataFrame or Series to the system clipboard (tab-separated, no index)."""
    return df.to_clipboard(index=False)


pd.DataFrame.to_clip = to_clip
pd.Series.to_clip = to_clip


def safe_reset_index(df: pd.DataFrame) -> pd.DataFrame:
    """Reset index, ensuring no index level name collides with existing columns."""
    if isinstance(df.index, pd.RangeIndex) and df.index.name is None:
        return df
    idx_names = list(df.index.names)
    col_names = [str(col) for col in df.columns]
    colliding = []
    used_names = set(col_names)
    for i, name in enumerate(idx_names):
        target_name = str(name) if name is not None else ("index" if len(idx_names) == 1 else f"level_{i}")
        if target_name in used_names:
            colliding.append(target_name)
    if colliding:
        raise ValueError(
            f"Resetting index failed: index level name(s) {colliding} collide with existing column name(s). "
            "Rename the column or index before resetting."
        )
    return df.reset_index()


def handle_missing(
    df: pd.DataFrame,
    fillna: str = ".",
    numeric_fill: Any = 0,
    cols: Sequence[str] | None = None,
    preserve_categories: bool = True,
) -> pd.DataFrame:
    """Fill missing values across columns with type-appropriate defaults.

    - String/object columns: filled with `fillna` (default '.') and stripped.
    - Categorical columns: category preserved (adds `fillna` category if needed) and filled.
    - Numeric columns: filled with `numeric_fill` (default 0). Can also be "mean", "median", or None.
    - Bool/datetime columns: left untouched by default to avoid corruption.

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame to process.
    fillna : str, default '.'
        Fill value for string, object, and categorical columns.
    numeric_fill : number, "mean", "median", or None, default 0
        Fill value for numeric columns. If None, numeric columns are not filled.
    cols : sequence of str, optional
        Specific columns to handle. If None, all applicable columns are processed.
    preserve_categories : bool, default True
        If True, categorical columns keep their categorical dtype instead of being
        converted to object.
    """
    df = df.copy()
    target_cols = set(cols) if cols is not None else set(df.columns)

    # Categorical columns
    cat_cols = [c for c in df.columns if c in target_cols and isinstance(df[c].dtype, pd.CategoricalDtype)]
    for c in cat_cols:
        if preserve_categories:
            if fillna not in df[c].cat.categories:
                df[c] = df[c].cat.add_categories([fillna])
            df[c] = df[c].fillna(fillna)
        else:
            df[c] = df[c].astype("object").fillna(fillna).str.strip()

    # String / object columns (excluding categorical)
    def _is_string_col(s: pd.Series) -> bool:
        if isinstance(s.dtype, pd.CategoricalDtype):
            return False
        if s.dtype == object:
            non_na = s.dropna()
            if len(non_na) == 0:
                return True
            return non_na.map(type).eq(str).all()
        return pd.api.types.is_string_dtype(s)

    str_cols = [c for c in df.columns if c in target_cols and _is_string_col(df[c])]
    for c in str_cols:
        df[c] = df[c].fillna(fillna)
        if hasattr(df[c], "str") and callable(getattr(df[c].str, "strip", None)):
            df[c] = df[c].str.strip()

    # Numeric columns
    if numeric_fill is not None:
        num_cols = [c for c in df.select_dtypes(include="number").columns if c in target_cols]
        for c in num_cols:
            if numeric_fill == "mean":
                df[c] = df[c].fillna(df[c].mean())
            elif numeric_fill == "median":
                df[c] = df[c].fillna(df[c].median())
            else:
                df[c] = df[c].fillna(numeric_fill)

    return df


def cols(df: pd.DataFrame, ascending: bool | None = True) -> list:
    '''
    Return the column names of the DataFrame sorted or in original order.
    
    Parameters:
    df (pd.DataFrame): The DataFrame whose columns are to be returned.
    ascending (bool or None, optional): 
        - True (default): Sort alphabetically A-Z.
        - False: Sort alphabetically Z-A.
        - None: Return columns in their original DataFrame order (unsorted).
    
    Returns:
    list: A list of column names in the specified order.
    
    Raises:
    ValueError: If an invalid ascending parameter is provided.
    '''
    columns = df.columns.to_list()
    
    if ascending is True:
        return sorted(columns)
    elif ascending is False:
        return sorted(columns, reverse=True)
    elif ascending is None:
        return columns
    else:
        raise ValueError(f"Invalid ascending value '{ascending}'. Must be True, False, or None")



def clean_column_names(
    names: Sequence[Hashable],
    *,
    strip: bool = False,
    strip_special: bool = False,
    squeeze: bool = False,
    fill: str | None = None,
    case: str | None = None,
    dedupe: bool = False,
) -> list:
    """Clean a list of header names, in a fixed order: strip -> strip_special
    -> squeeze -> fill -> case -> dedupe.

    strip: trim leading/trailing whitespace.
    strip_special: remove anything that isn't a letter, digit, underscore,
        whitespace, or the fill character (e.g. %, $, #, !, parentheses,
        quotes) — if fill is set, that character is kept in place rather
        than stripped, since it's the intended separator.
    squeeze: collapse runs of internal whitespace to a single space.
    fill: replace each individual whitespace character with this string (pair
        with squeeze=True to collapse multi-space runs to one separator first).
    case: 'lower', 'upper', or 'proper' (str.title()).
    dedupe: number collisions after cleaning (revenue, revenue_1, revenue_2, ...).
    """
    cleaned = list(names)
    if strip:
        cleaned = [str(name).strip() for name in cleaned]
    if strip_special:
        keep = "".join(re.escape(ch) for ch in fill) if fill else ""
        pattern = re.compile(rf"[^\w\s{keep}]")
        cleaned = [pattern.sub("", str(name)) for name in cleaned]
    if squeeze:
        cleaned = [re.sub(r"\s+", " ", str(name)) for name in cleaned]
    if fill is not None:
        cleaned = [re.sub(r"\s", fill, str(name)) for name in cleaned]
    if case == "lower":
        cleaned = [str(name).lower() for name in cleaned]
    elif case == "upper":
        cleaned = [str(name).upper() for name in cleaned]
    elif case == "proper":
        cleaned = [str(name).title() for name in cleaned]
    if dedupe:
        seen = {}
        used = set(cleaned)  # original names, so a generated suffix never collides with a real one
        deduped = []
        for name in cleaned:
            if name not in seen:
                seen[name] = 0
                deduped.append(name)
            else:
                seen[name] += 1
                candidate = f"{name}_{seen[name]}"
                while candidate in used:
                    seen[name] += 1
                    candidate = f"{name}_{seen[name]}"
                used.add(candidate)
                deduped.append(candidate)
        cleaned = deduped
    return cleaned


def clean_columns(
    df: pd.DataFrame,
    strip: bool = False,
    strip_special: bool = False,
    squeeze: bool = False,
    fill: str | None = None,
    case: str | None = None,
    dedupe: bool = False,
) -> pd.DataFrame:
    """Clean column header names (see clean_column_names() for the per-key behavior)."""
    df = df.copy()
    df.columns = clean_column_names(
        list(df.columns), strip=strip, strip_special=strip_special,
        squeeze=squeeze, fill=fill, case=case, dedupe=dedupe,
    )
    return df


def replace_values(
    df: pd.DataFrame,
    v: dict[Any, Any],
    c: str | Sequence[str] | None = None,
    exact: bool = True,
) -> pd.DataFrame:
    """Replace values (pandas replace()), optionally scoped to specific columns.
    Parameter names match the -replace_values CLI flag's v=/c=/exact= keys.

    v: dict of {old: new}.
    c: a column name, list of column names, or None for the whole DataFrame.
    exact: True (default) matches whole cell values; False matches a substring
        anywhere in the cell (v's keys are regex-escaped, so they're treated
        as literal text, not patterns).
    """
    df = df.copy()
    to_replace = v if exact else {re.escape(str(old)): new for old, new in v.items()}
    target_cols = ([c] if isinstance(c, str) else list(c)) if c is not None else list(df.columns)
    if not exact:
        non_string_cols = [col for col in target_cols if not pd.api.types.is_string_dtype(df[col])]
        if non_string_cols:
            raise ValueError(
                f"replace_values: exact=False (substring match) only works on string/object "
                f"columns; non-string column(s): {non_string_cols}"
            )
    if c is not None:
        df[target_cols] = df[target_cols].replace(to_replace, regex=not exact)
    else:
        df = df.replace(to_replace, regex=not exact)
    return df

################################################################################
# FILE: src/pytae/plotting.py
################################################################################
from __future__ import annotations

import io
import math
import warnings
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def _check_no_aggregation(kwargs: dict) -> None:
    if 'aggfunc' in kwargs or 'aggregate' in kwargs:
        raise ValueError(
            "Plotter does not perform aggregation. "
            "Aggregate your DataFrame first using df.pt.agg(...) or df.pt.pivot(...) before plotting."
        )


# pytae's own control kwargs -- these drive Plotter's behavior directly and are
# never forwarded to pandas.plot() (see _prepare_plot_kwargs/_filter_plot_kwargs).
_CONTROL_KWARGS = [
    'on', 'print_data', 'clip_data', 'secondary_y',
    'xlabel', 'ylabel', 'palette', 'annot', 'fmt',
]

# Per-kind config: `strip` = last_kwargs keys to remove before calling pandas'
# own .plot(kind=...) (pytae-only knobs that aren't real pandas.plot() kwargs for
# that kind); `unsupported` = {kwarg: warning message} for kwargs that don't
# apply to this kind at all (warned about, then dropped via `strip`).
_KIND_SPECS = {
    'scatter': {
        'strip': ['dropna', *_CONTROL_KWARGS],
        'unsupported': {
            'dropna': "The 'dropna' argument is not applicable to scatter plots and will be ignored.",
        },
    },
    'hexbin': {
        'strip': ['by', 'dropna', *_CONTROL_KWARGS],
        'unsupported': {
            'dropna': "The 'dropna' argument is not applicable to hex plots and will be ignored.",
            'by': "Use 'c' and 'cmap' to split the hex plot by a particular column. The 'by' argument will be ignored.",
        },
    },
    'pie': {
        'strip': ['x', 'by', 'dropna', *_CONTROL_KWARGS],
        'unsupported': {},
    },
    'line': {
        'strip': ['y', 'by', 'dropna', *_CONTROL_KWARGS],
        'unsupported': {},
    },
    'box': {
        'strip': ['x', 'y', 'by', 'dropna', 'column', *_CONTROL_KWARGS],
        'unsupported': {
            'dropna': "The 'dropna' argument is not applicable to box plots and will be ignored.",
        },
    },
    'heatmap': {
        'strip': ['x', 'y', 'by', 'dropna', 'column', *_CONTROL_KWARGS],
        'unsupported': {},
    },
    'other': {
        'strip': ['y', 'by', 'dropna', *_CONTROL_KWARGS],
        'unsupported': {},
    },
    'kde': {
        'strip': ['x', 'y', 'by', 'dropna', 'column', *_CONTROL_KWARGS],
        'unsupported': {},
    },
    'hist': {
        'strip': ['x', 'y', 'by', 'dropna', 'column', *_CONTROL_KWARGS],
        'unsupported': {},
    },
}
_KIND_SPECS['density'] = _KIND_SPECS['kde']  # 'kde' and 'density' are aliases

# Kwargs each kind needs set (via .data()/.plot(x=, y=, by=, column=, ...)) before
# it can render at all -- checked up front so a missing one raises a clear pytae
# error instead of a cryptic pandas KeyError deep inside pandas.
_REQUIRED_KWARGS = {
    'scatter': ['x', 'y'],
    'hexbin': ['x', 'y'],
    'line': ['x'],
    'other': ['x'],
    'box': ['y'],
    'pie': ['by', 'y'],
    'kde': ['column'],
    'density': ['column'],
    'hist': ['column'],
    'heatmap': [],
}


def _normalize_col_label(c: object) -> str:
    if c is None or (isinstance(c, float) and np.isnan(c)) or pd.isna(c):
        return "nan"
    if isinstance(c, (float, np.floating)) and float(c).is_integer():
        return str(int(c))
    return str(c)


def _holds_integer(cols: pd.Index) -> bool:
    return getattr(cols, 'dtype', None) is not None and getattr(cols.dtype, 'kind', None) in 'iu'


def _remap_coord_if_needed(cols: pd.Index, val: Any) -> Any:
    if not isinstance(val, bool) and isinstance(val, (int, np.integer)) and val in cols:
        if _holds_integer(cols):
            return val
        if hasattr(cols, '_holds_integer') and cols._holds_integer():
            return val
        loc = cols.get_loc(val)
        if isinstance(loc, slice):
            return loc.start
        if isinstance(loc, np.ndarray):
            return int(np.where(loc)[0][0])
        return loc
    return val


def _reshape_by_to_wide(df: pd.DataFrame, x: Any, by: Any, y: Any, kind_name: str, dropna: bool = False) -> pd.DataFrame:
    """Reshape tidy (x, by, y) into wide columns per 'by' category without aggregation."""
    sub = df[[x, by, y]].copy()
    if dropna:
        sub = sub.dropna(subset=[x, by])
    if sub[[x, by]].duplicated().any():
        raise ValueError(
            f"Plotter: cannot plot {kind_name} with by='{by}' because duplicate ({x}, {by}) pairs exist. "
            "Aggregate first using df.pt.pivot(...) or df.pt.agg(...) before plotting."
        )
    raw_piv = sub.pivot(index=x, columns=by, values=y)
    norm_x = _normalize_col_label(x)
    for col in raw_piv.columns:
        if _normalize_col_label(col) == norm_x:
            raise ValueError(
                f"Plotter: category '{col}' in 'by={by}' collides with x-axis column name '{x}'. "
                "Rename the category or x-axis column to avoid ambiguous plot data."
            )
    res = raw_piv.reset_index()
    res.columns.name = None
    return res


class Plotter:
    """
    Method chainable .data() .. .plot() .. .plot() .. .data() .. .plot() .. .finalize().. 
    All args from previous plot call are carried forward unless data is changed or plot kind is changed.

    All pandas.plot **kwargs are supported except secondary_y. This is because we are handling secondary axis ourselves (e.g. A^).
    All artists like color, style etc are supported explicitly via dict or built-in palette.
    """

    _last_active = None

    def __new__(cls, *args, **kwargs):
        df = args[0] if args else kwargs.get("df")
        by_col = kwargs.get("by") if kwargs.get("by") is not None else kwargs.get("col")
        has_facet_trigger = (
            kwargs.get("ncols") is not None
            or kwargs.get("facet") is True
            or "col" in kwargs
        )
        is_plot_call = bool(
            {"kind", "x", "y", "column"}.intersection(kwargs.keys())
        )

        if (
            not kwargs.get("_is_facet_grid")
            and isinstance(df, pd.DataFrame)
            and by_col is not None
            and by_col in df
            and has_facet_trigger
            and is_plot_call
        ):
            if "mosaic" in kwargs or (len(args) > 1 and args[1] is not None):
                raise ValueError("Cannot combine 'mosaic' with faceting ('by' with 'ncols'/'facet'/'col').")
            if "nrows" in kwargs:
                raise ValueError("Cannot combine 'nrows' with faceting ('by' with 'ncols'/'facet'/'col'). Use 'ncols' to control grid columns.")

            facet_kwargs = dict(kwargs)
            facet_kwargs.pop("df", None)
            facet_kwargs.pop("by", None)
            facet_kwargs.pop("col", None)
            facet_kwargs.pop("facet", None)
            ncols = facet_kwargs.pop("ncols", None)
            return cls.facet(df, by=by_col, ncols=ncols, **facet_kwargs)

        return super().__new__(cls)

    def __init__(self, df=None, mosaic=None, figsize=None, sharex=False, sharey=False, nrows=None, ncols=None, **kwargs):
        _check_no_aggregation(kwargs)
        if getattr(self, '_initialized', False):
            return
        self._initialized = True
        Plotter._last_active = self
        if df is not None and not isinstance(df, pd.DataFrame):
            # df was passed as mosaic string or grid list e.g. Plotter("AB")
            if mosaic is None:
                mosaic = df
                df = None

        if mosaic is None and (nrows is not None or ncols is not None):
            nrows = nrows or 1
            ncols = ncols or 1
            import string
            labels = string.ascii_uppercase
            grid = []
            idx = 0
            for r in range(nrows):
                row = []
                for c in range(ncols):
                    label = labels[idx] if idx < len(labels) else f"ax_{idx}"
                    row.append(label)
                    idx += 1
                grid.append(row)
            mosaic = grid

        if mosaic is None:  # If not provided, a single axis labeled 'A' is created. 'on' is optional in this case
            mosaic = """
            A
            """
        self.fig, self.axd = plt.subplot_mosaic(mosaic=mosaic, figsize=figsize, sharex=sharex, sharey=sharey)

        self.last_kwargs = {}
        self.df = None
        self.plot_kwargs_store = {}
        self.tables = {}
        if df is not None:
            self.data(df)
        plt.close()

        plot_keys = {"kind", "x", "y", "column"}
        plot_kwargs = {k: v for k, v in kwargs.items() if k != "_is_facet_grid"}
        if plot_keys.intersection(plot_kwargs.keys()):
            if self.df is not None:
                self.plot(**plot_kwargs)
        elif plot_kwargs:
            unexpected = ", ".join(repr(k) for k in plot_kwargs.keys())
            raise TypeError(f"Plotter.__init__() got unexpected keyword argument(s): {unexpected}")

    def _repr_png_(self):
        """Render figure as PNG bytes for automatic display in Jupyter notebooks."""
        buf = io.BytesIO()
        self.fig.savefig(buf, format="png", bbox_inches="tight")
        return buf.getvalue()

    @property
    def pt(self):
        """Allow consistent fluent chaining using the .pt namespace (e.g. df.pt.plot().pt.finalize())."""
        return self

    def data(self, df):
        Plotter._last_active = self
        self.df = df
        self.last_kwargs = {}  # Reset kwargs when new data is chained!
        return self

    def get_data(self, on="A"):
        """Retrieve the DataFrame plotted on the specified axis key."""
        if on in self.tables:
            return self.tables[on]
        for k, v in self.tables.items():
            if k == on or k.rstrip("^") == on:
                return v
        raise KeyError(f"No plotted data found for axis {on!r}; available: {list(self.tables)}")

    @classmethod
    def facet(cls, df, by, *, ncols=None, sort=True, titles=True, figsize=None, sharex=False, sharey=False, **plot_kwargs):
        """Build a small-multiples grid, one axis per distinct value of `by`, and plot
        the same chart (**plot_kwargs, e.g. x=/y=/kind=) on each facet's own subset of
        rows. A grid isn't always exactly fillable (e.g. 5 groups doesn't tile a 2x3
        rectangle) -- the leftover cell(s) are left blank rather than turned into
        empty/unused axes.

        Parameters:
        -----------
        df : pd.DataFrame
            The data to facet.
        by : str
            Column whose distinct values become one facet (axis) each.
        ncols : int, optional
            Columns in the grid. Defaults to a roughly square layout (ceil(sqrt(n))).
        sort : bool, optional
            Sort group values (default True). Ignored for a `category` dtype column,
            which always uses its own defined category order.
        titles : bool, optional
            Set each facet's axis title to its group value (default True).
        figsize :
            Passed straight through to the underlying Plotter().
        sharex, sharey : bool, optional
            Share x or y axis scales across facets (default False).
        **plot_kwargs :
            Passed straight through to .plot() for every facet, e.g. x=, y=, kind=.

        Returns:
        --------
        Plotter
            A regular Plotter instance (chain `.finalize()` etc. as usual); `.axd` is
            keyed by the stringified group values, e.g. `plotter.axd['Adelie']`.
        """
        _check_no_aggregation(plot_kwargs)
        if isinstance(df[by].dtype, pd.CategoricalDtype):
            present = set(df[by].dropna().unique())
            groups = [g for g in df[by].cat.categories if g in present]
        else:
            groups = list(pd.unique(df[by].dropna()))
            if sort:
                groups = sorted(groups)
        if not groups:
            raise ValueError(f"facet: no groups found in column '{by}'")

        keys = [str(g) for g in groups]
        if len(set(keys)) != len(keys):
            raise ValueError(f"facet: group values in '{by}' produce duplicate axis keys once stringified: {keys}")

        if ncols is not None and ncols < 1:
            raise ValueError(f"facet: ncols must be at least 1, got {ncols}")
        ncols = ncols or math.ceil(math.sqrt(len(groups)))
        nrows = math.ceil(len(groups) / ncols)
        cells = keys + ["."] * (nrows * ncols - len(keys))  # pad leftover cells blank -- grid may not tile exactly
        grid = [cells[i:i + ncols] for i in range(0, len(cells), ncols)]

        plot_kwargs = dict(plot_kwargs)
        if "mosaic" in plot_kwargs:
            raise ValueError("Cannot combine 'mosaic' with faceting. Faceting generates its own layout grid.")
        if "nrows" in plot_kwargs:
            raise ValueError("Cannot combine 'nrows' with faceting. Use 'ncols' to control grid columns.")
        plot_kwargs.pop('mosaic', None)
        plot_kwargs.pop('nrows', None)
        plot_kwargs.pop('on', None)
        suptitle = plot_kwargs.pop('suptitle', None) or plot_kwargs.pop('title', None)
        plotter = cls(mosaic=grid, figsize=figsize, sharex=sharex, sharey=sharey, _is_facet_grid=True)
        for group, key in zip(groups, keys):
            plotter.data(df[df[by] == group]).plot(on=key, **plot_kwargs)
            if titles:
                plotter.axd[key].set_title(str(group))
        if suptitle:
            plotter.fig.suptitle(str(suptitle))
        plotter.df = df  # restore the full input frame -- each facet call above narrowed it to one group
        plotter._initialized = True
        Plotter._last_active = plotter
        return plotter

    def plot(self, **kwargs):
        Plotter._last_active = self
        self._update_kwargs(kwargs)
        self._validate_required_kwargs()
        ax = self._get_target_axis()

        if self.kind == 'scatter':
            self._plot_scatter(ax)
        elif self.kind == 'hexbin':
            self._plot_hexbin(ax)
        elif self.kind in ['line']:
            self._plot_line(ax)
        elif self.kind in ['kde', 'density']:
            self._plot_density(ax)
        elif self.kind == 'pie':
            self._plot_pie(ax)
        elif self.kind == 'hist':
            self._plot_hist(ax)
        elif self.kind == 'box':
            self._plot_box(ax)
        elif self.kind == 'heatmap':
            self._plot_heatmap(ax)
        else:
            self._plot_other(ax)

        return self

    def _update_kwargs(self, kwargs):
        _check_no_aggregation(kwargs)
        # Extract print_data and clip_data before updating last_kwargs
        self.print_data = kwargs.get('print_data', False)
        self.clip_data = kwargs.get('clip_data', False)

        # Store the current kind and check if it has changed
        new_kind = kwargs.get('kind', self.current_kind if hasattr(self, 'current_kind') else 'line')
        
        if hasattr(self, 'current_kind') and self.current_kind != new_kind:
            self.last_kwargs = {}  # Reset kwargs if kind changes
        self.current_kind = new_kind

        # Combine kwargs but exclude print_data and clip_data from last_kwargs
        combined_kwargs = {**self.last_kwargs, **kwargs}
        self.last_kwargs = {k: v for k, v in combined_kwargs.items() if k not in ['print_data', 'clip_data']}
        self.last_kwargs['kind'] = new_kind
        
        self.x = combined_kwargs.get('x', None)
        self.y = combined_kwargs.get('y', None)
        self.by = combined_kwargs.get('by', None)
        col_arg = combined_kwargs.get('column', None)
        self.column = col_arg if col_arg is not None else self.x
        self.kind = new_kind

        if self.kind == 'box':
            self.y = self.y if self.y is not None else self.column
            self.x = self.x if self.x is not None else self.by

        self.dropna = combined_kwargs.get('dropna', False)

    def _get_target_axis(self):
        ax_key = self.last_kwargs.get('on', 'A')
        if self.last_kwargs.get('secondary_y') and not ax_key.endswith('^'):
            ax_key = f"{ax_key}^"
        if '^' in ax_key:
            base_key = ax_key.rstrip('^')
            if base_key not in self.axd:
                raise ValueError(f"Unknown mosaic key '{base_key}'; available: {list(self.axd)}")
            ax = self.axd[base_key]
            if not hasattr(ax, 'right_ax'):
                right_ax = ax.twinx()
                right_ax.set_label(ax_key)
                ax.right_ax = right_ax
                if ax_key not in self.axd:
                    self.axd[ax_key] = right_ax
                else:
                    warnings.warn(f"Axis '{ax_key}' already exists in axd; using existing axis instead of creating new one.")
                    right_ax = self.axd[ax_key]
            else:
                right_ax = ax.right_ax
            return right_ax
        else:
            if ax_key not in self.axd:
                raise ValueError(f"Unknown mosaic key '{ax_key}'; available: {list(self.axd)}")
            return self.axd[ax_key]

    def _validate_required_kwargs(self):
        """Raise a clear error if a kwarg this kind needs (x=/y=/by=/column=) was
        never set, instead of letting a cryptic pandas KeyError surface later
        from deep inside get_pivot_data()/etc."""
        spec_key = self.kind if self.kind in _REQUIRED_KWARGS else 'other'
        missing = [name for name in _REQUIRED_KWARGS[spec_key] if getattr(self, name, None) is None]
        if missing:
            needed = ', '.join(f"{name}=" for name in missing)
            raise ValueError(f"Plotter: kind='{self.kind}' needs {needed}")

    def _prepare_plot_kwargs(self, spec_key):
        """Warn about any kwarg unsupported by this kind (see _KIND_SPECS), then
        return last_kwargs with pytae's own control/unsupported keys stripped
        so only real pandas.plot() kwargs remain."""
        spec = _KIND_SPECS[spec_key]
        for key, message in spec['unsupported'].items():
            if key in self.last_kwargs:
                warnings.warn(message)
        return self._filter_plot_kwargs(spec['strip'])

    def _apply_axis_labels(self, ax):
        """Apply optional custom xlabel and ylabel."""
        xlabel = self.last_kwargs.get('xlabel')
        ylabel = self.last_kwargs.get('ylabel')
        if xlabel is not None:
            ax.set_xlabel(xlabel)
        if ylabel is not None:
            ax.set_ylabel(ylabel)

    def _get_palette_colors(self, categories, palette=None, explicit_colors=None):
        """Map category names to colors from an explicit dict, list, or named matplotlib colormap/palette."""
        if isinstance(explicit_colors, dict):
            color_map = {}
            for k, v in explicit_colors.items():
                color_map[k] = v
                color_map[str(k)] = v
            return color_map
        if isinstance(explicit_colors, str):
            return {cat: explicit_colors for cat in categories}
        if isinstance(explicit_colors, (list, tuple)):
            return {cat: explicit_colors[i % len(explicit_colors)] for i, cat in enumerate(categories)}

        palette = palette or "tab10"
        try:
            cmap = plt.get_cmap(palette)
        except (ValueError, AttributeError):
            warnings.warn(f"Unknown palette '{palette}', falling back to 'tab10'.")
            cmap = plt.get_cmap("tab10")

        n = len(categories)
        if hasattr(cmap, "colors") and len(cmap.colors) >= n:
            return {cat: cmap.colors[i] for i, cat in enumerate(categories)}
        elif n <= 1:
            return {categories[0]: cmap(0.5)}
        else:
            return {cat: cmap(i / max(n - 1, 1)) for i, cat in enumerate(categories)}

    def _plot_scatter(self, ax):
        """Plot a scatter chart with support for continuous 'c' or discrete 'by' categories."""
        plot_dict = self._prepare_plot_kwargs('scatter')
        self._store_plot_kwargs(ax, plot_dict)
        k = self.df.copy()

        if self.by is not None:
            if self.by not in k.columns:
                raise KeyError(f"Plotter: column '{self.by}' for 'by' not found in DataFrame")
            if isinstance(k[self.by].dtype, pd.CategoricalDtype):
                present = set(k[self.by].dropna().unique())
                groups = [g for g in k[self.by].cat.categories if g in present]
            else:
                groups = list(pd.unique(k[self.by].dropna()))

            palette_name = self.last_kwargs.get('palette')
            color_arg = self.last_kwargs.get('color')
            color_map = self._get_palette_colors(groups, palette=palette_name, explicit_colors=color_arg)

            scatter_kwargs = dict(plot_dict)
            for non_scatter in ['x', 'y', 'by', 'kind', 'title', 'color', 'c', 'legend', 'grid', 'logx', 'logy', 'rot', 'fontsize', 'colormap', 'colorbar', 'xlim', 'ylim']:
                scatter_kwargs.pop(non_scatter, None)

            s_arg = scatter_kwargs.pop('s', 20)
            marker = scatter_kwargs.pop('marker', 'o')
            alpha = scatter_kwargs.pop('alpha', 0.8)

            for group in groups:
                grp_mask = (k[self.by] == group)
                grp_df = k[grp_mask]
                grp_color = color_map.get(group) if isinstance(color_map, dict) else None
                if isinstance(s_arg, str) and s_arg in k.columns:
                    grp_s = k.loc[grp_mask, s_arg].to_numpy()
                elif hasattr(s_arg, '__len__') and not isinstance(s_arg, (str, bytes)):
                    grp_s = np.asarray(s_arg)[grp_mask.to_numpy()]
                else:
                    grp_s = s_arg

                ax.scatter(
                    grp_df[self.x],
                    grp_df[self.y],
                    label=str(group),
                    color=grp_color,
                    s=grp_s,
                    marker=marker,
                    alpha=alpha,
                    **scatter_kwargs,
                )
            if plot_dict.get('grid'):
                ax.grid(True)
            if plot_dict.get('logx'):
                ax.set_xscale('log')
            if plot_dict.get('logy'):
                ax.set_yscale('log')
            if 'xlim' in plot_dict:
                ax.set_xlim(plot_dict['xlim'])
            if 'ylim' in plot_dict:
                ax.set_ylim(plot_dict['ylim'])
            rot = plot_dict.get('rot')
            fontsize = plot_dict.get('fontsize')
            if rot is not None or fontsize is not None:
                tick_kw = {}
                if rot is not None:
                    tick_kw['labelrotation'] = rot
                if fontsize is not None:
                    tick_kw['labelsize'] = fontsize
                ax.tick_params(axis='x', **tick_kw)
            self.ax = ax
            if self.last_kwargs.get('title'):
                ax.set_title(self.last_kwargs['title'])
            if self.last_kwargs.get('xlabel') is None and self.x is not None:
                ax.set_xlabel(str(self.x))
            if self.last_kwargs.get('ylabel') is None and self.y is not None:
                ax.set_ylabel(str(self.y))
            self._apply_axis_labels(ax)
            self._handle_data_output(k, ax)
            return

        c = plot_dict.get('c', None)
        if c is not None:
            try:
                is_in_cols = (c in k.columns)
            except TypeError:
                is_in_cols = False
            if is_in_cols and not pd.api.types.is_numeric_dtype(k[c]):
                k[c] = k[c].astype('category')
        for coord in ('x', 'y'):
            if coord in plot_dict:
                plot_dict[coord] = _remap_coord_if_needed(k.columns, plot_dict[coord])
        self.ax = k.plot(ax=ax, **plot_dict)
        self._apply_axis_labels(ax)
        self._handle_data_output(k, ax)

    def _plot_pie(self, ax):
        """Plot a pie chart."""
        plot_dict = self._prepare_plot_kwargs('pie')
        self._store_plot_kwargs(ax, plot_dict)
        cat_col = self.by if self.by is not None else self.x
        if cat_col is not None and cat_col in self.df.columns:
            sub = self.df[[cat_col, self.y]].copy()
            if self.dropna:
                sub = sub.dropna(subset=[cat_col])
            if sub[cat_col].duplicated().any():
                raise ValueError(
                    f"Plotter: pie chart category '{cat_col}' has duplicate entries. "
                    "Aggregate first using df.pt.agg(...) or df.pt.freq(...) before plotting."
                )
            if not self.dropna and sub[cat_col].isna().any():
                sub[cat_col] = sub[cat_col].fillna("NaN")
            pie_df = sub.set_index(cat_col)[[self.y]]
        else:
            pie_df = self.df[[self.y]]

        if 'colors' in self.last_kwargs:
            color_arg = self.last_kwargs['colors']
            if isinstance(color_arg, (list, tuple)):
                plot_dict['colors'] = color_arg
            elif isinstance(color_arg, dict):
                plot_dict['colors'] = [color_arg.get(category, 'grey') for category in pie_df.index]
        elif 'palette' in self.last_kwargs:
            colors = self._get_palette_colors(list(pie_df.index), palette=self.last_kwargs['palette'])
            plot_dict['colors'] = [colors.get(category, 'grey') for category in pie_df.index]
        self.ax = pie_df.plot(ax=ax, **plot_dict)
        self._apply_axis_labels(ax)
        self._handle_data_output(pie_df, ax)

    def _plot_hexbin(self, ax):
        """Plot a hexbin chart."""
        plot_dict = self._prepare_plot_kwargs('hexbin')
        self._store_plot_kwargs(ax, plot_dict)
        for coord in ('x', 'y'):
            if coord in plot_dict:
                plot_dict[coord] = _remap_coord_if_needed(self.df.columns, plot_dict[coord])
        self.ax = self.df.plot(ax=ax, **plot_dict)
        self._apply_axis_labels(ax)
        self._handle_data_output(self.df, ax)

    def _plot_line(self, ax):
        """Plot a line chart."""
        plot_dict = self._prepare_plot_kwargs('line')
        self._store_plot_kwargs(ax, plot_dict)
        style = plot_dict.pop('style', None)
        width = plot_dict.pop('width', None)
        color_map = plot_dict.pop('color', None)
        palette = self.last_kwargs.get('palette')

        plot_data = self.df.copy()
        if self.by is not None:
            if self.by not in plot_data.columns:
                raise KeyError(f"Plotter: column '{self.by}' for 'by' not found in DataFrame")
            if self.y is None:
                raise ValueError(f"Plotter: kind='line' with by='{self.by}' needs y=")
            if self.x is not None:
                plot_data = _reshape_by_to_wide(plot_data, self.x, self.by, self.y, 'line', dropna=self.dropna)
        elif self.x is not None and self.y is not None:
            y_cols = list(self.y) if isinstance(self.y, (list, tuple)) else [self.y]
            cols = [self.x] + [c for c in y_cols if c != self.x and c in plot_data.columns]
            plot_data = plot_data[cols]
            if self.dropna and self.x in plot_data.columns:
                plot_data = plot_data.dropna(subset=[self.x])
        elif self.x is not None:
            if self.dropna and self.x in plot_data.columns:
                plot_data = plot_data.dropna(subset=[self.x])

        data_cols = [c for c in plot_data.columns if c != self.x]

        if color_map and isinstance(color_map, dict):
            color_list = [color_map.get(col, None) for col in data_cols]
            plot_dict['color'] = color_list
        elif palette:
            colors = self._get_palette_colors(data_cols, palette=palette)
            plot_dict['color'] = [colors.get(c) for c in data_cols]
        elif color_map is not None:
            plot_dict['color'] = color_map

        if style is not None and not isinstance(style, dict):
            plot_dict['style'] = style
        if width is not None and not isinstance(width, dict):
            plot_dict['linewidth'] = width

        if 'x' in plot_dict:
            plot_dict['x'] = _remap_coord_if_needed(plot_data.columns, plot_dict['x'])
        self.ax = plot_data.plot(ax=ax, **plot_dict)
        if style is not None and isinstance(style, dict):
            style_lookup = {str(k): v for k, v in style.items()}
            style_lookup.update(style)
            for line in self.ax.get_lines():
                lbl = line.get_label()
                if lbl in style_lookup:
                    line.set_linestyle(style_lookup[lbl])
        if width is not None:
            if isinstance(width, dict):
                width_lookup = {str(k): v for k, v in width.items()}
                width_lookup.update(width)
                for line in self.ax.get_lines():
                    lbl = line.get_label()
                    if lbl in width_lookup:
                        line.set_linewidth(width_lookup[lbl])
            else:
                for line in self.ax.get_lines():
                    line.set_linewidth(width)
        self._apply_axis_labels(ax)
        self._handle_data_output(plot_data, ax)

    def _plot_box(self, ax):
        """Plot a box chart showing unaggregated distributions per group or single column."""
        plot_dict = self._prepare_plot_kwargs('box')
        self._store_plot_kwargs(ax, plot_dict)
        val_col = self.y
        cat_col = self.x if self.x is not None else self.by

        if cat_col is not None:
            if cat_col not in self.df.columns:
                raise KeyError(f"Plotter: column '{cat_col}' not found in DataFrame")
            if isinstance(self.df[cat_col].dtype, pd.CategoricalDtype):
                present = set(self.df[cat_col].dropna().unique())
                categories = [c for c in self.df[cat_col].cat.categories if c in present]
            else:
                categories = list(pd.unique(self.df[cat_col].dropna()))
            cols = {str(cat): self.df.loc[self.df[cat_col] == cat, val_col].reset_index(drop=True) for cat in categories}
            box_df = pd.DataFrame(cols)
        else:
            box_df = self.df[[val_col]]

        palette = self.last_kwargs.get('palette')
        color_arg = self.last_kwargs.get('color')
        if palette or color_arg:
            plot_dict['patch_artist'] = True
            plot_dict.pop('color', None)

        self.ax = box_df.plot(ax=ax, **plot_dict)

        if palette or color_arg:
            colors = self._get_palette_colors(list(box_df.columns), palette=palette, explicit_colors=color_arg)
            new_patches = ax.patches[-len(box_df.columns):]
            for patch, col in zip(new_patches, box_df.columns):
                if col in colors:
                    patch.set_facecolor(colors[col])

        self._apply_axis_labels(ax)
        self._handle_data_output(box_df, ax)

    def _plot_heatmap(self, ax):
        """Plot a 2D heatmap matrix (correlation matrix or pivoted data)."""
        plot_dict = self._prepare_plot_kwargs('heatmap')
        self._store_plot_kwargs(ax, plot_dict)
        cmap = self.last_kwargs.get('cmap', self.last_kwargs.get('palette', 'viridis'))
        annot = self.last_kwargs.get('annot', False)
        fmt = self.last_kwargs.get('fmt', '.2f')

        imshow_kwargs = {}
        for key in ['vmin', 'vmax', 'interpolation', 'origin', 'norm', 'alpha']:
            if key in self.last_kwargs:
                imshow_kwargs[key] = self.last_kwargs[key]

        if self.x is not None and self.by is not None and self.y is not None:
            wide_df = _reshape_by_to_wide(self.df, self.x, self.by, self.y, 'heatmap', dropna=self.dropna)
            matrix = wide_df.set_index(self.x)
        elif self.x is not None and self.x in self.df.columns:
            sub = self.df.dropna(subset=[self.x]) if self.dropna else self.df
            matrix = sub.set_index(self.x).select_dtypes(include='number')
        else:
            matrix = self.df.select_dtypes(include='number')
            if self.dropna:
                matrix = matrix.dropna(how='all')

        if matrix.empty:
            raise ValueError("Plotter: heatmap requires numeric columns to plot.")

        im = ax.imshow(matrix.values, cmap=cmap, aspect='auto', **imshow_kwargs)
        ax.set_xticks(range(len(matrix.columns)))
        ax.set_xticklabels([str(c) for c in matrix.columns], rotation=45, ha='right')
        ax.set_yticks(range(len(matrix.index)))
        ax.set_yticklabels([str(idx) for idx in matrix.index])

        if annot:
            for i in range(len(matrix.index)):
                for j in range(len(matrix.columns)):
                    val = matrix.iloc[i, j]
                    is_num = isinstance(val, (int, float, np.integer, np.floating)) and not isinstance(val, (bool, np.bool_))
                    text_str = format(val, fmt) if is_num else str(val)
                    is_nan = is_num and (np.isnan(val) if isinstance(val, (float, np.floating)) else False)
                    if is_num and not is_nan:
                        rgba = im.cmap(im.norm(val))
                        luminance = 0.299 * rgba[0] + 0.587 * rgba[1] + 0.114 * rgba[2]
                        text_color = "white" if luminance < 0.5 else "black"
                    else:
                        text_color = "white"
                    ax.text(j, i, text_str, ha='center', va='center', color=text_color)

        self.fig.colorbar(im, ax=ax)
        if self.last_kwargs.get('title'):
            ax.set_title(self.last_kwargs['title'])
        self._apply_axis_labels(ax)
        self.ax = ax
        self._handle_data_output(matrix, ax)

    # bar, barh, area
    def _plot_other(self, ax):
        plot_dict = self._prepare_plot_kwargs('other')
        self._store_plot_kwargs(ax, plot_dict)
        plot_data = self.df.copy()
        if self.by is not None:
            if self.by not in plot_data.columns:
                raise KeyError(f"Plotter: column '{self.by}' for 'by' not found in DataFrame")
            if self.y is None:
                raise ValueError(f"Plotter: kind='{self.kind}' with by='{self.by}' needs y=")
            if self.x is not None:
                plot_data = _reshape_by_to_wide(plot_data, self.x, self.by, self.y, self.kind, dropna=self.dropna)
        elif self.x is not None and self.y is not None:
            y_cols = list(self.y) if isinstance(self.y, (list, tuple)) else [self.y]
            cols = [self.x] + [c for c in y_cols if c != self.x and c in plot_data.columns]
            plot_data = plot_data[cols]
            if self.dropna and self.x in plot_data.columns:
                plot_data = plot_data.dropna(subset=[self.x])
        elif self.x is not None:
            if self.dropna and self.x in plot_data.columns:
                plot_data = plot_data.dropna(subset=[self.x])

        if self.x is not None and self.x in plot_data.columns and self.kind in ['bar', 'barh']:
            if not pd.api.types.is_datetime64_any_dtype(plot_data[self.x]):
                plot_data[self.x] = plot_data[self.x].astype('object')

        palette = self.last_kwargs.get('palette')
        if palette and 'color' not in plot_dict:
            data_cols = [c for c in plot_data.columns if c != self.x]
            colors = self._get_palette_colors(data_cols, palette=palette)
            plot_dict['color'] = [colors.get(c) for c in data_cols]
        if 'x' in plot_dict:
            plot_dict['x'] = _remap_coord_if_needed(plot_data.columns, plot_dict['x'])
        self.ax = plot_data.plot(ax=ax, **plot_dict)
        self._apply_axis_labels(ax)
        self._handle_data_output(plot_data, ax)

    # kde, density
    def _plot_density(self, ax):
        """Plot a density or KDE chart."""
        plot_dict = self._prepare_plot_kwargs('density')
        self._store_plot_kwargs(ax, plot_dict)
        if self.by is not None:
            if self.by not in self.df.columns:
                raise KeyError(f"Plotter: column '{self.by}' for 'by' not found in DataFrame")
            sub = self.df[[self.by, self.column]].copy()
            if self.dropna:
                sub = sub.dropna(subset=[self.by])
            sub["__row__"] = sub.groupby(self.by, dropna=False).cumcount()
            k = sub.pivot(index="__row__", columns=self.by, values=self.column)
            k.index.name = None
            k.columns.name = None
            too_small = [col for col in k.columns if k[col].count() < 2]
            if too_small:
                raise ValueError(
                    f"Plotter: kde/density plot requires at least two non-null observations in each group. "
                    f"Group(s) with fewer observations: {too_small}"
                )
        else:
            k = self.df[[self.column]]
            if k[self.column].count() < 2:
                raise ValueError("Plotter: kde/density plot requires at least two non-null observations.")
        palette = self.last_kwargs.get('palette')
        if palette and 'color' not in plot_dict:
            colors = self._get_palette_colors(list(k.columns), palette=palette)
            plot_dict['color'] = [colors.get(c) for c in k.columns]
        self.ax = k.plot(ax=ax, **plot_dict)
        self._apply_axis_labels(ax)
        self._handle_data_output(k, ax)

    # hist
    def _plot_hist(self, ax):
        """Plot a histogram."""
        plot_dict = self._prepare_plot_kwargs('hist')
        self._store_plot_kwargs(ax, plot_dict)
        if self.by is not None:
            if self.by not in self.df.columns:
                raise KeyError(f"Plotter: column '{self.by}' for 'by' not found in DataFrame")
            sub = self.df[[self.by, self.column]].copy()
            if self.dropna:
                sub = sub.dropna(subset=[self.by])
            sub["__row__"] = sub.groupby(self.by, dropna=False).cumcount()
            k = sub.pivot(index="__row__", columns=self.by, values=self.column)
            k.index.name = None
            k.columns.name = None
        else:
            k = self.df[[self.column]]
        palette = self.last_kwargs.get('palette')
        if palette and 'color' not in plot_dict:
            colors = self._get_palette_colors(list(k.columns), palette=palette)
            plot_dict['color'] = [colors.get(c) for c in k.columns]
        self.ax = k.plot(ax=ax, **plot_dict)
        self._apply_axis_labels(ax)
        self._handle_data_output(k, ax)

    def _filter_plot_kwargs(self, keys_to_remove):
        return {k: v for k, v in self.last_kwargs.items() if k not in keys_to_remove}

    def _handle_data_output(self, data, ax=None):
        """Handle data output options (print, copy to clipboard, record table)."""
        if ax is not None:
            self.tables[ax.get_label()] = data
        if self.print_data:
            print(data)
        if self.clip_data:
            try:
                data.to_clipboard(index=False)
            except Exception as exc:
                warnings.warn(f"Plotter: failed to copy data to clipboard: {exc}")

    def _collect_handles_and_legends(self):
        """
        Collect unique handles and labels from all axes for legend creation.

        Returns:
            tuple: (handles, labels) for legend construction.
        """
        handles = []
        labels = []
        seen = set()
        
        for ax in self.fig.axes:
            ax_label = ax.get_label()
            
            # Skip this axis if legend=False is set in plot_kwargs_store
            if ax_label in self.plot_kwargs_store and not self.plot_kwargs_store[ax_label].get('legend', True):
                continue

            h, l = ax.get_legend_handles_labels()

            for handle, label in zip(h, l):
                label = label.replace(" (right)", "")
                identifier = (label, type(handle))
                if identifier not in seen:
                    handles.append(handle)
                    labels.append(label)
                    seen.add(identifier)

        return handles, labels

    def _store_plot_kwargs(self, ax, plot_dict):
        """Store plot kwargs for the given axis."""
        d = dict(plot_dict)
        d['kind'] = self.kind
        self.plot_kwargs_store[ax.get_label()] = d

    def finalize(self, consolidate_legends=False, bbox_to_anchor=(0.8, -0.05), ncols=10, hide_secondary_y=False, 
                 legend=True, legend_primary=True, legend_secondary=True, legend_loc='best', legend_frameon=False,
                 style=True, title=None, xlabel=None, ylabel=None, tight_layout=True, sharex=False, sharey=False, **kwargs):
        """
        Finalize the plot with layout adjustments and legend settings.
        """
        if kwargs:
            first_bad = next(iter(kwargs))
            raise ValueError(f"Plotter.finalize(): unexpected argument '{first_bad}'")

        Plotter._last_active = self
        self.consolidate_legends = consolidate_legends
        self.bbox_to_anchor = bbox_to_anchor
        self.ncols = ncols
        
        if title is not None:
            if len(self.fig.axes) == 1:
                self.fig.axes[0].set_title(title)
            else:
                self.fig.suptitle(title)

        if hasattr(self, 'ax') and self.ax is not None:
            if xlabel is not None:
                self.ax.set_xlabel(xlabel)
            if ylabel is not None:
                self.ax.set_ylabel(ylabel)

        if isinstance(style, str):
            try:
                plt.style.use(style)
            except (OSError, ValueError) as exc:
                raise ValueError(f"Plotter.finalize(): stylesheet '{style}' not recognized by matplotlib") from exc

        style_bool = style if isinstance(style, bool) else True
        if style_bool:
            self._hide_spines()
            self._adjust_ticks_and_spines()
            for ax in self.fig.axes:
                ax_label = ax.get_label()
                if '<colorbar>' in ax_label or ax_label == '':
                    continue
                stored = self.plot_kwargs_store.get(ax_label, {})
                kind = stored.get('kind', '')
                rot = stored.get('rot')
                fontsize = stored.get('fontsize')
                if rot is None and kind == 'bar':
                    rot = 90
                elif rot is None and kind == 'barh':
                    rot = 0
                target_axis = 'y' if kind == 'barh' else 'x'
                if fontsize is not None:
                    ax.tick_params(axis='both', labelsize=fontsize)
                if rot is not None:
                    ax.tick_params(axis=target_axis, labelrotation=rot)

        axes_list = list(self.axd.values()) if hasattr(self, 'axd') and self.axd else self.fig.axes
        plot_axes = [ax for ax in axes_list if '<colorbar>' not in ax.get_label() and not ax.get_label().endswith('^')]
        if sharex and len(plot_axes) > 1:
            base_ax = plot_axes[0]
            for other_ax in plot_axes[1:]:
                if getattr(other_ax, '_sharex', None) is not base_ax and getattr(base_ax, '_sharex', None) is not other_ax:
                    other_ax.sharex(base_ax)
            for ax in plot_axes:
                ax.autoscale()
        if sharey and len(plot_axes) > 1:
            base_ax = plot_axes[0]
            for other_ax in plot_axes[1:]:
                if getattr(other_ax, '_sharey', None) is not base_ax and getattr(base_ax, '_sharey', None) is not other_ax:
                    other_ax.sharey(base_ax)
            for ax in plot_axes:
                ax.autoscale()
        self._manage_legend(legend=legend, legend_primary=legend_primary, legend_secondary=legend_secondary, 
                            legend_loc=legend_loc, legend_frameon=legend_frameon)
        
        handles, labels = self._collect_handles_and_legends()
        
        if hide_secondary_y:
            for ax_key, ax in self.axd.items():
                if '^' in ax_key:
                    ax.tick_params(axis='y', labelright=False, right=False)
                    ax.spines['right'].set_visible(False)
        
        # Apply layout adjustments before drawing figure-level legend to avoid clipping
        if tight_layout:
            if self.consolidate_legends and self.bbox_to_anchor[1] < 0:
                self.fig.tight_layout(rect=[0, 0.08, 1, 1])
            else:
                self.fig.tight_layout()

        if legend and self.consolidate_legends and handles:
            self.fig.legend(handles, labels, bbox_to_anchor=self.bbox_to_anchor, ncol=self.ncols, frameon=legend_frameon)

        return self

    def show(self):
        """Display the figure interactively. Chainable."""
        try:
            import matplotlib._pylab_helpers as pylab_helpers
            manager = getattr(self.fig.canvas, 'manager', None)
            if manager is None or manager.num not in pylab_helpers.Gcf.figs:
                dummy = plt.figure()
                new_manager = dummy.canvas.manager
                if new_manager is not None:
                    new_manager.canvas.figure = self.fig
                    self.fig.set_canvas(new_manager.canvas)
        except Exception:
            pass
        plt.show()
        return self

    def save(self, path, **kwargs):
        """Save the figure (thin wrapper around fig.savefig()). Chainable, e.g.
        `Plotter().data(df).plot(...).finalize().save('out.png')`."""
        self.fig.savefig(path, **kwargs)
        return self

    @classmethod
    def supported_kwargs(cls, kind):
        """Describe which kwargs are ignored (with a warning) for a given plot
        `kind`, and which pytae-only control kwargs are never forwarded to
        pandas.plot() at all. Everything else in a .plot() call is passed
        straight through to the matching pandas.plot(kind=...) call.

        Returns:
        --------
        dict
            {'unsupported': [...], 'pytae_controls': [...]}
        """
        spec_key = 'density' if kind == 'kde' else kind
        spec = _KIND_SPECS.get(spec_key, _KIND_SPECS['other'])
        return {
            'unsupported': sorted(spec['unsupported']),
            'pytae_controls': sorted(set(spec['strip']) - set(spec['unsupported'])),
        }

    def _hide_spines(self):
        """Hide all spines by default, with adjustments for visibility."""
        for ax in self.fig.axes:
            ax_label = ax.get_label()
            if '<colorbar>' in ax_label or ax_label == '':
                continue
            
            ax.spines['top'].set_position(('outward', 5))
            ax.spines['bottom'].set_position(('outward', 5))
            ax.spines['left'].set_position(('outward', 5))
            ax.spines['right'].set_position(('outward', 5))
            
            for spine in ax.spines.values():
                spine.set_visible(False)
    
    def _adjust_ticks_and_spines(self):
        """Adjust tick visibility and spine settings based on whether anything was
        actually plotted on each axis."""
        for ax in self.fig.axes:
            ax_label = ax.get_label()
            if '<colorbar>' in ax_label or ax_label == '':
                continue

            blank = not ax.has_data()
            for spine, axis, label_position in [
                ('top', ax.xaxis, 'top'),
                ('bottom', ax.xaxis, 'bottom'),
                ('left', ax.yaxis, 'left'),
                ('right', ax.yaxis, 'right')
            ]:
                labels = ax.get_xticklabels() if axis == ax.xaxis else ax.get_yticklabels()
                labels = [label.get_text() for label in labels if label.get_text()]
                
                if (label_position == axis.get_label_position() and 
                    labels and 
                    not blank):
                    ax.spines[spine].set_visible(True)
                    
                if blank:
                    tick_params_position = 'labeltop' if label_position == 'top' else (
                        'labelbottom' if label_position == 'bottom' else (
                        'labelleft' if label_position == 'left' else 'labelright'
                    ))
                    ax.tick_params(
                        axis='x' if axis == ax.xaxis else 'y', 
                        which='both',
                        length=0,
                        **{tick_params_position: False}
                    )
    
    def _manage_legend(self, legend=True, legend_primary=True, legend_secondary=True, legend_loc='best', legend_frameon=False):
        """Manage legend display for each axis."""
        for ax in self.fig.axes:
            ax_label = ax.get_label()
            is_secondary = '^' in ax_label
            
            handles, labels = ax.get_legend_handles_labels()
            labels = [label.replace(" (right)", "") for label in labels]
            
            stored_legend = self.plot_kwargs_store.get(ax_label, {}).get('legend', True)
            should_display_legend = False
            if legend and stored_legend:
                if is_secondary and legend_secondary:
                    should_display_legend = True
                elif not is_secondary and legend_primary:
                    should_display_legend = True
            
            if should_display_legend and not self.consolidate_legends:
                ax.legend(handles, labels, frameon=legend_frameon, loc=legend_loc)
            
            if ax.get_legend() and (self.consolidate_legends or not should_display_legend):
                ax.get_legend().remove()
            
            if not self.consolidate_legends and ax.get_label() in self.plot_kwargs_store:
                if self.plot_kwargs_store[ax.get_label()].get('kind') == 'pie':
                    if ax.get_legend():
                        ax.get_legend().remove()


def plot(df: pd.DataFrame, *args, **kwargs) -> Plotter:
    """Create a pytae Plotter pre-loaded with `df`, optionally executing a plot.

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame to plot.
    *args :
        Positional arguments passed to Plotter constructor (e.g. mosaic).
    **kwargs :
        Plotter constructor arguments (mosaic, figsize, sharex, sharey)
        or .plot() parameters (kind, x, y, by, palette, etc.).

    Examples:
    ---------
    >>> import pytae as pt
    >>> tips = pt.sample("tips")
    >>> tips.pt.agg("day", total_bill="mean").pt.plot(kind="bar", x="day", y="total_bill").finalize()
    """
    return df.pt.plot(*args, **kwargs)


def finalize(plotter: Plotter | None = None, **kwargs) -> Plotter:
    """Finalize layout, spines, and legends for a Plotter.

    If `plotter` is omitted, finalizes the most recently active Plotter
    instance (similar to matplotlib's ``plt.tight_layout()``).

    Parameters:
    -----------
    plotter : Plotter, optional
        The Plotter instance to finalize. If None, finalizes the last active Plotter.
    **kwargs :
        Optional keyword arguments passed directly to ``Plotter.finalize()``:
        - `consolidate_legends` (bool): Consolidate all legends into one unified figure legend.
        - `bbox_to_anchor` (tuple): Position for consolidated legend (e.g. `(0.5, -0.05)`).
        - `ncols` (int): Number of columns in consolidated legend.
        - `hide_secondary_y` (bool): Hide secondary y-axis spines/ticks.
        - `legend` (bool): Master toggle to display or hide legends.
        - `style` (bool): Apply pytae's spine and tick cleanup.

    Returns:
    --------
    Plotter
        The finalized Plotter instance.

    Examples:
    ---------
    >>> import pytae as pt
    >>> tips = pt.sample("tips")
    >>> p = pt.plot(tips, kind="bar", x="day", y="total_bill")
    >>> pt.finalize(p)
    """
    if plotter is None:
        plotter = Plotter._last_active
        if plotter is None:
            raise ValueError("No active pytae Plotter to finalize.")
    elif not isinstance(plotter, Plotter):
        raise TypeError(f"Expected a Plotter instance, got {type(plotter).__name__}")
    return plotter.finalize(**kwargs)

################################################################################
# FILE: src/pytae/qry.py
################################################################################
from __future__ import annotations

import ast
import difflib
import operator
import re
from typing import Any

import pandas as pd

from pytae._text import unquote_name as _unquote_name
from pytae.cli_parsing import _split_qry_entries

# Dictionary mapping string operators to their corresponding functions
ops = {
    ">=": operator.ge, "<=": operator.le,
    ">": operator.gt, "<": operator.lt,
    "==": operator.eq, "!=": operator.ne
}

# String-accessor operators for tuple conditions, e.g. ('startswith', 'Ad')
str_ops = {
    "startswith": lambda s, v: s.str.startswith(v, na=False),
    "endswith": lambda s, v: s.str.endswith(v, na=False),
    "contains": lambda s, v: s.str.contains(v, regex=False, na=False),
    "regex": lambda s, v: s.str.contains(v, regex=True, na=False),
}

# No-value tuple operators, e.g. ('isna',)
unary_ops = {
    "isna": lambda s: s.isna(),
    "notna": lambda s: s.notna(),
}


def _parse_qry_string(raw: str) -> list[tuple[str, Any]]:
    stripped = raw.strip()
    if stripped.startswith("{") and stripped.endswith("}"):
        stripped = stripped[1:-1]
    try:
        entries = _split_qry_entries(stripped)
    except SystemExit as exc:
        raise ValueError(str(exc)) from exc
    pairs: list[tuple[str, Any]] = []
    for key_raw, value_raw in entries:
        key = _unquote_name(key_raw)
        if not key:
            raise ValueError("invalid qry conditions: empty column name")
        if not value_raw:
            raise ValueError(f"invalid qry conditions: '{key}' has no value")
        val_strip = value_raw.strip()
        m_int = re.match(r"^([\[(])\s*([^,()\[\]]+)\s*,\s*([^,()\[\]]+)\s*([\])])$", val_strip)
        if m_int:
            left, b1, b2, right = m_int.groups()
            b1_is_quoted = (b1.startswith("'") and b1.endswith("'")) or (b1.startswith('"') and b1.endswith('"'))
            b2_is_quoted = (b2.startswith("'") and b2.endswith("'")) or (b2.startswith('"') and b2.endswith('"'))
            if not (b1_is_quoted and b2_is_quoted and left == "[" and right == "]"):
                pairs.append((key, val_strip))
                continue

        op_match = re.match(r"^(>=|<=|!=|==|>|<)\s*(.+)$", val_strip)
        if op_match:
            op = op_match.group(1)
            sub_raw = op_match.group(2).strip()
            try:
                sub_val = ast.literal_eval(sub_raw)
            except (ValueError, SyntaxError):
                if len(sub_raw) >= 2 and sub_raw[0] in "'\"" and sub_raw[-1] == sub_raw[0]:
                    sub_val = sub_raw[1:-1]
                else:
                    sub_val = sub_raw
            pairs.append((key, (op, sub_val)))
            continue
        try:
            value = ast.literal_eval(val_strip)
        except (ValueError, SyntaxError):
            if len(val_strip) >= 2 and val_strip[0] in "'\"" and val_strip[-1] == val_strip[0]:
                value = val_strip[1:-1]
            else:
                value = val_strip
        pairs.append((key, value))
    if not pairs:
        raise ValueError("qry() expects at least one condition")
    return pairs

def qry(
    df: pd.DataFrame,
    *args: Any,
    **kwargs: Any,
) -> pd.DataFrame:
    """
    Filters a DataFrame based on string expressions, dictionary mappings, or keyword arguments.

    This method provides a clean, Pythonic way to filter rows in a DataFrame. Conditions can be
    passed as:
    - String expressions matching the CLI -qry syntax (e.g. df.pt.qry("body_mass_g > 5000"),
      df.pt.qry("bill length mm > 40"))
    - Plain dictionaries (e.g. df.pt.qry({"bill length mm": "> 40", "species": "Adelie"}))
    - Keyword arguments (e.g. df.pt.qry(species='Adelie', body_mass_g='> 5000'))
    - Any mix of the above (e.g. df.pt.qry("bill length mm > 40", species="Adelie"))

    Conditions can include direct values, lists of values, tuple-based comparisons (e.g., ('>', 100)),
    tuple-based list membership (e.g., ('in', ['a', 'b'])), or interval conditions (e.g., '(a,b)', '[a,b]').
    It supports both numeric and non-numeric columns. Index is not reset for the returned
    DataFrame since querying should not alter indexing.

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame to filter.
    *args : str or dict
        Positional filter conditions:
        - String expressions, e.g. "body_mass_g > 3500", "bill length mm > 40", or comma-separated
          "species = 'Adelie', body_mass_g > 3500".
        - Dictionaries mapping column names to conditions, e.g. {"bill length mm": "> 40"}.
    **kwargs : Any
        Filter conditions specified as keyword arguments where the keyword is the column name
        and the value is the condition to apply. Conditions can be:
        - A single value (e.g., 'Adelie'): Filters for rows where the column equals the value.
        - A list of values (e.g., ['Adelie', 'Gentoo']): Filters for rows where the column
          matches any value in the list.
        - A tuple with 'in' and list (e.g., ('in', ['Adelie', 'Gentoo'])): Filters for rows
          where the column matches any value in the list.
        - A tuple with 'not in' and list (e.g., ('not in', ['Adelie', 'Gentoo'])): Filters
          for rows where the column does not match any value in the list.
        - A tuple with an operator and value (e.g., ('>', 81500)): Filters for rows where
          the column satisfies the operator-based condition. Supported operators are
          >=, <=, >, <, ==, !=.
        - A tuple with a string operator and value (e.g., ('startswith', 'Ad')): Filters
          using pandas' `.str` accessor. Supported: 'startswith', 'endswith' (value may
          also be a list of prefixes/suffixes), 'contains', 'regex' (both search anywhere
          in the string, like re.search; 'regex' is 'contains' with regex=True). Missing
          values never match (na=False); requires a string-dtype column.
        - A one-element tuple ('isna',) or ('notna',): Filters for rows where the column
          is/isn't null. Takes no value since these are unary checks.
        - An interval condition (e.g., '(a,b)', '[a,b]'): Filters for rows where the column
          falls within the specified interval (parentheses for exclusive, brackets for inclusive).

    Returns:
    --------
    pd.DataFrame
        A filtered DataFrame containing only the rows that satisfy all conditions.

    Examples:
    ---------
    >>> import pandas as pd
    >>> import pytae as pt
    >>> data = {
    ...     'species': ['Adelie', 'Gentoo', 'Chinstrap', 'Adelie'],
    ...     'body mass g': [74125, 271425, 119925, 89100],
    ...     'code': ['A 1', 'B 2', 'C 3', 'D 4']
    ... }
    >>> df = pd.DataFrame(data)

    >>> # Filter using a string expression (columns with spaces work directly!)
    >>> df.pt.qry("body mass g > 81500")
         species  body mass g code
    1     Gentoo       271425  B 2
    2  Chinstrap       119925  C 3
    3     Adelie        89100  D 4

    >>> # Filter using a plain dictionary
    >>> df.pt.qry({"body mass g": "> 81500", "species": "Gentoo"})
      species  body mass g code
    1  Gentoo       271425  B 2

    >>> # Filter using keyword arguments
    >>> df.pt.qry(species='Adelie')
      species  body mass g code
    0  Adelie        74125  A 1
    3  Adelie        89100  D 4

    >>> # Mix string expressions and kwargs
    >>> df.pt.qry("body mass g > 80000", species="Adelie")
      species  body mass g code
    3  Adelie        89100  D 4

    Notes:
    ------
    - For numeric columns, tuple-based operator conditions (e.g., ('>', 81500)) will
      automatically convert the comparison value to a float. Since values are typically
      provided as literals, whitespace is handled by Python's parser prior to conversion.
    - For non-numeric columns, tuple-based operator conditions (e.g., ('==', 'A 1')) will
      treat the comparison value as-is, preserving any whitespace in string literals.
    - String-accessor operators ('startswith', 'endswith', 'contains', 'regex') use pandas'
      `.str` accessor and treat missing values as non-matching (`na=False`) rather than
      raising or propagating NaN. 'contains'/'regex' both do a search-anywhere match
      (like `re.search`); 'regex' is the same as 'contains' with `regex=True`.
    - 'isna'/'notna' are one-element tuples (e.g. ('isna',)) since they take no value.
    - Filtering does not modify the original DataFrame. Each condition is applied with
      `.loc[...]` and a new filtered frame is returned; the caller's object is unchanged.
    """
    cond_pairs: list[tuple[str, Any]] = []
    for arg in args:
        if isinstance(arg, str):
            cond_pairs.extend(_parse_qry_string(arg))
        elif isinstance(arg, dict):
            for col, cond in arg.items():
                cond_pairs.append((str(col), cond))
        elif isinstance(arg, (list, tuple)):
            for item in arg:
                if isinstance(item, (tuple, list)) and len(item) == 2:
                    cond_pairs.append((str(item[0]), item[1]))
                else:
                    raise TypeError(
                        f"qry() expects filter conditions as strings, dicts, or keyword arguments, got {type(arg).__name__}."
                    )
        else:
            raise TypeError(
                f"qry() expects filter conditions as strings, dicts, or keyword arguments, got {type(arg).__name__}."
            )
    for col, cond in kwargs.items():
        cond_pairs.append((col, cond))

    if not cond_pairs:
        raise ValueError("qry() expects at least one condition (e.g. df.pt.qry('col > 5000') or df.pt.qry(col='> 5000'))")

    normalized_conditions: list[tuple[str, Any]] = []
    for col, cond in cond_pairs:
        clean_col = _unquote_name(col)
        if isinstance(cond, str):
            op_match = re.match(r"^(>=|<=|!=|==|>|<)\s*(.+)$", cond.strip())
            if op_match:
                op = op_match.group(1)
                val_str = op_match.group(2).strip()
                try:
                    val = ast.literal_eval(val_str)
                except (ValueError, SyntaxError):
                    val = val_str.strip("'\"")
                normalized_conditions.append((clean_col, (op, val)))
                continue
        normalized_conditions.append((clean_col, cond))

    out = df
    available = list(df.columns)
    for col, cond in normalized_conditions:
        if col not in available:
            close = difflib.get_close_matches(col, available, n=1)
            hint = f" (did you mean '{close[0]}'?)" if close else ""
            raise KeyError(f"unknown column '{col}'{hint}")
        is_numeric = pd.api.types.is_numeric_dtype(out[col])

        if isinstance(cond, list):
            # Handle direct list conditions (e.g., ['Adelie', 'Gentoo'])
            out = out.loc[out[col].isin(cond)]
        elif isinstance(cond, tuple) and len(cond) == 1:
            op = cond[0]
            if op not in unary_ops:
                raise ValueError(f"Unsupported 1-element tuple operator '{op}' for '{col}'. Use one of {list(unary_ops.keys())}.")
            out = out.loc[unary_ops[op](out[col])]
        elif isinstance(cond, tuple) and len(cond) == 2:
            op, value = cond
            if op in ['in', 'not in']:
                if not isinstance(value, list):
                    raise ValueError(f"Second element of tuple for '{col}' with '{op}' must be a list, got {type(value)}")
                if op == 'in':
                    out = out.loc[out[col].isin(value)]
                elif op == 'not in':
                    out = out.loc[~out[col].isin(value)]
            elif op in str_ops:
                if op in ('startswith', 'endswith') and isinstance(value, list):
                    value = tuple(value)  # pandas' str.startswith()/endswith() take a tuple of prefixes, not a list
                try:
                    out = out.loc[str_ops[op](out[col], value)]
                except AttributeError as exc:
                    raise ValueError(
                        f"qry: '{op}' needs a string column; '{col}' is {out[col].dtype}. "
                        f"Cast it first, e.g. df.astype({{'{col}': str}}).qry(...)."
                    ) from exc
                except TypeError as exc:
                    raise ValueError(f"qry: '{op}' on '{col}': invalid value {value!r} ({exc})") from exc
            elif op in ops:
                if is_numeric and not isinstance(value, (int, float)):
                    try:
                        value = int(value)
                    except (ValueError, TypeError):
                        try:
                            value = float(value)
                        except (ValueError, TypeError):
                            pass
                out = out.loc[ops[op](out[col], value)]
            else:
                raise ValueError(
                    f"Unsupported tuple operator '{op}' for '{col}'. "
                    f"Use 'in', 'not in', or one of {list(ops.keys()) + list(str_ops.keys())}."
                )
        elif isinstance(cond, str) and re.match(r'^[\[(].*[\])]$', cond):
            # Handle interval conditions (e.g., '(a,b)', '[a,b]')
            interval_pattern = re.compile(r'^([\[(])([^,]+),([^,]+)([\])])$')
            match = interval_pattern.match(cond)
            if match:
                left_bracket, lower_raw, upper_raw, right_bracket = match.groups()
                lower_str = lower_raw.strip()
                upper_str = upper_raw.strip()
                lower: int | float | str
                upper: int | float | str
                if is_numeric:
                    try:
                        lower = int(lower_str)
                    except ValueError:
                        try:
                            lower = float(lower_str)
                        except ValueError as exc:
                            raise ValueError(
                                f"qry: interval lower bound for numeric column '{col}' must be a number, got {lower_str!r}"
                            ) from exc
                    try:
                        upper = int(upper_str)
                    except ValueError:
                        try:
                            upper = float(upper_str)
                        except ValueError as exc:
                            raise ValueError(
                                f"qry: interval upper bound for numeric column '{col}' must be a number, got {upper_str!r}"
                            ) from exc
                else:
                    lower = lower_str
                    upper = upper_str

                if left_bracket == '[':
                    lower_op = operator.ge
                else:
                    lower_op = operator.gt

                if right_bracket == ']':
                    upper_op = operator.le
                else:
                    upper_op = operator.lt

                out = out.loc[lower_op(out[col], lower) & upper_op(out[col], upper)]
            else:
                out = out.loc[out[col] == cond]
        else:
            # Handle single value equality (e.g., 'Adelie')
            out = out.loc[out[col] == cond]

    return out

################################################################################
# FILE: src/pytae/readers.py
################################################################################
"""Format-specific readers exposing a uniform metadata/inspection API.

Each reader avoids loading full row data for shape/columns/dtypes/head where
the file format allows it (parquet, sas7bdat carry that info in a header).
CSV/TXT have no embedded schema, so dtype inspection samples the first
DTYPE_SAMPLE_ROWS rows.
"""

from __future__ import annotations

import importlib.util
import io
import os
import re
import warnings
from pathlib import Path

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pa_parquet

from pytae.other_utilities import safe_reset_index

DEFAULT_CHUNK_SIZE = 200_000
CHUNK_SIZE = DEFAULT_CHUNK_SIZE
DTYPE_SAMPLE_ROWS = 10_000


def _split_path_suffixes(path: Path) -> tuple[str, str | None]:
    """Return (base_suffix, compression) e.g. ('.csv', 'gzip') or ('.parquet', None)."""
    suffixes = [s.lower() for s in path.suffixes]
    if not suffixes:
        return ("", None)
    if suffixes[-1] == ".gz":
        base = suffixes[-2] if len(suffixes) >= 2 else ""
        return (base, "gzip")
    return (suffixes[-1], None)


def _delimited_dtypes(path: Path, *, sep: str, encoding: str | None, compression: str | None = None) -> pd.Series:
    return pd.read_csv(
        path, sep=sep, encoding=encoding, nrows=DTYPE_SAMPLE_ROWS, low_memory=False, compression=compression
    ).dtypes


def _print_progress(done: int, total: int | None, label: str) -> None:
    if total:
        pct = min(100, int(done * 100 / total))
        print(f"\r{label}... {done}/{total} rows ({pct}%)", end="", flush=True)
    else:
        print(f"\r{label}... {done} rows", end="", flush=True)


class ParquetReader:
    def __init__(self, path: Path, chunk_size: int = DEFAULT_CHUNK_SIZE) -> None:
        self.path = path
        self.chunk_size = chunk_size
        self._pf = pa_parquet.ParquetFile(path)

    def shape(self) -> tuple[int, int]:
        md = self._pf.metadata
        return (md.num_rows, md.num_columns)

    def columns(self) -> list[str]:
        return list(self._pf.schema_arrow.names)

    def dtypes(self) -> pd.Series:
        return self._pf.schema_arrow.empty_table().to_pandas().dtypes

    def _empty(self) -> pd.DataFrame:
        return self._pf.schema_arrow.empty_table().to_pandas()

    def head(self, n: int) -> pd.DataFrame:
        if n <= 0 or self._pf.metadata.num_rows == 0:
            return self._empty()
        batch = next(self._pf.iter_batches(batch_size=n), None)
        if batch is None:
            return self._empty()
        return batch.to_pandas().head(n)

    def tail(self, n: int) -> pd.DataFrame:
        total = self._pf.metadata.num_rows
        if n <= 0 or total == 0:
            return self._empty()
        skip = max(total - n, 0)
        parts, seen = [], 0
        for batch in self._pf.iter_batches(batch_size=max(n, 1)):
            start_in_batch = max(0, skip - seen)
            if start_in_batch < batch.num_rows:
                parts.append(batch.to_pandas().iloc[start_in_batch:])
            seen += batch.num_rows
        df = pd.concat(parts, ignore_index=True) if parts else self._empty()
        return df.tail(n)

    def to_dataframe(self, columns: list[str] | None = None, progress: bool = False, nrows: int | None = None,
                     chunk_size: int | None = None) -> pd.DataFrame:
        if not progress and nrows is None:
            return pd.read_parquet(self.path, columns=columns)
        total = self._pf.metadata.num_rows
        limit = total if nrows is None else min(nrows, total)
        csize = chunk_size or self.chunk_size
        parts, done = [], 0
        for batch in self._pf.iter_batches(batch_size=csize, columns=columns):
            remaining = limit - done
            if remaining <= 0:
                break
            chunk = batch.to_pandas().iloc[:remaining]
            parts.append(chunk)
            done += len(chunk)
            if progress:
                _print_progress(done, limit, "reading")
        if progress:
            print()
        if parts:
            return pd.concat(parts, ignore_index=True)
        empty = self._empty()
        return empty[columns] if columns else empty


class CsvReader:
    def __init__(
        self,
        path: Path,
        sep: str = ",",
        encoding: str | None = None,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        compression: str | None = None,
    ) -> None:
        self.path = path
        self.sep = sep
        self.encoding = encoding
        self.chunk_size = chunk_size
        self.compression = compression

    def shape(self) -> tuple[int, int]:
        n_cols = len(self.columns())
        total = 0
        for chunk in pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            usecols=[0],
            chunksize=self.chunk_size,
            low_memory=False,
            compression=self.compression,
        ):
            total += len(chunk)
        return (total, n_cols)

    def columns(self) -> list[str]:
        try:
            return pd.read_csv(
                self.path,
                sep=self.sep,
                encoding=self.encoding,
                nrows=0,
                low_memory=False,
                compression=self.compression,
            ).columns.tolist()
        except pd.errors.EmptyDataError as exc:
            raise ValueError(f"'{self.path.name}' is empty or not a valid CSV file") from exc

    def dtypes(self) -> pd.Series:
        return _delimited_dtypes(self.path, sep=self.sep, encoding=self.encoding, compression=self.compression)

    def head(self, n: int) -> pd.DataFrame:
        return pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            nrows=n,
            low_memory=False,
            compression=self.compression,
        )

    def tail(self, n: int) -> pd.DataFrame:
        total = self.shape()[0]
        skip = range(1, max(total - n, 0) + 1)  # keep the header (row 0), skip everything before the tail
        return pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            skiprows=skip,
            low_memory=False,
            compression=self.compression,
        )

    def to_dataframe(
        self,
        columns: list[str] | None = None,
        progress: bool = False,
        nrows: int | None = None,
        chunk_size: int | None = None,
    ) -> pd.DataFrame:
        if not progress:
            return pd.read_csv(
                self.path,
                sep=self.sep,
                encoding=self.encoding,
                usecols=columns,
                nrows=nrows,
                low_memory=False,
                compression=self.compression,
            )
        total = self.shape()[0] if nrows is None else min(nrows, self.shape()[0])
        chunks, done = [], 0
        csize = chunk_size or self.chunk_size
        reader = pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            usecols=columns,
            chunksize=csize,
            nrows=nrows,
            low_memory=False,
            compression=self.compression,
        )
        for chunk in reader:
            chunks.append(chunk)
            done += len(chunk)
            _print_progress(done, total, "reading")
        print()
        return pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame(columns=columns or [])


class TxtReader:
    """Reads delimited .txt/.dat files (tab-delimited for .txt, pipe-delimited for .dat by default)."""

    def __init__(
        self,
        path: Path,
        sep: str = "\t",
        encoding: str | None = None,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        compression: str | None = None,
    ) -> None:
        self.path = path
        self.sep = sep
        self.encoding = encoding
        self.chunk_size = chunk_size
        self.compression = compression

    def shape(self) -> tuple[int, int]:
        n_cols = len(self.columns())
        total = 0
        for chunk in pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            usecols=[0],
            chunksize=self.chunk_size,
            low_memory=False,
            compression=self.compression,
        ):
            total += len(chunk)
        return (total, n_cols)

    def columns(self) -> list[str]:
        try:
            return pd.read_csv(
                self.path,
                sep=self.sep,
                encoding=self.encoding,
                nrows=0,
                low_memory=False,
                compression=self.compression,
            ).columns.tolist()
        except pd.errors.EmptyDataError as exc:
            raise ValueError(f"'{self.path.name}' is empty or not a valid delimited file") from exc

    def dtypes(self) -> pd.Series:
        return _delimited_dtypes(self.path, sep=self.sep, encoding=self.encoding, compression=self.compression)

    def head(self, n: int) -> pd.DataFrame:
        return pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            nrows=n,
            low_memory=False,
            compression=self.compression,
        )

    def tail(self, n: int) -> pd.DataFrame:
        total = self.shape()[0]
        skip = range(1, max(total - n, 0) + 1)  # keep the header (row 0), skip everything before the tail
        return pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            skiprows=skip,
            low_memory=False,
            compression=self.compression,
        )

    def to_dataframe(
        self,
        columns: list[str] | None = None,
        progress: bool = False,
        nrows: int | None = None,
        chunk_size: int | None = None,
    ) -> pd.DataFrame:
        if not progress:
            return pd.read_csv(
                self.path,
                sep=self.sep,
                encoding=self.encoding,
                usecols=columns,
                nrows=nrows,
                low_memory=False,
                compression=self.compression,
            )
        total = self.shape()[0] if nrows is None else min(nrows, self.shape()[0])
        chunks, done = [], 0
        csize = chunk_size or self.chunk_size
        reader = pd.read_csv(
            self.path,
            sep=self.sep,
            encoding=self.encoding,
            usecols=columns,
            chunksize=csize,
            nrows=nrows,
            low_memory=False,
            compression=self.compression,
        )
        for chunk in reader:
            chunks.append(chunk)
            done += len(chunk)
            _print_progress(done, total, "reading")
        print()
        return pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame(columns=columns or [])


class SasReader:
    """Reads .sas7bdat files. Row count/columns come from the file header,
    so shape/columns are cheap; dtype inspection reads a small sample."""

    def __init__(self, path: Path, encoding: str | None = None, chunk_size: int = DEFAULT_CHUNK_SIZE) -> None:
        self.path = path
        self.encoding = "utf-8" if encoding is None else encoding
        self.chunk_size = chunk_size

    def _open(self):
        return pd.read_sas(self.path, format="sas7bdat", encoding=self.encoding, iterator=True)

    def shape(self) -> tuple[int, int]:
        with self._open() as reader:
            return (reader.row_count, len(reader.column_names))

    def columns(self) -> list[str]:
        with self._open() as reader:
            return [str(name) for name in reader.column_names]

    def dtypes(self) -> pd.Series:
        with self._open() as reader:
            sample = reader.read(min(100, reader.row_count))
        return sample.dtypes

    def head(self, n: int) -> pd.DataFrame:
        with self._open() as reader:
            return reader.read(n)

    def tail(self, n: int) -> pd.DataFrame:
        with self._open() as reader:
            total = reader.row_count
            skip = max(total - n, 0)
            if skip:
                reader.read(skip)  # advance past rows we don't need
            return reader.read(min(n, total))

    def to_dataframe(self, columns: list[str] | None = None, progress: bool = False, nrows: int | None = None,
                     chunk_size: int | None = None) -> pd.DataFrame:
        if not progress and nrows is None:
            df = pd.read_sas(self.path, format="sas7bdat", encoding=self.encoding)
            return df[columns] if columns else df
        with self._open() as reader:
            total = reader.row_count if nrows is None else min(nrows, reader.row_count)
            chunks, done = [], 0
            csize = chunk_size or self.chunk_size
            while done < total:
                chunk = reader.read(min(csize, total - done))
                if chunk is None or len(chunk) == 0:
                    break
                chunks.append(chunk)
                done += len(chunk)
                if progress:
                    _print_progress(done, total, "reading")
        if progress:
            print()
        df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()
        return df[columns] if columns else df


class JsonlReader:
    """Reads line-delimited JSON (.jsonl / .ndjson) files."""

    def __init__(
        self,
        path: Path,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        compression: str | None = None,
    ) -> None:
        self.path = path
        self.chunk_size = chunk_size
        self.compression = compression

    def shape(self) -> tuple[int, int]:
        n_cols = len(self.columns())
        total = 0
        for chunk in pd.read_json(
            self.path,
            lines=True,
            chunksize=self.chunk_size,
            compression=self.compression,
        ):
            total += len(chunk)
        return (total, n_cols)

    def columns(self) -> list[str]:
        try:
            cols = pd.read_json(
                self.path,
                lines=True,
                nrows=1,
                compression=self.compression,
            ).columns.tolist()
            if not cols:
                raise ValueError(f"'{self.path.name}' is empty or not a valid JSON Lines file")
            return cols
        except Exception as exc:
            if isinstance(exc, ValueError) and "is empty or not a valid JSON Lines file" in str(exc):
                raise
            raise ValueError(f"'{self.path.name}' is empty or not a valid JSON Lines file") from exc

    def dtypes(self) -> pd.Series:
        return pd.read_json(
            self.path,
            lines=True,
            nrows=DTYPE_SAMPLE_ROWS,
            compression=self.compression,
        ).dtypes

    def head(self, n: int) -> pd.DataFrame:
        if n <= 0:
            return pd.DataFrame(columns=self.columns())
        return pd.read_json(
            self.path,
            lines=True,
            nrows=n,
            compression=self.compression,
        )

    def tail(self, n: int) -> pd.DataFrame:
        df = self.to_dataframe()
        return df.tail(n)

    def to_dataframe(
        self,
        columns: list[str] | None = None,
        progress: bool = False,
        nrows: int | None = None,
        chunk_size: int | None = None,
    ) -> pd.DataFrame:
        if not progress:
            df = pd.read_json(
                self.path,
                lines=True,
                nrows=nrows,
                compression=self.compression,
            )
            return df[columns] if columns else df
        total = self.shape()[0] if nrows is None else min(nrows, self.shape()[0])
        chunks, done = [], 0
        csize = chunk_size or self.chunk_size
        for chunk in pd.read_json(
            self.path,
            lines=True,
            chunksize=csize,
            nrows=nrows,
            compression=self.compression,
        ):
            chunks.append(chunk)
            done += len(chunk)
            _print_progress(done, total, "reading")
        print()
        df = pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame()
        return df[columns] if columns else df


class StdinReader:
    """Reader for data piped via STDIN, buffered in memory."""

    def __init__(
        self,
        data: bytes,
        *,
        fmt: str | None = None,
        sep: str | None = None,
        encoding: str | None = None,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
    ) -> None:
        self.data = data
        self.path = Path("<stdin>")
        self.chunk_size = chunk_size
        self.sep = sep
        self.encoding = encoding
        self.fmt = fmt
        self._df = self._load()

    def _load(self) -> pd.DataFrame:
        if not self.data:
            return pd.DataFrame()
        fmt = (self.fmt or "").lower().strip()
        if fmt in ("parquet", "pq") or (not fmt and self.data.startswith(b"PAR1")):
            return pd.read_parquet(io.BytesIO(self.data))
        if fmt in ("jsonl", "ndjson") or (not fmt and self.data.lstrip()[:1] in (b"{", b"[")):
            return pd.read_json(io.BytesIO(self.data), orient="records", lines=True)
        # Delimited text (CSV, TSV, etc.)
        sep = self.sep
        if sep is None:
            first_line = self.data.split(b"\n", 1)[0]
            if b"\t" in first_line and b"," not in first_line:
                sep = "\t"
            elif b"|" in first_line and b"," not in first_line:
                sep = "|"
            else:
                sep = ","
        compression = "gzip" if self.data.startswith(b"\x1f\x8b") else None
        return pd.read_csv(io.BytesIO(self.data), sep=sep, encoding=self.encoding, compression=compression)

    def shape(self) -> tuple[int, int]:
        return self._df.shape

    def columns(self) -> list[str]:
        return list(self._df.columns)

    def dtypes(self) -> pd.Series:
        return self._df.dtypes

    def head(self, n: int) -> pd.DataFrame:
        return self._df.head(n)

    def tail(self, n: int) -> pd.DataFrame:
        return self._df.tail(n)

    def to_dataframe(
        self,
        columns: list[str] | None = None,
        progress: bool = False,
        nrows: int | None = None,
        chunk_size: int | None = None,
    ) -> pd.DataFrame:
        out = self._df
        if columns is not None:
            out = out[columns]
        if nrows is not None:
            out = out.iloc[:nrows]
        return out


_READERS = {
    ".parquet": ParquetReader,
    ".pq": ParquetReader,
    ".csv": CsvReader,
    ".txt": TxtReader,
    ".dat": TxtReader,
    ".sas7bdat": SasReader,
    ".jsonl": JsonlReader,
    ".ndjson": JsonlReader,
}

# YAML connection configs (Databricks / SSH) are handled by the isolated pytae.connections
# module, loaded lazily so the core package never imports its optional dependencies.
_YAML_SUFFIXES = (".yaml", ".yml")
_EXTERNAL_CONNECTIONS_MODULE = None

# default field separator for TxtReader-backed suffixes (.txt/.dat)
_TXT_DEFAULT_SEP = {".txt": "\t", ".dat": "|"}

# default text encoding for TxtReader-backed suffixes; only .dat has one (.txt still
# falls through to pandas' own encoding inference, i.e. None)
_TXT_DEFAULT_ENCODING = {".dat": "latin-1"}


def get_reader(path: Path, *, sep: str | None = None, encoding: str | None = None, chunk_size: int = DEFAULT_CHUNK_SIZE):
    suffix, compression = _split_path_suffixes(path)
    if suffix in _YAML_SUFFIXES:
        global _EXTERNAL_CONNECTIONS_MODULE
        if _EXTERNAL_CONNECTIONS_MODULE is None:
            conn_env = os.environ.get("PYTAE_CONNECTIONS_PATH")
            if not conn_env:
                raise FileNotFoundError(
                    "connection reader not configured; set the PYTAE_CONNECTIONS_PATH environment "
                    "variable to a connections.py file"
                )
            module_path = Path(conn_env).expanduser()
            if not module_path.is_file():
                raise FileNotFoundError(
                    "connection reader not found; set PYTAE_CONNECTIONS_PATH to a connections.py file "
                    f"(looked for '{module_path}')"
                )
            module_spec = importlib.util.spec_from_file_location(
                "_pytae_external_connections", module_path
            )
            if module_spec is None or module_spec.loader is None:
                raise ImportError(f"cannot load connections module from '{module_path}'")
            _EXTERNAL_CONNECTIONS_MODULE = importlib.util.module_from_spec(module_spec)
            module_spec.loader.exec_module(_EXTERNAL_CONNECTIONS_MODULE)
        return _EXTERNAL_CONNECTIONS_MODULE.build_yaml_reader(path)
    try:
        cls = _READERS[suffix]
    except KeyError:
        supported = ", ".join(sorted((*_READERS, *_YAML_SUFFIXES, ".csv.gz", ".txt.gz", ".dat.gz", ".jsonl.gz", ".ndjson.gz")))
        raise ValueError(f"unsupported file type '{path.suffix or path.name}'; supported: {supported}") from None
    if cls is CsvReader:
        return CsvReader(path, sep=sep or ",", encoding=encoding, chunk_size=chunk_size, compression=compression)
    if cls is TxtReader:
        return TxtReader(
            path,
            sep=sep or _TXT_DEFAULT_SEP[suffix],
            encoding=encoding or _TXT_DEFAULT_ENCODING.get(suffix),
            chunk_size=chunk_size,
            compression=compression,
        )
    if cls is JsonlReader:
        return JsonlReader(path, chunk_size=chunk_size, compression=compression)
    if cls is SasReader:
        return SasReader(path, encoding=encoding, chunk_size=chunk_size)
    return cls(path, chunk_size=chunk_size)


# .sas7bdat is intentionally excluded: pandas has no writer for that format.
_WRITABLE_SUFFIXES = (
    ".parquet",
    ".pq",
    ".csv",
    ".txt",
    ".dat",
    ".jsonl",
    ".ndjson",
    ".csv.gz",
    ".txt.gz",
    ".dat.gz",
    ".jsonl.gz",
    ".ndjson.gz",
)


def write_dataframe(
    df: pd.DataFrame,
    dest: Path,
    *,
    sep: str | None = None,
    encoding: str | None = None,
    progress: bool = False,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    index: bool | None = None,
) -> None:
    if index is None:
        index = not (isinstance(df.index, pd.RangeIndex) and df.index.name is None)
    suffix, compression = _split_path_suffixes(dest)
    if suffix in (".parquet", ".pq"):
        _write_parquet(df, dest, progress=progress, chunk_size=chunk_size, index=index)
    elif suffix == ".csv":
        _write_delimited(
            df, dest, sep=sep or ",", encoding=encoding, progress=progress, chunk_size=chunk_size, index=index, compression=compression
        )
    elif suffix in _TXT_DEFAULT_SEP:
        _write_delimited(
            df, dest,
            sep=sep or _TXT_DEFAULT_SEP[suffix],
            encoding=encoding or _TXT_DEFAULT_ENCODING.get(suffix),
            progress=progress,
            chunk_size=chunk_size,
            index=index,
            compression=compression,
        )
    elif suffix in (".jsonl", ".ndjson"):
        _write_jsonl(df, dest, progress=progress, chunk_size=chunk_size, compression=compression)
    else:
        supported = ", ".join(_WRITABLE_SUFFIXES)
        raise ValueError(f"unsupported output type '{suffix or dest.name}'; supported: {supported}")


def _write_delimited(
    df: pd.DataFrame,
    dest: Path,
    *,
    sep: str,
    encoding: str | None,
    progress: bool,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    index: bool = False,
    compression: str | None = None,
) -> None:
    total = len(df)
    if not progress or total == 0:
        df.to_csv(dest, sep=sep, encoding=encoding, index=index, compression=compression)
        return
    done = 0
    for start in range(0, total, chunk_size):
        chunk = df.iloc[start:start + chunk_size]
        chunk.to_csv(
            dest,
            sep=sep,
            encoding=encoding,
            index=index,
            mode="w" if start == 0 else "a",
            header=(start == 0),
            compression=compression,
        )
        done += len(chunk)
        _print_progress(done, total, "writing")
    print()


def _write_jsonl(
    df: pd.DataFrame,
    dest: Path,
    *,
    progress: bool,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    compression: str | None = None,
) -> None:
    df = safe_reset_index(df)
    total = len(df)
    if not progress or total == 0:
        df.to_json(dest, orient="records", lines=True, compression=compression)
        return
    done = 0
    for start in range(0, total, chunk_size):
        chunk = df.iloc[start:start + chunk_size]
        chunk.to_json(
            dest,
            orient="records",
            lines=True,
            mode="w" if start == 0 else "a",
            compression=compression,
        )
        done += len(chunk)
        _print_progress(done, total, "writing")
    print()


def _coerce_failing_parquet_column(df: pd.DataFrame, error_msg: str) -> str | None:
    match = re.search(r"Conversion failed for column (.+?) with type object", error_msg)
    if match:
        bad_col = match.group(1).strip()
        if bad_col in df.columns:
            df[bad_col] = df[bad_col].map(lambda v: str(v) if pd.notna(v) else None)
            warnings.warn(
                f"Column '{bad_col}' contained mixed object types; coerced to string for Parquet serialization.",
                UserWarning,
                stacklevel=3,
            )
            return bad_col
    return None


def _write_parquet(
    df: pd.DataFrame,
    dest: Path,
    *,
    progress: bool,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    index: bool = False,
) -> None:
    df_to_write = df
    max_retries = max(len(df.columns), 1)

    if not progress or len(df) == 0:
        for _ in range(max_retries + 1):
            try:
                df_to_write.to_parquet(dest, index=index)
                return
            except Exception as e:
                if df_to_write is df:
                    df_to_write = df.copy()
                if _coerce_failing_parquet_column(df_to_write, str(e)):
                    continue
                raise
        return

    table = None
    for _ in range(max_retries + 1):
        try:
            table = pa.Table.from_pandas(df_to_write, preserve_index=index)
            break
        except Exception as e:
            if df_to_write is df:
                df_to_write = df.copy()
            if _coerce_failing_parquet_column(df_to_write, str(e)):
                continue
            raise

    if table is None:
        table = pa.Table.from_pandas(df_to_write, preserve_index=index)

    total = table.num_rows
    done = 0
    with pa_parquet.ParquetWriter(dest, table.schema) as writer:
        for batch in table.to_batches(max_chunksize=chunk_size):
            writer.write_batch(batch)
            done += batch.num_rows
            _print_progress(done, total, "writing")
    print()

################################################################################
# FILE: src/pytae/select.py
################################################################################
from __future__ import annotations

import difflib
import re
from collections.abc import Sequence
from typing import Any

import pandas as pd

from pytae._text import tokenize as _tokenize
from pytae._text import unquote_name as _unquote_name


class _EverythingMeta(type):
    def __repr__(cls) -> str:
        return "everything"


class everything(metaclass=_EverythingMeta):
    """Sentinel class/instance representing all remaining unselected columns in pt.select().

    Can be passed either as the class `pt.everything` or called as `pt.everything()`.
    """
    _instance: everything | None = None

    def __new__(cls) -> everything:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __repr__(self) -> str:
        return "everything()"


def select(
    df: pd.DataFrame,
    *args: Any,
    exclude: str | Sequence[str] | None = None,
    dtype: str | type | Sequence[str | type] | None = None,
    exclude_dtype: str | type | Sequence[str | type] | None = None,
    contains: str | Sequence[str] | None = None,
    startswith: str | Sequence[str] | None = None,
    endswith: str | Sequence[str] | None = None,
    regex: str | Sequence[str] | None = None,
) -> pd.DataFrame:
    """
    Select columns from a DataFrame based on names, regex patterns, slices, data types,
    string matching, or exclusions.

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame from which to select columns.
    *args : variable-length arguments
        Can be:
        - Exact column name(s) or slice (e.g. 'species', 'bill_length_mm:body_mass_g').
        - Negative column names to exclude (e.g. '-species', '~island', '-start:end').
        - `everything` or `everything()` to capture all remaining columns.
        - List of column names.
        - Callable `f(col_name) -> bool`.
    exclude : str or list of str, optional
        Column name(s) or slice to exclude from the selection.
    dtype : str, type, or list, optional
        Data type(s) to select (e.g., 'numeric', 'datetime').
    exclude_dtype : str, type, or list, optional
        Data type(s) to exclude by dtype.
    contains : str or list of str, optional
        Substring(s) to match in column names.
    startswith : str or list of str, optional
        Prefix(es) to match in column names.
    endswith : str or list of str, optional
        Suffix(es) to match in column names.
    regex : str or list of str, optional
        Regular expression(s) to match in column names.

    Returns:
    --------
    pd.DataFrame
        Selected columns, with explicit/regex/slice selections first, followed by everything()
        if specified, with any excluded columns removed.
    """
    if exclude_dtype is not None and (args or exclude or dtype or contains or startswith or endswith or regex):
        raise ValueError("exclude_dtype cannot be combined with other selection criteria.")

    all_cols = df.columns.tolist()  # List of all columns for slice positioning

    if exclude_dtype is not None:
        _shorthands = {
            "numeric": ["number"],
            "datetime": ["datetime", "datetimetz"],
            "category": ["category"],
            "bool": ["bool"],
        }
        if exclude_dtype == "non_numeric":
            exclude_cols = df.select_dtypes(include="number").columns.tolist()
        elif isinstance(exclude_dtype, str) and exclude_dtype in _shorthands:
            exclude_cols = df.select_dtypes(exclude=_shorthands[exclude_dtype]).columns.tolist()
        elif isinstance(exclude_dtype, (str, type)):
            exclude_cols = df.select_dtypes(exclude=[exclude_dtype]).columns.tolist()
        elif isinstance(exclude_dtype, (list, tuple)):
            exclude_cols = df.select_dtypes(exclude=list(exclude_dtype)).columns.tolist()
        else:
            raise TypeError("exclude_dtype must be a string, type, or list of strings/types")
        return df[exclude_cols]

    selected_cols: set[Any] = set()
    ordered_cols: list[Any] = []
    excluded_cols: set[Any] = set()
    everything_requested = False

    def _resolve_negation(neg_spec: str) -> list[str]:
        raw = _unquote_name(neg_spec.strip())
        if raw in all_cols:
            return [raw]
        elif ":" in raw:
            start_raw, end_raw = raw.split(":", 1)
            start: str | None = _unquote_name(start_raw.strip()) or None
            end: str | None = _unquote_name(end_raw.strip()) or None
            if start and start not in all_cols:
                raise KeyError(f"Start column '{start}' not found")
            if end and end not in all_cols:
                raise KeyError(f"End column '{end}' not found")
            start_idx = all_cols.index(start) if start in all_cols else 0
            end_idx = all_cols.index(end) if end in all_cols else len(all_cols) - 1
            return all_cols[start_idx : end_idx + 1]
        else:
            close = difflib.get_close_matches(raw, [str(c) for c in all_cols], n=1)
            hint = f" (did you mean '{close[0]}'?)" if close else ""
            raise KeyError(f"Column to exclude not found: '{raw}'{hint}")

    if exclude is not None:
        if isinstance(exclude, str):
            raw_exclude = exclude.strip()
            if raw_exclude.startswith("[") and raw_exclude.endswith("]"):
                raw_exclude = raw_exclude[1:-1].strip()
            for token in _tokenize(raw_exclude, ",", keep_quotes=True, track_brackets=True):
                token = _unquote_name(token.strip())
                if token:
                    excluded_cols.update(_resolve_negation(token))
        elif isinstance(exclude, (list, tuple, set)):
            for item in exclude:
                if isinstance(item, str):
                    clean_item = _unquote_name(item.strip())
                    if clean_item.startswith("[") and clean_item.endswith("]"):
                        clean_item = clean_item[1:-1].strip()
                    for sub in _tokenize(clean_item, ",", keep_quotes=True, track_brackets=True):
                        sub = _unquote_name(sub.strip())
                        if sub:
                            excluded_cols.update(_resolve_negation(sub))
                elif item in df.columns:
                    excluded_cols.add(item)
                else:
                    raise KeyError(f"Column to exclude not found: '{item}'")
        else:
            raise TypeError(f"exclude must be a string or sequence of strings, got {type(exclude)}")

    for arg in args:
        if arg is everything or isinstance(arg, everything):
            everything_requested = True
        elif isinstance(arg, list):
            for col in arg:
                clean_col = _unquote_name(col) if isinstance(col, str) else col
                if col in df.columns:
                    selected_cols.add(col)
                    if col not in ordered_cols:
                        ordered_cols.append(col)
                elif isinstance(clean_col, str) and clean_col in df.columns:
                    selected_cols.add(clean_col)
                    if clean_col not in ordered_cols:
                        ordered_cols.append(clean_col)
                elif isinstance(col, str) and col.startswith(("-", "~")):
                    excluded_cols.update(_resolve_negation(col[1:]))
                else:
                    close = difflib.get_close_matches(str(clean_col), [str(c) for c in all_cols], n=1)
                    hint = f" (did you mean '{close[0]}'?)" if close else ""
                    raise KeyError(f"Column not found: '{col}'{hint}")
        elif isinstance(arg, str):
            clean_arg = _unquote_name(arg)
            if arg in df.columns:  # Exact match first — a literal ':' or '-' in a real column name wins
                selected_cols.add(arg)
                if arg not in ordered_cols:
                    ordered_cols.append(arg)
            elif clean_arg in df.columns:
                selected_cols.add(clean_arg)
                if clean_arg not in ordered_cols:
                    ordered_cols.append(clean_arg)
            elif arg.startswith(("-", "~")):
                excluded_cols.update(_resolve_negation(arg[1:]))
            elif ":" in arg:  # Handle slice notation
                start_raw, end_raw = arg.split(":", 1)
                start = _unquote_name(start_raw.strip()) or None  # Empty start means from beginning
                end = _unquote_name(end_raw.strip()) or None     # Empty end means to end
                start_idx = all_cols.index(start) if start in all_cols else 0
                end_idx = all_cols.index(end) if end in all_cols else len(all_cols) - 1
                if start and start not in all_cols:
                    raise KeyError(f"Start column '{start}' not found")
                if end and end not in all_cols:
                    raise KeyError(f"End column '{end}' not found")
                slice_cols = all_cols[start_idx : end_idx + 1]
                selected_cols.update(slice_cols)
                ordered_cols.extend([col for col in slice_cols if col not in ordered_cols])
            elif "," in arg:
                for sub_token in _tokenize(arg, ",", track_brackets=True):
                    if not sub_token:
                        continue
                    clean_sub = _unquote_name(sub_token)
                    if sub_token in df.columns:
                        selected_cols.add(sub_token)
                        if sub_token not in ordered_cols:
                            ordered_cols.append(sub_token)
                    elif clean_sub in df.columns:
                        selected_cols.add(clean_sub)
                        if clean_sub not in ordered_cols:
                            ordered_cols.append(clean_sub)
                    elif sub_token.startswith(("-", "~")):
                        excluded_cols.update(_resolve_negation(sub_token[1:]))
                    elif ":" in sub_token:
                        s_raw, e_raw = sub_token.split(":", 1)
                        s = _unquote_name(s_raw.strip()) or None
                        e = _unquote_name(e_raw.strip()) or None
                        s_idx = all_cols.index(s) if s in all_cols else 0
                        e_idx = all_cols.index(e) if e in all_cols else len(all_cols) - 1
                        if s and s not in all_cols:
                            raise KeyError(f"Start column '{s}' not found")
                        if e and e not in all_cols:
                            raise KeyError(f"End column '{e}' not found")
                        slice_cols = all_cols[s_idx : e_idx + 1]
                        selected_cols.update(slice_cols)
                        ordered_cols.extend([col for col in slice_cols if col not in ordered_cols])
                    else:
                        close = difflib.get_close_matches(clean_sub, [str(c) for c in all_cols], n=1)
                        hint = f" (did you mean '{close[0]}'?)" if close else ""
                        raise KeyError(f"Column not found: '{sub_token}'{hint}")
            else:
                close = difflib.get_close_matches(clean_arg, [str(c) for c in all_cols], n=1)
                hint = f" (did you mean '{close[0]}'?)" if close else ""
                regex_hint = "; for a regex use select(regex=...)" if any(ch in arg for ch in "^$|*+?()") else ""
                raise KeyError(f"Column not found: '{arg}'{hint}{regex_hint}")

        elif callable(arg):
            func_cols = [col for col in df.columns if arg(col)]
            selected_cols.update(func_cols)
            ordered_cols.extend([col for col in func_cols if col not in ordered_cols])
        else:
            raise TypeError(f"Unsupported argument type: {type(arg)}.")

    if dtype is not None:
        if isinstance(dtype, str):
            if dtype == "numeric":
                dtype_cols = df.select_dtypes(include="number").columns.tolist()
            elif dtype == "non_numeric":
                dtype_cols = df.select_dtypes(exclude="number").columns.tolist()
            elif dtype == "datetime":
                dtype_cols = df.select_dtypes(include=["datetime", "datetimetz"]).columns.tolist()
            elif dtype == "category":
                dtype_cols = df.select_dtypes(include=["category"]).columns.tolist()
            elif dtype == "bool":
                dtype_cols = df.select_dtypes(include=["bool"]).columns.tolist()
            else:
                dtype_cols = df.select_dtypes(include=[dtype]).columns.tolist()
        elif isinstance(dtype, (type, list, tuple)):
            dtype_cols = df.select_dtypes(include=list(dtype) if isinstance(dtype, tuple) else dtype).columns.tolist()
        else:
            raise TypeError(f"dtype must be a string, type, list, or tuple, got {type(dtype)}")
        selected_cols.update(dtype_cols)
        ordered_cols.extend([col for col in dtype_cols if col not in ordered_cols])

    if contains is not None:
        if isinstance(contains, str):
            contains_cols = [col for col in df.columns if contains in str(col)]
        elif isinstance(contains, (list, tuple, Sequence)):
            contains_cols = [col for col in df.columns if any(sub in str(col) for sub in contains)]
        else:
            raise TypeError(f"contains must be a string or sequence of strings, got {type(contains)}")
        selected_cols.update(contains_cols)
        ordered_cols.extend([col for col in contains_cols if col not in ordered_cols])

    if startswith is not None:
        if isinstance(startswith, str):
            startswith_cols = [col for col in df.columns if str(col).startswith(startswith)]
        elif isinstance(startswith, (list, tuple, Sequence)):
            startswith_cols = [col for col in df.columns if any(str(col).startswith(sub) for sub in startswith)]
        else:
            raise TypeError(f"startswith must be a string or sequence of strings, got {type(startswith)}")
        selected_cols.update(startswith_cols)
        ordered_cols.extend([col for col in startswith_cols if col not in ordered_cols])

    if endswith is not None:
        if isinstance(endswith, str):
            endswith_cols = [col for col in df.columns if str(col).endswith(endswith)]
        elif isinstance(endswith, (list, tuple, Sequence)):
            endswith_cols = [col for col in df.columns if any(str(col).endswith(sub) for sub in endswith)]
        else:
            raise TypeError(f"endswith must be a string or sequence of strings, got {type(endswith)}")
        selected_cols.update(endswith_cols)
        ordered_cols.extend([col for col in endswith_cols if col not in ordered_cols])

    if regex is not None:
        patterns = [regex] if isinstance(regex, str) else list(regex)
        try:
            compiled = [re.compile(p) for p in patterns]
        except re.error as exc:
            raise ValueError(f"invalid regex: {exc}") from exc
        regex_cols = [col for col in df.columns if any(p.search(str(col)) for p in compiled)]
        selected_cols.update(regex_cols)
        ordered_cols.extend([col for col in regex_cols if col not in ordered_cols])

    # Final assembly:
    if not ordered_cols:
        if everything_requested or excluded_cols:
            final_cols = [col for col in all_cols if col not in excluded_cols]
            return df[final_cols]
        return df[[]]

    if everything_requested:
        remaining_cols = [col for col in all_cols if col not in selected_cols and col not in excluded_cols]
        ordered_cols.extend(remaining_cols)

    final_cols = [col for col in ordered_cols if col not in excluded_cols]
    return df[final_cols]

################################################################################
# FILE: src/pytae/shape.py
################################################################################
from __future__ import annotations

import difflib
from collections.abc import Sequence
from typing import Any

import numpy as np
import pandas as pd

from pytae._text import tokenize as _tokenize
from pytae._text import unquote_name as _unquote_name
from pytae.other_utilities import safe_reset_index

_safe_reset_index = safe_reset_index


def long(
    df: pd.DataFrame,
    cols: str | Sequence[str] | None = None,
    id_vars: str | Sequence[str] | None = None,
    c: str = "variable",
    v: str = "value",
    **kwargs: Any,
) -> pd.DataFrame:
    """Melt columns to rows (unpivot from wide to long format).

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame to melt.
    cols : str or sequence of str, optional
        Columns to unpivot into rows. Aliases: `values`, `value_vars`.
        If not specified and `id_vars` is provided, all remaining columns are melted.
        If neither `cols` nor `id_vars` is provided, all numeric columns are melted.
    id_vars : str or sequence of str, optional
        Identifier columns to keep fixed. Aliases: `by`, `id`.
        If not specified, all columns not in `cols` are kept as identifier variables.
    c : str, default 'variable'
        Name for the melted-variable column. Alias: `var_name`.
    v : str, default 'value'
        Name for the melted-value column. Alias: `value_name`.

    Returns:
    --------
    pd.DataFrame
        Unpivoted long DataFrame.
    """
    # Keyword argument aliases
    if cols is None:
        cols = kwargs.pop("values", kwargs.pop("value_vars", None))
    if id_vars is None:
        id_vars = kwargs.pop("by", kwargs.pop("id", None))
    if c == "variable":
        c = kwargs.pop("var_name", c)
    if v == "value":
        v = kwargs.pop("value_name", v)

    all_cols = list(df.columns)

    def _normalize_col_spec(spec: str | Sequence[str] | None, param_name: str) -> list[str] | None:
        if spec is None:
            return None
        if isinstance(spec, str):
            items = [_unquote_name(item.strip()) for item in _tokenize(spec, ",", keep_quotes=True, track_brackets=True) if item.strip()]
        elif isinstance(spec, (list, tuple, set)):
            items = [_unquote_name(item) if isinstance(item, str) else item for item in spec]
        else:
            raise TypeError(f"{param_name} must be a column name or sequence of names, got {type(spec).__name__}")
        for item in items:
            if item not in df.columns:
                close = difflib.get_close_matches(str(item), [str(col) for col in all_cols], n=1)
                hint = f" (did you mean '{close[0]}'?)" if close else ""
                raise KeyError(f"long(): {param_name} column '{item}' not found in DataFrame{hint}")
        return items

    melt_cols = _normalize_col_spec(cols, "cols")
    keep_ids = _normalize_col_spec(id_vars, "id_vars")

    if melt_cols is not None and keep_ids is not None:
        overlap = set(melt_cols).intersection(keep_ids)
        if overlap:
            raise ValueError(f"long(): column(s) {sorted(overlap)} cannot appear in both 'cols' and 'id_vars'")

    if melt_cols is not None:
        if keep_ids is None:
            keep_ids = [col for col in all_cols if col not in melt_cols]
    elif keep_ids is not None:
        melt_cols = [col for col in all_cols if col not in keep_ids]
        if not melt_cols:
            raise ValueError("long(): no columns left to melt after keeping id_vars")
    else:
        # Default: auto-detect numeric columns
        numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
        if not numeric_cols:
            raise ValueError(
                "long(): no numeric columns to melt. "
                "Specify columns to melt via cols= (e.g. df.pt.long(cols=['q1', 'q2'])) "
                "or id_vars= (e.g. df.pt.long(id_vars=['id']))."
            )
        melt_cols = numeric_cols
        keep_ids = [col for col in all_cols if col not in melt_cols]

    if c == v:
        raise ValueError(f"long(): output variable column 'c' ({c!r}) and value column 'v' ({v!r}) cannot have the same name")
    if c in keep_ids:
        raise ValueError(f"long(): output variable column 'c' ({c!r}) collides with kept id column(s)")
    if v in keep_ids:
        raise ValueError(f"long(): output value column 'v' ({v!r}) collides with kept id column(s)")
    if v in melt_cols:
        raise ValueError(f"long(): output value column 'v' ({v!r}) collides with melted column(s)")

    return pd.melt(
        df,
        id_vars=keep_ids,
        value_vars=melt_cols,
        var_name=c,
        value_name=v,
    )


def wide(
    df: pd.DataFrame,
    c: str = "variable",
    v: str = "value",
    index: str | Sequence[str] | None = None,
    **kwargs: Any,
) -> pd.DataFrame:
    """Pivot a long column into headers (pure 1-to-1 reshape).

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame to pivot.
    c : str, default 'variable'
        Column whose values become headers (columns).
    v : str, default 'value'
        Values column to populate table cells.
    index : str or sequence of str, optional
        Explicit index column(s) to use. Aliases: `r`, `rows`, `by`, `id_vars`.
        If omitted, all columns other than `c` and `v` are used.

    Returns:
    --------
    pd.DataFrame
        Pivoted wide DataFrame with standard RangeIndex.
    """
    for bad_arg in ("a", "agg", "aggfunc", "dropna"):
        if bad_arg in kwargs:
            raise ValueError(
                f"wide() does not support '{bad_arg}'. wide() is strictly for 1-to-1 reshaping "
                "without aggregation. Use pt.pivot() / -pivot for multi-dimensional aggregation."
            )

    if "cols" in kwargs and c == "variable":
        c = kwargs.pop("cols")
    elif "col" in kwargs and c == "variable":
        c = kwargs.pop("col")
    elif "columns" in kwargs and c == "variable":
        c = kwargs.pop("columns")

    if "values" in kwargs and v == "value":
        v = kwargs.pop("values")
    elif "val" in kwargs and v == "value":
        v = kwargs.pop("val")
    elif "vals" in kwargs and v == "value":
        v = kwargs.pop("vals")

    c = _unquote_name(c)
    if v is not None:
        v = _unquote_name(v)
    if c == v:
        raise ValueError(f"wide(): 'c' and 'v' cannot be the same column ('{c}')")
    if c not in df.columns:
        raise KeyError(f"wide(): columns 'c' column '{c}' not found in DataFrame")
    if v is not None and v not in df.columns:
        raise KeyError(f"wide(): values 'v' column '{v}' not found in DataFrame")

    if index is None:
        index = kwargs.pop("r", kwargs.pop("rows", kwargs.pop("by", kwargs.pop("id_vars", None))))

    if index is not None:
        if isinstance(index, str):
            index_cols = [_unquote_name(col.strip()) for col in _tokenize(index, ",", keep_quotes=True, track_brackets=True) if col.strip()]
        else:
            index_cols = [_unquote_name(col) if isinstance(col, str) else col for col in index]
        for col in index_cols:
            if col not in df.columns:
                raise KeyError(f"wide(): index column '{col}' not found in DataFrame")
        if c in index_cols:
            raise ValueError(f"wide(): column 'c' ('{c}') cannot also be in index columns")
        if v is not None and v in index_cols:
            raise ValueError(f"wide(): column 'v' ('{v}') cannot also be in index columns")
    else:
        index_cols = [col for col in df.columns if col not in [c, v]]

    try:
        if not index_cols:
            temp_df = df.copy()
            temp_df["__pt_wide_tmp__"] = 0
            pivoted = temp_df.pivot(index="__pt_wide_tmp__", columns=c, values=v)
        else:
            pivoted = df.pivot(index=index_cols, columns=c, values=v)
    except ValueError as exc:
        raise ValueError(
            f"wide() encountered duplicate entries for index {index_cols} and column '{c}'. "
            "wide() is strictly for 1-to-1 reshaping without aggregation. "
            "Use pt.pivot() / -pivot to summarize or aggregate duplicate entries."
        ) from exc

    wide_df = _safe_reset_index(pivoted)
    if not index_cols and "__pt_wide_tmp__" in wide_df.columns:
        wide_df = wide_df.drop(columns=["__pt_wide_tmp__"])
    wide_df.columns.name = None
    return wide_df


def pivot(
    df: pd.DataFrame,
    r: str | Sequence[str] | None = None,
    c: str | Sequence[str] | None = None,
    v: str | Sequence[str] | None = None,
    a: str = "sum",
    dropna: bool = False,
    fill_value: Any = None,
    **kwargs: Any,
) -> pd.DataFrame:
    """Summarize and aggregate a DataFrame across dimensions (Excel-style pivot table).

    Parameters:
    -----------
    df : pd.DataFrame
        The DataFrame to pivot.
    r : str or sequence of str, optional
        Row dimension(s) to group by. Aliases: `rows`, `row`, `index`, `by`.
    c : str or sequence of str, optional
        Column dimension(s) to spread as headers. Aliases: `cols`, `col`, `columns`.
    v : str or sequence of str
        Value column(s) to aggregate. REQUIRED. Aliases: `values`, `value`, `val`, `vals`.
    a : str, default 'sum'
        Aggregation function ('sum', 'mean', 'median', 'min', 'max', 'count', 'std', etc.).
        'n' is accepted as an alias for 'size' (group row count). Aliases: `agg`, `aggfunc`.
    dropna : bool, default False
        Whether to drop NA categories from grouping keys.
    fill_value : any, optional
        Value to replace missing grid intersections with (e.g. 0).

    Returns:
    --------
    pd.DataFrame
        Clean, flattened DataFrame with standard RangeIndex and 1D column names.
    """
    if r is None:
        r = kwargs.pop("rows", kwargs.pop("row", kwargs.pop("index", kwargs.pop("by", None))))
    if c is None:
        c = kwargs.pop("cols", kwargs.pop("col", kwargs.pop("columns", None)))
    if v is None:
        v = kwargs.pop("values", kwargs.pop("value", kwargs.pop("val", kwargs.pop("vals", None))))
    if a == "sum":
        a = kwargs.pop("agg", kwargs.pop("aggfunc", a))
    if fill_value is None:
        fill_value = kwargs.pop("fill", None)

    all_cols = list(df.columns)

    def _normalize_cols(spec: str | Sequence[str] | None, param_name: str) -> list[str]:
        if spec is None:
            return []
        if isinstance(spec, str):
            cols = [_unquote_name(item.strip()) for item in _tokenize(spec, ",", keep_quotes=True, track_brackets=True) if item.strip()]
        elif isinstance(spec, (list, tuple, set)):
            cols = [_unquote_name(item) if isinstance(item, str) else item for item in spec]
        else:
            raise TypeError(f"pivot(): {param_name} must be a column name or sequence of names, got {type(spec).__name__}")
        for col in cols:
            if col not in df.columns:
                close = difflib.get_close_matches(str(col), [str(x) for x in all_cols], n=1)
                hint = f" (did you mean '{close[0]}'?)" if close else ""
                raise KeyError(f"pivot(): {param_name} column '{col}' not found in DataFrame{hint}")
        return cols

    aggfunc = "size" if a == "n" else a
    if aggfunc == "size" and fill_value is None:
        fill_value = 0

    r_cols = _normalize_cols(r, "r (rows)")
    c_cols = _normalize_cols(c, "c (cols)")
    v_cols = _normalize_cols(v, "v (values)")

    dim_counts: dict[str, int] = {}
    for col in r_cols + c_cols:
        dim_counts[col] = dim_counts.get(col, 0) + 1
    dups = [col for col, count in dim_counts.items() if count > 1]
    if dups:
        raise ValueError(f"pivot(): column(s) {dups} cannot appear multiple times in grouping dimensions (r/c)")

    if len(v_cols) != len(set(v_cols)):
        v_dups = [col for col in dict.fromkeys(v_cols) if v_cols.count(col) > 1]
        raise ValueError(f"pivot(): value column(s) {v_dups} cannot appear multiple times in 'v'")

    if not v_cols and aggfunc != "size":
        raise ValueError("pivot(): value column 'v' is required (e.g. v='Sales')")

    overlap = [col for col in v_cols if (r_cols and col in r_cols) or (c_cols and col in c_cols)]
    if overlap:
        if aggfunc == "size":
            raise ValueError(
                f"pivot(): value column(s) {overlap} cannot also be in grouping dimensions (r/c). "
                "For row counts (a='n'), omit 'v' or specify a non-grouping column."
            )
        else:
            raise ValueError(f"pivot(): value column(s) {overlap} cannot also be in grouping dimensions (r/c)")

    observed = kwargs.get("observed", True)

    def _format_part(part: Any) -> str:
        if part is None or (isinstance(part, float) and np.isnan(part)) or pd.isna(part):
            return "nan"
        if isinstance(part, (float, np.floating)) and float(part).is_integer():
            return str(int(part))
        return str(part)

    def _flatten_cols(piv: pd.DataFrame) -> None:
        raw_names: list[str] = []
        if isinstance(piv.columns, pd.MultiIndex):
            for col in piv.columns:
                formatted_parts = [_format_part(part) for part in col if part is not None]
                raw_names.append("_".join(formatted_parts))
        else:
            raw_names = [_format_part(col) for col in piv.columns]

        duplicates = []
        seen = set()
        for name in raw_names:
            if name in seen and name not in duplicates:
                duplicates.append(name)
            seen.add(name)
        if duplicates:
            raise ValueError(
                f"pivot(): resulting column header(s) {duplicates} contain duplicate names. "
                "Ensure pivoted column keys have distinct string representations."
            )
        piv.columns = raw_names

    def _cast_size_ints(df_out: pd.DataFrame) -> pd.DataFrame:
        if aggfunc == "size":
            val_cols = [col for col in df_out.columns if not (r_cols and col in r_cols)]
            for col in val_cols:
                try:
                    if df_out[col].isna().any():
                        df_out[col] = df_out[col].astype("Int64")
                    else:
                        df_out[col] = df_out[col].astype("int64")
                except (TypeError, ValueError):
                    pass
        return df_out

    val_arg = v_cols[0] if len(v_cols) == 1 else v_cols

    if r_cols and c_cols:
        group_keys = list(r_cols) + list(c_cols)
        if aggfunc == "size":
            s = df.groupby(group_keys, dropna=dropna, observed=observed).size()
            pivoted = s.unstack(list(c_cols), fill_value=fill_value if fill_value is not None else 0)
        else:
            g = df.groupby(group_keys, dropna=dropna, observed=observed)[val_arg].agg(aggfunc)
            pivoted = g.unstack(list(c_cols), fill_value=fill_value)
        _flatten_cols(pivoted)
        res = _safe_reset_index(pivoted)
        res.columns.name = None
        return _cast_size_ints(res)
    elif r_cols:
        if aggfunc == "size":
            target_name = v_cols[0] if v_cols else "n"
            if target_name in r_cols:
                raise ValueError(
                    f"pivot(): count column name '{target_name}' collides with row grouping column. "
                    "For row counts (a='n'), omit 'v' or specify a non-colliding column name."
                )
            s = df.groupby(r_cols, dropna=dropna, observed=observed).size()
            res = s.to_frame(name=target_name).reset_index()
            res.columns.name = None
            return _cast_size_ints(res)
        else:
            g = df.groupby(r_cols, dropna=dropna, observed=observed)[val_arg].agg(aggfunc)
            pivoted = g if isinstance(g, pd.DataFrame) else g.to_frame()
            if fill_value is not None:
                pivoted = pivoted.fillna(fill_value)
            _flatten_cols(pivoted)
            res = _safe_reset_index(pivoted)
            res.columns.name = None
            return res
    elif c_cols:
        if aggfunc == "size":
            s = df.groupby(c_cols, dropna=dropna, observed=observed).size()
            pivoted = s.to_frame().T.reset_index(drop=True)
            _flatten_cols(pivoted)
            pivoted.columns.name = None
            return _cast_size_ints(pivoted)
        else:
            g = df.groupby(c_cols, dropna=dropna, observed=observed)[val_arg].agg(aggfunc)
            if len(v_cols) == 1:
                pivoted = g.to_frame().T.reset_index(drop=True)
                if fill_value is not None:
                    pivoted = pivoted.fillna(fill_value)
                _flatten_cols(pivoted)
                pivoted.columns.name = None
                return pivoted
            else:
                pivoted = g.T
                if fill_value is not None:
                    pivoted = pivoted.fillna(fill_value)
                _flatten_cols(pivoted)
                pivoted.index.name = "metric"
                res = _safe_reset_index(pivoted)
                res.columns.name = None
                return res
    else:
        if aggfunc == "size":
            target_cols = v_cols if v_cols else ["n"]
            res = pd.DataFrame({col: [len(df)] for col in target_cols}, dtype="int64")
        else:
            res = df[v_cols].agg(aggfunc).to_frame().T.reset_index(drop=True)
        if fill_value is not None:
            res = res.fillna(fill_value)
        _flatten_cols(res)
        res.columns.name = None
        return _cast_size_ints(res)

################################################################################
# FILE: src/pytae/sql.py
################################################################################
"""SQL over DataFrames via duckdb. Optional extra: pip install 'pytae[sql]'."""

from __future__ import annotations

import pandas as pd


def _require_duckdb():
    try:
        import duckdb
    except ImportError as exc:
        raise ImportError(
            "sql() requires duckdb. Install with: pip install 'pytae[sql]'"
        ) from exc
    return duckdb


def _normalize_sql(query: str) -> str:
    """Normalize query: read from @path if given, and convert [col] and `col` to \"col\"."""
    q = query.strip()
    if (q.startswith("'") and q.endswith("'")) or (q.startswith('"') and q.endswith('"')):
        q = q[1:-1].strip()
    if q.startswith("@"):
        path = q[1:].strip()
        if (path.startswith("'") and path.endswith("'")) or (path.startswith('"') and path.endswith('"')):
            path = path[1:-1].strip()
        from pathlib import Path
        file_path = Path(path)
        if not file_path.is_file():
            raise FileNotFoundError(f"sql: query file not found: '{path}'")
        q = file_path.read_text(encoding="utf-8").strip()
    out_chars: list[str] = []
    i, n = 0, len(q)
    quote = None
    while i < n:
        ch = q[i]
        if quote is not None:
            out_chars.append(ch)
            if ch == quote:
                if i + 1 < n and q[i + 1] == quote:
                    out_chars.append(q[i + 1])
                    i += 2
                    continue
                quote = None
            i += 1
            continue
        if ch in "'\"":
            quote = ch
            out_chars.append(ch)
            i += 1
            continue
        if ch == "`":
            j = q.find("`", i + 1)
            if j != -1:
                ident = q[i + 1 : j]
                out_chars.append(f'"{ident}"')
                i = j + 1
                continue
        if ch == "[":
            j = q.find("]", i + 1)
            if j != -1:
                inner = q[i + 1 : j]
                is_index = i > 0 and (q[i - 1].isalnum() or q[i - 1] in "_)]")
                is_list_literal = "," in inner or "'" in inner or '"' in inner
                has_ident = any(c.isalpha() or c == "_" for c in inner)
                all_valid = all(c.isalnum() or c in " _-" for c in inner)
                if not is_index and not is_list_literal and has_ident and all_valid:
                    out_chars.append(f'"{inner}"')
                    i = j + 1
                    continue

        out_chars.append(ch)
        i += 1
    return "".join(out_chars)


def sql(df: pd.DataFrame, query: str, /, **frames: pd.DataFrame) -> pd.DataFrame:
    """Run a SQL query via duckdb.

    The current DataFrame is table ``data`` (same as CLI ``-sql``). Extra keyword
    frames are registered under those names, like ``-file`` aliases.

    Parameters
    ----------
    df:
        Registered as table ``data``.
    query:
        DuckDB SQL, or '@query.txt' to load from a file. Spaced column names can
        use double quotes (\"col a\"), brackets ([col a]), or backticks (`col a`).
    **frames:
        Additional DataFrames registered as tables. The name ``data`` is reserved.

    Examples
    --------
    >>> import pytae as pt
    >>> penguins = pt.sample("penguins")
    >>> pt.sql(penguins, "select species, avg(body_mass_g) as avg_mass from data group by species order by species")
         species     avg_mass
    0     Adelie  3700.662252
    1  Chinstrap  3733.088235
    2     Gentoo  5076.016260
    >>> penguins.pt.sql("select * from data where species = 'Adelie'").shape
    (152, 7)
    >>> left = pd.DataFrame({"id": [1, 2], "x": ["a", "b"]})
    >>> right = pd.DataFrame({"id": [1, 3], "y": ["p", "q"]})
    >>> pt.sql(left, "select data.id, x, y from data join extra using (id)", extra=right)
       id  x  y
    0   1  a  p
    """
    if "data" in frames:
        raise ValueError("sql(): extra frame cannot be named 'data' — that's the current DataFrame")
    normalized_query = _normalize_sql(query)
    tables = {"data": df, **frames}
    duckdb = _require_duckdb()
    con = duckdb.connect()
    try:
        for name, frame in tables.items():
            con.register(name, frame)
        try:
            return con.sql(normalized_query).df()
        except Exception as exc:
            raise ValueError(f"sql: {exc}") from exc
    finally:
        con.close()

================================================================================
11. TEST SUITE IMPLEMENTATION SOURCE CODE (tests/)
================================================================================

################################################################################
# FILE: tests/__init__.py
################################################################################
# Makes `from tests.cli_helpers import ...` work when pytest's rootdir is the repo.

################################################################################
# FILE: tests/cli_helpers.py
################################################################################
"""Shared fixtures for CLI tests."""
from __future__ import annotations

import pandas as pd


def _write_csv(tmp_path, df: pd.DataFrame) -> str:
    path = tmp_path / "data.csv"
    df.to_csv(path, index=False)
    return str(path)


def _penguins_frame():
    return pd.DataFrame({
        "species": ["Adelie", "Adelie", "Adelie", "Chinstrap", "Chinstrap", "Gentoo", "Gentoo", "Gentoo"],
        "island": ["Biscoe", "Dream", "Dream", "Dream", "Dream", "Biscoe", "Biscoe", "Biscoe"],
        "sex": ["Male", "Female", "Male", "Male", "Female", "Male", "Female", "Male"],
        "body_mass_g": [3750, 3800, 4000, 3700, 3400, 5700, 4500, 5600],
    })


def _replace_frame():
    return pd.DataFrame({
        "col a": ["a magician", "not a magician exactly", "analphabet"],
        "colb": ["alpha", "alpha team", "x"],
    })


def _messy_headers_frame():
    return pd.DataFrame({
        "  Col A  ": [1],
        "col   b": [2],
        "Col A": [3],
        "100% Match!": [4],
    })


def _write_two_csvs(tmp_path):
    left = tmp_path / "left.csv"
    right = tmp_path / "right.csv"
    pd.DataFrame({"col a": [1, 2, 3], "val_l": ["a", "b", "c"]}).to_csv(left, index=False)
    pd.DataFrame({"cola": [1, 2, 4], "val_r": ["x", "y", "z"]}).to_csv(right, index=False)
    return str(left), str(right)


def _write_three_id_csvs(tmp_path):
    a = tmp_path / "a.csv"
    b = tmp_path / "b.csv"
    c = tmp_path / "c.csv"
    pd.DataFrame({"id": [1, 2, 3], "x": ["a", "b", "c"]}).to_csv(a, index=False)
    pd.DataFrame({"id": [1, 2, 4], "y": ["p", "q", "r"]}).to_csv(b, index=False)
    pd.DataFrame({"id": [1, 2], "z": ["m", "n"]}).to_csv(c, index=False)
    return str(a), str(b), str(c)

################################################################################
# FILE: tests/conftest.py
################################################################################
import os
import sys

import pytest

# Keep test imports consistent with existing test files.
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae


@pytest.fixture
def sample_csv(tmp_path):
    """Factory fixture: write a bundled pytae.sample_data dataset to CSV in tmp_path, return its path."""
    def _write(name: str) -> str:
        path = tmp_path / f"{name}.csv"
        pytae.sample_data[name].to_csv(path, index=False)
        return str(path)
    return _write


@pytest.fixture
def sample_parquet(tmp_path):
    """Factory fixture: write a bundled pytae.sample_data dataset to parquet in tmp_path, return its path."""
    def _write(name: str) -> str:
        path = tmp_path / f"{name}.parquet"
        pytae.sample_data[name].to_parquet(path)
        return str(path)
    return _write

################################################################################
# FILE: tests/test_accessor.py
################################################################################
import os
import sys

import pandas as pd

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt


def test_pt_accessor_select_then_agg_df():
    penguins = pt.sample("penguins")
    out = (
        penguins
        .pt.agg_df(["species", "island"], a=["mean", "n"])
    )
    assert list(out.columns)[:3] == ["species", "island", "n"]


def test_pandas_rename_then_pt_agg_df():
    df = pd.DataFrame({"g": ["a", "a", "b"], "v": [1, 2, 3]})
    out = df.rename(columns={"g": "grp", "v": "val"}).pt.agg_df("grp", a=["sum", "n"])
    assert list(out.columns) == ["grp", "n", "val"]
    assert out.loc[out["grp"] == "a", "val"].iloc[0] == 3


def test_pt_select_then_pandas_head():
    penguins = pt.sample("penguins")
    out = penguins.pt.select("species", "body_mass_g").head(3)
    assert list(out.columns) == ["species", "body_mass_g"]
    assert len(out) == 3

################################################################################
# FILE: tests/test_agg_df.py
################################################################################
import os
import sys

import numpy as np
import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt


def test_agg_df_sum():
    data = {
        "category": ["A", "A", "B", "B", "C"],
        "value": [10, 20, 30, 40, 50],
    }
    df = pd.DataFrame(data)

    result = pt.agg_df(df, "category", a=["sum"])
    expected_df = pd.DataFrame(data).groupby("category").sum().reset_index()
    pd.testing.assert_frame_equal(result, expected_df)


def test_all_agg():
    df = pd.DataFrame({
        "id": ["a", "b", "c", "d", "e", "", "f", "f"],
        "balance": [10, 20, 0, 21, 15, 10, 20, 25],
        "country": ["sg", "cn", "ca", "np", "in", "in", "in", "in"],
    })
    df["id"] = df["id"].replace("", np.nan)

    # Default dropna=False keeps NA groups
    result = pt.agg_df(df, ["id", "country"], a=["sum", "min", "mean", "min", "max", "n"])

    expected_df = (
        df.groupby(["id", "country"], dropna=False)
        .agg(
            n=("id", "size"),
            balance_sum=("balance", "sum"),
            balance_min=("balance", "min"),
            balance_mean=("balance", "mean"),
            balance_max=("balance", "max"),
        )
        .reset_index()
    )
    expected_df = expected_df[["id", "country", "n", "balance_sum", "balance_min", "balance_mean", "balance_max"]]
    pd.testing.assert_frame_equal(result, expected_df)


def test_all_agg_drop_na():
    df = pd.DataFrame({
        "id": ["a", "b", "c", "d", "e", "", "f", "f"],
        "balance": [10, 20, 0, 21, 15, 10, 20, 25],
        "country": ["sg", "cn", "ca", "np", "in", "in", "in", "in"],
    })
    df["id"] = df["id"].replace("", np.nan)

    result = pt.agg_df(df, ["id", "country"], a=["sum", "min", "mean", "min", "max", "n"], dropna=True)

    expected_df = (
        df.groupby(["id", "country"], dropna=True)
        .agg(
            n=("id", "size"),
            balance_sum=("balance", "sum"),
            balance_min=("balance", "min"),
            balance_mean=("balance", "mean"),
            balance_max=("balance", "max"),
        )
        .reset_index()
    )
    expected_df = expected_df[["id", "country", "n", "balance_sum", "balance_min", "balance_mean", "balance_max"]]
    pd.testing.assert_frame_equal(result, expected_df)


def test_agg_df_raises_when_n_collides_with_real_column():
    df = pd.DataFrame({"g": ["x", "x", "y"], "n": [1, 2, 3], "v": [10, 20, 30]})
    with pytest.raises(ValueError, match="already has a numeric column named 'n'"):
        pt.agg_df(df, "g", ["sum", "n"])


def test_agg_df_column_kwargs():
    df = pd.DataFrame({"grp": ["A", "A", "B"], "v1": [10, 20, 30], "v2": [1, 2, 3]})
    res = df.pt.agg_df("grp", v1="mean", v2="max", count="n")
    assert list(res.columns) == ["grp", "v1", "v2", "count"]
    assert list(res["v1"]) == [15.0, 30.0]
    assert list(res["v2"]) == [2, 3]
    assert list(res["count"]) == [2, 1]


def test_agg_df_rejects_dictionary():
    df = pd.DataFrame({"grp": ["A", "A", "B"], "v1": [10, 20, 30]})
    with pytest.raises(TypeError, match=r"agg\(\) no longer accepts dictionaries"):
        pt.agg_df(df, "grp", {"v1": "mean"})
    with pytest.raises(TypeError, match=r"agg\(\) no longer accepts dictionaries"):
        df.pt.agg_df("grp", a={"v1": "mean"})
    with pytest.raises(TypeError, match=r"agg\(\) no longer accepts dictionaries"):
        pt.agg_df(df, {"grp": "mean"})


def test_agg_df_missing_by_raises():
    df = pd.DataFrame({"grp": ["A", "A", "B"], "v1": [10, 20, 30]})
    with pytest.raises(TypeError, match="missing required argument: 'by'"):
        df.pt.agg_df()


def test_agg_df_aggfunc_as_by_raises_helpful_error():
    df = pd.DataFrame({"grp": ["A", "A", "B"], "v1": [10, 20, 30]})
    with pytest.raises(ValueError, match="is a known aggregation function"):
        df.pt.agg_df("mean")
    with pytest.raises(ValueError, match="looks like an aggregation list"):
        df.pt.agg_df(["mean", "n"])


def test_agg_df_unknown_group_column_raises():
    df = pd.DataFrame({"species": ["A", "B"], "val": [1, 2]})
    with pytest.raises(KeyError, match=r"group column 'speceis' not found in DataFrame \(did you mean 'species'\?\)"):
        df.pt.agg_df("speceis", "mean")


def test_agg_df_numeric_group_by():
    df = pd.DataFrame({
        "year": [2023, 2023, 2024],
        "sales": [100.0, 200.0, 300.0],
        "units": [10, 20, 30],
    })
    res = df.pt.agg_df("year", a=["sum", "n"])
    assert list(res.columns) == ["year", "n", "sales", "units"]
    assert list(res["year"]) == [2023, 2024]
    assert list(res["sales"]) == [300.0, 300.0]
    assert list(res["units"]) == [30, 30]
    assert list(res["n"]) == [2, 1]


def test_agg_df_whole_table_summary():
    df = pd.DataFrame({
        "sales": [100.0, 200.0, 300.0],
        "units": [10, 20, 30],
    })
    res = df.pt.agg_df(None, a=["sum", "mean", "n"])
    assert len(res) == 1
    assert list(res.columns) == ["n", "sales_sum", "sales_mean", "units_sum", "units_mean"]
    assert res["n"].iloc[0] == 3
    assert res["sales_sum"].iloc[0] == 600.0
    assert res["sales_mean"].iloc[0] == 200.0


def test_agg_df_string_spec_with_brackets():
    df = pd.DataFrame({
        "smoker status": ["Yes", "Yes", "No"],
        "tip amount": [2.0, 4.0, 3.0],
        "total bill": [10.0, 20.0, 30.0],
    })
    # 1. Single string spec with brackets
    res1 = df.pt.agg_df("[smoker status]", "[tip amount] = mean, [total bill] = mean, n = n")
    assert list(res1.columns) == ["smoker status", "tip amount", "total bill", "n"]
    assert list(res1["smoker status"]) == ["No", "Yes"]
    assert list(res1["tip amount"]) == [3.0, 3.0]
    assert list(res1["total bill"]) == [30.0, 15.0]
    assert list(res1["n"]) == [1, 2]

    # 2. Multiple positional string specs
    res2 = df.pt.agg_df("[smoker status]", "[tip amount] = mean", "[total bill] = mean", "n = n")
    pd.testing.assert_frame_equal(res1, res2)

    # 3. Passed via keyword a=
    res3 = df.pt.agg_df("[smoker status]", a="[tip amount] = mean, [total bill] = mean, n = n")
    pd.testing.assert_frame_equal(res1, res3)

    # 4. Mixed string spec and kwargs
    res4 = df.pt.agg_df("[smoker status]", "[total bill] = mean", **{"tip amount": "mean", "n": "n"})
    assert list(res4.columns) == ["smoker status", "total bill", "tip amount", "n"]
    assert list(res4["total bill"]) == [30.0, 15.0]

    # 5. Named aggregation with brackets in output and source
    res5 = df.pt.agg_df("[smoker status]", "[avg bill] = [total bill]:mean, [bill sum] = [total bill]:sum, [count] = n")
    assert list(res5.columns) == ["smoker status", "avg bill", "bill sum", "count"]
    assert list(res5["avg bill"]) == [30.0, 15.0]
    assert list(res5["bill sum"]) == [30.0, 30.0]
    assert list(res5["count"]) == [1, 2]

    # 6. Whole-table summary with string mapping
    res6 = df.pt.agg_df(None, "[avg bill] = [total bill]:mean, [count] = n")
    assert len(res6) == 1
    assert list(res6.columns) == ["avg bill", "count"]
    assert res6["avg bill"].iloc[0] == 20.0
    assert res6["count"].iloc[0] == 3


def test_agg_df_whole_frame_string_spec():
    df = pd.DataFrame({"g": ["a", "a", "b"], "v1": [1.0, 3.0, 5.0], "v2": [2.0, 4.0, 6.0]})
    res = df.pt.agg_df("g", "mean, n")
    assert list(res.columns) == ["g", "n", "v1", "v2"]
    assert list(res["n"]) == [2, 1]
    assert list(res["v1"]) == [2.0, 5.0]

    res_multi = df.pt.agg_df("g", "mean", "n")
    pd.testing.assert_frame_equal(res, res_multi)


def test_agg_df_user_pipeline_tips():
    tips = pt.sample("tips")
    res = (
        tips
        .pt.qry(day=["Sat", "Sun"], time="Dinner", size=">= 2", total_bill="> 10")
        .pt.select("smoker", "tip", "total_bill")
        .rename(columns={"total_bill": "total bill"})
        .pt.agg_df("smoker", "tip = mean, [total bill] = mean, n = n")
    )
    assert list(res.columns) == ["smoker", "tip", "total bill", "n"]
    assert len(res) == 2
    assert list(res["smoker"]) == ["Yes", "No"]
    assert list(res["n"]) == [57, 97]


def test_agg_df_string_spec_errors():
    df = pd.DataFrame({"g": ["a", "b"], "v": [1, 2]})
    with pytest.raises(ValueError, match="agg\\(\\):"):
        df.pt.agg_df("g", "invalid syntax without equals and not known agg")
    with pytest.raises(ValueError, match="cannot mix whole-frame aggregation"):
        df.pt.agg_df("g", "mean", v="sum")
    with pytest.raises(TypeError, match="multiple values for aggregation spec"):
        df.pt.agg_df("g", "mean", a="sum")


def test_agg_df_multi_agg_column_flattening_preserves_underscores():
    # Lib Issue 1: multi-agg flattening preserves leading/trailing underscores
    df = pd.DataFrame({"g": ["a", "b"], "_x": [1, 2], "y_": [3, 4]})
    res = df.pt.agg_df("g", ["mean", "sum"])
    assert "_x_mean" in res.columns
    assert "_x_sum" in res.columns
    assert "y__mean" in res.columns
    assert "y__sum" in res.columns


def test_agg_df_group_by_n_with_count_n_raises():
    # Lib Issue 2: grouping by 'n' and requesting 'n' raises collision error
    df = pd.DataFrame({"n": ["a", "a", "b"], "v": [1.0, 2.0, 4.0]})
    with pytest.raises(ValueError, match="collides with group column 'n'|already has a group column named 'n'"):
        df.pt.agg_df("n", ["mean", "n"])

    with pytest.raises(ValueError, match="collides with group column 'n'|already has a group column named 'n'"):
        df.pt.agg_df("n", v="mean", n="n")


def test_agg_df_output_name_collides_with_group_col():
    # Issue 1: named agg output matching a group column raises ValueError
    df = pd.DataFrame({"Region": ["East", "East", "West"], "Sales": [10, 30, 5], "Profit": [1, 2, 3]})
    with pytest.raises(ValueError, match="collides with group column"):
        df.pt.agg("Region", Region="Sales:sum")

    with pytest.raises(ValueError, match="collides with group column"):
        df.pt.agg("Region", "Region = Sales:sum")

def test_agg_df_duplicate_output_cols():
    # Issue 1: duplicate output column name in named_aggs / output_cols
    df = pd.DataFrame({"Region": ["East", "East", "West"], "Sales": [10, 30, 5], "Profit": [1, 2, 3]})
    with pytest.raises(ValueError, match="is duplicated in aggregation specification"):
        df.pt.agg("Region", Sales_sum="Profit:sum", Sales="Sales:sum,mean")


def test_agg_df_list_collision_with_group_col():
    # Issue 2: generated agg column collides with group column in _agg_df_list
    df = pd.DataFrame({"Sales_sum": ["A", "B"], "Sales": [10, 20]})
    with pytest.raises(ValueError, match="collides with group column 'Sales_sum'"):
        df.pt.agg("Sales_sum", ["sum", "mean"])


def test_agg_df_comma_separated_by_string():
    df = pd.DataFrame({
        "day": ["Thur", "Thur", "Fri"],
        "sex": ["Male", "Female", "Male"],
        "total_bill": [10.0, 20.0, 30.0],
    })
    res = df.pt.agg("day,sex", total_bill="mean")
    assert list(res.columns) == ["day", "sex", "total_bill"]
    assert len(res) == 3


if __name__ == "__main__":
    pytest.main()

################################################################################
# FILE: tests/test_cli_agg.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv


def test_order_agg_then_shape_uses_aggregated_frame(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", "x", "y", "y"],
            "v1": [1, 2, 3, 4],
            "v2": [10, 20, 30, 40],
        }
    )
    path = _write_csv(tmp_path, df)

    # With -by: groups by grp, v1/v2 summed -> 2 rows x 3 cols
    exit_code = cli.main([path, "-by", "grp", "-agg", "-shape"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert "(2, 3)" in out

    # Without -by: whole-table summary across numeric columns -> 1 row x 2 cols
    exit_code = cli.main([path, "-agg", "-shape"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert "(1, 2)" in out


def test_order_agg_dropna_default_and_flags(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", None],
            "v1": [1, 2],
        }
    )
    path = _write_csv(tmp_path, df)

    # Default is dropna=false: keeps the NA grouping key, so both groups remain (2 rows).
    exit_code = cli.main([path, "-by", "grp", "-agg", "sum", "-shape"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert "(2, 2)" in out

    # Explicit -dropna true drops the NA grouping key (1 row).
    exit_code_drop = cli.main([path, "-by", "grp", "-agg", "sum", "-dropna", "true", "-shape"])
    out_drop = capsys.readouterr().out
    assert exit_code_drop == 0
    assert "(1, 2)" in out_drop


def test_group_by_agg_bare_aggfunc_keeps_source_column_name(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", "x", "y", "y"],
            "v1": [1, 2, 3, 4],
            "v2": [10, 20, 30, 40],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-by", "grp", "-agg", "v1 = sum, v2 = mean"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "v1" in out and "v2" in out
    assert "3" in out and "15.0" in out
    assert "7" in out and "35.0" in out


def test_group_by_agg_named_output(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", "x", "y", "y"],
            "v1": [1, 2, 3, 4],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-by", "grp", "-agg", "total = v1:sum"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "total" in out
    assert "v1" not in out.split("\n")[0]  # header uses the custom output name, not the source column


def test_agg_same_func_on_multiple_columns(tmp_path, capsys):
    df = pd.DataFrame({"grp": ["x", "x"], "v1": [1, 2], "v2": [10, 20]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-by", "grp", "-agg", "sum"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "v1" in out and "v2" in out
    assert "3" in out and "30" in out


def test_agg_named_output_with_spaced_group(tmp_path, capsys):
    df = pd.DataFrame({"Scenario Name": ["A", "A", "B"], "value": [1, 2, 3]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main(
        [path, "-by", "Scenario Name", "-agg", "v = value:sum"]
    )

    out = capsys.readouterr().out
    assert exit_code == 0
    header = out.split("\n")[0]
    assert "v" in header and "value" not in header
    assert "3" in out


def test_agg_rejects_dict_literal(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"grp": ["x"], "v1": [1]}))

    with pytest.raises(SystemExit, match="use a name"):
        cli.main([path, "-by", "grp", "-agg", "{'v1':'sum'}"])


def test_agg_retired_syntax_raises_helpful_error(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"grp": ["x"], "v1": [1]}))

    with pytest.raises(SystemExit, match="syntax is retired"):
        cli.main([path, "-by", "grp", "-agg", "column=v1,aggfunc=sum"])


def test_by_requires_agg_or_mutate(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"grp": ["x"], "v1": [1]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-by", "grp"])
    assert exc_info.value.code == 2
    assert "requires -agg or -mutate" in capsys.readouterr().err


def test_agg_then_sort_by_prints_only_final_table(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["b", "a", "a"],
            "val": [2, 1, 3],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-by", "grp", "-agg", "sum", "-sort_by", "grp"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # Header should appear once: only final sorted aggregated table is printed.
    assert out.count("grp") == 1


def test_parse_agg_name_list_and_mapping():
    from pytae.cli_parsing import parse_agg
    assert parse_agg("mean") == "mean"
    assert parse_agg("mean,sum,n") == ["mean", "sum", "n"]
    assert parse_agg("val = sum, n = n") == {"val": "sum", "n": "n"}
    assert parse_agg("'body mass' = mean") == {"body mass": "mean"}
    assert parse_agg("total = v1:sum") == {"total": "v1:sum"}
    with pytest.raises(SystemExit, match="syntax is retired"):
        parse_agg("column=v1,aggfunc=sum")
    with pytest.raises(SystemExit, match="use '='"):
        parse_agg("val: sum, n: n")
    with pytest.raises(SystemExit, match="use a name"):
        parse_agg("['mean', 'sum']")
    with pytest.raises(SystemExit, match="mix of names"):
        parse_agg("mean, val = sum")


def test_agg_comma_list(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"grp": ["x", "x", "y"], "val": [1, 2, 3]}),
    )

    exit_code = cli.main([path, "-by", "grp", "-agg", "sum,n", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["grp", "n", "val"]


def test_agg_python_literal_is_rejected(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"grp": ["x"], "val": [1]}))
    with pytest.raises(SystemExit, match="use a name"):
        cli.main([path, "-agg", "{'val': 'sum'}"])


def test_agg_mapping_named_a_does_not_collide_with_kwargs(tmp_path, capsys):
    # CLI Issue 2: -agg "a = sum" passed positionally to agg_df
    df = pd.DataFrame({"grp": ["x", "x", "y"], "a": [1, 2, 3], "b": [10, 20, 30]})
    path = _write_csv(tmp_path, df)
    exit_code = cli.main([path, "-by", "grp", "-agg", "a = sum", "-cols"])
    assert exit_code == 0
    cols = capsys.readouterr().out.strip().splitlines()
    assert cols == ["grp", "a"]

################################################################################
# FILE: tests/test_cli_clean.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _messy_headers_frame, _replace_frame, _write_csv


def test_handle_missing_default_fill(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", None],
            "val": [1.0, None],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-handle_missing"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # object/category NA becomes '.', numeric NA becomes 0.
    assert "." in out
    assert "0.0" in out

def test_replace_whole_df_exact_match(tmp_path, capsys):
    path = _write_csv(tmp_path, _replace_frame())

    cli.main([path, "-replace_values", "v='a magician:the magic,alpha:bravo'"])

    out = capsys.readouterr().out
    assert "the magic" in out
    assert "bravo" in out
    assert "not a magician exactly" in out  # untouched: not an exact whole-cell match
    assert "analphabet" in out              # untouched: not an exact whole-cell match

def test_replace_scoped_columns_only(tmp_path, capsys):
    df = pd.DataFrame({"col a": ["alpha"], "colb": ["alpha"]})
    path = _write_csv(tmp_path, df)

    cli.main([path, "-replace_values", "c=col a,v=alpha:bravo"])

    out = capsys.readouterr().out.strip().splitlines()
    assert out[1].split() == ["bravo", "alpha"]

def test_replace_exact_false_matches_substring(tmp_path, capsys):
    path = _write_csv(tmp_path, _replace_frame())

    cli.main([path, "-replace_values", "v='a magician:the magic,alpha:bravo',exact=false"])

    out = capsys.readouterr().out
    assert "not the magic exactly" in out
    assert "anbravobet" in out

def test_replace_unknown_column_errors(tmp_path, capsys):
    path = _write_csv(tmp_path, _replace_frame())

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-replace_values", "c=missing,v=alpha:bravo"])
    assert exc_info.value.code == 2
    assert "-replace_values" in capsys.readouterr().err

def test_replace_requires_v(tmp_path):
    path = _write_csv(tmp_path, _replace_frame())

    with pytest.raises(SystemExit):
        cli.main([path, "-replace_values", "c=col a"])

def test_replace_values_mapping_quoting_is_optional(tmp_path, capsys):
    df = pd.DataFrame({"col a": ["alpha"], "colb": ["alpha"]})

    path = _write_csv(tmp_path, df)
    cli.main([path, "-replace_values", "v=alpha:bravo"])
    unquoted = capsys.readouterr().out.strip()

    path = _write_csv(tmp_path, df)
    cli.main([path, "-replace_values", "v=alpha:'bravo'"])
    quoted = capsys.readouterr().out.strip()

    assert unquoted == quoted
    assert "bravo" in unquoted

def test_replace_values_equals_raises_error(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2]}))
    with pytest.raises(SystemExit, match="use ':'"):
        cli.main([path, "-replace_values", "v=1=2"])

def test_rename_quoting_is_optional(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"old col": [1, 2], "b": [3, 4]}))
    out1 = tmp_path / "out1.csv"
    out2 = tmp_path / "out2.csv"

    cli.main([path, "-rename", "old col:new col", "-o", str(out1)])
    cli.main([path, "-rename", "'old col':'new col'", "-o", str(out2)])

    assert list(pd.read_csv(out1).columns) == ["new col", "b"]
    assert list(pd.read_csv(out2).columns) == ["new col", "b"]

def test_rename_equals_raises_error(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"old col": [1, 2], "b": [3, 4]}))
    with pytest.raises(SystemExit, match="use ':'"):
        cli.main([path, "-rename", "old col=new col"])

def test_rename_multiple_pairs(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"old col": [1, 2], "b": [3, 4]}))
    out1 = tmp_path / "out1.csv"
    out2 = tmp_path / "out2.csv"

    cli.main([path, "-rename", "old col:new col,b:beta", "-o", str(out1)])
    cli.main([path, "-rename", "'old col':'new col','b':'beta'", "-o", str(out2)])

    assert list(pd.read_csv(out1).columns) == ["new col", "beta"]
    assert list(pd.read_csv(out2).columns) == ["new col", "beta"]

def test_rename_pipeline_head(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"island": ["Torgersen"], "species": ["Adelie"]}))
    cli.main([path, "-rename", "island:kailand", "-head"])
    out = capsys.readouterr().out
    assert "kailand" in out
    assert "island" not in out

def test_rename_pipeline_show_all(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"island": ["Torgersen"], "species": ["Adelie"]}))
    cli.main([path, "-rename", "island:kailand"])
    out = capsys.readouterr().out
    assert "kailand" in out
    assert "island" not in out

def test_rename_pipeline_chained_with_select(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3]}))
    cli.main([path, "-rename", "a:alpha", "-select", "alpha,b", "-head"])
    out = capsys.readouterr().out
    assert "alpha" in out
    assert "b" in out
    assert "c" not in out

def test_rename_repeated_flags(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2]}))
    cli.main([path, "-rename", "a:alpha", "-rename", "b:beta", "-head"])
    out = capsys.readouterr().out
    assert "alpha" in out
    assert "beta" in out

def test_rename_unknown_column_raises_error(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2]}))
    with pytest.raises(SystemExit):
        cli.main([path, "-rename", "unknown_col:x", "-head"])

def test_clean_columns_strip_fill_case(tmp_path, capsys):
    path = _write_csv(tmp_path, _messy_headers_frame())

    cli.main([path, "-clean_columns", "strip,fill,case=lower"])

    header = capsys.readouterr().out.strip().splitlines()[0].split()
    assert header == ["col_a", "col___b", "col_a", "100%_match!"]

def test_clean_columns_custom_fill_char(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"col a": [1]}))

    cli.main([path, "-clean_columns", "fill=$"])

    header = capsys.readouterr().out.strip().splitlines()[0].split()
    assert header == ["col$a"]

def test_clean_columns_proper_case(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"col a": [1]}))

    cli.main([path, "-clean_columns", "case=proper"])

    header = capsys.readouterr().out.strip().splitlines()[0]
    assert header == "Col A"

def test_clean_columns_squeeze_and_strip_special_and_dedupe(tmp_path, capsys):
    path = _write_csv(tmp_path, _messy_headers_frame())

    cli.main([path, "-clean_columns", "strip,squeeze,strip_special,fill,case=lower,dedupe=true"])

    header = capsys.readouterr().out.strip().splitlines()[0].split()
    assert header == ["col_a", "col_b", "col_a_1", "100_match"]

def test_clean_columns_no_fill_leaves_whitespace(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"col a": [1]}))

    cli.main([path, "-clean_columns", "case=upper"])

    header = capsys.readouterr().out.strip().splitlines()[0]
    assert header == "COL A"

def test_clean_columns_unknown_key_errors(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"col a": [1]}))

    with pytest.raises(SystemExit, match="unknown key"):
        cli.main([path, "-clean_columns", "bogus=true"])

def test_clean_columns_case_without_value_errors(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"col a": [1]}))

    with pytest.raises(SystemExit):
        cli.main([path, "-clean_columns", "case"])

def test_clean_columns_strip_special_removes_quotes(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"'col a'": [1], '"col b"': [2]}))

    cli.main([path, "-clean_columns", "strip_special", "-cols"])

    header = capsys.readouterr().out.strip().splitlines()
    assert header == ["col a", "col b"]

def test_clean_columns_strip_special_keeps_fill_character(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"co-op's data": [1]}))

    cli.main([path, "-clean_columns", "strip_special,fill=-", "-cols"])

    header = capsys.readouterr().out.strip()
    assert header == "co-ops-data"

################################################################################
# FILE: tests/test_cli_enhancements.py
################################################################################
import io
import sys

import pandas as pd
import pytest

from pytae import cli
from pytae.readers import write_dataframe
from tests.cli_helpers import _write_csv


def test_parquet_mixed_object_coercion(tmp_path):
    df = pd.DataFrame({"Profit_Centre": [100.5, 200.0, "0828"], "val": [1, 2, 3]})
    dest = tmp_path / "mixed.parquet"
    with pytest.warns(UserWarning, match="contained mixed object types; coerced to string"):
        write_dataframe(df, dest)
    res = pd.read_parquet(dest)
    assert list(res["Profit_Centre"]) == ["100.5", "200.0", "0828"]
    assert res["val"].tolist() == [1, 2, 3]


class MockStdin:
    def __init__(self, raw_bytes: bytes) -> None:
        self.buffer = io.BytesIO(raw_bytes)


def test_stdin_csv_head_and_shape(monkeypatch, capsys):
    csv_bytes = b"a,b\n1,x\n2,y\n3,z\n"
    monkeypatch.setattr(sys, "stdin", MockStdin(csv_bytes))

    exit_code = cli.main(["-", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(3, 2)"

    monkeypatch.setattr(sys, "stdin", MockStdin(csv_bytes))
    exit_code = cli.main(["-", "-head", "2"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "1" in out and "2" in out and "3" not in out


def test_stdin_parquet_streaming(monkeypatch, capsys):
    df = pd.DataFrame({"col_x": [10, 20], "col_y": ["alpha", "beta"]})
    buf = io.BytesIO()
    df.to_parquet(buf, index=False)
    parquet_bytes = buf.getvalue()

    monkeypatch.setattr(sys, "stdin", MockStdin(parquet_bytes))
    exit_code = cli.main(["-", "-cols"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "col_x" in out and "col_y" in out


def test_stdin_with_sql_and_export(tmp_path, monkeypatch, capsys):
    csv_bytes = b"id,val\n1,10\n2,20\n3,30\n"
    monkeypatch.setattr(sys, "stdin", MockStdin(csv_bytes))
    dest = tmp_path / "filtered.csv"

    exit_code = cli.main(["-", "-sql", "select * from data where id > 1", "-o", str(dest)])
    assert exit_code == 0
    read_back = pd.read_csv(dest)
    assert len(read_back) == 2
    assert read_back["id"].tolist() == [2, 3]


def test_stdin_jsonl_streaming(monkeypatch, capsys):
    jsonl_bytes = b'{"name": "alice", "score": 95}\n{"name": "bob", "score": 88}\n'
    monkeypatch.setattr(sys, "stdin", MockStdin(jsonl_bytes))

    exit_code = cli.main(["-", "-fmt", "jsonl", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 2)"


def test_freq_categorical_bars(tmp_path, capsys):
    df = pd.DataFrame({"category": ["apple", "apple", "apple", "banana", "banana", "cherry"]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-freq", "category"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "category" in out and "Count" in out and "Distribution" in out
    assert "apple" in out and "3" in out and "(50.0%)" in out
    assert "banana" in out and "2" in out and "(33.3%)" in out
    assert "cherry" in out and "1" in out and "(16.7%)" in out
    assert "█" in out


def test_freq_dropna_default_and_flags(tmp_path, capsys):
    df = pd.DataFrame({"category": ["apple", None, "apple"]})
    path = _write_csv(tmp_path, df)

    # Default is dropna=false: retains NA category
    exit_code = cli.main([path, "-freq", "category"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "apple" in out
    assert "NaN" in out or "<NA>" in out or "None" in out or "nan" in out

    # Explicit -dropna true: drops NA category
    exit_code_drop = cli.main([path, "-freq", "category", "-dropna", "true"])
    assert exit_code_drop == 0
    out_drop = capsys.readouterr().out
    assert "apple" in out_drop
    assert "NaN" not in out_drop and "<NA>" not in out_drop and "None" not in out_drop


def test_hist_numeric_distribution(tmp_path, capsys):
    df = pd.DataFrame({"score": [10, 12, 14, 50, 52, 90]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-hist", "score:3"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "Range" in out and "Count" in out and "Distribution" in out
    assert "█" in out


def test_cli_plot_save_png(tmp_path, capsys):
    df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
    path = _write_csv(tmp_path, df)
    out_img = tmp_path / "chart.png"

    exit_code = cli.main([path, "-plot", "kind=scatter, x=x, y=y", "-o", str(out_img)])
    assert exit_code == 0
    assert out_img.exists()
    assert out_img.stat().st_size > 0
    assert f"Saved plot to {out_img}" in capsys.readouterr().out


def test_cli_plot_with_finalize(tmp_path, capsys):
    df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
    path = _write_csv(tmp_path, df)
    out_svg = tmp_path / "chart.svg"

    exit_code = cli.main([path, "-plot", "kind=line, x=x, y=y", "-finalize", "style=True", "-o", str(out_svg)])
    assert exit_code == 0
    assert out_svg.exists()
    assert out_svg.stat().st_size > 0


def test_library_plot_save_kwarg(tmp_path):
    df = pd.DataFrame({"x": [1, 2], "y": [3, 4]})
    out_file = tmp_path / "lib_plot.png"
    plotter = df.pt.plot(kind="scatter", x="x", y="y", save=str(out_file))
    assert plotter is not None
    assert out_file.exists()
    assert out_file.stat().st_size > 0


def test_cli_plot_forwards_dropna(tmp_path):
    df = pd.DataFrame({"cat": ["A", "B", None], "val": [10, 20, 30]})
    path = _write_csv(tmp_path, df)
    out_png = tmp_path / "chart.png"

    exit_code = cli.main([path, "-plot", "kind=bar, x=cat, y=val", "-dropna", "true", "-o", str(out_png)])
    assert exit_code == 0
    assert out_png.exists()


def test_cli_plot_error_clean_exit(tmp_path):
    # Issue 9: user error in -plot exits with code 2 cleanly instead of unhandled traceback
    df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
    path = _write_csv(tmp_path, df)
    with pytest.raises(SystemExit) as exc:
        cli.main([path, "-plot", "kind=scatter, x=nonexistent, y=y"])
    assert exc.value.code == 2

################################################################################
# FILE: tests/test_cli_filter.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv


def test_query_then_select_filters_like_pandas(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"keep": [1, 2, 3], "flt": ["A", "B", "A"], "val": [10, 20, 30]}),
    )

    exit_code = cli.main(
        [path, "-query", "flt == 'A'", "-select", "keep,val", "-shape"]
    )
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 2)"


def test_qry_without_braces_is_equivalent(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"keep": [1, 2, 3], "flt": ["A", "B", "A"], "val": [10, 20, 30]}),
    )

    exit_code = cli.main([path, "-qry", "'flt' = 'A'", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 3)"


def test_qry_column_key_quoting_is_optional(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"keep": [1, 2, 3], "flt": ["A", "B", "A"], "val": [10, 20, 30]}),
    )

    cli.main([path, "-qry", "flt='A'", "-shape"])
    unquoted = capsys.readouterr().out.strip()

    cli.main([path, "-qry", "'flt'='A'", "-shape"])
    quoted = capsys.readouterr().out.strip()

    assert unquoted == quoted == "(2, 3)"


def test_qry_unquoted_string_value_errors(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"flt": ["A", "B"]}))

    with pytest.raises(SystemExit, match="must be quoted"):
        cli.main([path, "-qry", "flt=A", "-shape"])


def test_qry_column_not_in_select_still_filters(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "keep": [1, 2, 3],
            "flt": ["A", "B", "A"],
            "val": [10, 20, 30],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main(
        [
            path,
            "-qry",
            "flt = 'A'",
            "-select",
            "'keep','val'",
            "-shape",
        ]
    )

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "(2, 2)" in out


def test_qry_clip_copies_filtered_frame_without_output_op(tmp_path, capsys, monkeypatch):
    path = _write_csv(tmp_path, pd.DataFrame({"month": [202606, 202607, 202607], "value": [1, 2, 3]}))

    copied = {}

    def _fake_to_clipboard(self, *args, **kwargs):
        copied["frame"] = self.copy()

    monkeypatch.setattr(pd.DataFrame, "to_clipboard", _fake_to_clipboard)

    exit_code = cli.main([path, "-qry", "month = 202607", "-o", "clip"])

    captured = capsys.readouterr()
    assert exit_code == 0
    assert captured.out == ""
    pd.testing.assert_frame_equal(
        copied["frame"].reset_index(drop=True),
        pd.DataFrame({"month": [202607, 202607], "value": [2, 3]}),
    )


def test_qry_direct_comparison(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000, 4000, 5000]}))

    exit_code = cli.main([path, "-qry", "body_mass_g > 3500", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 1)"


def test_qry_colon_raises_error(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000, 4000, 5000]}))

    with pytest.raises(SystemExit, match="use '='"):
        cli.main([path, "-qry", "body_mass_g: > 3500", "-shape"])


def test_qry_equals_separator(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000, 4000, 5000]}))

    # col= > val
    exit_code = cli.main([path, "-qry", "body_mass_g= > 3500", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 1)"

    # col=>val
    exit_code = cli.main([path, "-qry", "body_mass_g=>3500", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 1)"

    # col = > val
    exit_code = cli.main([path, "-qry", "body_mass_g = > 3500", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 1)"

    # col = val
    exit_code = cli.main([path, "-qry", "body_mass_g = 4000", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(1, 1)"


def test_qry_cli_bracket_interval(tmp_path, capsys):
    # CLI Issue 3: -qry "body_mass_g = [3000, 3200]" treats as interval
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000, 3100, 3200, 4000]}))
    exit_code = cli.main([path, "-qry", "body_mass_g = [3000, 3200]", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(3, 1)"

################################################################################
# FILE: tests/test_cli_inspect.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv


def test_order_head_then_shape(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(10), "b": range(10, 20)}))

    exit_code = cli.main([path, "-head", "3", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip() == "(3, 2)"

def test_sample_with_seed_is_reproducible(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(100), "b": range(100, 200)}))

    cli.main([path, "-sample", "10", "-seed", "42"])
    first = capsys.readouterr().out

    cli.main([path, "-sample", "10", "-seed", "42"])
    second = capsys.readouterr().out

    assert first == second

def test_sample_frac_selects_expected_row_count(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(100), "b": range(100, 200)}))

    exit_code = cli.main([path, "-sample", "-frac", "0.1", "-seed", "1", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip() == "(10, 2)"

def test_frac_requires_sample(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(10), "b": range(10, 20)}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-frac", "0.5", "-head"])
    assert exc_info.value.code == 2

def test_shape_cannot_be_followed_by_another_flag(tmp_path):
    # -shape returns a tuple in pandas terms (not a DataFrame), so nothing may
    # chain after it except -o clip.
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(10), "b": range(10, 20)}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-shape", "-head", "3"])
    assert exc_info.value.code == 2

@pytest.mark.parametrize("op", ["cols", "dtype", "nulls", "info"])
def test_non_df_ops_cannot_be_followed_by_another_flag(tmp_path, op):
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(3), "b": range(3, 6)}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, f"-{op}", "-head", "1"])
    assert exc_info.value.code == 2

def test_head_then_cols_prints_only_names(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(8), "b": range(8)}))

    exit_code = cli.main([path, "-head", "5", "-cols"])

    captured = capsys.readouterr()
    assert exit_code == 0, captured.err
    assert captured.out.strip().splitlines() == ["a", "b"]

def test_empty_parquet_head_keeps_columns(tmp_path, capsys):
    path = tmp_path / "empty.parquet"
    pd.DataFrame({"a": pd.Series(dtype="int64"), "b": pd.Series(dtype="object")}).to_parquet(path, index=False)

    exit_code = cli.main([str(path), "-head", "5"])

    captured = capsys.readouterr()
    assert exit_code == 0, captured.err
    assert "a" in captured.out and "b" in captured.out
    assert "Empty DataFrame" in captured.out

def test_info_prints_pandas_info(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}))

    exit_code = cli.main([path, "-info"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "DataFrame" in out
    assert "a" in out and "b" in out
    assert "2 entries" in out or "2 rows" in out or "RangeIndex: 2" in out

def test_head_then_info_prints_only_info(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": range(10), "b": range(10)}))

    exit_code = cli.main([path, "-head", "3", "-info"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "RangeIndex: 3" in out
    assert "DataFrame" in out

def test_describe_prints_pandas_summary(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2, 3]}))

    exit_code = cli.main([path, "-describe"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "count" in out and "mean" in out

def test_stats_flag_removed(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-stats"])
    assert exc_info.value.code == 2

def test_nonpositive_head_is_usage_error(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2, 3]}))

    for n in ("0", "-5"):
        with pytest.raises(SystemExit) as exc_info:
            cli.main([path, "-head", n])
        assert exc_info.value.code == 2

def test_cli_long_short_aliases(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"grp": ["a"], "n": [1], "x": [10]}))

    exit_code = cli.main([path, "-long", "c=feature,v=amount", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["grp", "feature", "amount"]

def test_cli_wide_short_aliases(tmp_path, capsys):
    df = pd.DataFrame({"id": ["a", "b"], "country": ["sg", "cn"], "balance": [10, 20]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-wide", "c=country,v=balance", "-cols"])

    captured = capsys.readouterr()
    assert exit_code == 0, captured.err
    assert captured.out.strip().splitlines() == ["id", "cn", "sg"]

def test_order_unique_then_shape_uses_deduplicated_frame(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", "x", "y"],
            "val": [1, 1, 2],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-unique", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # One duplicated row is removed by -unique, then -shape reflects the reduced frame.
    assert "(2, 2)" in out

def test_order_value_counts_then_shape_uses_count_table(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", "x", "y"],
            "val": [1, 2, 3],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "grp", "-value_counts", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # value_counts result has two columns: key + count.
    assert "(2, 2)" in out

def test_value_counts_dropna_default_and_flags(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", None, "x"],
        }
    )
    path = _write_csv(tmp_path, df)

    # Default is dropna=false: includes NA as its own key, yielding x + NA (2 rows).
    exit_code = cli.main([path, "-select", "grp", "-value_counts", "-shape"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert "(2, 2)" in out

    # Explicit -dropna true drops the NA key (1 row).
    exit_code_drop = cli.main([path, "-select", "grp", "-value_counts", "-dropna", "true", "-shape"])
    out_drop = capsys.readouterr().out
    assert exit_code_drop == 0
    assert "(1, 2)" in out_drop

def test_value_counts_uses_selected_multiple_columns(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "num_legs": [2, 4, 4, 6],
            "num_wings": [2, 0, 0, 0],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "num_legs,num_wings", "-value_counts", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # Two key columns + count, three unique key combinations.
    assert "(3, 3)" in out

def test_value_counts_without_argument_uses_selected_columns(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "num_legs": [2, 4, 4, 6],
            "num_wings": [2, 0, 0, 0],
            "animal": ["falcon", "dog", "cat", "ant"],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "num_legs,num_wings", "-value_counts", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # -select narrows working columns, so no-arg -value_counts uses those two.
    assert "(3, 3)" in out


def test_value_counts_collision_fails_cleanly(tmp_path, capsys):
    df = pd.DataFrame({"count": [1, 2], "n": [3, 4]})
    path = _write_csv(tmp_path, df)
    with pytest.raises(SystemExit) as exc:
        cli.main([path, "-value_counts"])
    assert exc.value.code == 2
    err = capsys.readouterr().err
    assert "columns contain both 'count' and 'n'" in err

def test_order_sort_by_then_shape_uses_sorted_frame(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["b", "a", "c"],
            "val": [2, 1, 3],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-sort_by", "grp", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "(3, 2)" in out

def test_sort_by_respects_select_and_descending(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["a", "b", "c"],
            "val": [1, 3, 2],
            "other": [10, 20, 30],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "grp,val", "-sort_by", "val desc", "-head", "1"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # After sorting by val desc, the first row should be grp=b, val=3.
    assert "b" in out
    assert "3" in out

def test_value_counts_then_sort_by_prints_only_final_table(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "grp": ["x", "x", "y"],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "grp", "-value_counts", "-sort_by", "count"])

    out = capsys.readouterr().out
    assert exit_code == 0
    # Header should appear once: only final sorted value_counts table is printed.
    assert out.count("grp") == 1

def test_cols_default_is_file_order(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"c": [1], "a": [2], "b": [3]}))

    exit_code = cli.main([path, "-cols"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip().splitlines() == ["c", "a", "b"]

def test_cols_asc_and_desc(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"c": [1], "a": [2], "b": [3]}))

    assert cli.main([path, "-cols", "asc"]) == 0
    assert capsys.readouterr().out.strip().splitlines() == ["a", "b", "c"]

    assert cli.main([path, "-cols", "desc"]) == 0
    assert capsys.readouterr().out.strip().splitlines() == ["c", "b", "a"]

def test_dtype_optional_name_order(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"c": [1], "a": ["x"]}))

    assert cli.main([path, "-dtype"]) == 0
    file_order = capsys.readouterr().out
    assert file_order.splitlines()[0].startswith("c")

    assert cli.main([path, "-dtype", "asc"]) == 0
    assert capsys.readouterr().out.splitlines()[0].startswith("a")

def test_sort_by_default_is_ascending(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"grp": ["b", "a", "c"], "val": [2, 1, 3]}))

    exit_code = cli.main([path, "-sort_by", "val", "-head", "1"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "a" in out
    assert "1" in out

def test_sort_by_explicit_asc(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"grp": ["b", "a"], "val": [2, 1]}))

    exit_code = cli.main([path, "-sort_by", "val asc", "-head", "1"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "a" in out
    assert "1" in out

def test_sort_flag_removed(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-sort", "desc"])
    assert exc_info.value.code == 2

def test_cols_rejects_invalid_order(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-cols", "foo"])
    assert exc_info.value.code == 2

def test_repeated_select_negative_subtracts_sequentially(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3], "d": [4]}))

    exit_code = cli.main([path, "-select", "-b", "-select", "-d", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["a", "c"]

def test_unquoted_spec_with_a_space_gets_a_quoting_hint(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"n": [1, 2]}))

    # simulates the shell splitting an unquoted "-select species,col name" into
    # separate argv tokens ("col" and "name" land as unrecognized leftovers)
    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-select", "n,col", "name"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "unrecognized arguments: col name" in err or "unrecognized arguments: name" in err
    assert "wrap the whole spec in quotes" in err

def test_typo_flag_does_not_get_the_quoting_hint(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"n": [1, 2]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-selecttypo", "n"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "unrecognized arguments" in err
    assert "wrap the whole spec in quotes" not in err

def test_describe_then_shape_is_describe_table(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))

    exit_code = cli.main([path, "-describe", "-shape"])

    assert exit_code == 0
    # pandas describe() of two numeric columns is 8 stats × 2
    assert capsys.readouterr().out.strip() == "(8, 2)"

def test_clip_suppresses_stdout_for_dataframe_ops(tmp_path, capsys, monkeypatch):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2], "b": [3, 4]}))

    copied = {"called": False}

    def _fake_to_clipboard(self, *args, **kwargs):
        copied["called"] = True

    monkeypatch.setattr(pd.DataFrame, "to_clipboard", _fake_to_clipboard)

    exit_code = cli.main([path, "-head", "2", "-o", "clip"])

    captured = capsys.readouterr()
    assert exit_code == 0
    assert captured.out == ""
    assert copied["called"] is True

def test_clip_shape_alone_succeeds(tmp_path, capsys, monkeypatch):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2], "b": [3, 4]}))
    copied = {}
    monkeypatch.setattr("pytae.cli_run._copy_to_clipboard", lambda s: copied.setdefault("text", s))

    exit_code = cli.main([path, "-shape", "-o", "clip"])

    captured = capsys.readouterr()
    assert exit_code == 0
    assert captured.out == ""
    assert copied["text"] == "(2, 2)"

def test_clip_shape_and_df_flag_errors(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2], "b": [3, 4]}))
    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-head", "1", "-shape", "-o", "clip"])
    assert exc_info.value.code == 2

def test_cli_convert_csv_to_parquet(tmp_path, capsys):
    src = tmp_path / "data.csv"
    pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}).to_csv(src, index=False)
    dest = tmp_path / "data.parquet"

    exit_code = cli.main([str(src), "-o", str(dest)])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert dest.exists()
    assert "Wrote 2 rows" in out
    result = pd.read_parquet(dest)
    assert list(result.columns) == ["a", "b"]
    assert len(result) == 2

def test_cli_output_bare_format(tmp_path, capsys):
    src = tmp_path / "input.csv"
    pd.DataFrame({"x": [10, 20]}).to_csv(src, index=False)

    exit_code = cli.main([str(src), "-o", "parquet"])
    assert exit_code == 0
    dest = tmp_path / "input.parquet"
    assert dest.exists()
    assert len(pd.read_parquet(dest)) == 2

def test_cli_output_shape_to_file_errors(tmp_path):
    src = tmp_path / "data.csv"
    pd.DataFrame({"a": [1]}).to_csv(src, index=False)
    dest = tmp_path / "out.parquet"

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(src), "-shape", "-o", str(dest)])
    assert exc_info.value.code == 2

def test_cli_removed_flags_migration_errors(tmp_path):
    src = tmp_path / "data.csv"
    pd.DataFrame({"a": [1]}).to_csv(src, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(src), "-to_clip"])
    assert exc_info.value.code == 2

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(src), "-convert"])
    assert exc_info.value.code == 2


def test_cli_multiple_positional_files_batch_convert(tmp_path, capsys):
    f1 = tmp_path / "data1.parquet"
    f2 = tmp_path / "data2.parquet"
    pd.DataFrame({"x": [1, 2]}).to_parquet(f1, index=False)
    pd.DataFrame({"y": [3, 4]}).to_parquet(f2, index=False)

    exit_code = cli.main([str(f1), str(f2), "-o", "csv"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert (tmp_path / "data1.csv").exists()
    assert (tmp_path / "data2.csv").exists()
    assert f"== {f1} ==" in out
    assert f"== {f2} ==" in out
    assert "Wrote 2 rows" in out


def test_cli_multiple_positional_files_inspect(tmp_path, capsys):
    f1 = tmp_path / "a.csv"
    f2 = tmp_path / "b.csv"
    pd.DataFrame({"col": [10, 20]}).to_csv(f1, index=False)
    pd.DataFrame({"col": [30, 40]}).to_csv(f2, index=False)

    exit_code = cli.main([str(f1), str(f2), "-head", "1"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert f"== {f1} ==" in out
    assert f"== {f2} ==" in out
    assert "10" in out
    assert "30" in out


def test_cli_multiple_files_single_destination_errors(tmp_path, capsys):
    f1 = tmp_path / "a.csv"
    f2 = tmp_path / "b.csv"
    pd.DataFrame({"col": [1]}).to_csv(f1, index=False)
    pd.DataFrame({"col": [2]}).to_csv(f2, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(f1), str(f2), "-o", str(tmp_path / "out.parquet")])
    assert exc_info.value.code == 2
    assert "requires a format" in capsys.readouterr().err


def test_cli_multiple_files_clip_errors(tmp_path, capsys):
    f1 = tmp_path / "a.csv"
    f2 = tmp_path / "b.csv"
    pd.DataFrame({"col": [1]}).to_csv(f1, index=False)
    pd.DataFrame({"col": [2]}).to_csv(f2, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(f1), str(f2), "-o", "clip"])
    assert exc_info.value.code == 2
    assert "-o clip cannot be used with multiple matched files" in capsys.readouterr().err


def test_cli_no_path_errors(capsys):
    with pytest.raises(SystemExit) as exc_info:
        cli.main([])
    assert exc_info.value.code == 2
    assert "the following arguments are required: path" in capsys.readouterr().err


def test_cli_positional_with_file_errors(tmp_path, capsys):
    f1 = tmp_path / "a.csv"
    f2 = tmp_path / "b.csv"
    pd.DataFrame({"id": [1]}).to_csv(f1, index=False)
    pd.DataFrame({"id": [2]}).to_csv(f2, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{f1}=a", str(f2), "-merge", "how=inner"])
    assert exc_info.value.code == 2
    assert "can't be combined with a positional path" in capsys.readouterr().err


def test_cli_progress_custom_chunk_size(tmp_path, capsys):
    src = tmp_path / "data.parquet"
    pd.DataFrame({"a": list(range(25))}).to_parquet(src, index=False)
    dest = tmp_path / "data.csv"

    exit_code = cli.main([str(src), "-o", str(dest), "-progress", "10"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert dest.exists()
    assert "reading... 10/25 rows (40%)" in out
    assert "writing... 10/25 rows (40%)" in out
    assert "Wrote 25 rows" in out


def test_cli_progress_bare_flag(tmp_path, capsys):
    src = tmp_path / "data.parquet"
    pd.DataFrame({"a": [1, 2]}).to_parquet(src, index=False)
    dest = tmp_path / "data.csv"

    exit_code = cli.main([str(src), "-o", str(dest), "-progress"])
    out = capsys.readouterr().out
    assert exit_code == 0
    assert dest.exists()
    assert "reading... 2/2 rows (100%)" in out
    assert "writing... 2/2 rows (100%)" in out
    assert "Wrote 2 rows" in out


def test_cli_invalid_progress_chunk_size_errors(tmp_path):
    src = tmp_path / "data.parquet"
    pd.DataFrame({"a": [1]}).to_parquet(src, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(src), "-progress", "0"])
    assert exc_info.value.code == 2


def test_describe_export_preserves_index(tmp_path):
    src = tmp_path / "data.parquet"
    df = pd.DataFrame({"sp": ["A", "B", "A"], "val": [10, 20, 30]})
    df.to_parquet(src, index=False)

    desc_csv = tmp_path / "desc.csv"
    assert cli.main([str(src), "-describe", "-o", str(desc_csv)]) == 0
    read_desc_csv = pd.read_csv(desc_csv, index_col=0)
    assert list(read_desc_csv.index) == ["count", "mean", "std", "min", "25%", "50%", "75%", "max"]

    desc_pq = tmp_path / "desc.parquet"
    assert cli.main([str(src), "-describe", "-o", str(desc_pq)]) == 0
    read_desc_pq = pd.read_parquet(desc_pq)
    assert list(read_desc_pq.index) == ["count", "mean", "std", "min", "25%", "50%", "75%", "max"]

    reg_csv = tmp_path / "reg.csv"
    assert cli.main([str(src), "-head", "2", "-o", str(reg_csv)]) == 0
    read_reg = pd.read_csv(reg_csv)
    assert list(read_reg.columns) == ["sp", "val"]


def test_batch_export_colliding_stems_error(tmp_path, capsys):
    f1 = tmp_path / "stem.parquet"
    f2 = tmp_path / "stem.pq"
    pd.DataFrame({"x": [1]}).to_parquet(f1, index=False)
    pd.DataFrame({"y": [2]}).to_parquet(f2, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(f1), str(f2), "-o", "csv"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "multiple input files resolve to the same destination path" in err


def test_cli_progress_preceding_path(tmp_path):
    src = tmp_path / "data.parquet"
    pd.DataFrame({"a": [1, 2]}).to_parquet(src, index=False)
    dest1 = tmp_path / "out1.csv"
    dest2 = tmp_path / "out2.csv"
    dest3 = tmp_path / "out3.csv"

    assert cli.main(["-progress", str(src), "-o", str(dest1)]) == 0
    assert dest1.exists()

    assert cli.main(["-o", str(dest2), "-progress", str(src)]) == 0
    assert dest2.exists()

    assert cli.main(["-progress", "1000", str(src), "-o", str(dest3)]) == 0
    assert dest3.exists()


def test_cli_intermixed_positional_args(tmp_path):
    f1 = tmp_path / "a.parquet"
    f2 = tmp_path / "b.parquet"
    pd.DataFrame({"a": [1]}).to_parquet(f1, index=False)
    pd.DataFrame({"b": [2]}).to_parquet(f2, index=False)

    exit_code = cli.main([str(f1), "-o", "csv", str(f2)])
    assert exit_code == 0
    assert (tmp_path / "a.csv").exists()
    assert (tmp_path / "b.csv").exists()


def test_in_place_overwrite_error_message(tmp_path, capsys):
    src = tmp_path / "file.csv"
    src.write_text("a,b\n1,2\n")

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(src), "-o", "csv"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "refusing to overwrite the source file" in err
    assert "specify a different format or an explicit path with -o/--output" in err


def test_non_df_terminal_op_error_message(tmp_path, capsys):
    src = tmp_path / "data.parquet"
    pd.DataFrame({"a": [1]}).to_parquet(src, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(src), "-shape", "-head", "5"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "-shape does not return a DataFrame/Series, so no flag may follow it (except '-o clip'" in err


def test_cli_out_dir_single_and_batch(tmp_path):
    df1 = pd.DataFrame({"col1": [1, 2], "col2": ["x", "y"]})
    df2 = pd.DataFrame({"col1": [3, 4], "col2": ["z", "w"]})
    f1 = tmp_path / "file1.csv"
    f2 = tmp_path / "file2.csv"
    df1.to_csv(f1, index=False)
    df2.to_csv(f2, index=False)

    out_directory = tmp_path / "exported"
    # Single file
    exit_code = cli.main([str(f1), "-o", "parquet", "-out_dir", str(out_directory)])
    assert exit_code == 0
    assert (out_directory / "file1.parquet").exists()

    # Batch files with -od alias
    exit_code2 = cli.main([str(f1), str(f2), "-o", "parquet", "-od", str(out_directory)])
    assert exit_code2 == 0
    assert (out_directory / "file1.parquet").exists()
    assert (out_directory / "file2.parquet").exists()


def test_cli_out_dir_validations(tmp_path, capsys):
    f = tmp_path / "data.csv"
    f.write_text("a,b\n1,2\n")
    out_directory = tmp_path / "exported"

    # Fails without -o
    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(f), "-out_dir", str(out_directory)])
    assert exc_info.value.code == 2
    assert "-out_dir requires -o/--output" in capsys.readouterr().err

    # Fails with -o clip
    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(f), "-o", "clip", "-out_dir", str(out_directory)])
    assert exc_info.value.code == 2
    assert "-out_dir cannot be used with '-o clip'" in capsys.readouterr().err


def test_cli_batch_out_dir_collision(tmp_path, capsys):
    sub1 = tmp_path / "sub1"
    sub2 = tmp_path / "sub2"
    sub1.mkdir()
    sub2.mkdir()
    f1 = sub1 / "data.csv"
    f2 = sub2 / "data.csv"
    f1.write_text("a,b\n1,2\n")
    f2.write_text("a,b\n3,4\n")

    out_directory = tmp_path / "exported"
    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(f1), str(f2), "-o", "parquet", "-out_dir", str(out_directory)])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "multiple input files resolve to the same destination path" in err


def test_cli_meta_parquet(tmp_path, capsys):
    p = tmp_path / "sample.parquet"
    pd.DataFrame({"id": range(10), "score": [float(i) * 1.5 for i in range(10)]}).to_parquet(p, index=False)

    exit_code = cli.main([str(p), "-meta"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "Format: Parquet" in out
    assert "Rows: 10" in out
    assert "Columns: 2" in out
    assert "Row groups: 1" in out
    assert "Schema:" in out
    assert "id" in out
    assert "score" in out


def test_cli_meta_restrictions(tmp_path, capsys):
    p = tmp_path / "sample.parquet"
    pd.DataFrame({"a": [1]}).to_parquet(p, index=False)

    # -meta cannot be chained before another flag
    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(p), "-meta", "-head", "5"])
    assert exc_info.value.code == 2
    assert "-meta does not return a DataFrame/Series, so no flag may follow it" in capsys.readouterr().err

    # -meta cannot export to file
    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(p), "-meta", "-o", "out.parquet"])
    assert exc_info.value.code == 2
    assert "-meta does not produce a tabular DataFrame, so it cannot be exported to a file" in capsys.readouterr().err


def test_cli_meta_clip(tmp_path, monkeypatch):
    p = tmp_path / "sample.parquet"
    pd.DataFrame({"a": [1, 2]}).to_parquet(p, index=False)

    copied = []
    monkeypatch.setattr(cli, "_process_path", cli._process_path)
    import pytae.cli_run as cli_run
    monkeypatch.setattr(cli_run, "_copy_to_clipboard", lambda text: copied.append(text))

    exit_code = cli.main([str(p), "-meta", "-o", "clip"])
    assert exit_code == 0
    assert len(copied) == 1
    assert "Format: Parquet" in copied[0]


def test_cli_diff_identical(tmp_path, capsys):
    df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
    f1 = tmp_path / "f1.csv"
    f2 = tmp_path / "f2.csv"
    df.to_csv(f1, index=False)
    df.to_csv(f2, index=False)

    exit_code = cli.main([str(f1), "-diff", str(f2)])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "Comparing:" in out
    assert "Rows: 3 vs 3" in out
    assert "Cols: 2 vs 2" in out
    assert "Identical: all cell values match exactly." in out


def test_cli_diff_mismatches_and_drift(tmp_path, capsys):
    df1 = pd.DataFrame({"a": [1, 2, 3], "b": ["x", None, "z"], "c": [10, 20, 30]})
    df2 = pd.DataFrame({"a": [1, 2, 4], "b": ["x", "y", "z"], "d": [1.0, 2.0, 3.0]})
    f1 = tmp_path / "left.csv"
    f2 = tmp_path / "right.csv"
    df1.to_csv(f1, index=False)
    df2.to_csv(f2, index=False)

    exit_code = cli.main([str(f1), "-diff", str(f2)])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "+ Added in left (1):   c" in out
    assert "- Removed in left (1): d" in out
    assert "Null Counts:" in out
    assert "Mismatches found" in out


def test_cli_diff_clip(tmp_path, monkeypatch):
    df = pd.DataFrame({"a": [1, 2]})
    f1 = tmp_path / "a.csv"
    f2 = tmp_path / "b.csv"
    df.to_csv(f1, index=False)
    df.to_csv(f2, index=False)

    copied = []
    import pytae.cli_run as cli_run
    monkeypatch.setattr(cli_run, "_copy_to_clipboard", lambda text: copied.append(text))

    exit_code = cli.main([str(f1), "-diff", str(f2), "-o", "clip"])
    assert exit_code == 0
    assert len(copied) == 1
    assert "Comparing:" in copied[0]


def test_cli_diff_restrictions(tmp_path, capsys):
    f = tmp_path / "a.csv"
    f.write_text("a\n1\n")

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(f), "-diff", str(f), "-o", "out.csv"])
    assert exc_info.value.code == 2
    assert "-diff does not produce a tabular DataFrame, so it cannot be exported to a file" in capsys.readouterr().err


def test_cli_jsonl_inspect_and_conversion(tmp_path, capsys):
    df = pd.DataFrame({"id": [1, 2, 3], "city": ["NYC", "SFO", "LON"]})
    jpath = tmp_path / "data.jsonl"
    df.to_json(jpath, orient="records", lines=True)

    # Inspect
    exit_code = cli.main([str(jpath), "-shape"])
    assert exit_code == 0
    assert "(3, 2)" in capsys.readouterr().out

    exit_code = cli.main([str(jpath), "-cols"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "id" in out
    assert "city" in out

    # Convert to Parquet
    exit_code2 = cli.main([str(jpath), "-o", "parquet"])
    assert exit_code2 == 0
    pq_path = tmp_path / "data.parquet"
    assert pq_path.exists()
    assert len(pd.read_parquet(pq_path)) == 3

    # Convert Parquet to ndjson
    exit_code3 = cli.main([str(pq_path), "-o", "ndjson"])
    assert exit_code3 == 0
    assert (tmp_path / "data.ndjson").exists()


def test_cli_transparent_compression(tmp_path):
    df = pd.DataFrame({"a": [1, 2, 3, 4], "b": ["p", "q", "r", "s"]})
    csv_path = tmp_path / "data.csv"
    df.to_csv(csv_path, index=False)

    # Convert csv to csv.gz
    exit_code = cli.main([str(csv_path), "-o", "csv.gz"])
    assert exit_code == 0
    gz_path = tmp_path / "data.csv.gz"
    assert gz_path.exists()

    # Convert csv.gz to parquet
    exit_code2 = cli.main([str(gz_path), "-o", "parquet"])
    assert exit_code2 == 0
    pq_path = tmp_path / "data.parquet"
    assert pq_path.exists()
    pd.testing.assert_frame_equal(pd.read_parquet(pq_path), df)

    # Convert parquet to jsonl.gz
    exit_code3 = cli.main([str(pq_path), "-o", "jsonl.gz"])
    assert exit_code3 == 0
    j_gz = tmp_path / "data.jsonl.gz"
    assert j_gz.exists()
    pd.testing.assert_frame_equal(pd.read_json(j_gz, lines=True), df)


def test_cli_pager(tmp_path, monkeypatch):
    df = pd.DataFrame({"col": range(5)})
    p = tmp_path / "data.csv"
    df.to_csv(p, index=False)

    paged = []
    import pydoc
    monkeypatch.setattr(pydoc, "pager", lambda text: paged.append(text))

    exit_code = cli.main([str(p), "-head", "3", "-pager"])
    assert exit_code == 0
    assert len(paged) == 1
    assert "col" in paged[0]


def test_cli_round_preserves_pipeline_precision(tmp_path, capsys):
    # CLI Issue 4: intermediate sort does not corrupt full precision for later agg
    df = pd.DataFrame({"g": ["x", "x"], "v": [0.4, 0.4]})
    path = _write_csv(tmp_path, df)
    exit_code = cli.main([path, "-round", "0", "-sort_by", "v", "-by", "g", "-agg", "sum"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "1" in out


def test_cli_repeated_flags_run_with_respective_values(tmp_path, capsys):
    # CLI Issue 7: -head 1 -head 3 runs head 1 then head 3 (result has 1 row)
    df = pd.DataFrame({"a": range(10)})
    path = _write_csv(tmp_path, df)
    exit_code = cli.main([path, "-head", "1", "-head", "3", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(1, 1)"


def test_cli_expand_paths_literal_bracket_file(tmp_path, capsys):
    # CLI Issue 9: data[1].csv read literally when file exists
    p = tmp_path / "data[1].csv"
    p.write_text("a,b\n1,2\n")
    exit_code = cli.main([str(p), "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(1, 2)"

################################################################################
# FILE: tests/test_cli_merge.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv, _write_three_id_csvs, _write_two_csvs


def test_merge_on_differing_column_names(tmp_path, capsys):
    left, right = _write_two_csvs(tmp_path)

    cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-merge", "left=df1,right=df2,on=col a:cola",
    ])

    out = capsys.readouterr().out.strip()
    assert "val_l" in out and "val_r" in out
    assert len(out.splitlines()) == 3  # header + 2 matching rows (inner join)

def test_merge_on_pair_quoting_is_optional(tmp_path, capsys):
    left, right = _write_two_csvs(tmp_path)

    cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-merge", "left=df1,right=df2,on=col a:'cola'",
    ])

    out = capsys.readouterr().out.strip()
    assert "val_l" in out and "val_r" in out
    assert len(out.splitlines()) == 3

def test_merge_shared_column_name_outer_join(tmp_path, capsys):
    left = tmp_path / "a.csv"
    right = tmp_path / "b.csv"
    pd.DataFrame({"id": [1, 2, 3], "x": ["a", "b", "c"]}).to_csv(left, index=False)
    pd.DataFrame({"id": [1, 2, 4], "y": ["p", "q", "r"]}).to_csv(right, index=False)

    cli.main([
        "-file", f"{left}=a;{right}=b",
        "-merge", "left=a,right=b,on=id,how=outer", "-shape",
    ])

    assert capsys.readouterr().out.strip() == "(4, 3)"

def test_merge_requires_file(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-merge", "left=df1,right=df2,on=id"])
    assert exc_info.value.code == 2

def test_file_requires_merge(tmp_path):
    left, right = _write_two_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{left}=df1;{right}=df2", "-shape"])
    assert exc_info.value.code == 2

def test_file_cannot_combine_with_positional_path(tmp_path):
    left, right = _write_two_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([left, "-file", f"{left}=df1;{right}=df2", "-merge", "left=df1,right=df2,on=col a:cola"])
    assert exc_info.value.code == 2

def test_merge_unknown_alias_errors(tmp_path):
    left, right = _write_two_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{left}=df1;{right}=df2", "-merge", "left=df1,right=bogus,on=col a:cola"])
    assert exc_info.value.code == 2

def test_merge_output_format_requires_concrete_path(tmp_path):
    left, right = _write_two_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{left}=df1;{right}=df2", "-merge", "left=df1,right=df2,on=col a:cola", "-o", "csv"])
    assert exc_info.value.code == 2

def test_merge_must_be_first_op_with_file(tmp_path):
    left, right = _write_two_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([
            "-file", f"{left}=df1;{right}=df2",
            "-shape",
            "-merge", "left=df1,right=df2,on=col a:cola",
        ])
    assert exc_info.value.code == 2

def test_merge_chained_via_data_alias(tmp_path, capsys):
    a, b, c = _write_three_id_csvs(tmp_path)

    cli.main([
        "-file", f"{a}=a;{b}=b;{c}=c",
        "-merge", "left=a,right=b,on=id",
        "-merge", "left=data,right=c,on=id",
    ])

    out = capsys.readouterr().out
    assert all(col in out for col in ("x", "y", "z"))

def test_merge_data_not_available_before_anything_produced(tmp_path, capsys):
    a, b, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{a}=a;{b}=b", "-merge", "left=data,right=b,on=id"])
    assert exc_info.value.code == 2
    assert "'data' isn't available yet" in capsys.readouterr().err

def test_concat_stacks_three_frames_with_reset_index(tmp_path, capsys):
    a, b, c = _write_three_id_csvs(tmp_path)

    cli.main(["-file", f"{a}=a;{b}=b;{c}=c", "-concat", "frames='a,b,c'"])

    lines = capsys.readouterr().out.strip().splitlines()
    assert len(lines) == 1 + 3 + 3 + 2  # header + 3 rows from each of a, b, c

def test_concat_resets_index(tmp_path, capsys, monkeypatch):
    a, b, _ = _write_three_id_csvs(tmp_path)
    copied = {}

    def _fake_to_clipboard(self, *args, **kwargs):
        copied["frame"] = self.copy()

    monkeypatch.setattr(pd.DataFrame, "to_clipboard", _fake_to_clipboard)

    cli.main(["-file", f"{a}=a;{b}=b", "-concat", "frames='a,b'", "-o", "clip"])

    result = copied["frame"]
    assert list(result.index) == list(range(len(result)))

def test_concat_chained_via_data_alias(tmp_path, capsys):
    a, b, c = _write_three_id_csvs(tmp_path)

    cli.main([
        "-file", f"{a}=a;{b}=b;{c}=c",
        "-concat", "frames='a,b'",
        "-concat", "frames='data,c'",
        "-shape",
    ])

    assert capsys.readouterr().out.strip() == "(8, 4)"

def test_concat_requires_file(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-concat", "frames='a,b'"])
    assert exc_info.value.code == 2

def test_concat_unknown_alias_errors(tmp_path, capsys):
    a, b, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{a}=a;{b}=b", "-concat", "frames='a,bogus'"])
    assert exc_info.value.code == 2
    assert "unknown -file alias" in capsys.readouterr().err

def test_file_dlim_and_encoding_overrides_applied(tmp_path, capsys):
    pipe_path = tmp_path / "pipe.csv"
    pipe_path.write_bytes("id|name\n1|caf\xe9\n".encode("latin-1"))
    plain_path = tmp_path / "plain.csv"
    pd.DataFrame({"id": [2], "name": ["bravo"]}).to_csv(plain_path, index=False)

    cli.main([
        "-file", f"{pipe_path}=p,dlim=|,encoding=latin-1; {plain_path}=q",
        "-concat", "frames='p,q'",
    ])

    out = capsys.readouterr().out
    assert "café" in out
    assert "bravo" in out

def test_file_requires_at_least_two_entries(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1]}))

    with pytest.raises(SystemExit, match="at least two"):
        cli.main(["-file", f"{path}=a", "-concat", "frames='a,a'"])

def test_concat_requires_at_least_two_frame_names(tmp_path):
    a, b, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit, match="frames= needs at least two names"):
        cli.main(["-file", f"{a}=a;{b}=b", "-concat", "frames=a"])

def test_file_duplicate_alias_errors(tmp_path):
    a, b, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit, match="duplicate alias"):
        cli.main(["-file", f"{a}=x;{b}=x", "-concat", "frames='x,x'"])

def test_merge_on_mixed_style_errors(tmp_path):
    a, b, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit, match="mixes plain columns and left:right pairs"):
        cli.main(["-file", f"{a}=a;{b}=b", "-merge", "left=a,right=b,on='id:id,x'"])

def test_merge_missing_required_keys_errors(tmp_path):
    a, b, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit, match="missing required key"):
        cli.main(["-file", f"{a}=a;{b}=b", "-merge", "left=a"])

def test_merge_how_cross_does_not_require_on(tmp_path, capsys):
    a, b, _ = _write_three_id_csvs(tmp_path)

    cli.main(["-file", f"{a}=a;{b}=b", "-merge", "left=a,right=b,how=cross"])
    out = capsys.readouterr().out.strip()
    assert len(out.splitlines()) == 1 + 3 * 3  # header + 3x3 cross join rows

def test_merge_how_cross_rejects_on(tmp_path):
    a, b, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit, match="on= cannot be used with how=cross"):
        cli.main(["-file", f"{a}=a;{b}=b", "-merge", "left=a,right=b,how=cross,on=id"])

def test_merge_validate_failure_errors(tmp_path, capsys):
    left = tmp_path / "left.csv"
    right = tmp_path / "right.csv"
    pd.DataFrame({"id": [1, 1, 2], "x": ["a", "b", "c"]}).to_csv(left, index=False)
    pd.DataFrame({"id": [1, 2], "y": ["p", "q"]}).to_csv(right, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([
            "-file", f"{left}=l;{right}=r",
            "-merge", "left=l,right=r,on=id,validate=one_to_one",
        ])
    assert exc_info.value.code == 2
    assert "-merge:" in capsys.readouterr().err

def test_output_succeeds_after_merge(tmp_path):
    left, right = _write_two_csvs(tmp_path)
    dest = tmp_path / "merged.parquet"

    exit_code = cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-merge", "left=df1,right=df2,on=col a:cola",
        "-o", str(dest),
    ])

    assert exit_code == 0
    assert dest.exists()
    result = pd.read_parquet(dest)
    assert list(result.columns) == ["col a", "val_l", "cola", "val_r"]
    assert len(result) == 2

def test_file_nonexistent_path_errors(tmp_path, capsys):
    a, _, _ = _write_three_id_csvs(tmp_path)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{a}=a;{tmp_path / 'missing.csv'}=b", "-concat", "frames='a,b'"])
    assert exc_info.value.code == 2
    assert "file not found" in capsys.readouterr().err

def test_file_bad_encoding_errors(tmp_path, capsys):
    bad_path = tmp_path / "bad.csv"
    bad_path.write_bytes("id,name\n1,caf\xe9\n".encode("latin-1"))
    ok_path = tmp_path / "ok.csv"
    pd.DataFrame({"id": [2], "name": ["bravo"]}).to_csv(ok_path, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main(["-file", f"{bad_path}=a;{ok_path}=b", "-concat", "frames='a,b'"])
    assert exc_info.value.code == 2
    assert "try -encoding" in capsys.readouterr().err


def test_concat_unquoted_frames(tmp_path, capsys):
    a, b, c = _write_three_id_csvs(tmp_path)

    # unquoted frames=a,b,c without inner quotes
    exit_code = cli.main(["-file", f"{a}=a;{b}=b;{c}=c", "-concat", "frames=a,b,c", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(8, 4)"


def test_merge_unquoted_on_pairs(tmp_path, capsys):
    left = tmp_path / "l.csv"
    right = tmp_path / "r.csv"
    pd.DataFrame({"id": [1, 2], "code": ["x", "y"], "v1": [10, 20]}).to_csv(left, index=False)
    pd.DataFrame({"id": [1, 2], "code": ["x", "y"], "v2": [100, 200]}).to_csv(right, index=False)

    exit_code = cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-merge", "left=df1,right=df2,on=id:id,code:code",
        "-shape",
    ])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 4)"

################################################################################
# FILE: tests/test_cli_mutate.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv


def test_mutate_then_select_new_column(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"body_mass_g": [3000.0, 4000.0], "bill_length_mm": [30.0, 40.0]}),
    )

    exit_code = cli.main(
        [path, "-mutate", "bmi = body_mass_g / bill_length_mm ** 2", "-select", "bmi", "-shape"]
    )
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 1)"


def test_mutate_multiple_entries_and_chained_reference(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000.0, 4000.0]}))

    cli.main(
        [path, "-mutate", "mass_kg = body_mass_g / 1000, mass_lb = mass_kg * 2.20462",
         "-select", "mass_kg,mass_lb", "-head"]
    )
    out = capsys.readouterr().out
    assert "mass_kg" in out and "mass_lb" in out


def test_mutate_key_quoting_is_optional(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000.0, 4000.0]}))

    cli.main([path, "-mutate", "mass_kg = body_mass_g / 1000", "-select", "mass_kg", "-shape"])
    unquoted = capsys.readouterr().out.strip()

    cli.main([path, "-mutate", "'mass_kg' = body_mass_g / 1000", "-select", "mass_kg", "-shape"])
    quoted = capsys.readouterr().out.strip()

    assert unquoted == quoted == "(2, 1)"


def test_mutate_unknown_column_reference_errors(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000.0, 4000.0]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-mutate", "bmi = bod_mass_g / 1000", "-shape"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "-mutate" in err
    assert "body_mass_g" in err


def test_mutate_at_local_var_gives_cli_specific_error(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"n": [5, 15]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-mutate", "heavy = n >= @threshold"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "-mutate" in err
    assert "library-only" in err


def test_mutate_at_inside_quoted_string_is_not_flagged_as_local_var(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"s": ["a@b", "c"]}))

    exit_code = cli.main([path, "-mutate", "flag = s == 'a@b'"])

    assert exit_code == 0
    out = capsys.readouterr().out
    assert "True" in out
    assert "False" in out


def test_mutate_load_from_file_cli(tmp_path, capsys):
    data_path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000.0, 4000.0]}))
    spec_path = tmp_path / "specs.txt"
    spec_path.write_text("mass_kg = body_mass_g / 1000\n# comment\nmass_lb = mass_kg * 2.2\n")

    exit_code = cli.main([data_path, "-mutate", f"@{spec_path}", "-select", "mass_kg,mass_lb", "-head"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "mass_kg" in out and "mass_lb" in out


def test_mutate_bracketed_column_names_cli(tmp_path, capsys):
    data_path = _write_csv(tmp_path, pd.DataFrame({"col a": [10, 20], "col b": [1, 2]}))

    exit_code = cli.main([data_path, "-mutate", "total = [col a] + [col b]", "-select", "total", "-head"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "11" in out and "22" in out


def test_mutate_coalesce_cli(tmp_path, capsys):
    data_path = _write_csv(tmp_path, pd.DataFrame({"a": [1.0, None], "b": [None, 2.0]}))

    exit_code = cli.main([data_path, "-mutate", "c = coalesce(a, b, 0.0)", "-select", "c", "-head"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "1.0" in out and "2.0" in out


def test_mutate_colon_raises_cli_error(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"body_mass_g": [3000.0, 4000.0]}))
    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-mutate", "bmi: body_mass_g / 1000"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "use '='" in err


def test_mutate_grouped_cli(tmp_path, capsys):
    df = pd.DataFrame({"grp": ["x", "x", "y", "y"], "val": [10.0, 20.0, 30.0, 40.0]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([
        path, "-by", "grp",
        "-mutate", "avg = mean(val), diff = val - avg, cnt = n",
        "-select", "grp,val,avg,diff,cnt",
        "-head",
    ])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "15.0" in out and "35.0" in out
    assert "-5.0" in out and "5.0" in out


def test_mutate_grouped_with_spaced_columns_cli(tmp_path, capsys):
    df = pd.DataFrame({"grp": ["a", "a"], "total bill": [10.0, 30.0]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([
        path, "-by", "grp",
        "-mutate", "diff = [total bill] - mean([total bill])",
        "-select", "diff",
        "-head",
    ])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "-10.0" in out and "10.0" in out


def test_mutate_by_unknown_column_cli(tmp_path, capsys):
    df = pd.DataFrame({"grp": ["a"], "val": [1.0]})
    path = _write_csv(tmp_path, df)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-by", "nonexistent", "-mutate", "x = mean(val)"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "-by: unknown column" in err

################################################################################
# FILE: tests/test_cli_parse.py
################################################################################
import os
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))



def test_parse_sort_by_trailing_direction():
    from pytae.cli_parsing import parse_sort_by
    assert parse_sort_by("val") == (["val"], "asc")
    assert parse_sort_by("val desc") == (["val"], "desc")
    assert parse_sort_by("species,body_mass_g desc") == (["species", "body_mass_g"], "desc")
    assert parse_sort_by("bill length mm desc") == (["bill length mm"], "desc")


def test_expand_paths_single_and_list(tmp_path):
    import pytest

    from pytae.cli_parsing import expand_paths

    f1 = tmp_path / "a.csv"
    f2 = tmp_path / "b.csv"
    f1.write_text("a,b\n1,2\n")
    f2.write_text("a,b\n3,4\n")

    assert expand_paths(str(f1)) == [f1]
    assert expand_paths([str(f1), str(f2)]) == [f1, f2]
    assert expand_paths(str(tmp_path / "*.csv")) == [f1, f2]
    assert expand_paths([str(f1), str(f1), str(f2)]) == [f1, f2]

    with pytest.raises(SystemExit) as exc_info:
        expand_paths(str(tmp_path / "*.parquet"))
    assert "no files matched pattern" in str(exc_info.value)


def test_parse_rename_conflicting_mappings():
    import pytest

    from pytae.cli_parsing import parse_rename

    assert parse_rename("old:new, other:target") == {"old": "new", "other": "target"}
    with pytest.raises(SystemExit) as exc_info:
        parse_rename("col:new1, col:new2")
    assert "conflicting mappings for 'col'" in str(exc_info.value)

################################################################################
# FILE: tests/test_cli_reshape.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv


def test_long_melts_numeric_columns(tmp_path, capsys):
    df = pd.DataFrame({"grp": ["a", "b"], "n": [1, 2], "x": [10, 20]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-long", "-cols"])

    captured = capsys.readouterr()
    assert exit_code == 0, captured.err
    assert captured.out.strip().splitlines() == ["grp", "variable", "value"]

def test_long_renames_melt_columns(tmp_path, capsys):
    df = pd.DataFrame({"grp": ["a"], "n": [1], "x": [10]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-long", "c=feature,v=amount", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["grp", "feature", "amount"]

def test_wide_pivots_long_frame(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "id": ["a", "b", "c"],
            "country": ["sg", "cn", "ca"],
            "balance": [10, 20, 0],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-wide", "c=country,v=balance", "-cols"])

    captured = capsys.readouterr()
    assert exit_code == 0, captured.err
    assert captured.out.strip().splitlines() == ["id", "ca", "cn", "sg"]

def test_long_quoted_names_with_spaces(tmp_path, capsys):
    df = pd.DataFrame({"grp": ["a"], "n": [1], "x": [10]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-long", "c=feature name,v=amount col", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["grp", "feature name", "amount col"]

def test_wide_quoted_names_with_spaces(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "id": ["a", "b"],
            "country name": ["sg", "cn"],
            "body mass": [10, 20],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-wide", "c=country name,v=body mass", "-cols"])

    captured = capsys.readouterr()
    assert exit_code == 0, captured.err
    assert captured.out.strip().splitlines() == ["id", "cn", "sg"]

def test_wide_rejects_a_arg(tmp_path):
    df = pd.DataFrame(
        {
            "id": ["a", "b"],
            "country": ["sg", "cn"],
            "balance": [10, 20],
        }
    )
    path = _write_csv(tmp_path, df)

    with pytest.raises(SystemExit, match=r"-wide: -wide is strictly for 1-to-1 reshaping without aggregation"):
        cli.main([path, "-wide", "c=country,v=balance,a=mean"])


def test_wide_duplicate_keys_fail_cleanly(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "id": ["a", "a"],
            "country": ["sg", "sg"],
            "balance": [10, 20],
        }
    )
    path = _write_csv(tmp_path, df)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-wide", "c=country,v=balance"])
    assert exc_info.value.code == 2
    captured = capsys.readouterr()
    assert "duplicate entries" in captured.err.lower()
    assert "-pivot" in captured.err

def test_wide_unknown_column_errors(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"id": ["a"], "balance": [1]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-wide", "c=country,v=balance"])
    assert exc_info.value.code == 2

def test_wide_rejects_dropna_kwarg(tmp_path):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"country": ["a"], "variable": ["x"], "value": [1]}),
    )
    with pytest.raises(SystemExit, match="unknown key 'dropna'"):
        cli.main([path, "-wide", "dropna=false"])


def test_long_unknown_column_validation(tmp_path):
    # CLI Issue 11: -long validates columns
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    p = _write_csv(tmp_path, df)
    with pytest.raises(SystemExit) as exc_info:
        cli.main([p, "-long", "cols=nope"])
    assert exc_info.value.code == 2


def test_crosstab_unrecognized_argument(tmp_path):
    df = pd.DataFrame({"species": ["Adelie", "Gentoo"], "island": ["Biscoe", "Dream"]})
    p = _write_csv(tmp_path, df)
    with pytest.raises(SystemExit) as exc_info:
        cli.main([p, "-crosstab", "index=species,columns=island"])
    assert exc_info.value.code == 2

################################################################################
# FILE: tests/test_cli_sample_datasets.py
################################################################################
"""Integration tests against pytae's bundled real datasets (pytae.sample_data).

Unlike the tiny hand-crafted frames in test_cli_order.py (kept deliberately
minimal for exact, easy-to-verify assertions), these exercise the CLI against
real, messier data: multiple categories, real NaNs, larger row counts. This is
also where docs/cli.md's "Sample datasets" section examples are regression
tested, so a broken documented command fails CI instead of only being caught
by manually re-running the docs.
"""

import os
import sys

import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli

_SAMPLE_DATASET_NAMES = ("penguins", "tips", "titanic", "diamonds", "mpg", "flights")


@pytest.fixture(scope="module")
def sample_dataset_dir(tmp_path_factory):
    """Write every bundled sample dataset to parquet once, shared across this module's tests."""
    import pytae
    directory = tmp_path_factory.mktemp("sample_datasets")
    for name in _SAMPLE_DATASET_NAMES:
        pytae.sample_data[name].to_parquet(directory / f"{name}.parquet")
    return directory


# Mirrors docs/cli.md's "Sample datasets" section — keep in sync with those examples.
_DOC_EXAMPLES = [
    ("penguins.parquet", ["-qry", "species = 'Adelie'", "-by", "species", "-agg", "mean"]),
    ("penguins.parquet", ["-pivot", "r=species,c=island,v=sex,a=n"]),
    ("tips.parquet", ["-select", "day,total_bill,tip", "-by", "day", "-mutate", "avg_tip = mean(tip)"]),
    ("titanic.parquet", ["-freq", "survived"]),
    ("diamonds.parquet", ["-select", "cut,price", "-by", "cut", "-agg", "mean"]),
    ("mpg.parquet", ["-select", "origin,mpg", "-sort_by", "mpg desc", "-head", "5"]),
    ("flights.parquet", ["-wide", "c=month,v=passengers,r=year", "-head", "5"]),
]


@pytest.mark.parametrize("filename, extra_args", _DOC_EXAMPLES)
def test_docs_sample_dataset_examples_run_cleanly(sample_dataset_dir, filename, extra_args, capsys):
    exit_code = cli.main([str(sample_dataset_dir / filename), *extra_args])
    assert exit_code == 0, capsys.readouterr().err


def test_pivot_a_n_counts_on_real_penguins(sample_parquet, capsys):
    path = sample_parquet("penguins")

    exit_code = cli.main([path, "-pivot", "r=species,c=island,v=sex,a=n"])
    assert exit_code == 0
    wide_out = capsys.readouterr().out

    def _row(line):
        return [int(tok) for tok in line.split()[1:]]

    wide_rows = {line.split()[0]: _row(line) for line in wide_out.strip().splitlines()[1:]}
    assert wide_rows["Adelie"] == [44, 56, 52]
    assert wide_rows["Chinstrap"] == [0, 68, 0]
    assert wide_rows["Gentoo"] == [124, 0, 0]


def test_wide_pure_reshape_on_real_flights(sample_parquet, capsys):
    path = sample_parquet("flights")

    exit_code = cli.main([path, "-wide", "c=month,v=passengers,r=year", "-cols"])
    assert exit_code == 0
    cols = capsys.readouterr().out.strip().splitlines()
    assert cols[0] == "year"
    assert "Jan" in cols
    assert "Dec" in cols


def test_agg_groups_by_diamonds_cut_categories(sample_parquet, capsys):
    path = sample_parquet("diamonds")

    exit_code = cli.main([path, "-select", "cut,price", "-by", "cut", "-agg", "mean", "-shape"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip() == "(5, 2)"  # 5 known cut categories: Fair/Good/Very Good/Premium/Ideal


def test_group_by_agg_sum_matches_known_flights_totals(sample_parquet, capsys):
    path = sample_parquet("flights")

    exit_code = cli.main([
        path, "-by", "year", "-agg", "passengers = sum", "-head", "3",
    ])

    out = capsys.readouterr().out
    assert exit_code == 0
    rows = [line.split() for line in out.strip().splitlines()[1:]]
    assert rows == [["1949", "1520"], ["1950", "1676"], ["1951", "2042"]]


def test_handle_missing_fills_real_titanic_nans(sample_parquet, capsys):
    # titanic has genuine NaNs in numeric (age) and object (deck) columns.
    path = sample_parquet("titanic")

    exit_code = cli.main([path, "-handle_missing", "-select", "age,deck", "-nulls"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip().splitlines() == ["age     0", "deck    0"]

################################################################################
# FILE: tests/test_cli_select.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv, _write_two_csvs


def test_select_regex_matches_column_names(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "species": ["A", "B"],
            "bill_length_mm": [1, 2],
            "bill_depth_mm": [3, 4],
            "body_mass_g": [5, 6],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "regex=^bill", "-cols"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip().splitlines() == ["bill_length_mm", "bill_depth_mm"]

def test_select_regex_combines_with_explicit_select(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "species": ["A"],
            "island": ["Torgersen"],
            "bill_length_mm": [1],
            "body_mass_g": [2],
        }
    )
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "species,regex=bill|body", "-cols"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip().splitlines() == ["species", "bill_length_mm", "body_mass_g"]

def test_parse_select_spec_names_and_kwargs():
    names, kwargs = cli.parse_select_spec("species,contains=bill,dtype=numeric")
    assert names == ["species"]
    assert kwargs == {"contains": "bill", "dtype": "numeric"}

def test_parse_select_spec_quoted_name_and_repeated_key():
    names, kwargs = cli.parse_select_spec("'bill length mm',contains=bill,contains=body")
    assert names == ["bill length mm"]
    assert kwargs == {"contains": ["bill", "body"]}

def test_select_dtype_unions_with_explicit_names(tmp_path, capsys):
    df = pd.DataFrame({"cola": ["a"], "colb": [1], "n": [2.0], "label": ["x"]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "cola,colb,dtype=numeric", "-cols"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip().splitlines() == ["cola", "colb", "n"]

def test_select_exclude_dtype_standalone(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"name": ["a"], "n": [1]}))

    exit_code = cli.main([path, "-select", "exclude_dtype=numeric", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["name"]

def test_select_exclude_dtype_non_numeric_keeps_numbers(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"name": ["a"], "n": [1], "x": [2.0]}))

    exit_code = cli.main([path, "-select", "exclude_dtype=non_numeric", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["n", "x"]

def test_select_exclude_dtype_cannot_combine(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"name": ["a"], "n": [1]}))

    with pytest.raises(SystemExit, match="exclude_dtype cannot be combined"):
        cli.main([path, "-select", "name,exclude_dtype=numeric", "-cols"])

def test_select_slice(tmp_path, capsys):
    df = pd.DataFrame({"a": [1], "b": [2], "c": [3], "d": [4]})
    path = _write_csv(tmp_path, df)

    exit_code = cli.main([path, "-select", "a:c", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["a", "b", "c"]

def test_second_select_filters_remaining_names(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3], "d": [4]}))

    exit_code = cli.main([path, "-select", "a,b,c", "-select", "a,c", "-cols"])

    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["a", "c"]

def test_second_select_rejects_name_not_kept_by_first(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3]}))

    # last-wins would yield b,c; chaining rejects c because the first spec dropped it
    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-select", "a,b", "-select", "b,c", "-cols"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "unknown column" in err
    assert "'c'" in err

def test_select_only_prints_selected_table(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3]}))

    exit_code = cli.main([path, "-select", "a,b"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "a" in out and "b" in out
    assert "c" not in out.split()

def test_select_unknown_exact_name_errors(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-select", "d,a,b", "-cols"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "unknown column" in err
    assert "'d'" in err

def test_select_after_agg_df_sees_agg_columns(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"grp": ["x", "x", "y"], "val": [1, 2, 3], "z": [9, 8, 7]}),
    )

    # n is created by agg; a starting-view -select would reject it
    exit_code = cli.main(
        [path, "-by", "grp", "-agg", "val = sum, n = n", "-select", "grp,n", "-cols"]
    )
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["grp", "n"]

def test_select_agg_df_select_shape_chains(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"grp": ["x", "x", "y"], "val": [1, 2, 3], "z": [9, 8, 7]}),
    )

    # starting-view chaining would collapse to val-only before agg → (1, 1)
    exit_code = cli.main(
        [path, "-select", "grp,val", "-by", "grp", "-agg", "sum", "-select", "val", "-shape"]
    )
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 1)"

def test_second_select_intersects_dtype_and_contains(tmp_path, capsys):
    df = pd.DataFrame(
        {
            "species": ["A"],
            "bill_length_mm": [1],
            "bill_depth_mm": [2],
            "body_mass_g": [3],
            "island": ["Torgersen"],
        }
    )
    path = _write_csv(tmp_path, df)

    # last-wins of dtype=numeric would keep body_mass_g too
    exit_code = cli.main([path, "-select", "contains=bill", "-select", "dtype=numeric", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["bill_length_mm", "bill_depth_mm"]

    exit_code = cli.main([path, "-select", "dtype=numeric", "-select", "contains=bill", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["bill_length_mm", "bill_depth_mm"]

def test_removed_select_star_flags_are_unknown(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-select-regex", "^a", "-cols"])
    assert exc_info.value.code == 2

def test_select_caret_pattern_is_not_implicit_regex(tmp_path):
    path = _write_csv(tmp_path, pd.DataFrame({"bill_length_mm": [1], "species": ["A"]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-select", "^bill", "-cols"])
    assert exc_info.value.code == 2

def test_select_then_qry_on_dropped_column_errors(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"keep": [1, 2], "flt": ["A", "B"], "val": [10, 20]}),
    )

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-select", "keep,val", "-qry", "flt = 'A'", "-shape"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "-qry" in err
    assert "flt" in err

def test_select_runs_after_merge(tmp_path, capsys):
    left, right = _write_two_csvs(tmp_path)

    cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-merge", "left=df1,right=df2,on=col a:cola",
        "-select", "val_l,val_r", "-shape",
    ])

    assert capsys.readouterr().out.strip() == "(2, 2)"


def test_select_cli_negative_column_syntax(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"species": ["A"], "island": ["B"], "val": [1]}),
    )

    exit_code = cli.main([path, "-select", "-species", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["island", "val"]

    exit_code2 = cli.main([path, "-select", "~island", "-cols"])
    assert exit_code2 == 0
    assert capsys.readouterr().out.strip().splitlines() == ["species", "val"]


def test_select_cli_exclude_kwarg(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"species": ["A"], "island": ["B"], "val": [1]}),
    )

    exit_code = cli.main([path, "-select", "exclude=species", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["island", "val"]


def test_select_cli_multiple_negative_columns(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "c": [2], "b": [3], "d": [4]}))

    exit_code = cli.main([path, "-select", "-c,-d", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["a", "b"]


def test_select_cli_negative_slice(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3], "d": [4]}))

    exit_code = cli.main([path, "-select", "-b:c", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["a", "d"]


def test_select_cli_negative_quoted_name_with_space(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"bill length mm": [1], "body mass g": [2], "keep": [3]}))

    exit_code = cli.main([path, "-select", "-'bill length mm'", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["body mass g", "keep"]


def test_select_cli_chained_negative(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1], "b": [2], "c": [3], "d": [4]}))

    exit_code = cli.main([path, "-select", "a,b,c", "-select", "-b", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["a", "c"]


def test_select_cli_negative_after_agg(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"grp": ["x", "x", "y"], "val": [1, 2, 3], "z": [9, 8, 7]}),
    )

    exit_code = cli.main(
        [path, "-by", "grp", "-agg", "val = sum, n = n", "-select", "-val", "-cols"]
    )
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["grp", "n"]


def test_select_cli_negative_after_qry(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"keep": [1, 2], "flt": ["A", "B"], "val": [10, 20]}),
    )

    exit_code = cli.main([path, "-qry", "flt = 'A'", "-select", "-flt", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["keep", "val"]


def test_select_cli_negative_unknown_name_hint(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"species": [1], "island": [2]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-select", "-speces", "-cols"])
    assert exc_info.value.code == 2
    err = capsys.readouterr().err
    assert "unknown column" in err
    assert "species" in err


def test_select_cli_bracketed_columns(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"species": ["A"], "island name": ["B"], "bill length mm": [1.0]}),
    )

    exit_code = cli.main([path, "-select", "[island name], [bill length mm]", "-cols"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip().splitlines() == ["island name", "bill length mm"]

    exit_code2 = cli.main([path, "-select", "species, [bill length mm]", "-cols"])
    assert exit_code2 == 0
    assert capsys.readouterr().out.strip().splitlines() == ["species", "bill length mm"]

    exit_code3 = cli.main([path, "-select", "-[island name]", "-cols"])
    assert exit_code3 == 0
    assert capsys.readouterr().out.strip().splitlines() == ["species", "bill length mm"]

################################################################################
# FILE: tests/test_cli_sql.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

from pytae import cli
from tests.cli_helpers import _write_csv, _write_two_csvs


def test_sql_query_via_data_alias(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"col a": [1, 20, 3], "col b": ["x", "y", "z"]}))

    exit_code = cli.main([path, "-sql", 'select "col b" from data where "col a" > 10'])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "y" in out
    assert "x" not in out

def test_sql_spaced_column_name_needs_double_quotes(tmp_path, capsys):
    # Standard SQL identifier quoting: double quotes for names with spaces, not single quotes.
    path = _write_csv(tmp_path, pd.DataFrame({"bill length mm": [1, 20, 3], "species": ["a", "b", "c"]}))

    exit_code = cli.main([path, "-sql", 'select species from data where "bill length mm" > 10'])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert out.strip().splitlines()[-1].strip() == "b"

def test_sql_only_data_is_registered_no_file_derived_alias(tmp_path, capsys):
    # The file is already named on the command line; -sql does not also register a
    # file-stem alias (e.g. "my_data" for my_data.csv) -- only `data` is queryable.
    path = tmp_path / "my_data.csv"
    pd.DataFrame({"a": [1, 2]}).to_csv(path, index=False)

    with pytest.raises(SystemExit) as exc_info:
        cli.main([str(path), "-sql", "select * from my_data"])
    assert exc_info.value.code == 2
    assert "-sql" in capsys.readouterr().err

def test_sql_reserved_word_table_is_not_registered(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-sql", "select * from table"])
    assert exc_info.value.code == 2
    assert "-sql" in capsys.readouterr().err

def test_sql_chains_after_select(tmp_path, capsys):
    path = _write_csv(
        tmp_path,
        pd.DataFrame({"keep": [1, 2, 3], "flt": ["A", "B", "A"], "val": [10, 20, 30]}),
    )

    exit_code = cli.main(
        [path, "-select", "keep,flt", "-sql", 'select "keep" from data where "flt" = \'A\'', "-shape"]
    )

    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 1)"

def test_sql_invalid_query_errors(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2]}))

    with pytest.raises(SystemExit) as exc_info:
        cli.main([path, "-sql", "not valid sql"])
    assert exc_info.value.code == 2
    assert "-sql" in capsys.readouterr().err

def test_sql_first_op_scans_parquet_directly(tmp_path, capsys):
    # -sql as the very first op should give the same result whether duckdb scans the
    # source file directly (fast path) or pandas materializes it first (fallback).
    path = tmp_path / "data.parquet"
    pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}).to_parquet(path, index=False)

    exit_code = cli.main([str(path), "-sql", "select b from data where a > 1"])

    out = capsys.readouterr().out.strip().splitlines()
    assert exit_code == 0
    assert out[1:] == ["y", "z"]

def test_sql_first_op_respects_nrows(tmp_path, capsys):
    path = tmp_path / "data.parquet"
    pd.DataFrame({"a": range(10)}).to_parquet(path, index=False)

    exit_code = cli.main([str(path), "-nrows", "3", "-sql", "select * from data"])

    out = capsys.readouterr().out.strip().splitlines()
    assert exit_code == 0
    assert len(out) - 1 == 3  # header + 3 rows

def test_sql_first_op_with_progress_still_works(tmp_path, capsys):
    # -progress forces the pandas-materialize fallback instead of the direct duckdb scan.
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2, 3]}))

    exit_code = cli.main([path, "-progress", "-sql", "select * from data where a > 1"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "reading" in out  # progress line proves the pandas fallback path ran
    table_lines = [line for line in out.strip().splitlines() if "reading" not in line]
    assert len(table_lines) - 1 == 2

def test_sql_first_op_with_custom_encoding_still_works(tmp_path, capsys):
    # A non-utf8 encoding forces the pandas-materialize fallback (duckdb's CSV reader
    # doesn't support arbitrary encodings the way pandas does).
    path = tmp_path / "data.csv"
    pd.DataFrame({"a": [1, 2], "name": ["café", "naïve"]}).to_csv(path, index=False, encoding="latin-1")

    exit_code = cli.main([str(path), "-encoding", "latin-1", "-sql", "select * from data"])

    out = capsys.readouterr().out
    assert exit_code == 0
    assert "café" in out

def test_sql_can_join_file_aliases_directly(tmp_path, capsys):
    left, right = _write_two_csvs(tmp_path)

    cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-sql", 'select * from df1 inner join df2 on df1."col a" = df2.cola',
    ])

    out = capsys.readouterr().out
    assert "val_l" in out and "val_r" in out

def test_sql_after_merge_can_query_data(tmp_path, capsys):
    left, right = _write_two_csvs(tmp_path)

    cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-merge", "left=df1,right=df2,on=col a:cola",
        "-sql", "select count(*) as n from data",
    ])

    assert capsys.readouterr().out.strip().splitlines()[-1].strip() == "2"

def test_sql_as_first_op_satisfies_file_requirement(tmp_path):
    left, right = _write_two_csvs(tmp_path)

    # no -merge at all — -sql alone should be accepted as the founding op
    exit_code = cli.main([
        "-file", f"{left}=df1;{right}=df2",
        "-sql", "select * from df1",
    ])
    assert exit_code == 0


def test_sql_load_from_file_cli(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2], "b": [10, 20]}))
    qfile = tmp_path / "query.txt"
    qfile.write_text("select a, b from data where a > 1")

    exit_code = cli.main([str(path), "-sql", f"@{qfile}", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(1, 2)"


def test_sql_load_from_quoted_file_cli(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"a": [1, 2], "b": [10, 20]}))
    sub = tmp_path / "subdir"
    sub.mkdir()
    qfile = sub / "query.txt"
    qfile.write_text("select a, b from data where a > 1")

    # Test with outer quotes around '@path'
    exit_code = cli.main([str(path), "-sql", f"'@{qfile}'", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(1, 2)"

    # Test with quotes around the path itself '@"path"'
    exit_code = cli.main([str(path), "-sql", f'@"{qfile}"', "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(1, 2)"


def test_sql_bracketed_identifiers_cli(tmp_path, capsys):
    path = _write_csv(tmp_path, pd.DataFrame({"col a": [1, 2], "col b": [10, 20]}))

    exit_code = cli.main([str(path), "-sql", "select [col a], `col b` from data", "-shape"])
    assert exit_code == 0
    assert capsys.readouterr().out.strip() == "(2, 2)"

################################################################################
# FILE: tests/test_import.py
################################################################################
import os
import subprocess
import sys

_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))


def test_import_does_not_load_plotting_shape_or_samples():
    code = r"""
import sys
sys.path.insert(0, %r)
import pytae

assert "pytae.plotting" not in sys.modules
assert "pytae.shape" not in sys.modules
assert "matplotlib" not in sys.modules
assert "duckdb" not in sys.modules
assert pytae._cache == {}
assert "penguins" in pytae.sample_data
assert list(pytae.sample_data.keys())
assert pytae._cache == {}

# keys() returns a plain readable tuple, not a bare KeysView object
keys = pytae.sample_data.keys()
assert isinstance(keys, tuple)
assert "penguins" in keys

df = pytae.sample("penguins")
assert len(df) > 0
# sample() returns a copy each time, so mutating it never poisons the cache
assert pytae.sample_data["penguins"] is not df
df.iloc[0, 0] = "MUTATED"
assert pytae.sample_data["penguins"].iloc[0, 0] != "MUTATED"

assert "pytae.shape" not in sys.modules
long = pytae.long
assert "pytae.shape" in sys.modules
assert long is not None

assert callable(pytae.mutate)
assert hasattr(__import__("pandas").DataFrame, "pt")
assert not hasattr(__import__("pandas").DataFrame, "mutate")

Plotter = pytae.Plotter
assert "pytae.plotting" in sys.modules
assert Plotter is not None
""" % _SRC
    subprocess.check_call([sys.executable, "-c", code])


def test_cli_import_does_not_load_shape_or_plotting():
    code = r"""
import sys
sys.path.insert(0, %r)
import pytae.cli

assert "pytae.plotting" not in sys.modules
assert "pytae.shape" not in sys.modules
assert "matplotlib" not in sys.modules
assert callable(__import__("pytae", fromlist=["qry"]).qry)
assert callable(__import__("pytae", fromlist=["select"]).select)
DF = sys.modules["pandas"].DataFrame
assert hasattr(DF, "pt")
assert not hasattr(DF, "qry")
assert not hasattr(DF, "select")
assert not hasattr(DF, "mutate")
""" % _SRC
    subprocess.check_call([sys.executable, "-c", code])


def test_plotter_import_error_mentions_plot_extra():
    code = r"""
import builtins
import sys
sys.path.insert(0, %r)
real_import = builtins.__import__

def blocked(name, *args, **kwargs):
    if name == "matplotlib" or name.startswith("matplotlib."):
        raise ImportError("No module named matplotlib")
    return real_import(name, *args, **kwargs)

builtins.__import__ = blocked
import pytae
try:
    pytae.Plotter
except ImportError as exc:
    msg = str(exc)
    assert "matplotlib" in msg
    assert "pytae[plot]" in msg
else:
    raise AssertionError("expected ImportError")
""" % _SRC
    subprocess.check_call([sys.executable, "-c", code])


def test_sample_unknown_dataset():
    import pytae

    try:
        pytae.sample("not_a_real_dataset")
    except KeyError as exc:
        assert "not_a_real_dataset" in str(exc)
        assert "penguins" in str(exc)
    else:
        raise AssertionError("expected KeyError")


def test_sql_import_error_mentions_sql_extra(tmp_path):
    csv_path = tmp_path / "data.csv"
    csv_path.write_text("a,b\n1,2\n")
    code = r"""
import builtins
import sys
sys.path.insert(0, %r)
real_import = builtins.__import__

def blocked(name, *args, **kwargs):
    if name == "duckdb" or name.startswith("duckdb."):
        raise ImportError("No module named duckdb")
    return real_import(name, *args, **kwargs)

builtins.__import__ = blocked
import pytae.cli as cli
try:
    cli.main([%r, "-sql", "select * from df"])
except SystemExit as exc:
    assert exc.code == 2
else:
    raise AssertionError("expected SystemExit")
""" % (_SRC, str(csv_path))
    result = subprocess.run(
        [sys.executable, "-c", code], capture_output=True, text=True
    )
    assert result.returncode == 0, result.stderr
    assert "pytae[sql]" in result.stderr

################################################################################
# FILE: tests/test_mutate.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt
from pytae.mutate import parse_mutate_spec


def _df():
    return pd.DataFrame(
        {
            "species": ["Adelie", "Gentoo"],
            "body_mass_g": [3000.0, 4000.0],
            "bill_length_mm": [30.0, 40.0],
        }
    )


def test_mutate_basic_expression():
    result = pt.mutate(_df(), bmi="body_mass_g / bill_length_mm ** 2")
    assert list(result["bmi"]) == [pytest.approx(3.333333, rel=1e-4), pytest.approx(2.5)]


def test_mutate_multiple_entries_in_one_call():
    result = pt.mutate(_df(), heavy="body_mass_g > 3500", mass_kg="body_mass_g / 1000")
    assert list(result["heavy"]) == [False, True]
    assert list(result["mass_kg"]) == [3.0, 4.0]


def test_mutate_later_entry_references_earlier_one():
    result = pt.mutate(_df(), mass_kg="body_mass_g / 1000", mass_lb="mass_kg * 2.20462")
    assert result["mass_lb"].iloc[0] == pytest.approx(6.61386)


def test_mutate_string_comparison_needs_quotes():
    result = pt.mutate(_df(), is_adelie="species == 'Adelie'")
    assert list(result["is_adelie"]) == [True, False]


def test_mutate_does_not_modify_original():
    df = _df()
    pt.mutate(df, bmi="body_mass_g / bill_length_mm ** 2")
    assert "bmi" not in df.columns


def test_mutate_overwrites_existing_column():
    result = pt.mutate(_df(), body_mass_g="body_mass_g / 1000")
    assert list(result["body_mass_g"]) == [3.0, 4.0]


def test_mutate_unknown_column_suggests_typo():
    with pytest.raises(KeyError, match="did you mean 'body_mass_g'"):
        pt.mutate(_df(), bmi="bod_mass_g / bill_length_mm")


def test_mutate_quoting_column_reference_breaks_arithmetic():
    with pytest.raises(TypeError):
        pt.mutate(_df(), bmi="'body_mass_g' / 'bill_length_mm' ** 2")


def test_mutate_empty_spec_errors():
    with pytest.raises(ValueError, match="mutate\\(\\) expects at least one"):
        pt.mutate(_df())


def test_mutate_if_else_basic():
    result = pt.mutate(_df(), size="if_else(body_mass_g >= 3500, 'heavy', 'light')")
    assert list(result["size"]) == ["light", "heavy"]


def test_mutate_if_else_numeric_values():
    result = pt.mutate(_df(), bonus="if_else(body_mass_g >= 3500, 100, 10)")
    assert list(result["bonus"]) == [10, 100]


def test_mutate_if_else_value_can_be_column_expression():
    result = pt.mutate(_df(), adjusted="if_else(body_mass_g >= 3500, body_mass_g / 1000, body_mass_g)")
    assert list(result["adjusted"]) == [3000.0, 4.0]


def test_mutate_if_else_wrong_arg_count_errors():
    with pytest.raises(TypeError):
        pt.mutate(_df(), size="if_else(body_mass_g >= 3500, 'heavy')")


def test_mutate_if_else_mixed_dtype_branches():
    result = pt.mutate(_df(), g="if_else(body_mass_g >= 3500, 'heavy', 0)")
    assert list(result["g"]) == [0, "heavy"]


def test_mutate_case_when_mixed_dtype_choices():
    result = pt.mutate(_df(), g="case_when((body_mass_g >= 3500, 'heavy'), 0)")
    assert list(result["g"]) == [0, "heavy"]


def test_mutate_case_when_first_match_wins():
    result = pt.mutate(
        _df(),
        grade="case_when((body_mass_g >= 3800, 'A'), (body_mass_g >= 3200, 'B'), 'C')",
    )
    assert list(result["grade"]) == ["C", "A"]


def test_mutate_case_when_no_default_gives_nan_for_unmatched():
    result = pt.mutate(_df(), grade="case_when((body_mass_g >= 3800, 'A'))")
    assert result["grade"].iloc[1] == "A"
    assert pd.isna(result["grade"].iloc[0])


def test_mutate_case_when_default_must_be_last():
    with pytest.raises(ValueError, match="case_when default .*must be the last argument"):
        pt.mutate(_df(), grade="case_when('C', (body_mass_g >= 3800, 'A'))")


def test_mutate_case_when_requires_at_least_one_entry():
    with pytest.raises(ValueError, match="case_when expects at least one"):
        pt.mutate(_df(), grade="case_when()")


def test_mutate_if_else_and_case_when_chain_with_other_entries():
    result = pt.mutate(
        _df(),
        mass_kg="body_mass_g / 1000",
        size="if_else(mass_kg >= 3.5, 'heavy', 'light')",
    )
    assert list(result["size"]) == ["light", "heavy"]


def test_mutate_local_var_reference():
    threshold = 3500.0
    result = pt.mutate(_df(), heavy="body_mass_g >= @threshold")
    assert list(result["heavy"]) == [False, True]


def test_mutate_local_var_reference_via_accessor():
    threshold = 3500.0
    result = _df().pt.mutate(heavy="body_mass_g >= @threshold")
    assert list(result["heavy"]) == [False, True]


def test_mutate_local_var_reference_in_if_else():
    threshold = 3500.0
    result = pt.mutate(_df(), size="if_else(body_mass_g >= @threshold, 'heavy', 'light')")
    assert list(result["size"]) == ["light", "heavy"]


def test_mutate_local_var_reference_in_case_when():
    low, high = 3200.0, 3800.0
    result = pt.mutate(
        _df(),
        grade="case_when((body_mass_g >= @high, 'A'), (body_mass_g >= @low, 'B'), 'C')",
    )
    assert list(result["grade"]) == ["C", "A"]


def test_mutate_unknown_local_var_suggests_typo_free_error():
    with pytest.raises(KeyError, match="threshol"):
        pt.mutate(_df(), heavy="body_mass_g >= @threshol")


def test_mutate_kwargs():
    df = _df()
    res = pt.mutate(df, bmi="body_mass_g / bill_length_mm ** 2", mass_kg="body_mass_g / 1000")
    assert "bmi" in res.columns
    assert "mass_kg" in res.columns
    assert list(res["mass_kg"]) == [3.0, 4.0]


def test_mutate_rejects_unsupported_positional_types():
    df = _df()
    with pytest.raises(TypeError, match="unexpected positional argument"):
        df.pt.mutate(123)
    with pytest.raises(TypeError, match="unexpected positional argument"):
        df.pt.mutate([1, 2, 3])
    with pytest.raises(TypeError, match="callables as keyword arguments"):
        df.pt.mutate(lambda d: d)


def test_mutate_callable_lambda():
    df = _df()
    res = pt.mutate(df, bmi=lambda d: d["body_mass_g"] / d["bill_length_mm"] ** 2)
    assert "bmi" in res.columns
    assert list(res["bmi"]) == [pytest.approx(3.333333, rel=1e-4), pytest.approx(2.5)]


def test_mutate_constants():
    df = _df()
    res = pt.mutate(df, status="'active'", count=42)
    assert list(res["status"]) == ["active", "active"]
    assert list(res["count"]) == [42, 42]


def test_mutate_case_when_default_kwarg():
    df = _df()
    res = pt.mutate(df, g="case_when((body_mass_g >= 3500, 'A'), default='B')")
    assert list(res["g"]) == ["B", "A"]


def test_mutate_case_when_flat_alternating_pairs():
    df = _df()
    res = pt.mutate(df, g="case_when(body_mass_g >= 3800, 'A', body_mass_g >= 3200, 'B', default='C')")
    assert list(res["g"]) == ["C", "A"]


def test_mutate_coalesce_helper():
    df = pd.DataFrame({"a": [1.0, None, None], "b": [None, 2.0, None]})
    res = pt.mutate(df, c="coalesce(a, b, 0.0)")
    assert list(res["c"]) == [1.0, 2.0, 0.0]


def test_mutate_bracketed_and_backtick_columns_in_fallback():
    df = pd.DataFrame({"col a": [10, 20], "col b": [1, 2]})
    res1 = pt.mutate(df, c="if_else([col a] > 15, 'high', 'low')")
    assert list(res1["c"]) == ["low", "high"]
    res2 = pt.mutate(df, c="if_else(`col a` > 15, 'high', 'low')")
    assert list(res2["c"]) == ["low", "high"]
    res3 = pt.mutate(df, total="[col a] + [col b]")
    assert list(res3["total"]) == [11, 22]


def test_parse_mutate_spec_equals_and_comments():
    spec = """
    # calculate weight in kg
    mass_kg = body_mass_g / 1000
    # calculate weight in lbs
    mass_lb = mass_kg * 2.2
    """
    res = parse_mutate_spec(spec)
    assert "mass_kg" in res
    assert "mass_lb" in res
    assert res["mass_kg"] == "body_mass_g / 1000"


def test_parse_mutate_spec_rejects_colon():
    with pytest.raises(ValueError, match="use '=' for assignment"):
        parse_mutate_spec("mass_kg: body_mass_g / 1000")


def test_mutate_load_from_file(tmp_path):
    df = _df()
    file_path = tmp_path / "specs.txt"
    file_path.write_text("mass_kg = body_mass_g / 1000\n# comment\nmass_lb = mass_kg * 2.2\n")
    res = pt.mutate(df, f"@{file_path}")
    assert "mass_kg" in res.columns
    assert "mass_lb" in res.columns


def test_mutate_explicit_params():
    df = _df()
    res = pt.mutate(df, flag="body_mass_g >= @thresh", params={"thresh": 3500.0})
    assert list(res["flag"]) == [False, True]


def test_mutate_unquoted_string_literal_hint():
    df = _df()
    with pytest.raises(KeyError, match="if you intended a string literal, quote it like 'heavy'"):
        pt.mutate(df, status="if_else(body_mass_g > 3500, heavy, light)")


def test_mutate_positional_string_expr():
    df = _df()
    res = df.pt.mutate("mass_kg = body_mass_g / 1000")
    assert "mass_kg" in res.columns
    assert list(res["mass_kg"]) == [3.0, 4.0]

    # Multiple positional expressions
    res2 = df.pt.mutate("mass_kg = body_mass_g / 1000", "mass_lb = mass_kg * 2.2")
    assert "mass_kg" in res2.columns
    assert "mass_lb" in res2.columns


def test_mutate_positional_dict():
    df = _df()
    res = df.pt.mutate({"mass_kg": "body_mass_g / 1000"})
    assert list(res["mass_kg"]) == [3.0, 4.0]


def test_mutate_params_column_creation():
    df = _df()
    # Integer literal assigned to column named 'params'
    res1 = df.pt.mutate(params=10)
    assert "params" in res1.columns
    assert list(res1["params"]) == [10, 10]

    # String expression assigned to column named 'params'
    res2 = df.pt.mutate(params="body_mass_g / 1000")
    assert "params" in res2.columns
    assert list(res2["params"]) == [3.0, 4.0]


def test_mutate_explicit_underscore_params():
    df = _df()
    res = df.pt.mutate(flag="body_mass_g >= @thresh", _params={"thresh": 3500.0})
    assert list(res["flag"]) == [False, True]


def test_mutate_bracketed_columns():
    df = pd.DataFrame({"body mass g": [3000.0, 4000.0], "height": [1.5, 2.0]})
    # Target column has brackets, source has brackets
    res1 = df.pt.mutate("[mass kg] = [body mass g] / 1000")
    assert "mass kg" in res1.columns
    assert list(res1["mass kg"]) == [3.0, 4.0]

    # Mixed source columns (one with space, one without)
    res2 = df.pt.mutate("[bmi score] = [body mass g] / (height ** 2)")
    assert "bmi score" in res2.columns


def test_mutate_grouped_single_col():
    df = pd.DataFrame({"grp": ["a", "a", "b"], "val": [10.0, 20.0, 60.0]})
    res = df.pt.mutate(
        "avg = mean(val), diff = val - avg, total = sum(val), cnt = n, cnt2 = n()",
        by="grp",
    )
    assert list(res["avg"]) == [15.0, 15.0, 60.0]
    assert list(res["diff"]) == [-5.0, 5.0, 0.0]
    assert list(res["total"]) == [30.0, 30.0, 60.0]
    assert list(res["cnt"]) == [2, 2, 1]
    assert list(res["cnt2"]) == [2, 2, 1]


def test_mutate_grouped_multi_col():
    df = pd.DataFrame({
        "g1": ["a", "a", "a", "b"],
        "g2": ["x", "x", "y", "x"],
        "val": [10.0, 30.0, 5.0, 100.0],
    })
    res1 = df.pt.mutate("m = mean(val), cnt = n", by=["g1", "g2"])
    assert list(res1["m"]) == [20.0, 20.0, 5.0, 100.0]
    assert list(res1["cnt"]) == [2, 2, 1, 1]

    # Comma-separated string by="g1, g2"
    res2 = df.pt.mutate("m = mean(val)", by="g1, g2")
    assert list(res2["m"]) == [20.0, 20.0, 5.0, 100.0]


def test_mutate_grouped_spaced_brackets():
    df = pd.DataFrame({"group code": ["a", "a", "b"], "total bill": [10.0, 30.0, 50.0]})
    res = df.pt.mutate("[diff bill] = [total bill] - mean([total bill])", by="group code")
    assert list(res["diff bill"]) == [-10.0, 10.0, 0.0]


def test_mutate_grouped_dropna_observed():
    df = pd.DataFrame({"grp": ["a", None, "a"], "val": [10.0, 50.0, 30.0]})
    # Default is dropna=False: computes on the NA group
    res_default = df.pt.mutate("m = mean(val)", by="grp")
    assert res_default["m"].iloc[1] == 50.0
    assert res_default["m"].iloc[0] == 20.0
    assert res_default["m"].iloc[2] == 20.0

    res_dropna = df.pt.mutate("m = mean(val)", by="grp", dropna=True)
    assert pd.isna(res_dropna["m"].iloc[1])
    assert res_dropna["m"].iloc[0] == 20.0
    assert res_dropna["m"].iloc[2] == 20.0

    res_keepna = df.pt.mutate("m = mean(val)", by="grp", dropna=False)
    assert res_keepna["m"].iloc[1] == 50.0


def test_mutate_grouped_duplicate_index():
    df = pd.DataFrame({"grp": ["a", "b", "a"], "val": [10.0, 50.0, 30.0]}, index=[0, 1, 0])
    res = df.pt.mutate("m = mean(val)", by="grp")
    assert list(res["m"]) == [20.0, 50.0, 20.0]
    assert list(res.index) == [0, 1, 0]


def test_mutate_grouped_functional_and_callable():
    df = pd.DataFrame({"grp": ["a", "a", "b"], "val": [10.0, 20.0, 60.0]})
    # Functional pt.mutate
    res1 = pt.mutate(df, "avg = mean(val)", by="grp")
    assert list(res1["avg"]) == [15.0, 15.0, 60.0]

    # Callable inside grouped mutate
    res2 = df.pt.mutate(half_mean=lambda g: g["val"].mean() / 2, by="grp")
    assert list(res2["half_mean"]) == [7.5, 7.5, 30.0]


def test_mutate_grouped_unknown_col():
    df = pd.DataFrame({"grp": ["a", "b"], "val": [1, 2]})
    with pytest.raises(KeyError, match="grouping column 'grp_typo' not found"):
        df.pt.mutate("m = mean(val)", by="grp_typo")


def test_mutate_coalesce_scalar_broadcast():
    # Lib Issue 5
    df = pd.DataFrame({"a": [None, 2.0, None]}, index=[5, 6, 7])
    res1 = pt.mutate(df, c="coalesce(0, a)")
    assert list(res1["c"]) == [0.0, 0.0, 0.0]

    res2 = pt.mutate(df, c="coalesce(1, 2)")
    assert list(res2["c"]) == [1, 1, 1]


def test_mutate_column_named_sum_overrides_helper():
    # Lib Issue 7
    df = pd.DataFrame({"sum": [1, 2], "x": [10, 20]})
    res = df.pt.mutate(y="if_else(x > 15, sum, 0)")
    assert list(res["y"]) == [0, 2]


def test_mutate_map_preserves_explicit_null():
    # Lib Issue 11
    df = pd.DataFrame({"a": ["x", "y", "z"]})
    res = pt.mutate(df, b="map(a, {'x': 1, 'y': None}, -1)")
    assert res["b"].iloc[0] == 1.0
    assert pd.isna(res["b"].iloc[1])
    assert res["b"].iloc[2] == -1.0


def test_mutate_if_else_and_case_when_nullable_boolean():
    # Lib Issue 12
    df = pd.DataFrame({"a": pd.Series([1, pd.NA, 3], dtype="Int64")})
    res_if = pt.mutate(df, b="if_else(a > 2, 'y', 'n')")
    assert list(res_if["b"]) == ["n", "n", "y"]

    res_cw = pt.mutate(df, b="case_when((a > 2, 'y'), 'n')")
    assert list(res_cw["b"]) == ["n", "n", "y"]


def test_mutate_n_row_count_vs_col_and_helper():
    # Lib Issue 16: caller local 'n' does not shadow row count 'n' unless referenced as '@n'
    n = 999  # noqa: F841
    df = pd.DataFrame({"val": [10, 20, 30]})
    res = df.pt.mutate("cnt = n, caller_n = @n")
    assert list(res["cnt"]) == [3, 3, 3]
    assert list(res["caller_n"]) == [999, 999, 999]


def test_mutate_grouped_callable_sees_original_index():
    # Lib Issue 17: callable sees original index in grouped mutate
    s = pd.Series([100, 200, 300], index=[10, 20, 30])
    df = pd.DataFrame({"g": ["a", "b", "a"], "x": [1, 2, 3]}, index=[10, 20, 30])
    res = df.pt.mutate(y=lambda d: d["x"] + s.reindex(d.index), by="g")
    assert list(res["y"]) == [101, 202, 303]
    assert list(res.index) == [10, 20, 30]


def test_mutate_scalar_if_else_and_case_when():
    # Lib Issue 1 (Review 2ec5bbba): scalar condition in if_else and case_when broadcasts across rows and groups
    df = pd.DataFrame({"val": [10, 20, 30]})
    assert df.pt.mutate(y="if_else(n > 1, 1, 0)")["y"].tolist() == [1, 1, 1]
    assert df.pt.mutate(y="case_when((n > 1, 'big'), 'small')")["y"].tolist() == ["big", "big", "big"]

    df2 = pd.DataFrame({"val": [10, 20, 30]}, index=[5, 6, 7])
    assert df2.pt.mutate(y="if_else(n > 1, 1, 0)")["y"].tolist() == [1, 1, 1]

    df3 = pd.DataFrame({"g": ["a", "a", "b"], "val": [10, 20, 30]})
    assert df3.pt.mutate(y="if_else(n > 1, 1, 0)", by="g")["y"].tolist() == [1, 1, 0]

################################################################################
# FILE: tests/test_pivot.py
################################################################################
import numpy as np
import pandas as pd
import pytest

import pytae as pt
from pytae import cli


@pytest.fixture
def sales_df():
    return pd.DataFrame({
        "Region": ["East", "East", "West", "West", "East"],
        "Store": ["S1", "S2", "S1", "S2", "S1"],
        "Year": [2023, 2023, 2023, 2024, 2024],
        "Quarter": ["Q1", "Q2", "Q1", "Q1", "Q2"],
        "Sales": [100.0, 200.0, 300.0, 400.0, 150.0],
        "Profit": [10.0, 25.0, 35.0, 50.0, 20.0],
    })


def test_pivot_both_r_and_c(sales_df):
    res = sales_df.pt.pivot(r="Region", c="Year", v="Sales", a="sum")
    assert isinstance(res, pd.DataFrame)
    assert list(res.columns) == ["Region", "2023", "2024"]
    assert res.columns.name is None
    assert list(res.index) == [0, 1]

    east_row = res[res["Region"] == "East"].iloc[0]
    assert east_row["2023"] == 300.0  # 100 + 200
    assert east_row["2024"] == 150.0

    west_row = res[res["Region"] == "West"].iloc[0]
    assert west_row["2023"] == 300.0
    assert west_row["2024"] == 400.0


def test_pivot_r_only(sales_df):
    res = pt.pivot(sales_df, r="Region", v="Sales", a="sum")
    assert list(res.columns) == ["Region", "Sales"]
    assert res.columns.name is None
    assert len(res) == 2
    assert res.loc[res["Region"] == "East", "Sales"].iloc[0] == 450.0
    assert res.loc[res["Region"] == "West", "Sales"].iloc[0] == 700.0


def test_pivot_c_only(sales_df):
    res = sales_df.pt.pivot(c="Year", v="Sales", a="sum")
    assert list(res.columns) == ["2023", "2024"]
    assert len(res) == 1
    assert res["2023"].iloc[0] == 600.0
    assert res["2024"].iloc[0] == 550.0


def test_pivot_neither_r_nor_c(sales_df):
    res = sales_df.pt.pivot(v="Sales", a="sum")
    assert list(res.columns) == ["Sales"]
    assert len(res) == 1
    assert res["Sales"].iloc[0] == 1150.0


def test_pivot_multiple_rows(sales_df):
    res = sales_df.pt.pivot(r=["Region", "Store"], c="Year", v="Sales", a="sum")
    assert list(res.columns[:2]) == ["Region", "Store"]
    assert "2023" in res.columns
    assert "2024" in res.columns
    assert len(res) == 4  # East-S1, East-S2, West-S1, West-S2 -> 4 combinations
    assert res.columns.name is None


def test_pivot_multiple_values(sales_df):
    res = sales_df.pt.pivot(r="Region", c="Year", v=["Sales", "Profit"], a="sum")
    assert res.columns[0] == "Region"
    assert "Sales_2023" in res.columns
    assert "Profit_2023" in res.columns
    assert "Sales_2024" in res.columns
    assert "Profit_2024" in res.columns
    assert not isinstance(res.columns, pd.MultiIndex)


def test_pivot_multiple_columns(sales_df):
    res = sales_df.pt.pivot(r="Region", c=["Year", "Quarter"], v="Sales", a="sum")
    assert res.columns[0] == "Region"
    # Column names are flattened with underscore
    assert any("2023_Q1" in col for col in res.columns)
    assert not isinstance(res.columns, pd.MultiIndex)


def test_pivot_agg_n_alias(sales_df):
    res = sales_df.pt.pivot(r="Region", c="Year", v="Sales", a="n")
    assert list(res.columns) == ["Region", "2023", "2024"]
    assert res["2023"].dtype == "int64"
    assert res["2024"].dtype == "int64"
    east_row = res[res["Region"] == "East"].iloc[0]
    assert east_row["2023"] == 2  # two rows for East in 2023
    assert east_row["2024"] == 1


def test_pivot_agg_n_fill_zero_by_default():
    # If a combination has 0 occurrences, it should be 0 and int64
    df = pd.DataFrame({
        "Island": ["Biscoe", "Biscoe", "Dream"],
        "Species": ["Adelie", "Gentoo", "Chinstrap"],
        "Sex": ["M", "F", "M"],
    })
    res = df.pt.pivot(r="Island", c="Species", v="Sex", a="n")
    assert list(res.columns) == ["Island", "Adelie", "Chinstrap", "Gentoo"]
    assert res["Adelie"].dtype == "int64"
    assert res["Chinstrap"].dtype == "int64"
    assert res["Gentoo"].dtype == "int64"

    biscoe = res[res["Island"] == "Biscoe"].iloc[0]
    assert biscoe["Adelie"] == 1
    assert biscoe["Chinstrap"] == 0
    assert biscoe["Gentoo"] == 1

    dream = res[res["Island"] == "Dream"].iloc[0]
    assert dream["Adelie"] == 0
    assert dream["Chinstrap"] == 1
    assert dream["Gentoo"] == 0


def test_pivot_agg_n_r_only(sales_df):
    res = sales_df.pt.pivot(r="Region", v="Sales", a="n")
    assert list(res.columns) == ["Region", "Sales"]
    assert res["Sales"].dtype == "int64"
    assert res.loc[res["Region"] == "East", "Sales"].iloc[0] == 3
    assert res.loc[res["Region"] == "West", "Sales"].iloc[0] == 2


def test_pivot_agg_n_c_only(sales_df):
    res = sales_df.pt.pivot(c="Year", v="Sales", a="n")
    assert list(res.columns) == ["2023", "2024"]
    assert res["2023"].dtype == "int64"
    assert res["2024"].dtype == "int64"
    assert res["2023"].iloc[0] == 3
    assert res["2024"].iloc[0] == 2


def test_pivot_agg_n_neither_r_nor_c(sales_df):
    res = sales_df.pt.pivot(v="Sales", a="n")
    assert list(res.columns) == ["Sales"]
    assert res["Sales"].dtype == "int64"
    assert res["Sales"].iloc[0] == 5


def test_pivot_agg_n_explicit_nan_fill():
    df = pd.DataFrame({
        "Island": ["Biscoe", "Dream"],
        "Species": ["Adelie", "Chinstrap"],
        "Sex": ["M", "F"],
    })
    res = df.pt.pivot(r="Island", c="Species", v="Sex", a="n", fill_value=np.nan)
    # When explicit NaN fill is requested, missing combinations are <NA> of Int64 dtype
    assert res["Adelie"].dtype == "Int64"
    assert res["Chinstrap"].dtype == "Int64"
    assert pd.isna(res.loc[res["Island"] == "Dream", "Adelie"].iloc[0])
    assert res.loc[res["Island"] == "Dream", "Chinstrap"].iloc[0] == 1


def test_pivot_fill_value(sales_df):
    res = sales_df.pt.pivot(r=["Region", "Store"], c="Year", v="Sales", a="sum", fill_value=0.0)
    # Check that missing combinations have 0.0 instead of NaN
    assert not res.isna().any().any()


def test_pivot_dropna_false_default():
    df_na = pd.DataFrame({
        "Region": ["East", np.nan, "West"],
        "Year": [2023, 2023, np.nan],
        "Sales": [10.0, 20.0, 30.0],
    })
    res = df_na.pt.pivot(r="Region", c="Year", v="Sales", a="sum")
    # Region has NaN, so there should be a row with NaN in Region
    assert res["Region"].isna().any()
    # Year has NaN, so there should be a column representing nan
    col_names = [str(c) for c in res.columns]
    assert any("nan" in c.lower() for c in col_names)


def test_pivot_aliases(sales_df):
    res = sales_df.pt.pivot(rows="Region", cols="Year", values="Sales", aggfunc="mean", fill=0.0)
    assert list(res.columns) == ["Region", "2023", "2024"]


def test_pivot_missing_v_error(sales_df):
    with pytest.raises(ValueError, match="value column 'v' is required"):
        sales_df.pt.pivot(r="Region", c="Year")


def test_pivot_unknown_column_error(sales_df):
    with pytest.raises(KeyError, match="not found in DataFrame"):
        sales_df.pt.pivot(r="NonExistent", c="Year", v="Sales")


def test_wide_accepts_r_alias(sales_df):
    # wide accepts r as alias for index/by
    tall = sales_df.pt.long(cols=["Sales", "Profit"], id_vars=["Region", "Store", "Year"])
    back = tall.pt.wide(r=["Region", "Store", "Year"], c="variable", v="value")
    assert "Sales" in back.columns
    assert "Profit" in back.columns


# ---------------- CLI Tests ----------------


def test_cli_pivot_basic(tmp_path, capsys):
    df = pd.DataFrame({
        "Region": ["East", "East", "West"],
        "Year": [2023, 2024, 2023],
        "Sales": [100, 150, 200],
    })
    path = str(tmp_path / "data.csv")
    df.to_csv(path, index=False)

    exit_code = cli.main([path, "-pivot", "r=Region,c=Year,v=Sales,a=sum"])
    assert exit_code == 0
    captured = capsys.readouterr()
    assert "East" in captured.out
    assert "West" in captured.out
    assert "2023" in captured.out
    assert "2024" in captured.out


def test_cli_pivot_r_only(tmp_path, capsys):
    df = pd.DataFrame({
        "Region": ["East", "East", "West"],
        "Sales": [100, 150, 200],
    })
    path = str(tmp_path / "data.csv")
    df.to_csv(path, index=False)

    exit_code = cli.main([path, "-pivot", "r=Region,v=Sales,a=sum"])
    assert exit_code == 0
    captured = capsys.readouterr()
    assert "East" in captured.out
    assert "250" in captured.out
    assert "200" in captured.out


def test_cli_pivot_c_only(tmp_path, capsys):
    df = pd.DataFrame({
        "Year": [2023, 2024, 2023],
        "Sales": [100, 150, 200],
    })
    path = str(tmp_path / "data.csv")
    df.to_csv(path, index=False)

    exit_code = cli.main([path, "-pivot", "c=Year,v=Sales,a=sum"])
    assert exit_code == 0
    captured = capsys.readouterr()
    assert "300" in captured.out
    assert "150" in captured.out


def test_cli_pivot_unknown_column(tmp_path):
    df = pd.DataFrame({"Region": ["East"], "Sales": [100]})
    path = str(tmp_path / "data.csv")
    df.to_csv(path, index=False)

    with pytest.raises(SystemExit) as exc:
        cli.main([path, "-pivot", "r=NoSuchCol,v=Sales"])
    assert exc.value.code == 2


def test_cli_pivot_export_parquet(tmp_path):
    df = pd.DataFrame({
        "Region": ["East", "West"],
        "Year": [2023, 2024],
        "Sales": [100, 200],
    })
    in_path = str(tmp_path / "in.csv")
    out_path = str(tmp_path / "out.parquet")
    df.to_csv(in_path, index=False)

    exit_code = cli.main([in_path, "-pivot", "r=Region,c=Year,v=Sales", "-o", out_path])
    assert exit_code == 0

    pivoted_df = pd.read_parquet(out_path)
    assert "Region" in pivoted_df.columns
    assert "2023" in pivoted_df.columns
    assert "2024" in pivoted_df.columns


def test_pivot_no_cartesian_product_unobserved_combinations():
    doc = pd.DataFrame({
        "Region": ["East", "East", "East", "East", "West", "West"],
        "Store": ["Store A", "Store A", "Store B", "Store B", "Store C", "Store C"],
        "Year": [2023, 2024, 2023, 2024, 2023, 2024],
        "Sales": [700.0, 850.0, 800.0, 950.0, 900.0, 1100.0],
    })
    res_sum = doc.pt.pivot(r=["Region", "Store"], c="Year", v="Sales", a="sum")
    # Only observed Region-Store pairs should exist (3 rows), not 6
    assert len(res_sum) == 3
    assert set(res_sum["Store"]) == {"Store A", "Store B", "Store C"}

    res_n = doc.pt.pivot(r=["Region", "Store"], c="Year", v="Sales", a="n")
    assert len(res_n) == 3


def test_pivot_multi_c_cols_a_n_preserves_outer_keys():
    df = pd.DataFrame({
        "Region": ["East", "East", "East", "East"],
        "Year": [2023, 2023, 2023, 2024],
        "Quarter": ["Q1", "Q1", "Q2", "Q1"],
        "Sales": [1, 1, 1, 1],
    })
    out = df.pt.pivot(r="Region", c=["Year", "Quarter"], v="Sales", a="n")
    assert "2023_Q1" in out.columns
    assert "2023_Q2" in out.columns
    assert "2024_Q1" in out.columns
    assert out["2023_Q1"].iloc[0] == 2
    assert out["2023_Q2"].iloc[0] == 1
    assert out["2024_Q1"].iloc[0] == 1


def test_pivot_header_deduplication():
    df1 = pd.DataFrame({"r": ["a", "a"], "c": [1, "1"], "v": [10.0, 20.0]})
    with pytest.raises(ValueError, match="contain duplicate names"):
        df1.pt.pivot(r="r", c="c", v="v", a="sum")

    df2 = pd.DataFrame({"r": ["a", "a", "a"], "c": [np.nan, "nan", "ok"], "v": [1.0, 2.0, 3.0]})
    with pytest.raises(ValueError, match="contain duplicate names"):
        df2.pt.pivot(r="r", c="c", v="v", a="sum")


def test_pivot_c_only_multi_values_keeps_metric_index():
    df = pd.DataFrame({
        "Year": [2023, 2024, 2023],
        "Sales": [100.0, 150.0, 200.0],
        "Profit": [10.0, 15.0, 20.0],
    })
    out = df.pt.pivot(c="Year", v=["Sales", "Profit"], a="sum")
    assert "metric" in out.columns
    assert set(out["metric"]) == {"Sales", "Profit"}


def test_pivot_row_only_a_n_count_key_collision():
    df1 = pd.DataFrame({"island": ["A", "A", "B"]})
    out1 = df1.pt.pivot(r="island", a="n")
    assert "island" in out1.columns
    assert "n" in out1.columns
    assert list(out1["n"]) == [2, 1]

    # If v is specified and overlaps with r, raise ValueError
    with pytest.raises(ValueError, match="cannot also be in grouping dimensions"):
        df1.pt.pivot(r="island", v="island", a="n")

    # If row grouping column is named "n", omitting v will default to "n" which collides
    df_colliding = pd.DataFrame({"n": ["a", "a", "b"], "val": [1, 2, 3]})
    with pytest.raises(ValueError, match="collides with row grouping column"):
        df_colliding.pt.pivot(r="n", a="n")

    # Providing a non-colliding column for v works
    out_custom = df_colliding.pt.pivot(r="n", v="val", a="n")
    assert list(out_custom.columns) == ["n", "val"]


def test_pivot_v_in_dimensions_validation():
    df = pd.DataFrame({"Region": ["East", "West"], "Year": [2023, 2024], "Sales": [10, 20]})
    with pytest.raises(ValueError, match="cannot also be in grouping dimensions"):
        df.pt.pivot(r="Region", c="Year", v="Region", a="sum")


def test_cli_pivot_fill_numeric(tmp_path, capsys):
    df = pd.DataFrame({
        "Region": ["East", "West"],
        "Store": ["A", "B"],
        "Sales": [4.0, 2.0],
    })
    path = str(tmp_path / "sales.csv")
    df.to_csv(path, index=False)
    exit_code = cli.main([path, "-pivot", "r=Region,c=Store,v=Sales,a=sum,fill=0", "-round", "1"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "0" in out


def test_cli_pivot_unknown_agg_fails_cleanly(tmp_path):
    df = pd.DataFrame({"Region": ["East"], "Sales": [100]})
    path = str(tmp_path / "sales.csv")
    df.to_csv(path, index=False)
    with pytest.raises(SystemExit) as exc:
        cli.main([path, "-pivot", "r=Region,v=Sales,a=nope"])
    assert exc.value.code == 2


def test_wide_cols_and_values_aliases(tmp_path):
    df = pd.DataFrame({
        "id": ["a", "b"],
        "country": ["sg", "cn"],
        "balance": [10, 20],
    })
    w1 = df.pt.wide(cols="country", values="balance")
    assert "cn" in w1.columns and "sg" in w1.columns

    path = str(tmp_path / "data.csv")
    df.to_csv(path, index=False)
    exit_code = cli.main([path, "-wide", "cols=country,values=balance"])
    assert exit_code == 0


def test_cli_value_counts_with_existing_count_column(tmp_path, capsys):
    df = pd.DataFrame({"count": ["a", "a", "b"]})
    path = str(tmp_path / "counts.csv")
    df.to_csv(path, index=False)
    exit_code = cli.main([path, "-value_counts"])
    assert exit_code == 0
    out = capsys.readouterr().out
    assert "count" in out


def test_pivot_multi_key_empty_string_component():
    # Issue 3: multi-key pivot keeps empty string level rather than dropping it
    df = pd.DataFrame({
        "r": ["r1", "r1", "r1"],
        "y": [2023, 2023, 2024],
        "q": ["", "Q1", "Q2"],
        "v": [15, 10, 20],
    })
    piv = df.pt.pivot(r="r", c=["y", "q"], v="v", a="sum")
    assert "2023_" in piv.columns
    assert "2023_Q1" in piv.columns
    assert piv.loc[piv["r"] == "r1", "2023_"].iloc[0] == 15


def test_pivot_repeated_columns_in_dimensions_and_values():
    # Issue 3: Repeated column in pivot dimensions or values
    df = pd.DataFrame({
        "Region": ["East", "West"],
        "Year": [2023, 2024],
        "Sales": [100, 200],
    })
    with pytest.raises(ValueError, match="cannot appear multiple times in grouping dimensions"):
        df.pt.pivot(r="Region", c="Region", v="Sales")

    with pytest.raises(ValueError, match="cannot appear multiple times in grouping dimensions"):
        df.pt.pivot(r=["Region", "Region"], v="Sales")

    with pytest.raises(ValueError, match="cannot appear multiple times in grouping dimensions"):
        df.pt.pivot(c=["Year", "Year"], v="Sales")

    with pytest.raises(ValueError, match="cannot appear multiple times in 'v'"):
        df.pt.pivot(r="Region", v=["Sales", "Sales"])

################################################################################
# FILE: tests/test_plotter.py
################################################################################
import os
import sys

import matplotlib
import pandas as pd
import pytest

matplotlib.use("Agg")

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt
from pytae.plotting import Plotter


def test_plotter_constructs_default_axis():
    plotter = Plotter()
    assert "A" in plotter.axd
    assert plotter.fig is not None


def test_pie_plot_does_not_overwrite_plotter_df():
    df = pd.DataFrame({"cat": ["a", "b"], "val": [1, 2]})
    plotter = Plotter().data(df)
    plotter.plot(x="cat", y="val", kind="pie", by="cat")
    pd.testing.assert_frame_equal(plotter.df, df)


def test_unknown_mosaic_on_key_raises():
    df = pd.DataFrame({"cat": ["a", "b"], "val": [1, 2]})
    plotter = Plotter(mosaic="AB").data(df)
    with pytest.raises(ValueError, match="Unknown mosaic key"):
        plotter.plot(x="cat", y="val", kind="bar", on="Z")


def _facet_frame():
    return pd.DataFrame(
        {
            "species": ["Adelie"] * 3 + ["Gentoo"] * 3 + ["Chinstrap"] * 3 + ["Emperor"] * 3 + ["King"] * 3,
            "x": range(15),
            "y": [v * 2 for v in range(15)],
        }
    )


def test_facet_leaves_unfillable_grid_cells_blank():
    # 5 groups in a 2-column grid needs 3 rows (6 cells) -- the leftover cell must
    # stay blank, not become an empty/unused 6th axis.
    plotter = Plotter.facet(_facet_frame(), by="species", ncols=2, x="x", y="y", kind="line")
    assert len(plotter.axd) == 5
    assert len(plotter.fig.axes) == 5


def test_facet_one_axis_per_group_with_own_data():
    df = _facet_frame()
    plotter = Plotter.facet(df, by="species", ncols=2, x="x", y="y", kind="line")
    assert set(plotter.axd.keys()) == {"Adelie", "Gentoo", "Chinstrap", "Emperor", "King"}


def test_facet_sets_axis_titles_by_default():
    plotter = Plotter.facet(_facet_frame(), by="species", ncols=2, x="x", y="y", kind="line")
    assert plotter.axd["Adelie"].get_title() == "Adelie"


def test_facet_titles_can_be_disabled():
    plotter = Plotter.facet(_facet_frame(), by="species", ncols=2, x="x", y="y", kind="line", titles=False)
    assert plotter.axd["Adelie"].get_title() == ""


def test_facet_respects_categorical_order():
    df = _facet_frame()
    df["species"] = pd.Categorical(df["species"], categories=["King", "Emperor", "Gentoo", "Chinstrap", "Adelie"])
    plotter = Plotter.facet(df, by="species", x="x", y="y", kind="line")
    assert list(plotter.axd.keys()) == ["King", "Emperor", "Gentoo", "Chinstrap", "Adelie"]


def test_facet_default_ncols_is_roughly_square():
    plotter = Plotter.facet(_facet_frame(), by="species", x="x", y="y", kind="line")
    assert len(plotter.axd) == 5


def test_facet_no_groups_raises():
    df = _facet_frame()
    with pytest.raises(ValueError, match="no groups found"):
        Plotter.facet(df[0:0], by="species", x="x", y="y", kind="line")


def test_facet_chains_into_finalize():
    plotter = Plotter.facet(_facet_frame(), by="species", ncols=2, x="x", y="y", kind="line")
    assert plotter.finalize() is plotter


def test_facet_restores_full_df_not_just_last_group():
    df = _facet_frame()
    plotter = Plotter.facet(df, by="species", ncols=2, x="x", y="y", kind="line")
    pd.testing.assert_frame_equal(plotter.df, df)


def test_facet_rejects_ncols_less_than_one():
    with pytest.raises(ValueError, match="ncols must be at least 1"):
        Plotter.facet(_facet_frame(), by="species", ncols=0, x="x", y="y", kind="line")


def test_kde_column_without_by_works():
    df = pd.DataFrame({"a": [1.0, 2.0, 3.0, 4.0]})
    plotter = Plotter().data(df).plot(kind="kde", column="a")
    assert plotter.ax is not None


def test_scatter_missing_required_kwargs_raises():
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    with pytest.raises(ValueError, match="kind='scatter' needs y="):
        Plotter().data(df).plot(kind="scatter", x="a")


def test_pie_missing_required_kwargs_raises():
    df = pd.DataFrame({"cat": ["a", "b"], "val": [1, 2]})
    with pytest.raises(ValueError, match="kind='pie' needs by="):
        Plotter().data(df).plot(kind="pie", y="val")


def test_hist_missing_required_kwarg_raises():
    df = pd.DataFrame({"a": [1, 2]})
    with pytest.raises(ValueError, match="kind='hist' needs column="):
        Plotter().data(df).plot(kind="hist")


def test_scatter_groups_by_category_with_labels():
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4], "cat": ["x", "y"]})
    plotter = Plotter().data(df).plot(kind="scatter", x="a", y="b", by="cat")
    handles, labels = plotter.ax.get_legend_handles_labels()
    assert set(labels) == {"x", "y"}


def test_aggfunc_raises_helpful_error():
    df = pd.DataFrame({"a": [1, 2, 3]})
    with pytest.raises(ValueError, match="Plotter does not perform aggregation"):
        Plotter().data(df).plot(kind="hist", column="a", aggfunc="mean")


def test_supported_kwargs_lists_unsupported_and_controls():
    info = Plotter.supported_kwargs("scatter")
    assert set(info["unsupported"]) == {"dropna"}
    assert "on" in info["pytae_controls"]
    assert "palette" in info["pytae_controls"]


def test_supported_kwargs_kde_alias_matches_density():
    assert Plotter.supported_kwargs("kde") == Plotter.supported_kwargs("density")



def test_save_writes_figure(tmp_path):
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    out = tmp_path / "out.png"
    plotter = Plotter().data(df).plot(kind="scatter", x="a", y="b").finalize().save(str(out))
    assert out.exists()
    assert isinstance(plotter, Plotter)


def test_finalize_style_false_skips_spine_and_tick_cleanup():
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    plotter = Plotter(mosaic="AB").data(df)
    plotter.plot(kind="scatter", x="a", y="b", on="A")
    plotter.finalize(style=False)
    unused_ax = plotter.axd["B"]
    # style=False means the blank second axis keeps matplotlib's own default
    # ticks/spines instead of pytae hiding them.
    assert unused_ax.spines["bottom"].get_visible()


def test_finalize_style_true_hides_blank_axis_ticks():
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    plotter = Plotter(mosaic="AB").data(df)
    plotter.plot(kind="scatter", x="a", y="b", on="A")
    plotter.finalize()
    unused_ax = plotter.axd["B"]
    assert not unused_ax.has_data()
    assert unused_ax.xaxis.get_tick_params()["length"] == 0


def test_plotter_init_with_dataframe():
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    p1 = Plotter(df)
    assert p1.df is df
    assert "A" in p1.axd

    p2 = Plotter(df, mosaic="AB")
    assert p2.df is df
    assert "B" in p2.axd


def test_df_pt_plot_accessor():
    import pytae  # noqa: F401 registers df.pt
    df = pd.DataFrame({"x": ["a", "b", "c"], "y": [10, 20, 30]})
    p = df.pt.plot(kind="bar", x="x", y="y").finalize()
    assert isinstance(p, Plotter)
    assert p.ax is not None


def test_box_plot_distribution_not_aggregated():
    df = pd.DataFrame({
        "group": ["A", "A", "A", "B", "B", "B"],
        "val": [1, 2, 9, 10, 11, 20],
    })
    p = Plotter(df).plot(kind="box", x="group", y="val")
    box_data = p.get_data("A")
    assert set(box_data.columns) == {"A", "B"}
    assert len(box_data["A"].dropna()) == 3
    assert len(box_data["B"].dropna()) == 3


def test_kde_hist_accepts_x_as_column_alias():
    df = pd.DataFrame({"val": [1.0, 2.0, 3.0, 4.0, 5.0]})
    # hist with x= instead of column=
    p_hist = Plotter(df).plot(kind="hist", x="val")
    assert p_hist.ax is not None

    # kde with x= instead of column=
    p_kde = Plotter(df).plot(kind="kde", x="val")
    assert p_kde.ax is not None


def test_palette_and_axis_labels():
    df = pd.DataFrame({"cat": ["x", "y"], "v1": [1, 2], "v2": [3, 4]})
    p = Plotter(df).plot(kind="bar", x="cat", y="v1", palette="tab10", xlabel="Category", ylabel="Value")
    assert p.ax.get_xlabel() == "Category"
    assert p.ax.get_ylabel() == "Value"


def test_heatmap_rendering():
    df = pd.DataFrame({"x": ["A", "A", "B", "B"], "y": ["C", "D", "C", "D"], "v": [1, 2, 3, 4]})
    p = Plotter(df).plot(kind="heatmap", x="x", y="v", by="y", annot=True)
    assert p.ax is not None
    table = p.get_data("A")
    assert table.shape == (2, 2)


def test_repr_png_returns_bytes():
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    p = Plotter(df).plot(kind="scatter", x="a", y="b").finalize()
    png_bytes = p._repr_png_()
    assert isinstance(png_bytes, bytes)
    assert png_bytes.startswith(b"\x89PNG")


def test_facet_sharex_sharey():
    df = pd.DataFrame({"group": ["g1", "g2"], "x": [1, 2], "y": [3, 4]})
    p = Plotter.facet(df, by="group", x="x", y="y", kind="line", sharex=True, sharey=True)
    assert len(p.axd) == 2


def test_pt_plot_top_level_function():
    import pytae as pt
    df = pd.DataFrame({"day": ["Thur", "Fri"], "bill": [10, 20]})
    p = pt.plot(df, kind="bar", x="day", y="bill").finalize()
    assert isinstance(p, Plotter)
    assert p.ax is not None


def test_plotter_nrows_ncols_without_mosaic():
    df = pd.DataFrame({"x": [1, 2], "y": [3, 4]})
    p = Plotter(df, nrows=1, ncols=2)
    assert set(p.axd.keys()) == {"A", "B"}
    p.plot(on="A", kind="scatter", x="x", y="y")
    p.plot(on="B", kind="line", x="x", y="y")
    p.finalize()
    assert len(p.axd) == 2


def test_scatter_continuous_c_preserves_numeric_ticks():
    df = pd.DataFrame({
        "x": list(range(50)),
        "y": list(range(50)),
        "val": [float(i * 1.5) for i in range(50)],
    })
    p = Plotter(df).plot(kind="scatter", x="x", y="y", c="val", cmap="viridis").finalize()
    # Colorbar axis should have normal small number of ticks, not 50 category ticks
    cb_axes = [ax for ax in p.fig.axes if ax != p.axd["A"]]
    assert len(cb_axes) == 1
    assert len(cb_axes[0].get_yticks()) < 15


def test_pt_finalize_function_explicit_and_implicit():
    import pytae as pt
    df = pd.DataFrame({"day": ["Thur", "Fri"], "bill": [10, 20]})

    # 1. Explicit plotter passed
    p1 = pt.plot(df, kind="bar", x="day", y="bill")
    res1 = pt.finalize(p1)
    assert res1 is p1
    assert not p1.ax.spines["top"].get_visible()

    # 2. Implicit active plotter (zero args)
    p2 = pt.plot(df, kind="bar", x="day", y="bill")
    res2 = pt.finalize()
    assert res2 is p2
    assert not p2.ax.spines["top"].get_visible()

    # 3. Type error when invalid object passed
    with pytest.raises(TypeError, match="Expected a Plotter instance"):
        pt.finalize("not a plotter")

    # 4. Error when no active plotter exists
    Plotter._last_active = None
    with pytest.raises(ValueError, match="No active pytae Plotter"):
        pt.finalize()


def test_plotter_pt_namespace_chaining():
    import pytae as pt
    df = pd.DataFrame({"day": ["Thur", "Fri"], "bill": [10, 20]})

    # Fluent .pt namespace throughout the entire pipeline
    p = df.pt.plot(kind="bar", x="day", y="bill").pt.finalize()
    assert isinstance(p, pt.Plotter)
    assert not p.ax.spines["top"].get_visible()


def test_control_kwargs_secondary_y_and_print_clip(capsys):
    import pytae as pt
    df = pd.DataFrame({"day": ["Thur", "Fri"], "bill": [10, 20], "tip": [2, 4]})

    # secondary_y=True should automatically route to A^
    p = (
        pt.Plotter(df, figsize=(8, 4))
        .plot(on="A", kind="bar", x="day", y="bill")
        .plot(kind="line", x="day", y="tip", secondary_y=True, print_data=True, clip_data=True)
        .finalize()
    )
    assert "A^" in p.axd
    captured = capsys.readouterr()
    assert "tip" in captured.out


def test_auto_facet_by_and_ncols():
    import pytae as pt
    df = pd.DataFrame({
        "grp": ["g1", "g1", "g2", "g2", "g3", "g3"],
        "x": [1, 2, 1, 2, 1, 2],
        "y": [10, 20, 15, 25, 12, 22],
    })

    # 1. Top-level pt.plot with by and ncols auto-facets
    p1 = pt.plot(df, by="grp", ncols=3, kind="scatter", x="x", y="y", title="Facet Title").finalize()
    assert set(p1.axd.keys()) == {"g1", "g2", "g3"}
    assert p1.fig._suptitle.get_text() == "Facet Title"

    # 2. DataFrame accessor df.pt.plot with by and ncols auto-facets
    p2 = df.pt.plot(by="grp", ncols=3, kind="scatter", x="x", y="y").pt.finalize()
    assert set(p2.axd.keys()) == {"g1", "g2", "g3"}

    # 3. Class constructor Plotter(df, by=..., ncols=...) auto-facets
    p3 = pt.Plotter(df, by="grp", ncols=3, kind="scatter", x="x", y="y").finalize()
    assert set(p3.axd.keys()) == {"g1", "g2", "g3"}

    # 4. col="grp" alias auto-facets
    p4 = pt.plot(df, col="grp", ncols=3, kind="scatter", x="x", y="y").finalize()
    assert set(p4.axd.keys()) == {"g1", "g2", "g3"}

    # 5. Non-faceted subplots (no by) still creates standard panels A, B
    p5 = pt.Plotter(df, nrows=1, ncols=2)
    assert set(p5.axd.keys()) == {"A", "B"}


def test_plot_line_color_string_and_list():
    import matplotlib.colors as mcolors

    import pytae as pt
    df = pd.DataFrame({"day": ["Thur", "Fri", "Sat"], "bill": [10, 20, 30], "tip": [2, 4, 6]})

    # String color
    p = pt.Plotter(df).plot(kind="line", x="day", y="bill", color="crimson").finalize()
    line_color = p.axd["A"].get_lines()[0].get_color()
    assert mcolors.to_hex(line_color) == mcolors.to_hex("crimson")

    # List of colors for multiple lines
    df_multi = pd.DataFrame({"day": ["Thur", "Fri"], "A": [1, 2], "B": [3, 4]})
    p2 = pt.Plotter(df_multi).plot(kind="line", x="day", y=["A", "B"], color=["crimson", "navy"]).finalize()
    lines = p2.axd["A"].get_lines()
    assert mcolors.to_hex(lines[0].get_color()) == mcolors.to_hex("crimson")
    assert mcolors.to_hex(lines[1].get_color()) == mcolors.to_hex("navy")


def test_plot_scatter_by_title_labels_kwargs():
    import pytae as pt
    df = pd.DataFrame({
        "species": ["Adelie", "Gentoo", "Adelie", "Gentoo"],
        "bill_len": [39.1, 46.5, 40.3, 48.0],
        "bill_dep": [18.7, 14.5, 18.0, 15.0],
        "size_col": [20, 40, 25, 45],
    })

    p = (
        pt.Plotter(df)
        .plot(
            kind="scatter",
            x="bill_len",
            y="bill_dep",
            by="species",
            title="Penguin Bill Dimensions",
            s=df["size_col"],
            edgecolors="black",
        )
        .finalize()
    )
    ax = p.axd["A"]
    assert ax.get_title() == "Penguin Bill Dimensions"
    assert ax.get_xlabel() == "bill_len"
    assert ax.get_ylabel() == "bill_dep"
    # Check collections have edgecolors set
    for coll in ax.collections:
        assert len(coll.get_edgecolors()) > 0


def test_heatmap_title_vmin_vmax_and_contrast():
    import pytae as pt
    df = pd.DataFrame({
        "A": [1.0, 0.0, -1.0],
        "B": [0.0, 1.0, 0.5],
        "C": [-1.0, 0.5, 1.0],
    })
    p = pt.Plotter(df).plot(
        kind="heatmap",
        annot=True,
        fmt=".2f",
        cmap="coolwarm",
        title="Correlation Heatmap",
        vmin=-1,
        vmax=1,
    ).finalize()
    ax = p.axd["A"]
    assert ax.get_title() == "Correlation Heatmap"
    images = ax.get_images()
    assert len(images) == 1
    assert images[0].get_clim() == (-1.0, 1.0)
    # Check text contrast: text annotations exist and contain black or white
    texts = ax.texts
    assert len(texts) == 9
    colors = {t.get_color() for t in texts}
    assert "white" in colors or "black" in colors


def test_unknown_palette_warning():
    import pytest

    import pytae as pt
    df = pd.DataFrame({"day": ["Thur", "Fri"], "bill": [10, 20]})
    with pytest.warns(UserWarning, match="Unknown palette 'non_existent_palette'"):
        pt.Plotter(df).plot(kind="bar", x="day", y="bill", palette="non_existent_palette").finalize()


def test_facet_layout_kwargs_conflict_error():
    import pytest

    import pytae as pt
    df = pd.DataFrame({"grp": ["g1", "g2"], "x": [1, 2], "y": [3, 4]})

    with pytest.raises(ValueError, match="Cannot combine 'mosaic' with faceting"):
        pt.Plotter(df, mosaic="AB", ncols=2, by="grp", kind="line", x="x", y="y")

    with pytest.raises(ValueError, match="Cannot combine 'nrows' with faceting"):
        df.pt.plot(nrows=1, ncols=2, by="grp", kind="bar", x="x", y="y")

    with pytest.raises(ValueError, match="Cannot combine 'mosaic' with faceting"):
        pt.Plotter(df, "AB", by="grp", ncols=2, kind="scatter", x="x", y="y")

    with pytest.raises(ValueError, match="Cannot combine 'mosaic' with faceting"):
        pt.Plotter.facet(df, by="grp", ncols=2, mosaic="AB", kind="line", x="x", y="y")

    with pytest.raises(ValueError, match="Cannot combine 'nrows' with faceting"):
        pt.Plotter.facet(df, by="grp", ncols=2, nrows=1, kind="line", x="x", y="y")


def test_plotter_init_immediate_plot_and_unexpected_kwargs():
    import pytest

    import pytae as pt
    df = pd.DataFrame({"day": ["Thur", "Fri"], "bill": [10, 20]})

    # Passing plot kwargs to Plotter constructor draws immediately
    p = pt.Plotter(df, kind="bar", x="day", y="bill")
    assert p.axd["A"].has_data()

    # Unexpected kwargs raise TypeError
    with pytest.raises(TypeError, match="unexpected keyword argument"):
        pt.Plotter(df, nonexistent_kwarg="invalid")


def test_box_palette_colors():
    import pytae as pt
    df = pd.DataFrame({
        "species": ["Adelie", "Gentoo", "Adelie", "Gentoo"],
        "body_mass": [3000, 4500, 3200, 4700],
    })
    p = pt.Plotter(df).plot(kind="box", x="species", y="body_mass", palette="Set1").finalize()
    ax = p.axd["A"]
    assert len(ax.patches) == 2
    # Verify patches have facecolors assigned from palette
    c0 = ax.patches[0].get_facecolor()
    c1 = ax.patches[1].get_facecolor()
    assert c0 != c1


def test_plotter_finalize_reapplies_rot_and_fontsize():
    df = pd.DataFrame({"day": ["Thursday", "Friday", "Saturday"], "bill": [10, 20, 30]})
    p = pt.Plotter(df).plot(kind="bar", x="day", y="bill", rot=45, fontsize=8).finalize()
    ax = p.axd["A"]
    labels = ax.get_xticklabels()
    assert labels[0].get_rotation() == 45
    assert labels[0].get_fontsize() == 8


def test_plotter_manage_legend_respects_legend_false():
    df = pd.DataFrame({"day": ["Thur", "Fri"], "bill": [10, 20], "sex": ["M", "F"]})
    p = pt.Plotter(df).plot(kind="bar", x="day", y="bill", by="sex", legend=False).finalize()
    ax = p.axd["A"]
    assert ax.get_legend() is None


def test_plotter_default_kind_finalize_does_not_raise_keyerror():
    df = pd.DataFrame({"a": [1, 2, 3], "b": [3, 1, 2]})
    p = pt.Plotter(df).plot(x="a", y="b").finalize()
    assert p.axd["A"].has_data()

    df_facet = pd.DataFrame({"g": ["x", "y"], "x": [1, 2], "y": [3, 4]})
    pf = pt.Plotter.facet(df_facet, by="g", x="x", y="y").finalize()
    assert len(pf.axd) == 2


def test_plotter_numeric_x_dtype_preserved_for_line_and_area():
    df = pd.DataFrame({"x": [1, 2, 10], "y": [1, 4, 2]})
    p_line = pt.Plotter(df).plot(kind="line", x="x", y="y")
    data_line = p_line.get_data("A")
    assert pd.api.types.is_numeric_dtype(data_line["x"])

    p_area = pt.Plotter(df).plot(kind="area", x="x", y="y")
    data_area = p_area.get_data("A")
    assert pd.api.types.is_numeric_dtype(data_area["x"])


def test_plotter_line_style_and_width_handling():
    df = pd.DataFrame({"x": [1, 2], "A": [10, 20], "B": [30, 40]})
    # String style and int width
    p1 = pt.Plotter(df).plot(kind="line", x="x", y=["A", "B"], style="--", width=3)
    lines1 = p1.axd["A"].get_lines()
    for line_item in lines1:
        assert line_item.get_linestyle() == "--"
        assert line_item.get_linewidth() == 3

    # Dict style and dict width mapped by label
    p2 = pt.Plotter(df).plot(kind="line", x="x", y=["A", "B"], style={"B": ":", "A": "-."}, width={"B": 5, "A": 1})
    lines2 = p2.axd["A"].get_lines()
    for line_item in lines2:
        if line_item.get_label() == "A":
            assert line_item.get_linestyle() == "-."
            assert line_item.get_linewidth() == 1
        elif line_item.get_label() == "B":
            assert line_item.get_linestyle() == ":"
            assert line_item.get_linewidth() == 5


def test_plotter_grouped_scatter_xlim_ylim_rot_fontsize_s_col():
    df = pd.DataFrame({
        "x": [1, 2, 3, 4],
        "y": [10, 20, 30, 40],
        "g": ["a", "a", "b", "b"],
        "sz": [15, 25, 35, 45],
    })
    p = pt.Plotter(df).plot(kind="scatter", x="x", y="y", by="g", s="sz", xlim=(0, 10), ylim=(0, 50), rot=45, fontsize=9)
    ax = p.axd["A"]
    assert ax.get_xlim() == (0, 10)
    assert ax.get_ylim() == (0, 50)
    assert len(ax.collections) == 2


def test_plotter_color_string_and_box_color_list():
    df = pd.DataFrame({
        "g": ["a", "b", "a", "b"],
        "v": [10, 20, 30, 40],
    })
    # Grouped scatter with color="crimson" string
    p1 = pt.Plotter(df).plot(kind="scatter", x="v", y="v", by="g", color="crimson")
    assert len(p1.axd["A"].collections) == 2

    # Box plot with color list
    p2 = pt.Plotter(df).plot(kind="box", x="g", y="v", color=["crimson", "navy"])
    assert len(p2.axd["A"].patches) == 2


def test_plotter_heatmap_numpy_int_annotation_formatting():
    import numpy as np
    df = pd.DataFrame({"a": [np.int64(0), np.int64(100)], "b": [np.int64(100), np.int64(0)]})
    p = pt.Plotter(df).plot(kind="heatmap", annot=True, fmt=".2f", cmap="viridis")
    texts = [t.get_text() for t in p.axd["A"].texts]
    assert "100.00" in texts
    assert "0.00" in texts


def test_plotter_ungrouped_scatter_color_and_array_c():
    import numpy as np
    df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
    # c is string color name not in columns
    p1 = pt.Plotter(df).plot(kind="scatter", x="x", y="y", c="red")
    assert p1.axd["A"].has_data()

    # c is numpy array
    c_arr = np.array([0.1, 0.5, 0.9])
    p2 = pt.Plotter(df).plot(kind="scatter", x="x", y="y", c=c_arr)
    assert p2.axd["A"].has_data()


def test_plotter_pie_colors_list():
    df = pd.DataFrame({"cat": ["a", "b"], "val": [1, 2]})
    p = pt.Plotter(df).plot(kind="pie", by="cat", y="val", colors=["red", "blue"])
    assert p.axd["A"].has_data()


def test_plotter_list_y_order_and_aggregate_error():
    df = pd.DataFrame({"q": ["Q1", "Q2"], "rev": [100, 200], "cost": [40, 80]})
    p1 = pt.Plotter(df).plot(kind="bar", x="q", y=["rev", "cost"])
    assert list(p1.get_data("A").columns) == ["q", "rev", "cost"]

    # Passing aggregate raises ValueError
    with pytest.raises(ValueError, match="Plotter does not perform aggregation"):
        pt.Plotter(df).plot(kind="bar", x="q", y=["rev", "cost"], aggregate=True)


def test_plotter_hist_and_kde_with_by_on_duplicate_index():
    df = pd.DataFrame({"g": ["a", "a", "b", "b"], "v": [1.0, 2.0, 3.0, 4.0]}, index=[0, 0, 1, 1])
    # hist with duplicate index
    p_hist = pt.Plotter(df).plot(kind="hist", column="v", by="g")
    assert p_hist.axd["A"].has_data()

    # kde with duplicate index
    p_kde = pt.Plotter(df).plot(kind="kde", column="v", by="g")
    assert p_kde.axd["A"].has_data()


def test_plotter_finalize_fontsize_both_axes_and_rot():
    # Plot Issue 1 (Review 2ec5bbba): fontsize applies to both axes, rot to category axis
    df = pd.DataFrame({"day": ["Thursday", "Friday", "Saturday"], "bill": [10, 20, 30]})
    p = pt.Plotter(df).plot(kind="bar", x="day", y="bill", rot=45, fontsize=8).finalize()
    ax = p.axd["A"]
    yticks = ax.yaxis.get_major_ticks()
    if yticks:
        assert yticks[0].label1.get_size() == 8
    xticks = ax.xaxis.get_major_ticks()
    if xticks:
        assert xticks[0].label1.get_size() == 8
        assert xticks[0].label1.get_rotation() == 45


def test_plotter_line_style_width_integer_keys():
    # Plot Issue 2 (Review 2ec5bbba): style and width dicts with integer keys match stringified line labels
    df = pd.DataFrame({"x": [1, 2], 1: [10, 20], 2: [30, 40]})
    p = pt.Plotter(df).plot(kind="line", x="x", y=[1, 2], style={2: ":", 1: "-."}, width={2: 5, 1: 1})
    lines = p.axd["A"].get_lines()
    for line in lines:
        if line.get_label() == "1":
            assert line.get_linestyle() == "-."
            assert line.get_linewidth() == 1
        elif line.get_label() == "2":
            assert line.get_linestyle() == ":"
            assert line.get_linewidth() == 5


def test_plotter_hist_and_kde_with_by_having_missing_values():
    penguins = pt.sample("penguins")
    p_hist = pt.Plotter(penguins).plot(kind="hist", column="bill_length_mm", by="sex")
    assert p_hist.axd["A"].has_data()

    p_kde = pt.Plotter(penguins).plot(kind="kde", column="body_mass_g", by="sex")
    assert p_kde.axd["A"].has_data()


def test_plotter_pie_with_dropna():
    df = pd.DataFrame({"cat": ["a", "b", None], "val": [1, 2, 5]})
    p1 = pt.Plotter(df).plot(kind="pie", by="cat", y="val", dropna=False)
    assert p1.axd["A"].has_data()
    assert len(p1.get_data("A")) == 3
    p2 = pt.Plotter(df).plot(kind="pie", by="cat", y="val", dropna=True)
    assert p2.axd["A"].has_data()
    assert len(p2.get_data("A")) == 2


def test_plotter_finalize_with_title_and_tight_layout():
    df = pd.DataFrame({"x": [1, 2], "y": [3, 4]})
    p = pt.Plotter(df).plot(kind="scatter", x="x", y="y")
    p.finalize(title="Custom Title", tight_layout=True, xlabel="X Axis", ylabel="Y Axis")
    assert p.axd["A"].get_xlabel() == "X Axis"
    assert p.axd["A"].get_ylabel() == "Y Axis"


def test_plotter_heatmap_matrix_honors_annot_and_numeric():
    df = pd.DataFrame({"x": ["a", "b"], "c1": [1.0, 2.0], "c2": [3.0, 4.0]})
    p = pt.Plotter(df).plot(kind="heatmap", x="x", annot=True)
    assert p.axd["A"].has_data()
    assert p.get_data("A").shape == (2, 2)


def test_plotter_grouped_bar_by_value_equals_x_name():
    df = pd.DataFrame({
        "year": [2020, 2020, 2021],
        "metric": ["year", "sales", "sales"],
        "value": [1, 10, 12],
    })
    with pytest.raises(ValueError, match="collides with x-axis column name"):
        pt.Plotter(df).plot(kind="bar", x="year", y="value", by="metric")


def test_plotter_box_color_dict_with_integer_keys():
    df = pd.DataFrame({"g": [1, 1, 2, 2], "v": [1, 2, 3, 4]})
    p = pt.Plotter(df).plot(kind="box", x="g", y="v", color={1: "crimson", 2: "navy"})
    assert p.axd["A"].has_data()


def test_plotter_falsy_column_name_zero():
    # String column "0" works cleanly
    df_str = pd.DataFrame({"x": [1, 2], "y": [3, 4], "0": ["g1", "g2"]})
    p_bar_str = pt.Plotter(df_str).plot(kind="bar", x="x", y="y", by="0")
    assert p_bar_str.axd["A"].has_data()

    # Integer column 0 works cleanly
    df_int = pd.DataFrame({"x": [1, 2], "y": [3, 4], 0: ["g1", "g2"]})
    p_bar_int = pt.Plotter(df_int).plot(kind="bar", x="x", y="y", by=0)
    assert p_bar_int.axd["A"].has_data()


def test_plotter_finalize_invalid_style():
    df = pd.DataFrame({"x": [1, 2], "y": [3, 4]})
    p = pt.Plotter(df).plot(kind="line", x="x", y="y")
    with pytest.raises(ValueError, match="stylesheet 'non_existent_style_xyz' not recognized"):
        p.finalize(style="non_existent_style_xyz")


def test_plotter_normalized_by_category_collision_with_x():
    # Issue 4: stringified collision between by category and x column name
    df1 = pd.DataFrame({"1": ["a", "b", "c"], "g": [1, 1, 2], "v": [10, 20, 30]})
    with pytest.raises(ValueError, match="collides with x-axis column name"):
        pt.Plotter(df1).plot(kind="bar", x="1", y="v", by="g")

    df2 = pd.DataFrame({"nan": [1, 2, 3], "g": [None, None, "a"], "v": [10, 20, 30]})
    with pytest.raises(ValueError, match="collides with x-axis column name"):
        pt.Plotter(df2).plot(kind="bar", x="nan", y="v", by="g")


def test_plotter_integer_x_labels():
    # Issue 5: integer x column label is treated as column, not out-of-bounds positional index
    df = pd.DataFrame({1: ["a", "a", "b"], "g": ["x", "y", "x"], "v": [10, 20, 30]})
    p1 = pt.Plotter(df).plot(kind="bar", x=1, y="v", by="g")
    assert p1.axd["A"].has_data()

    df_2020 = pd.DataFrame({2020: ["a", "b", "c"], "v": [10, 20, 30]})
    p2 = pt.Plotter(df_2020).plot(kind="bar", x=2020, y="v")
    assert p2.axd["A"].has_data()


def test_plotter_integer_column_scatter_and_hexbin():
    # Issue 6: integer column label in scatter and hexbin
    df = pd.DataFrame({"x": [1, 2, 3], 0: [4, 5, 6], "g": ["a", "a", "b"]})
    p_scat = pt.Plotter(df).plot(kind="scatter", x="x", y=0)
    assert p_scat.axd["A"].has_data()

    p_hex = pt.Plotter(df).plot(kind="hexbin", x="x", y=0)
    assert p_hex.axd["A"].has_data()


def test_plotter_box_plot_with_cat_col_zero():
    # Issue 7: box grouping on column named 0
    df = pd.DataFrame({0: [1, 1, 2, 2], "v": [1.0, 2.0, 100.0, 200.0]})
    p_box_x = pt.Plotter(df).plot(kind="box", x=0, y="v")
    assert len(p_box_x.tables["A"].columns) == 2

    p_box_by = pt.Plotter(df).plot(kind="box", by=0, y="v")
    assert len(p_box_by.tables["A"].columns) == 2


def test_plotter_faceting_with_by_zero():
    # Issue 8: faceting on by=0
    df = pd.DataFrame({0: ["a", "a", "b", "b"], "x": [1, 2, 3, 4], "y": [10, 20, 30, 40]})
    p = pt.Plotter(df, by=0, ncols=2, kind="scatter", x="x", y="y")
    assert "a" in p.axd and "b" in p.axd
    assert p.axd["a"].has_data()
    assert p.axd["b"].has_data()

    p_acc = df.pt.plot(by=0, ncols=2, kind="scatter", x="x", y="y")
    assert "a" in p_acc.axd and "b" in p_acc.axd
    assert p_acc.axd["a"].has_data()
    assert p_acc.axd["b"].has_data()


def test_plotter_mixed_y_aggregate_true():
    # Issue 10: mixed integer and string column names in y with aggregate=True
    df = pd.DataFrame({"x": [1, 2], 0: [4, 5], "b": [7, 8]})
    p = pt.Plotter(df).plot(kind="bar", x="x", y=[0, "b"])
    assert p.axd["A"].has_data()


def test_plotter_kde_minimum_count_check():
    # Issue 11: kde raises ValueError when observation count per group < 2
    df = pd.DataFrame({"g": ["a", "b"], "v": [1.0, 2.0]})
    with pytest.raises(ValueError, match="kde/density plot requires at least two"):
        pt.Plotter(df).plot(kind="kde", column="v", by="g")


def test_plotter_finalize_sharex_and_sharey():
    # Issue 12: finalize(sharex=True, sharey=True) links axes
    df1 = pd.DataFrame({"x": [1, 2], "y": [10, 20]})
    df2 = pd.DataFrame({"x": [1, 2], "y": [100, 200]})
    p = pt.Plotter(mosaic="AB")
    p.data(df1).plot(on="A", kind="bar", x="x", y="y")
    p.data(df2).plot(on="B", kind="bar", x="x", y="y")
    p.finalize(sharey=True)
    ax_a = p.axd["A"]
    ax_b = p.axd["B"]
    assert ax_a.get_shared_y_axes().joined(ax_a, ax_b)


def test_plotter_integer_coords_on_integer_columns():
    # Issue 5: integer coordinates x and y on DataFrame with integer column index
    df = pd.DataFrame([[10, 20], [30, 40]], columns=[1, 2])
    p1 = pt.Plotter(df).plot(kind="scatter", x=1, y=2)
    assert p1.axd["A"].has_data()

    p2 = pt.Plotter(df).plot(kind="line", x=1, y=2)
    assert p2.axd["A"].has_data()

    p3 = pt.Plotter(df).plot(kind="bar", x=1, y=2)
    assert p3.axd["A"].has_data()


def test_plotter_dropna_in_line_bar_heatmap():
    # Issue 6: dropna dropped NA values before reshaping/plotting
    df = pd.DataFrame({
        "x": ["a", "b", None, "c"],
        "g": ["g1", "g1", "g2", "g2"],
        "y": [1.0, 2.0, 3.0, 4.0],
    })
    # Line with dropna=True drops None row in x
    p_line = pt.Plotter(df).plot(kind="line", x="x", y="y", by="g", dropna=True)
    table_line = p_line.get_data("A")
    assert None not in table_line["x"].values

    # Bar with dropna=True
    p_bar = pt.Plotter(df).plot(kind="bar", x="x", y="y", by="g", dropna=True)
    table_bar = p_bar.get_data("A")
    assert None not in table_bar["x"].values

    # Heatmap with dropna=True
    p_heat = pt.Plotter(df).plot(kind="heatmap", x="x", y="y", by="g", dropna=True)
    table_heat = p_heat.get_data("A")
    assert None not in table_heat.index


def test_plotter_grouped_kde_singleton_raises_value_error():
    # Issue 7: Grouped kde requires at least 2 non-null observations in each group
    df = pd.DataFrame({
        "g": ["A", "A", "B"],  # B has only 1 observation
        "v": [10.0, 20.0, 30.0],
    })
    with pytest.raises(ValueError, match="at least two non-null observations in each group.*'B'"):
        pt.Plotter(df).plot(kind="kde", column="v", by="g")


def test_plotter_finalize_sharex_and_sharey_three_axes():
    # Issue 8: finalize(sharex=True, sharey=True) with 3+ axes
    df = pd.DataFrame({"x": [1, 2], "y": [10, 20]})
    p = pt.Plotter(mosaic="ABC")
    p.data(df).plot(on="A", kind="line", x="x", y="y")
    p.data(df).plot(on="B", kind="line", x="x", y="y")
    p.data(df).plot(on="C", kind="line", x="x", y="y")
    p.finalize(sharex=True, sharey=True)
    ax_a = p.axd["A"]
    ax_b = p.axd["B"]
    ax_c = p.axd["C"]
    assert ax_a.get_shared_x_axes().joined(ax_a, ax_b)
    assert ax_a.get_shared_x_axes().joined(ax_a, ax_c)
    assert ax_a.get_shared_y_axes().joined(ax_a, ax_b)
    assert ax_a.get_shared_y_axes().joined(ax_a, ax_c)


def test_plotter_show_method():
    # Issue 10: Plotter.show() exists and is chainable
    df = pd.DataFrame({"x": [1, 2], "y": [10, 20]})
    p = pt.Plotter(df).plot(kind="line", x="x", y="y")
    ret = p.show()
    assert ret is p


def test_plotter_by_column_validation():
    # Issue 12: Validate that 'by' column exists in DataFrame
    df = pd.DataFrame({"x": [1, 2], "y": [10, 20]})
    with pytest.raises(KeyError, match="column 'nonexistent' for 'by' not found"):
        pt.Plotter(df).plot(kind="line", x="x", y="y", by="nonexistent")

    with pytest.raises(KeyError, match="column 'nonexistent' for 'by' not found"):
        pt.Plotter(df).plot(kind="bar", x="x", y="y", by="nonexistent")

    with pytest.raises(KeyError, match="column 'nonexistent' for 'by' not found"):
        pt.Plotter(df).plot(kind="scatter", x="x", y="y", by="nonexistent")

    with pytest.raises(KeyError, match="column 'nonexistent' for 'by' not found"):
        pt.Plotter(df).plot(kind="kde", column="y", by="nonexistent")

    with pytest.raises(KeyError, match="column 'nonexistent' for 'by' not found"):
        pt.Plotter(df).plot(kind="hist", column="y", by="nonexistent")


def test_plotter_hist_and_kde_no_dropna_warning():
    # Issue 13: dropna is supported in hist and kde without warning
    import warnings
    df = pd.DataFrame({"g": ["a", "a", None], "v": [1.0, 2.0, 3.0]})
    with warnings.catch_warnings(record=True) as record:
        warnings.simplefilter("always")
        pt.Plotter(df).plot(kind="hist", column="v", by="g", dropna=True)
    dropna_warnings = [w for w in record if "dropna" in str(w.message).lower()]
    assert len(dropna_warnings) == 0


def test_plotter_wide_data_optional_y():
    # When plotting wide data (e.g. from pivot), y is optional
    df = pd.DataFrame({
        "day": ["Thur", "Fri"],
        "Lunch": [10.0, 15.0],
        "Dinner": [20.0, 25.0],
    })
    p_barh = df.pt.plot(kind="barh", x="day")
    assert p_barh.axd["A"].has_data()

    p_line = df.pt.plot(kind="line", x="day")
    assert p_line.axd["A"].has_data()

def test_plotter_remap_coord_pandas2_and_3():
    from unittest.mock import MagicMock

    from pytae.plotting import _remap_coord_if_needed

    # Simulating pandas 2.x Index with _holds_integer() -> True
    mock_cols_pd2 = MagicMock()
    mock_cols_pd2.dtype.kind = "O"
    mock_cols_pd2._holds_integer.return_value = True
    mock_cols_pd2.__contains__.side_effect = lambda k: k in ["x", 0, "g"]
    assert _remap_coord_if_needed(mock_cols_pd2, 0) == 0
    assert _remap_coord_if_needed(mock_cols_pd2, "x") == "x"

    # Current pandas (e.g. pandas 3.x) where Index does not have _holds_integer
    cols_pd3 = pd.Index(["x", 0, "g"])
    assert _remap_coord_if_needed(cols_pd3, 0) == 1
    assert _remap_coord_if_needed(cols_pd3, "x") == "x"

################################################################################
# FILE: tests/test_qry.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt
from pytae.cli_parsing import parse_qry


def _df():
    return pd.DataFrame(
        {
            "species": ["Adelie", "Gentoo", "Chinstrap", "Adelie"],
            "body_mass_g": [74125, 271425, 119925, 89100],
            "code": ["A 1", "B 2", "C 3", "D 4"],
        }
    )


def test_qry_equality():
    result = pt.qry(_df(), species="Adelie")
    assert list(result.index) == [0, 3]


def test_qry_list_membership():
    result = pt.qry(_df(), species=["Adelie", "Gentoo"])
    assert list(result["species"]) == ["Adelie", "Gentoo", "Adelie"]


def test_qry_in_and_not_in():
    df = _df()
    assert list(pt.qry(df, species=("in", ["Gentoo"]))["species"]) == ["Gentoo"]
    assert list(pt.qry(df, species=("not in", ["Adelie", "Gentoo"]))["species"]) == ["Chinstrap"]


def test_qry_comparison():
    result = pt.qry(_df(), body_mass_g="> 81500")
    assert list(result["species"]) == ["Gentoo", "Chinstrap", "Adelie"]


def test_qry_interval():
    result = pt.qry(_df(), body_mass_g="(85000,100000)")
    assert list(result["species"]) == ["Adelie"]
    assert result["body_mass_g"].iloc[0] == 89100


def test_qry_does_not_mutate_original():
    df = _df()
    pt.qry(df, species="Adelie")
    assert len(df) == 4


def test_qry_in_requires_list():
    with pytest.raises(ValueError, match="must be a list"):
        pt.qry(_df(), species=("in", "Adelie"))


def test_qry_unknown_column_suggests_typo():
    with pytest.raises(KeyError, match=r"unknown column 'speceis' \(did you mean 'species'\?\)"):
        pt.qry(_df(), speceis="Adelie")


def test_qry_startswith():
    result = pt.qry(_df(), species=("startswith", "Ad"))
    assert list(result["species"]) == ["Adelie", "Adelie"]


def test_qry_endswith():
    result = pt.qry(_df(), code=("endswith", "3"))
    assert list(result["species"]) == ["Chinstrap"]


def test_qry_contains():
    result = pt.qry(_df(), species=("contains", "in"))
    assert list(result["species"]) == ["Chinstrap"]


def test_qry_regex():
    result = pt.qry(_df(), code=("regex", r"^[AB]"))
    assert list(result["species"]) == ["Adelie", "Gentoo"]


def test_qry_startswith_accepts_list_of_prefixes():
    result = pt.qry(_df(), species=("startswith", ["Ad", "Ge"]))
    assert list(result["species"]) == ["Adelie", "Gentoo", "Adelie"]


def test_qry_string_op_on_non_string_column_raises_clear_error():
    with pytest.raises(ValueError, match="needs a string column"):
        pt.qry(_df(), body_mass_g=("startswith", "1"))


def test_qry_isna():
    df = _df()
    df.loc[0, "species"] = None
    result = pt.qry(df, species=("isna",))
    assert len(result) == 1
    assert result.index[0] == 0


def test_qry_notna():
    df = _df()
    df.loc[0, "species"] = None
    result = pt.qry(df, species=("notna",))
    assert list(result["species"]) == ["Gentoo", "Chinstrap", "Adelie"]


def test_qry_string_op_treats_missing_as_no_match():
    df = _df()
    df.loc[0, "species"] = None
    result = pt.qry(df, species=("startswith", "Ad"))
    assert list(result["species"]) == ["Adelie"]


def test_qry_unsupported_tuple_operator_lists_string_ops():
    with pytest.raises(ValueError, match="startswith"):
        pt.qry(_df(), species=("badop", "x"))


def test_qry_unsupported_unary_operator():
    with pytest.raises(ValueError, match="Unsupported 1-element tuple operator"):
        pt.qry(_df(), species=("badop",))


def test_qry_accepts_positional_dict_and_string():
    df = _df()
    res_dict = pt.qry(df, {"species": "Adelie"})
    assert list(res_dict.index) == [0, 3]

    res_str = df.pt.qry("body_mass_g > 100000")
    assert list(res_str["species"]) == ["Gentoo", "Chinstrap"]


def test_qry_rejects_invalid_positional_type():
    df = _df()
    with pytest.raises(TypeError, match=r"qry\(\) expects filter conditions as strings, dicts, or keyword arguments"):
        pt.qry(df, 123)
    with pytest.raises(TypeError, match=r"qry\(\) expects filter conditions as strings, dicts, or keyword arguments"):
        df.pt.qry(["species = Adelie"])


def test_qry_empty_conditions_raises():
    df = _df()
    with pytest.raises(ValueError, match=r"qry\(\) expects at least one condition"):
        df.pt.qry()


def test_qry_spaced_column_name_string_and_dict():
    df = pd.DataFrame(
        {
            "bill length mm": [35.0, 42.0, 51.0],
            "body mass g": [3000, 4500, 3200],
            "species": ["Adelie", "Gentoo", "Chinstrap"],
        }
    )
    # String condition with spaces in column name (no dict unpacking!)
    res_str = df.pt.qry("bill length mm > 40")
    assert list(res_str["species"]) == ["Gentoo", "Chinstrap"]

    # Plain dict with spaces in column name (no **kwargs unpacking!)
    res_dict = df.pt.qry({"bill length mm": "> 40"})
    assert list(res_dict["species"]) == ["Gentoo", "Chinstrap"]


def test_qry_mixed_strings_dicts_and_kwargs():
    df = pd.DataFrame(
        {
            "bill length mm": [35.0, 42.0, 51.0],
            "body mass g": [3000, 4500, 3200],
            "species": ["Adelie", "Gentoo", "Chinstrap"],
        }
    )
    res = df.pt.qry("bill length mm > 40", {"body mass g": "> 4000"}, species="Gentoo")
    assert len(res) == 1
    assert res.iloc[0]["species"] == "Gentoo"


def test_qry_multiple_conditions_on_same_column():
    df = pd.DataFrame(
        {
            "mass": [2000, 3500, 4500, 6000],
        }
    )
    # Comma-separated in single string
    res1 = df.pt.qry("mass > 3000, mass < 5000")
    assert list(res1["mass"]) == [3500, 4500]

    # Multiple string args
    res2 = df.pt.qry("mass > 3000", "mass < 5000")
    assert list(res2["mass"]) == [3500, 4500]


def test_qry_kwargs_operator():
    df = _df()
    res = df.pt.qry(body_mass_g="> 100000")
    assert list(res["species"]) == ["Gentoo", "Chinstrap"]


def test_qry_kwargs_tuple():
    df = _df()
    res = df.pt.qry(body_mass_g=(">", 100000))
    assert list(res["species"]) == ["Gentoo", "Chinstrap"]


def test_qry_kwargs_equality():
    df = _df()
    res = df.pt.qry(species="Adelie")
    assert list(res.index) == [0, 3]


def test_parse_qry_cli():
    res1 = parse_qry("body_mass_g= > 100000")
    res2 = parse_qry("body_mass_g=>100000")
    res3 = parse_qry("body_mass_g = > 100000")
    res4 = parse_qry("body_mass_g = 74125")
    assert res1 == [("body_mass_g", (">", 100000))]
    assert res2 == [("body_mass_g", (">", 100000))]
    assert res3 == [("body_mass_g", (">", 100000))]
    assert res4 == [("body_mass_g", 74125)]


def test_qry_bracketed_columns():
    df = pd.DataFrame({
        "species": ["Adelie", "Gentoo"],
        "body mass g": [3500, 5000],
    })
    res1 = df.pt.qry("[body mass g] > 4000")
    assert list(res1["species"]) == ["Gentoo"]

    res2 = df.pt.qry("[species] == 'Adelie'")
    assert list(res2["species"]) == ["Adelie"]

    res3 = df.pt.qry({"[body mass g]": "> 4000"})
    assert list(res3["species"]) == ["Gentoo"]

    parsed = parse_qry("[body mass g] > 4000, [species] = 'Gentoo'")
    assert parsed == [("body mass g", (">", 4000)), ("species", "Gentoo")]


def test_qry_interval_string_spec():
    # Lib Issue 3
    df = pd.DataFrame({"mass": [2000, 3500, 4000, 4500, 6000]})
    res_bracket = df.pt.qry("mass = [3500, 4500]")
    assert list(res_bracket["mass"]) == [3500, 4000, 4500]

    res_paren = df.pt.qry("mass = (3500, 4500)")
    assert list(res_paren["mass"]) == [4000]


def test_qry_large_int_precision():
    # Lib Issue 4: integers > 2**53 do not get collapsed by float cast
    val = 2**53 + 1
    df = pd.DataFrame({"id": [2**53, val, 2**53 + 2]})
    res_eq = df.pt.qry(id=("==", val))
    assert list(res_eq["id"]) == [val]

    res_str = df.pt.qry(f"id == {val}")
    assert list(res_str["id"]) == [val]

    res_gt = df.pt.qry(id=(">", 2**53))
    assert list(res_gt["id"]) == [2**53 + 1, 2**53 + 2]


def test_qry_bare_string_comparison():
    # Lib Issue 8
    df = pd.DataFrame({"species": ["Adelie", "Gentoo", "Chinstrap"]})
    res = df.pt.qry("species == Adelie")
    assert list(res["species"]) == ["Adelie"]


def test_qry_string_interval_spaces():
    # Lib Issue 9: spaces in interval bounds are stripped
    df = pd.DataFrame({"name": ["a", "b", "c", "m", "aa"]})
    res1 = df.pt.qry(name="[a, c)")
    assert list(res1["name"]) == ["a", "b", "aa"]

    res2 = df.pt.qry(name="(a, c]")
    assert list(res2["name"]) == ["b", "c", "aa"]


def test_qry_parenthesized_string_without_comma_is_equality():
    # Lib Issue 10: parenthesized/bracketed string without comma falls through to equality
    df = pd.DataFrame({"x": ["(abc)", "def", "[xyz]"]})
    res1 = df.pt.qry("x = (abc)")
    assert list(res1["x"]) == ["(abc)"]

    res2 = df.pt.qry("x = [xyz]")
    assert list(res2["x"]) == ["[xyz]"]


def test_parse_qry_multiple_conditions_on_same_column():
    # CLI Issue 1: multiple predicates on same column kept in order
    conditions = parse_qry("body_mass_g >= 3000, body_mass_g <= 4500")
    assert conditions == [("body_mass_g", (">=", 3000)), ("body_mass_g", ("<=", 4500))]


def test_qry_numeric_interval_invalid_bounds():
    df = pd.DataFrame({"mass": [2000, 3000, 4000]})
    with pytest.raises(ValueError, match="interval lower bound for numeric column 'mass' must be a number"):
        df.pt.qry("mass = [abc, 4000]")

    with pytest.raises(ValueError, match="interval upper bound for numeric column 'mass' must be a number"):
        df.pt.qry("mass = [2000, xyz]")

################################################################################
# FILE: tests/test_readers.py
################################################################################
import os
import sys
from pathlib import Path

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))


import pytae.readers as readers
from pytae.readers import get_reader, write_dataframe


def _frame():
    return pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})


def test_csv_reader_shape_head_tail_and_load(tmp_path):
    path = tmp_path / "t.csv"
    df = _frame()
    df.to_csv(path, index=False)
    reader = get_reader(path)

    assert reader.shape() == (3, 2)
    assert reader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(reader.head(2), df.head(2))
    pd.testing.assert_frame_equal(reader.tail(1).reset_index(drop=True), df.tail(1).reset_index(drop=True))
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)


def test_parquet_reader_roundtrip(tmp_path):
    path = tmp_path / "t.parquet"
    df = _frame()
    df.to_parquet(path, index=False)
    reader = get_reader(path)

    assert reader.shape() == (3, 2)
    assert reader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(reader.head(1), df.head(1))
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)


def test_parquet_head_and_tail_keep_schema_when_empty(tmp_path):
    path = tmp_path / "empty.parquet"
    pd.DataFrame({"a": pd.Series(dtype="int64"), "b": pd.Series(dtype="object")}).to_parquet(path, index=False)
    reader = get_reader(path)

    for frame in (reader.head(5), reader.tail(5)):
        assert list(frame.columns) == ["a", "b"]
        assert len(frame) == 0


def test_txt_reader_tab_delimited(tmp_path):
    path = tmp_path / "t.txt"
    df = _frame()
    df.to_csv(path, sep="\t", index=False)
    reader = get_reader(path)

    assert reader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)


def test_dat_reader_pipe_delimited(tmp_path):
    path = tmp_path / "t.dat"
    df = _frame()
    df.to_csv(path, sep="|", index=False)
    reader = get_reader(path)

    assert reader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)


def test_dat_reader_defaults_to_latin1_encoding(tmp_path):
    path = tmp_path / "t.dat"
    df = pd.DataFrame({"name": ["café", "naïve"]})
    df.to_csv(path, sep="|", index=False, encoding="latin-1")

    reader = get_reader(path)
    assert reader.encoding == "latin-1"
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)


def test_dat_reader_encoding_override():
    path = Path("dummy.dat")
    assert get_reader(path).encoding == "latin-1"
    assert get_reader(path, encoding="utf-8").encoding == "utf-8"


def test_txt_reader_encoding_defaults_to_none(tmp_path):
    path = tmp_path / "t.txt"
    df = _frame()
    df.to_csv(path, sep="\t", index=False)
    reader = get_reader(path)
    assert reader.encoding is None


def test_write_dataframe_csv_and_parquet(tmp_path):
    df = _frame()
    csv_path = tmp_path / "out.csv"
    pq_path = tmp_path / "out.parquet"
    write_dataframe(df, csv_path)
    write_dataframe(df, pq_path)
    pd.testing.assert_frame_equal(pd.read_csv(csv_path), df)
    pd.testing.assert_frame_equal(pd.read_parquet(pq_path), df)


def test_parquet_to_dataframe_keeps_schema_when_empty_with_nrows_or_progress(tmp_path):
    path = tmp_path / "empty.parquet"
    pd.DataFrame({"a": pd.Series(dtype="int64"), "b": pd.Series(dtype="int64")}).to_parquet(path)
    reader = get_reader(path)

    assert list(reader.to_dataframe().columns) == ["a", "b"]
    assert list(reader.to_dataframe(nrows=5).columns) == ["a", "b"]
    assert list(reader.to_dataframe(progress=True).columns) == ["a", "b"]


def test_csv_shape_counts_quoted_newline_as_one_row(tmp_path):
    path = tmp_path / "q.csv"
    path.write_text('a,b\n"hello\nworld",2\nfoo,3\n')
    reader = get_reader(path)

    assert reader.shape() == (2, 2)
    pd.testing.assert_frame_equal(reader.to_dataframe(), pd.read_csv(path))


def test_csv_tail_correct_after_quoted_newline(tmp_path):
    path = tmp_path / "q.csv"
    path.write_text('a,b\n"hello\nworld",2\nfoo,3\n')
    reader = get_reader(path)

    tail = reader.tail(1)
    assert tail["a"].tolist() == ["foo"]


def test_empty_csv_raises_friendly_error(tmp_path):
    path = tmp_path / "empty.csv"
    path.write_text("")
    reader = get_reader(path)

    with pytest.raises(ValueError, match="empty or not a valid"):
        reader.shape()


def test_nrows_caps_head_and_tail(tmp_path):
    from pytae.cli_pipeline import _Pipeline

    path = tmp_path / "x.csv"
    pd.DataFrame({"a": range(20), "b": range(20)}).to_csv(path, index=False)

    pipe = _Pipeline(get_reader(path), nrows=5)
    assert len(pipe.head(15)) == 5

    pipe2 = _Pipeline(get_reader(path), nrows=5)
    assert pipe2.tail(3)["a"].tolist() == [2, 3, 4]


def test_write_dataframe_dat_uses_pipe_delimiter(tmp_path):
    df = _frame()
    dat_path = tmp_path / "out.dat"
    write_dataframe(df, dat_path)
    pd.testing.assert_frame_equal(pd.read_csv(dat_path, sep="|"), df)


def test_write_dataframe_dat_defaults_to_latin1(tmp_path):
    df = pd.DataFrame({"name": ["café", "naïve"]})
    dat_path = tmp_path / "out.dat"
    write_dataframe(df, dat_path)
    pd.testing.assert_frame_equal(pd.read_csv(dat_path, sep="|", encoding="latin-1"), df)
    pd.testing.assert_frame_equal(get_reader(dat_path).to_dataframe(), df)


def test_unsupported_reader_and_writer(tmp_path):
    bad = tmp_path / "x.xlsx"
    bad.write_text("nope")
    with pytest.raises(ValueError, match="unsupported"):
        get_reader(bad)
    with pytest.raises(ValueError, match="unsupported"):
        write_dataframe(_frame(), tmp_path / "x.sas7bdat")


def test_csv_and_txt_dtypes_sample_not_full_file(tmp_path, monkeypatch):
    monkeypatch.setattr(readers, "DTYPE_SAMPLE_ROWS", 3)
    rows = ["n,label"] + [f"{i},a" for i in range(3)] + ["not_a_number,z"]
    csv_path = tmp_path / "t.csv"
    txt_path = tmp_path / "t.txt"
    csv_path.write_text("\n".join(rows) + "\n")
    txt_path.write_text("\n".join(row.replace(",", "\t") for row in rows) + "\n")

    csv_dtypes = get_reader(csv_path).dtypes()
    txt_dtypes = get_reader(txt_path).dtypes()
    assert pd.api.types.is_integer_dtype(csv_dtypes["n"])
    assert pd.api.types.is_integer_dtype(txt_dtypes["n"])
    full = pd.read_csv(csv_path)
    assert not pd.api.types.is_integer_dtype(full["n"])


def test_sas_reader_defaults_to_utf8():
    path = Path("dummy.sas7bdat")
    assert get_reader(path).encoding == "utf-8"
    assert get_reader(path, encoding="latin-1").encoding == "latin-1"


def test_get_reader_missing_suffix(tmp_path):
    path = Path(tmp_path / "noext")
    path.write_text("a,b\n1,2\n")
    with pytest.raises(ValueError, match="unsupported"):
        get_reader(path)


def test_jsonl_reader_and_writer(tmp_path):
    df = _frame()
    jsonl_path = tmp_path / "test.jsonl"
    write_dataframe(df, jsonl_path)
    assert jsonl_path.exists()

    reader = get_reader(jsonl_path)
    assert reader.shape() == (3, 2)
    assert reader.columns() == ["a", "b"]
    assert len(reader.dtypes()) == 2
    pd.testing.assert_frame_equal(reader.head(2), df.head(2))
    pd.testing.assert_frame_equal(reader.tail(1).reset_index(drop=True), df.tail(1).reset_index(drop=True))
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)
    pd.testing.assert_frame_equal(reader.to_dataframe(columns=["a"]), df[["a"]])

    # Test ndjson alias
    ndjson_path = tmp_path / "test.ndjson"
    write_dataframe(df, ndjson_path)
    nd_reader = get_reader(ndjson_path)
    assert nd_reader.shape() == (3, 2)
    pd.testing.assert_frame_equal(nd_reader.to_dataframe(), df)


def test_jsonl_reader_empty_error(tmp_path):
    empty_file = tmp_path / "empty.jsonl"
    empty_file.write_text("")
    reader = get_reader(empty_file)
    with pytest.raises(ValueError, match="is empty or not a valid JSON Lines file"):
        reader.columns()


def test_jsonl_chunked_progress(tmp_path):
    df = pd.DataFrame({"x": range(15), "y": [f"val_{i}" for i in range(15)]})
    path = tmp_path / "chunked.jsonl"
    write_dataframe(df, path, progress=True, chunk_size=5)

    reader = get_reader(path, chunk_size=5)
    read_df = reader.to_dataframe(progress=True, chunk_size=5)
    pd.testing.assert_frame_equal(read_df, df)


def test_compressed_csv_and_jsonl(tmp_path):
    df = _frame()
    csv_gz = tmp_path / "data.csv.gz"
    write_dataframe(df, csv_gz, progress=True, chunk_size=2)
    assert csv_gz.exists()

    reader = get_reader(csv_gz)
    assert reader.shape() == (3, 2)
    assert reader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)

    jsonl_gz = tmp_path / "data.jsonl.gz"
    write_dataframe(df, jsonl_gz, progress=True, chunk_size=2)
    assert jsonl_gz.exists()

    jreader = get_reader(jsonl_gz)
    assert jreader.shape() == (3, 2)
    assert jreader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(jreader.to_dataframe(), df)


def test_compressed_txt_and_dat(tmp_path):
    df = _frame()
    txt_gz = tmp_path / "data.txt.gz"
    write_dataframe(df, txt_gz)
    assert txt_gz.exists()
    reader = get_reader(txt_gz)
    assert reader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(reader.to_dataframe(), df)

    dat_gz = tmp_path / "data.dat.gz"
    write_dataframe(df, dat_gz)
    assert dat_gz.exists()
    dreader = get_reader(dat_gz)
    assert dreader.columns() == ["a", "b"]
    pd.testing.assert_frame_equal(dreader.to_dataframe(), df)

################################################################################
# FILE: tests/test_select.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt


def test_select_regex_kwarg():
    df = pd.DataFrame(
        {
            "species": [1],
            "bill_length_mm": [2],
            "bill_depth_mm": [3],
            "body_mass_g": [4],
        }
    )

    result = pt.select(df, regex="^bill")
    assert list(result.columns) == ["bill_length_mm", "bill_depth_mm"]


def test_select_regex_additive_with_explicit_name():
    df = pd.DataFrame(
        {
            "species": [1],
            "bill_length_mm": [2],
            "body_mass_g": [3],
        }
    )

    result = pt.select(df, "species", regex="bill|body")
    assert list(result.columns) == ["species", "bill_length_mm", "body_mass_g"]


def test_select_regex_invalid_pattern():
    df = pd.DataFrame({"a": [1]})
    with pytest.raises(ValueError, match="invalid regex"):
        pt.select(df, regex="[")


def test_select_dtype_numeric_includes_uint():
    df = pd.DataFrame({"i": pd.array([1, 2], dtype="int64"), "u": pd.array([1, 2], dtype="uint8")})
    assert pt.select(df, dtype="numeric").columns.tolist() == ["i", "u"]


def test_select_dtype_datetime_includes_tz_aware():
    df = pd.DataFrame(
        {
            "d": pd.to_datetime(["2020-01-01"]),
            "dtz": pd.to_datetime(["2020-01-01"]).tz_localize("UTC"),
        }
    )
    assert pt.select(df, dtype="datetime").columns.tolist() == ["d", "dtz"]


def test_select_dtype_accepts_tuple():
    df = pd.DataFrame({"i": [1], "u": pd.array([1], dtype="uint8"), "s": ["x"]})
    assert pt.select(df, dtype=("int64", "uint8")).columns.tolist() == ["i", "u"]


def test_select_colon_column_name_exact_match_wins():
    df = pd.DataFrame({"a:b": [1, 2], "c": [3, 4]})
    assert pt.select(df, "a:b").columns.tolist() == ["a:b"]


def test_select_integer_column_names_with_contains_startswith():
    df = pd.DataFrame({0: [1, 2], "b": [3, 4]})
    assert pt.select(df, contains="0").columns.tolist() == [0]
    assert pt.select(df, startswith="0").columns.tolist() == [0]


def _wide():
    return pd.DataFrame(
        {
            "species": [1],
            "island": [2],
            "bill_length_mm": [3],
            "bill_depth_mm": [4],
            "body_mass_g": [5],
        }
    )


def test_select_contains_startswith_endswith():
    df = _wide()
    assert list(pt.select(df, contains="bill").columns) == ["bill_length_mm", "bill_depth_mm"]
    assert list(pt.select(df, startswith="bill").columns) == ["bill_length_mm", "bill_depth_mm"]
    assert list(pt.select(df, endswith="_mm").columns) == ["bill_length_mm", "bill_depth_mm"]


def test_select_dtype_numeric_and_non_numeric():
    df = pd.DataFrame({"name": ["a"], "n": [1], "x": [1.5]})
    assert list(pt.select(df, dtype="numeric").columns) == ["n", "x"]
    assert list(pt.select(df, dtype="non_numeric").columns) == ["name"]


def test_select_exclude_dtype_numeric_vs_non_numeric():
    """exclude_dtype='numeric' keeps strings; exclude_dtype='non_numeric' keeps numbers."""
    df = pd.DataFrame({"name": ["a"], "n": [1], "x": [1.5]})
    assert list(pt.select(df, exclude_dtype="numeric").columns) == ["name"]
    assert list(pt.select(df, exclude_dtype="non_numeric").columns) == ["n", "x"]
    assert list(pt.select(df, exclude_dtype="numeric").columns) != list(
        pt.select(df, exclude_dtype="non_numeric").columns
    )


def test_select_exclude_dtype_cannot_combine():
    df = pd.DataFrame({"name": ["a"], "n": [1]})
    with pytest.raises(ValueError, match="exclude_dtype cannot be combined"):
        pt.select(df, "name", exclude_dtype="numeric")


def test_select_slice_and_missing_list():
    df = _wide()
    assert list(pt.select(df, "species:bill_length_mm").columns) == [
        "species",
        "island",
        "bill_length_mm",
    ]
    with pytest.raises(KeyError, match="not found"):
        pt.select(df, ["nope"])
    with pytest.raises(KeyError, match=r"Column not found: '\^bill'"):
        pt.select(df, "^bill")


def test_select_unknown_exact_name_does_not_skip():
    df = pd.DataFrame({"a": [1], "b": [2], "c": [3]})
    with pytest.raises(KeyError, match=r"Column not found: 'd'"):
        pt.select(df, "d", "a", "b")


def test_select_everything():
    from pytae.select import everything

    df = _wide()
    # Both everything() instance and everything class
    assert list(pt.select(df, "species", everything()).columns) == list(df.columns)
    assert list(pt.select(df, "species", everything).columns) == list(df.columns)
    assert list(pt.select(df, "species", pt.everything).columns) == list(df.columns)


def test_select_negative_column_syntax():
    df = _wide()
    # Single negative column
    assert list(pt.select(df, "-species").columns) == [
        "island",
        "bill_length_mm",
        "bill_depth_mm",
        "body_mass_g",
    ]
    # Tilde syntax
    assert list(pt.select(df, "~island").columns) == [
        "species",
        "bill_length_mm",
        "bill_depth_mm",
        "body_mass_g",
    ]
    # Multiple negative columns
    assert list(pt.select(df, "-species", "-island").columns) == [
        "bill_length_mm",
        "bill_depth_mm",
        "body_mass_g",
    ]
    # Negative column slice
    assert list(pt.select(df, "-bill_length_mm:body_mass_g").columns) == [
        "species",
        "island",
    ]


def test_select_exclude_kwarg():
    df = _wide()
    assert list(pt.select(df, exclude="species").columns) == [
        "island",
        "bill_length_mm",
        "bill_depth_mm",
        "body_mass_g",
    ]
    assert list(pt.select(df, exclude=["species", "island"]).columns) == [
        "bill_length_mm",
        "bill_depth_mm",
        "body_mass_g",
    ]


def test_select_positive_and_negative_combination():
    df = _wide()
    # Positive slice minus one column
    res = pt.select(df, "species:body_mass_g", "-island")
    assert list(res.columns) == [
        "species",
        "bill_length_mm",
        "bill_depth_mm",
        "body_mass_g",
    ]


def test_select_negative_typo_raises_hint():
    df = _wide()
    with pytest.raises(KeyError, match=r"Column to exclude not found: 'speceis' \(did you mean 'species'\?\)"):
        pt.select(df, "-speceis")


def test_select_bracketed_names_and_slices():
    df = pd.DataFrame({
        "species": [1],
        "island name": [2],
        "bill length mm": [3],
        "body mass g": [4],
    })
    # Single bracketed name
    assert list(pt.select(df, "[island name]").columns) == ["island name"]
    # Multiple bracketed names
    assert list(pt.select(df, "[island name]", "[body mass g]").columns) == ["island name", "body mass g"]
    # Comma-separated bracketed string
    assert list(pt.select(df, "species, [bill length mm]").columns) == ["species", "bill length mm"]
    # Bracketed slice
    assert list(pt.select(df, "[island name]:[body mass g]").columns) == ["island name", "bill length mm", "body mass g"]
    # Bracketed negation
    assert list(pt.select(df, "-[island name]").columns) == ["species", "bill length mm", "body mass g"]


def test_select_contains_startswith_endswith_tuples():
    # Lib Issue 14: contains, startswith, endswith accept tuples
    df = pd.DataFrame({"bill_length_mm": [1], "bill_depth_mm": [2], "body_mass_g": [3]})
    res_c = pt.select(df, contains=("bill", "body"))
    assert list(res_c.columns) == ["bill_length_mm", "bill_depth_mm", "body_mass_g"]

    res_sw = pt.select(df, startswith=("bill", "body"))
    assert list(res_sw.columns) == ["bill_length_mm", "bill_depth_mm", "body_mass_g"]

    res_ew = pt.select(df, endswith=("mm", "g"))
    assert list(res_ew.columns) == ["bill_length_mm", "bill_depth_mm", "body_mass_g"]


def test_select_exclude_bracketed_string():
    # Issue 8: exclude with bracketed names
    df = pd.DataFrame({"col a": [1], "col b": [2], "keep": [3]})
    res = pt.select(df, exclude="[col a, col b]")
    assert list(res.columns) == ["keep"]

################################################################################
# FILE: tests/test_shape.py
################################################################################
import os
import sys

import pandas as pd
import pytest

# Assuming the current working directory is where the project root is
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))

import pytae as pt
from pytae import sample


def test_long():

    penguins = sample("penguins")

    
    result = pt.long(penguins, c='features')
    

    numeric_cols = penguins.select_dtypes(include=['number']).columns.tolist()
    expected_df = pd.melt(penguins, id_vars=[col for col in penguins.columns if col not in numeric_cols],
                        value_vars=numeric_cols, var_name='features', 
                                                 value_name='value')
    
    pd.testing.assert_frame_equal(result, expected_df)


def test_wide_raises_on_duplicate_keys():
    """Duplicate (id, country) raises ValueError directing user to pt.pivot()."""
    df = pd.DataFrame(
        {
            "id": ["a", "b", "c", "d", "e", "", "f", "f"],
            "balance": [10, 20, 0, 21, 15, 10, 20, 25],
            "country": ["sg", "cn", "ca", "np", "in", "in", "in", "in"],
        }
    )

    with pytest.raises(ValueError, match=r"wide\(\) encountered duplicate entries.*Use pt\.pivot\(\)"):
        pt.wide(df, c="country", v="balance")


def test_wide_uses_pivot_when_keys_are_unique():
    df = pd.DataFrame(
        {
            "id": ["a", "b", "c", "d", "e", "", "f"],
            "balance": [10, 20, 0, 21, 15, 10, 20],
            "country": ["sg", "cn", "ca", "np", "in", "in", "in"],
        }
    )

    result = pt.wide(df, c="country", v="balance")

    expected_df = df.pivot(index="id", columns="country", values="balance").reset_index()
    expected_df.columns.name = None
    pd.testing.assert_frame_equal(result, expected_df)


def test_wide_rejects_aggregation_arguments():
    """wide() is strictly for 1-to-1 reshaping and rejects aggregation 'a'."""
    df = pd.DataFrame(
        {
            "id": ["a", "b"],
            "balance": [10, 20],
            "country": ["sg", "cn"],
        }
    )

    with pytest.raises(ValueError, match=r"wide\(\) does not support 'a'.*Use pt\.pivot\(\)"):
        pt.wide(df, c="country", v="balance", a="sum")

    with pytest.raises(ValueError, match=r"wide\(\) does not support 'agg'.*Use pt\.pivot\(\)"):
        df.pt.wide(c="country", v="balance", agg="mean")


def test_long_raises_when_no_numeric_columns():
    df = pd.DataFrame({"a": ["x", "y"], "b": ["p", "q"]})
    with pytest.raises(ValueError, match="no numeric columns to melt"):
        pt.long(df)


def test_long_with_non_numeric_cols():
    df = pd.DataFrame({
        "id": ["1", "2"],
        "phone_home": ["123", "456"],
        "phone_work": ["789", "012"],
    })
    res = df.pt.long(cols=["phone_home", "phone_work"], c="phone_type", v="number")
    assert list(res.columns) == ["id", "phone_type", "number"]
    assert len(res) == 4
    assert set(res["phone_type"]) == {"phone_home", "phone_work"}


def test_long_with_id_vars():
    df = pd.DataFrame({
        "user": ["alice", "bob"],
        "color": ["red", "blue"],
        "size": ["M", "L"],
    })
    res = df.pt.long(id_vars=["user"])
    assert list(res.columns) == ["user", "variable", "value"]
    assert len(res) == 4
    assert set(res["variable"]) == {"color", "size"}


def test_long_unknown_col_raises_helpful_error():
    df = pd.DataFrame({"a": [1], "b": [2]})
    with pytest.raises(KeyError, match=r"long\(\): cols column 'c' not found in DataFrame"):
        pt.long(df, cols=["c"])


def test_wide_with_explicit_index():
    df = pd.DataFrame({
        "region": ["East", "West"],
        "rep": ["Alice", "Bob"],
        "quarter": ["Q1", "Q2"],
        "sales": [100, 200],
    })
    res = df.pt.wide(index="region", c="quarter", v="sales")
    assert "region" in res.columns
    assert "Q1" in res.columns
    assert "Q2" in res.columns


def test_wide_bracketed_names_and_column_validation():
    # Lib Issue 6: wide unquotes bracketed names and checks column existence
    df = pd.DataFrame({
        "region name": ["East", "West"],
        "quarter": ["Q1", "Q2"],
        "sales": [1, 2],
    })
    res = df.pt.wide(index="[region name]", c="quarter", v="sales")
    assert "region name" in res.columns
    assert "Q1" in res.columns
    assert "Q2" in res.columns

    with pytest.raises(KeyError, match=r"wide\(\): index column 'nope' not found in DataFrame"):
        df.pt.wide(index="nope", c="quarter", v="sales")

    with pytest.raises(KeyError, match=r"wide\(\): columns 'c' column 'nope' not found in DataFrame"):
        df.pt.wide(index="region name", c="nope", v="sales")


def test_safe_reset_index_name_collision():
    # Lib Issue 13: index name collision does not raise ValueError
    df = pd.DataFrame({"region": ["East", "West"], "c": ["Q1", "Q2"], "v": [1, 2]})
    df.index.name = "Q1"  # index name collides with pivoted column name "Q1"
    res = df.pt.wide(index="region", c="c", v="v")
    assert "region" in res.columns
    assert "Q1" in res.columns


def test_safe_reset_index_disambiguates_collision():
    # When pivoted column name matches index name, raise ValueError rather than magical renaming
    df = pd.DataFrame({"id": [1, 2], "k": ["id", "a"], "v": [10, 20]})
    with pytest.raises(ValueError, match="collide with existing column name"):
        df.pt.wide(index="id", c="k", v="v")


def test_long_cols_and_id_vars_overlap():
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]})
    with pytest.raises(ValueError, match="cannot appear in both 'cols' and 'id_vars'"):
        df.pt.long(cols=["a", "b"], id_vars=["b", "c"])


def test_wide_c_and_v_collisions():
    df = pd.DataFrame({"x": ["A", "B"], "y": [1, 2], "z": [3, 4]})
    with pytest.raises(ValueError, match="'c' and 'v' cannot be the same column"):
        df.pt.wide(index="x", c="y", v="y")

    with pytest.raises(ValueError, match="cannot also be in index columns"):
        df.pt.wide(index=["x", "y"], c="y", v="z")


def test_long_output_name_collisions():
    # Issue 2: long() output names colliding with id_vars, cols, or each other
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    with pytest.raises(ValueError, match="cannot have the same name"):
        df.pt.long(cols=["a"], id_vars=["b"], c="val", v="val")

    with pytest.raises(ValueError, match="collides with kept id column"):
        df.pt.long(cols=["a"], id_vars=["b"], c="b", v="val")

    with pytest.raises(ValueError, match="collides with kept id column"):
        df.pt.long(cols=["a"], id_vars=["b"], c="var", v="b")

def test_wide_no_id_columns():
    # Issue 4: wide() where index_cols is empty
    df = pd.DataFrame({"metric": ["Sales", "Profit"], "val": [100, 20]})
    res = df.pt.wide(c="metric", v="val")
    assert list(res.columns) == ["Profit", "Sales"]
    assert len(res) == 1
    assert res.iloc[0]["Sales"] == 100
    assert res.iloc[0]["Profit"] == 20

    # Duplicates should raise clean ValueError
    df_dup = pd.DataFrame({"metric": ["Sales", "Sales"], "val": [100, 200]})
    with pytest.raises(ValueError, match="encountered duplicate entries"):
        df_dup.pt.wide(c="metric", v="val")


def test_long_bracketed_column_names():
    # Issue 11: long() bracket tokenization for cols and id_vars
    df = pd.DataFrame({
        "order id": [1, 2],
        "total bill": [10.5, 20.0],
        "tip amount": [2.0, 3.5],
    })
    res = df.pt.long(id_vars="[order id]", cols="[total bill], [tip amount]")
    assert list(res.columns) == ["order id", "variable", "value"]
    assert len(res) == 4


if __name__ == '__main__':
    pytest.main()

################################################################################
# FILE: tests/test_sql.py
################################################################################
import os
import sys

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt

duckdb = pytest.importorskip("duckdb")


def test_sql_query_via_data_alias():
    df = pd.DataFrame({"col a": [1, 20, 3], "col b": ["x", "y", "z"]})
    out = pt.sql(df, 'select "col b" from data where "col a" > 10')
    assert list(out["col b"]) == ["y"]


def test_sql_via_accessor():
    penguins = pt.sample("penguins")
    out = penguins.pt.sql(
        "select species, avg(body_mass_g) as avg_mass from data group by species order by species"
    )
    assert list(out.columns) == ["species", "avg_mass"]
    assert set(out["species"]) == {"Adelie", "Chinstrap", "Gentoo"}


def test_sql_extra_frames():
    left = pd.DataFrame({"id": [1, 2], "x": ["a", "b"]})
    right = pd.DataFrame({"id": [1, 3], "y": ["p", "q"]})
    out = pt.sql(left, "select data.id, x, y from data inner join extra using (id)", extra=right)
    assert list(out["id"]) == [1]
    assert list(out["x"]) == ["a"]
    assert list(out["y"]) == ["p"]


def test_sql_rejects_data_as_extra_name():
    df = pd.DataFrame({"a": [1]})
    with pytest.raises(ValueError, match="cannot be named 'data'"):
        pt.sql(df, "select * from data", data=df)


def test_sql_invalid_query_errors():
    df = pd.DataFrame({"a": [1]})
    with pytest.raises(ValueError, match="sql:"):
        pt.sql(df, "not valid sql")


def test_sql_then_pandas_head():
    df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
    out = df.pt.sql("select b from data where a > 1").head(1)
    assert list(out["b"]) == ["y"]


def test_sql_load_from_file(tmp_path):
    df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
    qfile = tmp_path / "query.txt"
    qfile.write_text("select b from data where a > 1")
    out = pt.sql(df, f"@{qfile}")
    assert list(out["b"]) == ["y", "z"]


def test_sql_bracketed_identifiers():
    df = pd.DataFrame({"col a": [1, 2], "col b": [10, 20]})
    out = pt.sql(df, "select [col a], `col b` from data")
    assert list(out["col a"]) == [1, 2]
    assert list(out["col b"]) == [10, 20]


def test_sql_preserves_duckdb_lists_and_indexing():
    df = pd.DataFrame({"a": [1, 2, 3], "col a": [10, 20, 30], "arr": [[100, 200], [300, 400], [500, 600]]})
    # List literal in WHERE clause with spaces
    out1 = pt.sql(df, "select [col a] from data where a in [1, 2]")
    assert list(out1["col a"]) == [10, 20]

    # Array indexing arr[1]
    out2 = pt.sql(df, "select arr[1] as first_elem from data where a = 1")
    assert list(out2["first_elem"]) == [100]

    # String list literal
    out3 = pt.sql(df, "select ['x', 'y'] as pair from data where a = 1")
    assert list(out3["pair"][0]) == ["x", "y"]

################################################################################
# FILE: tests/test_utilities.py
################################################################################
import os
import sys
from typing import Any

import pandas as pd
import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))

import pytae as pt
from pytae.other_utilities import clean_column_names


def test_to_clip_copies_as_method_and_does_not_shadow_pandas_clip(monkeypatch):
    df = pd.DataFrame({"a": [-1, 2]})
    copied: dict[str, Any] = {}

    def _fake_to_clipboard(self, *args, **kwargs):
        copied["called"] = True
        copied["index"] = kwargs.get("index")

    monkeypatch.setattr(pd.DataFrame, "to_clipboard", _fake_to_clipboard)

    # hasattr does not trigger clipboard copy
    assert hasattr(df, "to_clip") is True
    assert copied.get("called") is None

    df.to_clip()
    assert copied["called"] is True
    assert copied["index"] is False

    assert not hasattr(pt, "to_clip")
    assert not hasattr(df.pt, "to_clip")

    copied_series: dict[str, Any] = {}
    monkeypatch.setattr(pd.Series, "to_clipboard", lambda self, *args, **kwargs: copied_series.setdefault("called", True))
    df["a"].to_clip()
    assert copied_series.get("called") is True

    clipped = df.clip(lower=0)
    pd.testing.assert_series_equal(clipped["a"], pd.Series([0, 2], name="a"))


def test_handle_missing_fills_object_and_numeric():
    df = pd.DataFrame({"grp": pd.Series(["x", None], dtype=object), "val": [1.0, None]})
    result = pt.handle_missing(df)
    assert result["grp"].tolist() == ["x", "."]
    assert result["val"].tolist() == [1.0, 0.0]
    assert df["grp"].isna().any()


def test_handle_missing_preserves_categorical():
    s = pd.Series(pd.Categorical(["A", None, "B"]))
    df = pd.DataFrame({"cat": s, "val": [1.0, None, 3.0]})
    result = pt.handle_missing(df, fillna="Missing")
    assert isinstance(result["cat"].dtype, pd.CategoricalDtype)
    assert result["cat"].tolist() == ["A", "Missing", "B"]
    assert "Missing" in result["cat"].cat.categories


def test_handle_missing_custom_numeric_fill_and_cols():
    df = pd.DataFrame({"a": [10.0, None], "b": [100.0, None], "txt": ["hello", None]})
    # Only fill 'a' with mean, leave 'b' and 'txt' untouched
    result = pt.handle_missing(df, numeric_fill="mean", cols=["a"])
    assert result["a"].tolist() == [10.0, 10.0]
    assert pd.isna(result["b"].iloc[1])
    assert pd.isna(result["txt"].iloc[1])

    # numeric_fill=None leaves numerics alone
    result2 = pt.handle_missing(df, fillna="N/A", numeric_fill=None)
    assert pd.isna(result2["a"].iloc[1])
    assert result2["txt"].tolist() == ["hello", "N/A"]



def test_cols_sort_orders():
    df = pd.DataFrame({"c": [1], "a": [2], "b": [3]})
    assert pt.cols(df) == ["a", "b", "c"]
    assert pt.cols(df, ascending=False) == ["c", "b", "a"]
    assert pt.cols(df, ascending=None) == ["c", "a", "b"]
    with pytest.raises(ValueError, match="Invalid ascending"):
        pt.cols(df, ascending="nope")



def test_clean_columns_strip_fill_case():
    df = pd.DataFrame({"  Col A  ": [1], "col   b": [2]})
    result = pt.clean_columns(df, strip=True, fill="_", case="lower")
    assert list(result.columns) == ["col_a", "col___b"]
    assert list(df.columns) == ["  Col A  ", "col   b"]  # original untouched


def test_clean_columns_squeeze_strip_special_dedupe():
    df = pd.DataFrame({"Col A": [1], "col  a": [2], "100% Match!": [3]})
    result = pt.clean_columns(df, strip=True, squeeze=True, strip_special=True, fill="_", case="lower", dedupe=True)
    assert list(result.columns) == ["col_a", "col_a_1", "100_match"]


def test_clean_columns_no_fill_leaves_whitespace():
    df = pd.DataFrame({"col a": [1]})
    result = pt.clean_columns(df, case="upper")
    assert list(result.columns) == ["COL A"]


def test_clean_columns_strip_special_removes_quotes():
    df = pd.DataFrame({"'col a'": [1], '"col b"': [2]})
    result = pt.clean_columns(df, strip_special=True)
    assert list(result.columns) == ["col a", "col b"]


def test_clean_columns_strip_special_keeps_fill_character():
    df = pd.DataFrame({"co-op's data": [1]})
    result = pt.clean_columns(df, strip_special=True, fill="-")
    assert list(result.columns) == ["co-ops-data"]


def test_replace_values_exact_whole_df():
    df = pd.DataFrame({"a": ["x", "not x exactly"], "b": ["x", "y"]})
    result = pt.replace_values(df, {"x": "z"})
    assert result["a"].tolist() == ["z", "not x exactly"]
    assert result["b"].tolist() == ["z", "y"]
    assert df["a"].tolist() == ["x", "not x exactly"]  # original untouched


def test_replace_values_scoped_to_columns():
    df = pd.DataFrame({"a": ["x"], "b": ["x"]})
    result = pt.replace_values(df, {"x": "z"}, c="a")
    assert result["a"].tolist() == ["z"]
    assert result["b"].tolist() == ["x"]


def test_replace_values_substring_match():
    df = pd.DataFrame({"a": ["not a magician exactly"]})
    result = pt.replace_values(df, {"a magician": "the magic"}, exact=False)
    assert result["a"].tolist() == ["not the magic exactly"]


def test_replace_values_substring_match_on_numeric_column_raises():
    df = pd.DataFrame({"a": [1, 2, 3]})
    with pytest.raises(ValueError, match="only works on string/object columns"):
        pt.replace_values(df, {1: 9}, exact=False)


def test_handle_missing_leaves_bool_column_untouched():
    df = pd.DataFrame({"b": [True, None]})
    result = pt.handle_missing(df)
    assert result["b"].iloc[0] is True
    assert pd.isna(result["b"].iloc[1])


def test_clean_columns_int_names_are_stringified():
    df = pd.DataFrame({0: [1, 2], "b": [3, 4]})
    result = pt.clean_columns(df, strip=True)
    assert list(result.columns) == ["0", "b"]


def test_clean_columns_dedupe_avoids_new_collisions():
    result = clean_column_names(["revenue", "revenue", "revenue_1"], dedupe=True)
    assert result == ["revenue", "revenue_2", "revenue_1"]
    assert len(set(result)) == len(result)


def test_handle_missing_all_null_object_column():
    # Lib Issue 15: all-null object column is filled by handle_missing
    df = pd.DataFrame({"a": [None, None]}, dtype=object)
    result = pt.handle_missing(df)
    assert list(result["a"]) == [".", "."]
