#!/usr/bin/env python3
"""Repair Yahoo adj-close / event-date misalignment (see
scripts/scan_adj_misalign.py).

For a hit date D (adj spikes, close flat, price drops later) the repair is
purely arithmetic and leaves cumulative returns unchanged:

  r(t) = adj(t)/close(t);  k = r(D)/r(D-1)          # the bogus spike
  f   = 1 - 1/k                                      # dist fraction Yahoo applied
  d   = f * close(D-1)                               # implied per-share dist
  E   = first t > D with close(t) < close(t-1) - 0.5*d   # true ex-div date
  for t in [D, E):  adj(t) *= (1 - f)                # restore pre-event ratio

For t before D and from E on the series is already consistent, so only the
[D, E) window is patched. The patch is written as a 'history' correction
(overrides/corrections/{SYM}.json, merged, with a note), so it survives
re-downloads that rewrite the raw CSVs with the same Yahoo values.

Usage:
  python3 scripts/fix_adj_misalign.py [--only SYM[,SYM...]] [--dry-run]
"""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from scripts.scan_adj_misalign import scan  # noqa: E402

ROOT_DEFAULT = Path.home() / "prog/fin/stocks"
CORR = Path(__file__).resolve().parent.parent / "overrides" / "corrections"
NOTE = ("adj-close repair: Yahoo dated a distribution on {d} (record date) "
        "but the price went ex-div on {e} (close {c0:.2f} -> {c1:.2f}, drop "
        "~{d_amt:.3f}); the raw Adj Close column spiked {spike:+.1%} on {d} "
        "and snapped back later. Cells in [{d}, {e}) rescaled by (1-f) so "
        "the adjusted path is smooth; cumulative returns are unchanged. "
        "Mechanism: scripts/fix_adj_misalign.py, "
        "scan: reports/adj_misalign/scan.md.")


def load_rows(sym: str, root: Path):
    import data as D
    fr = D.FROZEN_DIR / f"{sym}-history.csv"
    f = fr if (D.FROZEN_DIR.is_dir() and fr.exists()) else root / f"{sym}-history.csv"
    rows = []
    for l in f.read_text(errors="replace").splitlines()[1:]:
        if not l.strip():
            continue
        a = l.split(",")
        if len(a) < 6:
            continue
        try:
            rows.append((a[0][:10], float(a[4]), float(a[5])))
        except ValueError:
            continue
    return rows


def repair(sym: str, date: str, rows) -> dict:
    """Return {date: {"Adj Close": value}} patch, or {} if not repairable."""
    i = next((k for k, r in enumerate(rows) if r[0] == date), None)
    if i is None or i == 0:
        return {}
    dates, closes, adjs = (list(x) for x in zip(*rows))
    r = [a / c for _, c, a in rows]
    if r[i - 1] <= 0 or r[i] <= 0 or r[i] <= r[i - 1]:
        return {}
    f = 1.0 - r[i - 1] / r[i]
    if not (0.005 < f < 0.95):
        return {}
    d_amt = f * closes[i - 1]
    E = None
    for t in range(i + 1, min(i + 11, len(rows))):
        if closes[t] < closes[t - 1] - 0.5 * d_amt:
            E = t
            break
    if E is None:
        return {}
    k = 1.0 - f
    patch = {}
    for t in range(i, E):
        patch[dates[t]] = {"Adj Close": round(adjs[t] * k, 6)}
    return patch, f, d_amt, E


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", type=Path, default=ROOT_DEFAULT)
    ap.add_argument("--only", default="", help="comma-separated symbols")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()
    only = {s.strip().lower() for s in args.only.split(",") if s.strip()}

    hits = scan(args.data, 0.05, 0.02, 3)
    if only:
        hits = [h for h in hits if h["symbol"] in only]
    by_sym: dict[str, list] = {}
    for h in hits:
        by_sym.setdefault(h["symbol"], []).append(h)
    applied = 0
    for sym, hs in sorted(by_sym.items()):
        rows = load_rows(sym, args.data)
        cells, notes = {}, []
        for h in hs:
            try:
                patch, f, d_amt, E = repair(sym, h["date"], rows)
            except Exception as ex:
                print(f"  !! {sym} {h['date']}: {ex}")
                continue
            if not patch:
                print(f"  -- {sym} {h['date']}: not repairable (skipped)")
                continue
            cells.update(patch)
            dates = [x[0] for x in rows]
            di = next(k for k, r in enumerate(rows) if r[0] == h["date"])
            notes.append(NOTE.format(d=h["date"], e=dates[E], c0=rows[di - 1][1],
                                     c1=rows[E][1], d_amt=d_amt,
                                     spike=h["adj_ret"]))
            print(f"  {sym} {h['date']}: patch {len(patch)} cell(s), "
                  f"implied dist {d_amt:.3f}, ex-div {dates[E]}")
        if not cells:
            continue
        pf = CORR / f"{sym.upper()}.json"
        d = json.loads(pf.read_text()) if pf.exists() else {"symbol": sym}
        d.setdefault("symbol", sym)
        d.setdefault("note", "")
        if notes:
            d["note"] = (d["note"] + " " if d["note"] else "") + \
                " ".join(notes).strip()
        d.setdefault("history", {})
        for dt, cells_ in cells.items():
            d["history"].setdefault(dt, {}).update(cells_)
        if not args.dry_run:
            pf.write_text(json.dumps(d, indent=2) + "\n")
        applied += 1
        print(f"{'(dry) ' if args.dry_run else ''}wrote {pf.name}: "
              f"{sum(len(v) for v in d['history'].values())} cell(s)")
    print(f"{'dry-run: ' if args.dry_run else ''}{applied} symbol(s) repaired "
          f"of {len(by_sym)} candidate(s)")
    return 0


if __name__ == "__main__":
    main()