"""kit.py: one-file toolkit for a timed dataset exercise.

Paste this into the first cell of a notebook (or save it next to the notebook
and `from kit import *`). Every function prints what it did, so the printout is
your audit trail for the data slide.

Conventions: T = time column name, E = entity column name (None if the file is
a single series). Nothing here fits a model; it scores signals honestly.
"""
import numpy as np
import pandas as pd


# ---------- 1. load and look -------------------------------------------------

def load(path, T, E=None, **read_kw):
    """Read, parse time, sort by (entity, time). Sorting first makes every later shift/diff correct."""
    df = pd.read_csv(path, **read_kw)
    col = df[T]
    if pd.api.types.is_integer_dtype(col) and col.between(1000, 3000).all():
        df[T] = pd.to_datetime(col.astype(str), format="%Y")      # plain years, not nanoseconds
    else:
        df[T] = pd.to_datetime(col)
        if pd.api.types.is_numeric_dtype(col):
            print("WARNING: %s was numeric; check the parsed dates are real (epoch seconds? pass unit=)" % T)
    keys = [E, T] if E else [T]
    df = df.sort_values(keys).reset_index(drop=True)
    print("rows %d, cols %d | %s -> %s" % (len(df), df.shape[1], df[T].min(), df[T].max()))
    if E:
        print("entities: %d" % df[E].nunique())
    return df


def bar_size(df, T, E=None):
    """Typical gap between consecutive rows of one entity: the unit of every return you compute."""
    gap = df.groupby(E)[T].diff() if E else df[T].diff()
    return gap.median()


def triage(df, T, E=None, ohlc=("open", "high", "low", "close"), volume="volume", max_gap=None):
    """Rows, key, missing, wrong. Prints one screen; returns the numbers as a dict for the data slide.
    max_gap: longest normal gap (daily equities: pd.Timedelta(days=4), so weekends are not flagged)."""
    out = {}
    keys = [T, E] if E else [T]
    print("missing per column:", {k: int(v) for k, v in df.isna().sum().items() if v})
    out["dup_keys"] = int(df.duplicated(keys).sum())
    out["exact_dups"] = int(df.duplicated().sum())
    print("duplicate keys: %d | exact duplicate rows: %d%s" % (
        out["dup_keys"], out["exact_dups"],
        "  <- some copies DISAGREE" if out["dup_keys"] > out["exact_dups"] else ""))
    bar = bar_size(df, T, E)
    out["bar"] = bar
    print("bar size:", bar)
    if E:
        span = df.groupby(E)[T].agg(["min", "max", "count"])
        late = int((span["min"] > df[T].min()).sum())
        early = int((span["max"] < df[T].max()).sum())
        out["start_late"], out["end_early"] = late, early
        print("entities starting late: %d | ending early: %d%s" % (
            late, early, "  <- survivors only? (say so on the last slide)" if late and not early else ""))
        print(span.sort_values("count").head(5))
        per_t = df.groupby(T)[E].nunique()
        print("entities per timestamp: min %d, median %d, max %d" % (per_t.min(), per_t.median(), per_t.max()))
        gap = df.groupby(E)[T].diff()
    else:
        gap = df[T].diff()
    big = gap > (max_gap if max_gap is not None else bar * 1.5)   # 1.5x tolerates leap years, month lengths
    out["gaps"] = int(big.sum())
    if big.any():
        top = df.loc[big, T].value_counts().head(3)
        print("gaps longer than one bar: %d; most common timestamps after a gap:" % big.sum())
        print(top.to_string())
    have = [c for c in ohlc if c in df.columns]
    if len(have) == 4:
        o, h, l, c = [df[x] for x in ohlc]
        bad = (c > h) | (c < l) | (h < l)
        flat = (o == h) & (h == l) & (l == c)
        out["bad_prints"], out["flat"] = int(bad.sum()), int(flat.sum())
        msg = "close outside [low, high]: %d | flat bars: %d" % (bad.sum(), flat.sum())
        if volume in df.columns:
            out["zero_vol"] = int((df[volume] == 0).sum())
            msg += " | zero volume: %d (flat and zero: %d)" % (out["zero_vol"], (flat & (df[volume] == 0)).sum())
        print(msg)
    num = df.select_dtypes("number")
    q = num.quantile([0.001, 0.5, 0.999]).T
    q["max"] = num.max()
    print("numeric tails (look for a max far beyond the 99.9%):")
    print(q.round(4).to_string())
    return out


def show_rows(df, mask, keys):
    """When a count looks odd: print every row that shares a key with the odd rows."""
    odd = df.loc[mask, keys].drop_duplicates()
    rows = df.merge(odd, on=keys)
    print("%d odd rows -> %d rows sharing their keys" % (mask.sum(), len(rows)))
    return rows


# ---------- 2. clean, with a log ---------------------------------------------

def clean(df, T, E=None, dedupe="last", drop_bad_prints=True, drop_zero_volume=True,
          ohlc=("open", "high", "low", "close"), volume="volume"):
    """Apply the stated rules and print what each removed. Returns (df, log)."""
    log = []
    keys = [T, E] if E else [T]
    n = len(df)
    if dedupe:
        df = df.drop_duplicates(keys, keep=dedupe)
        log.append("duplicate keys removed (keep=%s): %d" % (dedupe, n - len(df)))
    if drop_bad_prints and all(c in df.columns for c in ohlc):
        o, h, l, c = [df[x] for x in ohlc]
        bad = (c > h) | (c < l) | (h < l)
        df = df[~bad]
        log.append("bad prints removed (close outside [low, high]): %d" % bad.sum())
    if drop_zero_volume and volume in df.columns:
        z = df[volume] == 0
        df = df[~z]
        log.append("zero-volume bars removed: %d" % z.sum())
    keys_sorted = [E, T] if E else [T]
    df = df.sort_values(keys_sorted).reset_index(drop=True)
    log.append("rows left: %d of %d" % (len(df), n))
    for line in log:
        print(line)
    return df, log


# ---------- 3. targets and lags that cannot cheat ----------------------------

def _gap_ok(df, T, E, bar, max_gap):
    gap = df.groupby(E)[T].diff() if E else df[T].diff()
    return gap <= (max_gap if max_gap is not None else bar * 1.5)


def lag(df, col, T, E=None, k=1, max_gap=None):
    """col shifted one row back within each entity, blanked when the previous row is more than
    1.5 bars earlier (outage, halt, late listing). Daily equities: max_gap=pd.Timedelta(days=4).
    For k > 1 no gap guard is applied; say so if you use it."""
    bar = bar_size(df, T, E)
    s = df.groupby(E)[col].shift(k) if E else df[col].shift(k)
    if k == 1:
        bad = ~_gap_ok(df, T, E, bar, max_gap) & s.notna()
        s = s.mask(bad)
        print("lag %s: %d values blanked across gaps" % (col, int(bad.sum())))
    return s


def forward(df, col, T, E=None, k=1, max_gap=None):
    """col k rows AHEAD within each entity (a next-period target). For k == 1 it is blanked
    when the next row is more than 1.5 bars later."""
    bar = bar_size(df, T, E)
    s = df.groupby(E)[col].shift(-k) if E else df[col].shift(-k)
    nxt = (df.groupby(E)[T].shift(-1) if E else df[T].shift(-1)) - df[T]
    if k == 1:
        s = s.where(nxt <= (max_gap if max_gap is not None else bar * 1.5))
    return s


def when_known(columns):
    """Print a checklist: fill in before / at / after for each column, relative to the moment you trade."""
    print("For each column: known BEFORE, AT, or only AFTER the moment you trade?")
    for c in columns:
        print("  %-20s  before / at / after" % c)
    print("Only 'before' or 'at' can be a signal. 'After' is a leak, however good it scores.")


# ---------- 4. split -------------------------------------------------------------

def split(df, T, cut, embargo="1D"):
    """Build = before cut; check = from cut + embargo. The gap between them is thrown away."""
    cut = pd.Timestamp(cut)
    emb = pd.Timedelta(embargo)
    build = df[df[T] < cut]
    check = df[df[T] >= cut + emb]
    for name, d in (("build", build), ("check", check)):
        print("%s: %s -> %s  (%d rows, %d days)" % (
            name, d[T].min(), d[T].max(), len(d), d[T].dt.normalize().nunique()))
    return build, check


# ---------- 5. score -------------------------------------------------------------

def ic_series(df, signal, target, T, method="spearman", min_n=5):
    """One cross-sectional correlation per timestamp (needs an entity column)."""
    def one(g):
        g = g[[signal, target]].dropna()
        return g[signal].corr(g[target], method=method) if len(g) >= min_n else np.nan
    return df.groupby(T)[[signal, target]].apply(one).dropna()


def summary(x):
    """Mean, t (mean over its standard error), hit rate, count."""
    x = pd.Series(x).dropna()
    t = x.mean() / (x.std() / np.sqrt(len(x))) if len(x) > 1 else np.nan
    return pd.Series({"mean": x.mean(), "t": t, "hit": (x > 0).mean(), "n": len(x)})


def ic_table(parts, signals, target, T, method="spearman"):
    """parts = {"build": build, "check": check}. The slide-5 table, one row per signal."""
    rows = {}
    for s in signals:
        r = {}
        for name, d in parts.items():
            m = summary(ic_series(d, s, target, T, method))
            r[name + " IC"] = round(m["mean"], 4)
            r[name + " t"] = round(m["t"], 2)
        rows[s] = r
    tab = pd.DataFrame(rows).T
    print(tab.to_string())
    return tab


def ts_score(df, signal, target, method="spearman"):
    """Single series (no entity): one correlation over time, with a rough t. Check overlap first."""
    d = df[[signal, target]].dropna()
    r = d[signal].corr(d[target], method=method)
    t = r * np.sqrt((len(d) - 2) / max(1e-12, 1 - r * r))
    print("%s vs %s: r %.4f, t %.2f, n %d (t assumes independent rows)" % (signal, target, r, t, len(d)))
    return r, t


def incremental_t(df, signal, controls, target, T, min_n=10):
    """Does signal add anything beyond the controls? Per-timestamp rank regression; t of signal's slope."""
    cols = [signal] + list(controls)
    b = []
    for _, g in df.dropna(subset=cols + [target]).groupby(T):
        if len(g) < min_n:
            continue
        X = np.column_stack([np.ones(len(g))] + [g[c].rank().values for c in cols])
        b.append(np.linalg.lstsq(X, g[target].rank().values, rcond=None)[0][1])
    m = summary(b)
    print("%s beyond %s: slope t %.2f over %d timestamps" % (signal, ", ".join(controls), m["t"], m["n"]))
    return m


def by_period(ics, freq="M"):
    """IC averaged per month (or 'W'): is the edge stable or decaying?"""
    p = ics.groupby(ics.index.to_period(freq)).mean().round(4)
    print(p.to_string())
    return p


# ---------- 6. can you trade it --------------------------------------------------

def quantile_returns(df, signal, target, T, q=10, scale=1e4):
    """Mean target per signal quantile (in bps by default), low to high, plus top minus bottom."""
    d = df.dropna(subset=[signal, target]).copy()
    d["_q"] = d.groupby(T)[signal].transform(
        lambda s: pd.qcut(s.rank(method="first"), q, labels=False) if s.size >= q else np.nan)
    m = d.groupby("_q")[target].mean() * scale
    print("per quantile:", m.round(2).tolist())
    print("top minus bottom: %.2f" % (m.iloc[-1] - m.iloc[0]))
    return m


def top_overlap(df, signal, T, E, q=10):
    """Share of the top quantile that is still there next timestamp. 0.1 means ~90% turnover."""
    d = df.dropna(subset=[signal]).copy()
    d["_r"] = d.groupby(T)[signal].rank(pct=True)
    tops = d[d["_r"] > 1 - 1.0 / q].groupby(T)[E].apply(set)
    ov = [len(a & b) / len(a) for a, b in zip(tops.iloc[:-1], tops.iloc[1:]) if len(a)]
    m = float(np.mean(ov)) if ov else np.nan
    print("top-quantile overlap step to step: %.2f" % m)
    return m


# ---------- 7. figures for slides -----------------------------------------------

def save_rolling_ic(ics, window, path, title=""):
    """Rolling mean of the IC series, saved as a PNG for slide 6."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(8, 3.5))
    ics.rolling(window).mean().plot(ax=ax)
    ax.axhline(0, color="grey", lw=0.8)
    ax.set_title(title or "rolling %d-period mean IC" % window)
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print("saved", path)


def save_bars(series, path, title="", ylabel=""):
    """Bar chart of a small Series (e.g. quantile returns) as a PNG."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(6, 3.5))
    series.plot.bar(ax=ax)
    ax.axhline(0, color="grey", lw=0.8)
    ax.set_title(title)
    ax.set_ylabel(ylabel)
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)
    print("saved", path)