import numpy as np
import pandas as pd
from pandas.api import types as ptypes
from collections import OrderedDict

class SkimResult:
    """
    Holds header info and one DataFrame per variable type.
    Nicely prints sections similar to skimr::skim().
    """
    def __init__(self, header: pd.DataFrame, sections: OrderedDict):
        self.header = header
        self.sections = sections  # OrderedDict[type_name -> DataFrame]

    def __repr__(self):
        parts = []

        # Header
        parts.append("Skim summary\n")
        parts.append(self.header.to_string(index=False))
        parts.append("")

        # Sections
        for t, df in self.sections.items():
            if df.empty:
                continue
            bar = "─" * 20
            parts.append(f"{bar} Variable type: {t} {bar}")
            # Use a compact string repr
            parts.append(df.to_string(index=False))
            parts.append("")

        return "\n".join(parts)


def _numeric_section(df: pd.DataFrame, cols):
    if not cols:
        return pd.DataFrame(columns=[
            "name","n","missing","complete_rate","mean","sd","p0","p25","p50","p75","p100"
        ])

    rows = []
    for name in cols:
        s = pd.to_numeric(df[name], errors="coerce")
        n = len(s)
        miss = int(s.isna().sum())
        comp = (n - miss) / n if n else np.nan
        valid = s.dropna()
        if valid.empty:
            q = [np.nan]*5
            mean = np.nan
            sd = np.nan
        else:
            q = np.percentile(valid, [0, 25, 50, 75, 100])
            mean = valid.mean()
            sd = valid.std(ddof=1) if len(valid) > 1 else 0.0

        rows.append(dict(
            name=name, n=n, missing=miss, complete_rate=round(comp, 3),
            mean=round(float(mean), 6) if pd.notna(mean) else np.nan,
            sd=round(float(sd), 6) if pd.notna(sd) else np.nan,
            p0=q[0], p25=q[1], p50=q[2], p75=q[3], p100=q[4]
        ))
    return pd.DataFrame(rows)[[
        "name","n","missing","complete_rate","mean","sd","p0","p25","p50","p75","p100"
    ]].sort_values("name").reset_index(drop=True)


def _bool_section(df: pd.DataFrame, cols):
    if not cols:
        return pd.DataFrame(columns=["name","n","missing","complete_rate","n_true","n_false"])
    rows = []
    for name in cols:
        s = df[name].astype("boolean")  # preserves <NA>
        n = len(s)
        miss = int(s.isna().sum())
        comp = (n - miss) / n if n else np.nan
        n_true = int((s == True).sum(skipna=True))
        n_false = int((s == False).sum(skipna=True))
        rows.append(dict(
            name=name, n=n, missing=miss, complete_rate=round(comp, 3),
            n_true=n_true, n_false=n_false
        ))
    return pd.DataFrame(rows)[["name","n","missing","complete_rate","n_true","n_false"]].sort_values("name").reset_index(drop=True)


def _categorical_section(df: pd.DataFrame, cols, top_k=3):
    if not cols:
        return pd.DataFrame(columns=["name","n","missing","complete_rate","n_unique","top_counts"])
    rows = []
    for name in cols:
        s = df[name]
        n = len(s)
        miss = int(s.isna().sum())
        comp = (n - miss) / n if n else np.nan
        n_unique = int(s.nunique(dropna=True))
        vc = s.value_counts(dropna=False)
        items = []
        for k, v in vc.head(top_k).items():
            label = "<NA>" if pd.isna(k) else str(k)
            items.append(f"{label}: {int(v)}")
        rows.append(dict(
            name=name, n=n, missing=miss, complete_rate=round(comp, 3),
            n_unique=n_unique, top_counts=", ".join(items)
        ))
    return pd.DataFrame(rows)[["name","n","missing","complete_rate","n_unique","top_counts"]].sort_values("name").reset_index(drop=True)


def _datetime_section(df: pd.DataFrame, cols):
    if not cols:
        return pd.DataFrame(columns=["name","n","missing","complete_rate","min","max"])
    rows = []
    for name in cols:
        s = pd.to_datetime(df[name], errors="coerce")
        n = len(s)
        miss = int(s.isna().sum())
        comp = (n - miss) / n if n else np.nan
        mn = s.min(skipna=True)
        mx = s.max(skipna=True)
        rows.append(dict(
            name=name, n=n, missing=miss, complete_rate=round(comp, 3),
            min=mn if pd.notna(mn) else pd.NaT,
            max=mx if pd.notna(mx) else pd.NaT
        ))
    return pd.DataFrame(rows)[["name","n","missing","complete_rate","min","max"]].sort_values("name").reset_index(drop=True)


def _timedelta_section(df: pd.DataFrame, cols):
    if not cols:
        return pd.DataFrame(columns=["name","n","missing","complete_rate","min","max","mean"])
    rows = []
    for name in cols:
        s = pd.to_timedelta(df[name], errors="coerce")
        n = len(s)
        miss = int(s.isna().sum())
        comp = (n - miss) / n if n else np.nan
        mn = s.min(skipna=True)
        mx = s.max(skipna=True)
        mean = s.mean(skipna=True)
        rows.append(dict(
            name=name, n=n, missing=miss, complete_rate=round(comp, 3),
            min=mn if pd.notna(mn) else pd.NaT,
            max=mx if pd.notna(mx) else pd.NaT,
            mean=mean if pd.notna(mean) else pd.NaT
        ))
    return pd.DataFrame(rows)[["name","n","missing","complete_rate","min","max","mean"]].sort_values("name").reset_index(drop=True)


def skim(df: pd.DataFrame, by=None, top_k=3) -> SkimResult:
    """
    skim(df, by=None, top_k=3) -> SkimResult

    - Produces a skimr-like summary with separate sections per variable type.
    - If `by` is provided (str or list[str]), returns one section set per group value
      by stacking group labels into the variable names, e.g., 'col | group=A'.

    Sections produced (if present): numeric, bool, categorical, datetime, timedelta.
    """
    if by is not None:
        # Grouped skim: concatenate skims across groups, tagging names
        if isinstance(by, str):
            by = [by]
        pieces = []
        for keys, sub in df.groupby(by, dropna=False, sort=True):
            if not isinstance(keys, tuple):
                keys = (keys,)
            label = ", ".join(f"{b}={k}" for b, k in zip(by, keys))
            # Recursively skim subgroup (without grouping), then tag names
            res = skim(sub, by=None, top_k=top_k)
            tagged = OrderedDict()
            for t, sec in res.sections.items():
                if sec.empty:
                    tagged[t] = sec.copy()
                else:
                    sec2 = sec.copy()
                    sec2["name"] = sec2["name"].astype(str) + f" | {label}"
                    tagged[t] = sec2
            # Build a header-like stub carrying group sizes (useful if inspected)
            pieces.append((label, tagged))

        # Merge sections across groups by row-binding per type
        all_types = ["numeric","bool","categorical","datetime","timedelta"]
        merged = OrderedDict()
        for t in all_types:
            merged[t] = pd.concat(
                [sec_map.get(t, pd.DataFrame(columns=[])) for _, sec_map in pieces],
                axis=0, ignore_index=True
            ) if any(t in sec_map and not sec_map[t].empty for _, sec_map in pieces) else pd.DataFrame()

        # Header for grouped output (overall df, plus group counts)
        header = _make_header(df)
        return SkimResult(header=header, sections=OrderedDict((k, v) for k, v in merged.items() if not v.empty))

    # ---- Ungrouped path ----
    # Detect columns per type (order matters; bool before numeric)
    bool_cols = [c for c in df.columns if ptypes.is_bool_dtype(df[c])]
    dt_cols = [c for c in df.columns if ptypes.is_datetime64_any_dtype(df[c])]
    td_cols = [c for c in df.columns if ptypes.is_timedelta64_dtype(df[c])]
    numeric_cols = [c for c in df.columns
                    if ptypes.is_numeric_dtype(df[c]) and c not in bool_cols and c not in td_cols]
    cat_cols = []
    for c in df.columns:
        if c in bool_cols or c in dt_cols or c in td_cols or c in numeric_cols:
            continue
        # New-style categorical detection (fixes deprecation)
        if isinstance(df[c].dtype, pd.CategoricalDtype):
            cat_cols.append(c)
        else:
            # Treat the rest (object, string, mixed) as categorical-like
            cat_cols.append(c)

    # Build sections
    sections = OrderedDict()
    num_sec = _numeric_section(df, numeric_cols)
    if not num_sec.empty:
        sections["numeric"] = num_sec

    bool_sec = _bool_section(df, bool_cols)
    if not bool_sec.empty:
        sections["bool"] = bool_sec

    cat_sec = _categorical_section(df, cat_cols, top_k=top_k)
    if not cat_sec.empty:
        sections["categorical"] = cat_sec

    dt_sec = _datetime_section(df, dt_cols)
    if not dt_sec.empty:
        sections["datetime"] = dt_sec

    td_sec = _timedelta_section(df, td_cols)
    if not td_sec.empty:
        sections["timedelta"] = td_sec

    header = _make_header(df)

    return SkimResult(header=header, sections=sections)


def _make_header(df: pd.DataFrame) -> pd.DataFrame:
    n_rows = len(df)
    n_cols = df.shape[1]
    # Overall missingness and completeness
    total_cells = n_rows * max(n_cols, 1)
    missing = int(df.isna().sum().sum())
    complete_rate = (total_cells - missing) / total_cells if total_cells else np.nan

    # Type counts
    type_counts = {
        "numeric": 0, "bool": 0, "categorical": 0, "datetime": 0, "timedelta": 0
    }
    for c in df.columns:
        if ptypes.is_bool_dtype(df[c]):
            type_counts["bool"] += 1
        elif ptypes.is_datetime64_any_dtype(df[c]):
            type_counts["datetime"] += 1
        elif ptypes.is_timedelta64_dtype(df[c]):
            type_counts["timedelta"] += 1
        elif ptypes.is_numeric_dtype(df[c]):
            type_counts["numeric"] += 1
        elif isinstance(df[c].dtype, pd.CategoricalDtype):
            type_counts["categorical"] += 1
        else:
            type_counts["categorical"] += 1

    header = pd.DataFrame([{
        "n_rows": n_rows,
        "n_cols": n_cols,
        "complete_rate": round(complete_rate, 3),
        "numeric": type_counts["numeric"],
        "bool": type_counts["bool"],
        "categorical": type_counts["categorical"],
        "datetime": type_counts["datetime"],
        "timedelta": type_counts["timedelta"]
    }])
    return header