"""Return-of-capital (ROC) detector from N-CSR annual reports.

A mutual fund's per-share "Financial Highlights" table (item 7 of the
annual report, filed as N-CSR) lists distributions by character:

    Distributions from:
        Net investment income ...
        (Realized) capital gains ...
        (Tax) return of capital ...
        Total distributions ...

The "return of capital" line is the definitive answer to "is this
fund paying me back my own money?" - the price-data heuristic
(NAV falling while paying out) only SUGGESTS it, and for
long-duration funds NAV declines are usually rate damage, not ROC.

Facts used:
- The table shows up to 5-10 fiscal years, MOST RECENT LEFT.
- Values are parenthesized outflows: "(1.16)" = $1.16 distributed.
- A row with fewer tokens than the NII row has blanks dropped; the
  arithmetic check (NII + gains + ROC == total per column) is used
  to place a lone value (verified on PGSIX: the lone (0.35) only
  reconciles in the oldest column: 0.14 + 0.35 = 0.49).

Run:  python -m fundlab.rocdetect [SYM ...]
      (no args = the 16 shortlist + 22 cross-checked funds)
Output: fundlab/roc_results.json  (cached per fund in roc_cache/)
"""
from __future__ import annotations

import json
import re
import time
from pathlib import Path

from fundlab import edgar

HERE = Path(__file__).parent
CACHE = HERE / "roc_cache"
RESULTS = HERE / "roc_results.json"
NUM = re.compile(r"\$?\(?[\d,]+(?:\.\d+)?\)?:?")


def _nums(tokens: list[str]) -> list[float | None]:
    out = []
    for t in tokens:
        t = t.strip().rstrip(":")
        if not re.fullmatch(r"\$?\(?\d[\d,]*(?:\.\d+)?\)?", t):
            out.append(None)
            continue
        neg = t.startswith("(")
        v = float(t.strip("$()").replace(",", ""))
        out.append(-v if neg else v)
    return out


NUMTOK = re.compile(r"\$?\(?\d[\d,]*(?:\.\d+)?\)?")


def _parse_table(block: str) -> dict | None:
    """Parse one per-share distribution block."""

    def row(label_rx: str) -> list[float | None]:
        rm = re.search(label_rx, block, re.I)
        if not rm:
            return []
        tail = re.sub(r"\.{2,}", " ", block[rm.end(): rm.end() + 600])
        # some reports print "(1.01 )" / "( 1.01 )" - rejoin the parens
        tail = re.sub(r"\(\s+", "(", tail)
        tail = re.sub(r"\s+\)", ")", tail)
        vals: list[str] = []
        for t in re.findall(r"\S+", tail):
            if t == "$":
                continue         # "NAV ... $ 15.88 $ 15.23 ..."
            if NUMTOK.fullmatch(t):
                vals.append(t)
            elif vals:
                break            # numbers over - next label reached
        return _nums(vals[:10])

    nii = row(r"net investment income")
    gains = row(r"(?:realized |unrealized )?capital gains?")
    roc = row(r"(?:tax )?return of capital")
    tot = row(r"total distributions")
    nav = row(r"net asset value, end of year")
    if not nii and not tot:
        return None
    ncols = max(len(x) for x in (nii, gains, roc, tot, nav))
    if ncols == 0:
        return None
    # place short rows by arithmetic: find the column where the
    # components sum to the total (within rounding)
    def _get(row: list[float | None], i: int) -> float:
        if i < len(row) and row[i] is not None:
            return row[i]
        return 0.0

    def place(vals: list[float | None]) -> list[float | None]:
        if len(vals) == ncols:
            return vals
        from itertools import combinations
        nz = [v for v in vals if v is not None]
        if not nz:
            return [0.0] * ncols

        def reconcile(assign: list[int], nz_: list[float]) -> bool:
            put = {c: v for c, v in zip(assign, nz_)}
            for c in range(ncols):
                if not tot or c >= len(tot) or tot[c] is None:
                    continue
                s = (_get(nii, c) + _get(gains, c) + put.get(c, 0.0))
                if abs(s - tot[c]) > 0.011 * max(1.0, abs(tot[c])) + 1e-9:
                    return False
            return True

        # blanks are dropped when the table flattens, so a short row's
        # values may sit in ANY columns - try placements and keep the
        # one where every column reconciles (NII+gains+ROC == total)
        if len(nz) <= 6 and tot:
            for assign in combinations(range(ncols), len(nz)):
                if reconcile(list(assign), nz):
                    put = {c: v for c, v in zip(assign, nz)}
                    return [put.get(c, 0.0) for c in range(ncols)]
        full = [0.0] * ncols
        for i, v in enumerate(vals):
            if v is not None:
                full[i] = v
        return full
    roc = place(roc)
    # sanity: the per-share table has small values.  Dollar-basis
    # statements ("distributions to shareholders: $459,612,118") and
    # fiscal-year date columns ("2025") leak in from other sections.
    def _bad(vals: list[float | None]) -> bool:
        for v in vals:
            if v is None:
                continue
            if abs(v) > 1000:
                return True
            if v == int(v) and 1900 <= v <= 2100:
                return True
        return False
    if _bad(tot) or _bad(nav) or _bad(nii):
        return None
    if tot and all(v == 0 for v in tot if v is not None):
        return None     # all-zero total = wrong table (these all pay)
    # table-level validity: every column with a total must reconcile
    # (NII + gains + ROC == total).  Catches rows that leaked in from
    # the operations section or prose ("...capital gains of $3.00...").
    if tot:
        for c in range(ncols):
            t = tot[c] if c < len(tot) else None
            if t is None:
                continue
            s = _get(nii, c) + _get(gains, c) + (roc[c] if c < len(roc) else 0.0)
            if abs(s - t) > 0.011 * max(1.0, abs(t)) + 0.005:
                return None
    return {"nii": nii, "gains": gains, "roc": roc, "tot": tot,
            "nav": nav, "ncols": ncols}


def parse_highlights(txt: str, nav_now: float | None = None) -> dict | None:
    """Extract the per-share distribution table for ONE fund.

    Anchors = 'distributions from / declared to / to shareholders' lines
    immediately followed by 'net investment income' (the per-share
    table, not the dollar-basis statement).  Family annual reports
    carry one table PER FUND (NexPoint's report has three) - when the
    fund's current NAV is known, pick the table whose end-of-year NAV
    is closest to it; otherwise the first.
    """
    txt = re.sub(r"\s+", " ", txt)
    tables: list[dict] = []
    for m in re.finditer(
            r"distributions (?:declared )?(?:from|to shareholders)",
            txt, re.I):
        if not re.search(r"net investment income",
                         txt[m.end(): m.end() + 300], re.I):
            continue
        t = _parse_table(txt[m.start(): m.start() + 3000])
        if t:
            tables.append(t)
    if not tables:
        return None
    if nav_now:
        def navdist(t: dict) -> float:
            v = t["nav"][0] if t["nav"] and t["nav"][0] else None
            return abs(v - nav_now) / nav_now if v else 99.0
        tables.sort(key=navdist)
        if navdist(tables[0]) > 0.15:
            return None     # no table matches this fund's NAV
    return tables[0]


def roc_share(h: dict) -> float | None:
    tot = sum(v or 0 for v in h["tot"])
    if not tot:
        return None
    return sum(v or 0 for v in h["roc"]) / tot


def _nav_now(sym: str) -> float | None:
    """Latest raw NAV from the price file (to disambiguate family
    reports that carry one per-share table per fund)."""
    try:
        import pandas as pd
        p = Path.home() / "prog" / "fin" / "stocks" / f"{sym.lower()}-history.csv"
        d = pd.read_csv(p, parse_dates=["Date"])
        c = d["Close"].dropna()
        return float(c.iloc[-1]) if len(c) else None
    except Exception:
        return None


def detect(sym: str, force: bool = False) -> dict | None:
    """Full pipeline for one fund: CIK -> latest N-CSR -> highlights."""
    CACHE.mkdir(exist_ok=True)
    cpath = CACHE / f"{sym.lower()}.json"
    if cpath.exists() and not force:
        return json.loads(cpath.read_text())
    tc = edgar.ticker_to_company(sym)
    if not tc:
        return {"sym": sym, "error": "CIK not found"}
    cik, name = tc
    fil = edgar.cik_recent_filings(cik, "N-CSR", count=1)
    if not fil:
        return {"sym": sym, "error": "no N-CSR on file", "cik": cik}
    f0 = fil[0]
    txt = edgar.to_text(edgar.sec_get(
        edgar.doc_url(cik, f0["accession"], f0["doc"]), timeout=120))
    h = parse_highlights(txt, nav_now=_nav_now(sym))
    rs = roc_share(h) if h else None
    out = {
        "sym": sym, "cik": cik, "trust": name,
        "filing": f0["filed"],
        "table": h,
        "roc_share_5y": (round(rs, 3) if rs is not None else None),
        "nav_change_5y": (round(h["nav"][0] / h["nav"][-1] - 1, 3)
                          if h and len(h["nav"]) >= 2
                          and h["nav"][0] and h["nav"][-1] else None),
    }
    cpath.write_text(json.dumps(out, indent=1))
    return out


def _default_syms() -> list[str]:
    dr = json.loads((HERE / "decompose_results.json").read_text())
    xc = json.loads((HERE / "xcheck_report.json").read_text())
    return sorted(set(s.upper() for s in dr)
                  | set(s.upper() for s in xc))


def _print(res: dict) -> None:
    print(f"  {'fund':<7} {'ROC 5y':>7} {'NAV 5y':>8}  {'dist/yr (most recent first)':<44} filing")
    for s in sorted(res, key=lambda x: -(res[x].get("roc_share_5y")
                                         or -1)):
        r = res[s]
        if "error" in r:
            print(f"  {s:<7}   ({r['error']})")
            continue
        h = r["table"]
        if not h:
            print(f"  {s:<7}   (no per-share table found)")
            continue
        yrs = []
        for i in range(h["ncols"]):
            t = h["tot"][i] if i < len(h["tot"]) else None
            ro = h["roc"][i] if i < len(h["roc"]) else None
            yrs.append(f"tot {t if t is not None else 0:.2f}"
                       + (f"/roc {ro:.2f}" if ro else ""))
        rc = r.get("roc_share_5y")
        nv = r.get("nav_change_5y")
        print(f"  {s:<7} {(f'{rc*100:5.1f}%' if rc is not None else '   -'):>7} "
              f"{(f'{nv*100:+6.1f}%' if nv is not None else '    -'):>8}  "
              f"{'  '.join(yrs)[:44]:<44} {r['filing']}")


def run(syms: list[str] | None = None, force: bool = False) -> dict:
    syms = syms or _default_syms()
    res: dict = {}
    for s in syms:
        for attempt in range(3):
            try:
                r = detect(s, force=force)
                break
            except Exception as e:  # SEC 503/429 throttling - back off
                if attempt == 2:
                    r = {"sym": s, "error": f"{type(e).__name__}: {e}"}
                else:
                    time.sleep(15 * (attempt + 1))
                    r = None
        if r:
            res[s.upper()] = r
        time.sleep(0.5)
    RESULTS.write_text(json.dumps(res, indent=1))
    _print(res)
    return res


if __name__ == "__main__":
    import sys
    args = [a for a in sys.argv[1:] if a != "--force"]
    run(args or None, force="--force" in sys.argv)