"""
"Lazy Prices" text similarity (Cohen, Malloy & Nguyen 2020, JF) for every
S&P 500 ever-member in the US research panel, from SEC EDGAR.

    PYTHONPATH=. uv run python -m research.lab.fetch_lazy_prices              # index, 10-K, 10-Q
    PYTHONPATH=. uv run python -m research.lab.fetch_lazy_prices --forms 10-K # 10-Ks only
    nohup env PYTHONPATH=. uv run python -m research.lab.fetch_lazy_prices \
        > logs/fetch_lazy_prices.log 2>&1 &

Stages, each resumable (a rerun skips what is already on disk):
  1. index   submissions JSON of each CIK (+ its older pages from 2004 on) →
             data/research/lazy_prices_index.csv (original 10-K/10-K405/10-Q),
             data/research/sec_sic.csv (ticker, cik, sic, sicDescription, name),
             data/research/sec_8k_earnings.csv (8-K and 8-K/A with item 2.02).
  2. 10-K    each 10-K vs the previous 10-K (period end 300-430 days earlier).
  3. 10-Q    each 10-Q vs the 10-Q of the same fiscal quarter a year earlier
             (period end 350-380 days earlier).

Text: the filing's primary document only (no exhibits); HTML, hidden inline
XBRL, scripts and tables whose characters are >15% digits are dropped, then
lines that are >25% digits and "table of contents" headers; lowercase words
of 2+ letters minus English stop words. Similarities: cosine of term-count
vectors and Jaccard of word sets, for the whole document and (10-K only) for
Item 1A Risk Factors when a section of 200+ words can be located.

Output: data/research/lazy_prices.csv, one row per (ticker, filing), appended
as filings finish: ticker, cik, form, filed, period, accession, ref_accession,
cos_sim, jac_sim, n_words, cos_rf, jac_rf, n_words_rf. No documents are kept;
data/research/lazy_prices_vec/ holds the term-count vectors of the filings a
later filing still needs as its reference (plus the latest 10-K / four 10-Qs).

SEC fair access: one User-Agent, at most 5 requests/second across all worker
processes (a shared limiter), direct connection (proxy variables ignored).
Amendments (10-K/A, 10-Q/A) are left out as in the paper: they are often
partial documents and would read as "changers".
"""

from __future__ import annotations

import argparse
import csv
import fcntl
import html
import multiprocessing as mp
import os
import pickle
import re
import time
import traceback
import zlib
from collections import Counter
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import pandas as pd
import requests

from research.lab.fetch_us_fundamentals import cik_map
from research.lab.panel import US_ETFS

DIR = os.path.join("data", "research")
INDEX = os.path.join(DIR, "lazy_prices_index.csv")
OUT = os.path.join(DIR, "lazy_prices.csv")
SIC = os.path.join(DIR, "sec_sic.csv")
EARN = os.path.join(DIR, "sec_8k_earnings.csv")
VEC = os.path.join(DIR, "lazy_prices_vec")
COLS = ["ticker", "cik", "form", "filed", "period", "accession", "ref_accession",
        "cos_sim", "jac_sim", "n_words", "cos_rf", "jac_rf", "n_words_rf"]

UA = {"User-Agent": "quant-research/0.1 gahow@example.com", "Accept-Encoding": "gzip, deflate"}
MIN_INTERVAL = 0.21            # seconds between request starts, all processes (≈4.8/s)
START = "2004-01-01"
GROUPS = {"10-K": ("10-K", "10-K405"), "10-Q": ("10-Q",)}
# ever-members missing from today's company_tickers.json; CIKs checked by name
EXTRA_CIKS = {"AVB": 915912, "EQR": 906107, "EA": 712515, "GENZ": 732485, "GR": 42542}

STOP = frozenset("""
a about above after again against all am an and any are as at be because been before being
below between both but by can cannot could did do does doing down during each few for from
further had has have having he her here hers herself him himself his how i if in into is it
its itself let me more most my myself no nor not of off on once only or other ought our ours
ourselves out over own same she should so some such than that the their theirs them themselves
then there these they this those through to too under until up very was we were what when where
which while who whom why with would you your yours yourself yourselves will shall
""".split())

# ---------------------------------------------------------------------------
# HTTP: a limiter shared by threads and worker processes
# ---------------------------------------------------------------------------

_LOCK = None
_NEXT = None
_SESSION: requests.Session | None = None


def _init_worker(lock, nxt) -> None:
    global _LOCK, _NEXT, _SESSION
    _LOCK, _NEXT = lock, nxt
    _SESSION = requests.Session()
    _SESSION.trust_env = False            # direct connection, never the shell's proxy


def _throttle() -> None:
    with _LOCK:
        t = max(time.time(), _NEXT.value)
        _NEXT.value = t + MIN_INTERVAL
    wait = t - time.time()
    if wait > 0:
        time.sleep(wait)


def _get(url: str) -> requests.Response | None:
    """GET with the shared limiter; None for a missing document."""
    for attempt in range(7):
        _throttle()
        try:
            r = _SESSION.get(url, headers=UA, timeout=120)
        except requests.RequestException as exc:
            print(f"  ! {url}: {exc.__class__.__name__}; retry {attempt + 1}", flush=True)
            time.sleep(5 * (attempt + 1))
            continue
        if r.status_code == 200:
            return r
        if r.status_code == 404:
            return None
        rate = r.status_code == 429 or (r.status_code == 403
                                        and b"Request Rate Threshold" in r.content[:5000])
        if r.status_code == 403 and not rate:
            return None
        # rate threshold (SEC blocks for ~10 minutes) or server error: back off
        wait = 60 * (attempt + 1) if rate else 10 * (attempt + 1)
        print(f"  ! {url}: HTTP {r.status_code}; sleeping {wait}s", flush=True)
        time.sleep(wait)
    return None


# ---------------------------------------------------------------------------
# stage 1: filing index, SIC codes, 8-K earnings dates
# ---------------------------------------------------------------------------

def _columns(block: dict) -> pd.DataFrame:
    keys = ["accessionNumber", "filingDate", "reportDate", "acceptanceDateTime", "form",
            "items", "primaryDocument"]
    return pd.DataFrame({k: block.get(k, [""] * len(block["accessionNumber"])) for k in keys})


def _submissions(cik: int) -> tuple[dict, pd.DataFrame] | None:
    r = _get(f"https://data.sec.gov/submissions/CIK{cik:010d}.json")
    if r is None:
        return None
    js = r.json()
    parts = [_columns(js["filings"]["recent"])]
    for f in js["filings"].get("files", []):
        if f.get("filingTo", "9999") < START:
            continue
        r2 = _get(f"https://data.sec.gov/submissions/{f['name']}")
        if r2 is not None:
            parts.append(_columns(r2.json()))
    df = pd.concat(parts, ignore_index=True).drop_duplicates("accessionNumber")
    meta = {"sic": js.get("sic", ""), "sicDescription": js.get("sicDescription", ""),
            "name": js.get("name", "")}
    return meta, df


def ticker_ciks(tickers: list[str]) -> dict[str, int]:
    cm = cik_map()
    out = {t: cm[t] for t in tickers if t in cm}
    for t, c in EXTRA_CIKS.items():
        if t in tickers and t not in out:
            out[t] = c
    return out


def build_index(tickers: list[str]) -> None:
    tc = ticker_ciks(tickers)
    print(f"{len(tc)}/{len(tickers)} panel tickers map to a CIK "
          f"({len(set(tc.values()))} CIKs; {len(tc) - sum(t in EXTRA_CIKS for t in tc)} "
          f"from company_tickers.json)", flush=True)
    by_cik: dict[int, list[str]] = {}
    for t, c in tc.items():
        by_cik.setdefault(c, []).append(t)
    res: dict[int, tuple] = {}
    t0 = time.time()
    with ThreadPoolExecutor(max_workers=5) as ex:
        for k, (cik, out) in enumerate(zip(by_cik, ex.map(_submissions, by_cik)), 1):
            if out is not None:
                res[cik] = out
            if k % 50 == 0:
                print(f"  submissions {k}/{len(by_cik)} [{time.time() - t0:.0f}s]", flush=True)
    sic, idx, earn = [], [], []
    for cik, (meta, df) in res.items():
        for t in sorted(by_cik[cik]):
            sic.append({"ticker": t, "cik": cik, **meta})
            e = df[df["form"].isin(["8-K", "8-K/A"])
                   & df["items"].fillna("").str.split(",").apply(lambda xs: "2.02" in [x.strip() for x in xs])]
            if len(e):
                earn.append(e.assign(ticker=t, cik=cik)[
                    ["ticker", "cik", "form", "filingDate", "acceptanceDateTime", "accessionNumber", "items"]])
        f = df[df["form"].isin(GROUPS["10-K"] + GROUPS["10-Q"]) & (df["filingDate"] >= START)]
        idx.append(pd.DataFrame({"cik": cik, "tickers": "|".join(sorted(by_cik[cik])), "form": f["form"],
                                 "filed": f["filingDate"], "period": f["reportDate"],
                                 "accession": f["accessionNumber"], "doc": f["primaryDocument"]}))
    pd.DataFrame(sic).sort_values("ticker").to_csv(SIC, index=False)
    ear = pd.concat(earn, ignore_index=True).sort_values(["ticker", "filingDate"])
    ear.to_csv(EARN, index=False)
    ix = pd.concat(idx, ignore_index=True).sort_values(["cik", "filed"])
    ix.to_csv(INDEX, index=False)
    print(f"index: {len(ix)} filings ({(ix['form'] != '10-Q').sum()} 10-K) for {ix['cik'].nunique()} CIKs; "
          f"{len(ear)} 8-K item 2.02 rows ({ear['filingDate'].min()} → {ear['filingDate'].max()}); "
          f"sic for {len(sic)} tickers [{time.time() - t0:.0f}s]", flush=True)


def with_refs(ix: pd.DataFrame, group: str) -> pd.DataFrame:
    """Filings of one form group with their comparable prior filing."""
    f = ix[ix["form"].isin(GROUPS[group])].copy()
    f["filed_d"] = pd.to_datetime(f["filed"])
    f["period_d"] = pd.to_datetime(f["period"], errors="coerce")
    # a period filed twice as an original (re-filing) keeps the first copy
    f = f.sort_values("filed_d")
    f = pd.concat([f[f["period_d"].isna()],
                   f[f["period_d"].notna()].drop_duplicates(["cik", "period_d"], keep="first")])
    lo, hi = (300, 430) if group == "10-K" else (350, 380)
    refs = []
    for _, g in f.sort_values(["cik", "filed_d"]).groupby("cik", sort=False):
        pe, fe, acc = g["period_d"].values, g["filed_d"].values, g["accession"].values
        for i in range(len(g)):
            best, best_gap = "", None
            for j in range(i - 1, -1, -1):
                if not pd.isna(pe[i]) and not pd.isna(pe[j]):
                    gap = (pe[i] - pe[j]).astype("timedelta64[D]").astype(int)
                    ok = lo <= gap <= hi
                else:                       # no period on one side: filing dates, wider window
                    gap = (fe[i] - fe[j]).astype("timedelta64[D]").astype(int)
                    ok = lo - 20 <= gap <= hi + 20
                if ok and (best_gap is None or abs(gap - 365) < abs(best_gap - 365)):
                    best, best_gap = acc[j], gap
                if gap > hi + 60:
                    break
            refs.append((g.index[i], best))
    f.loc[[i for i, _ in refs], "ref"] = [r for _, r in refs]
    return f.sort_values(["cik", "filed_d"]).drop(columns=["filed_d", "period_d"])


# ---------------------------------------------------------------------------
# text → term counts
# ---------------------------------------------------------------------------

_BLOCK = ("p", "div", "br", "tr", "li", "h1", "h2", "h3", "h4", "h5", "h6", "table", "pre",
          "center", "ul", "ol", "dl", "dt", "dd", "blockquote", "hr")
_DROP_TAGS = {"script", "style", "head", "title", "ix:header", "ix:hidden", "xbrli:context"}
_DIG = re.compile(r"\d")
_ALPHA = re.compile(r"[A-Za-z]")
_WORD = re.compile(r"[a-z]{2,}")
_TOC = re.compile(r"^\s*(table\s+of\s+contents|index|page)\s*$", re.I)
_RF_START = re.compile(r"^[ \t]*item[\s.]*1a\b[\s.:\-–—]*risk\s+factors", re.I | re.M)
_RF_END = re.compile(r"^[ \t]*item[\s.]*(?:1b|2)\b", re.I | re.M)
_HTMLISH = re.compile(rb"<(html|body|div|p|font|br|table)[\s>/]", re.I)


def _digit_share(s: str) -> float:
    d = len(_DIG.findall(s))
    a = len(_ALPHA.findall(s))
    return d / (a + d) if a + d else 0.0


def _html_text(raw: bytes) -> str:
    import lxml.html
    parser = lxml.html.HTMLParser(huge_tree=True, remove_comments=True, remove_pis=True)
    root = lxml.html.document_fromstring(raw, parser=parser)
    drop = []
    for el in root.iter():
        if not isinstance(el.tag, str):
            continue
        if el.tag.lower() in _DROP_TAGS:
            drop.append(el)
            continue
        st = el.get("style")
        if st and "display:none" in st.replace(" ", "").lower():
            drop.append(el)
    for el in drop:
        el.drop_tree()
    for tb in reversed(list(root.iter("table"))):        # innermost first
        if _digit_share(tb.text_content()) > 0.15:
            tb.drop_tree()
    for el in root.iter(*_BLOCK):
        el.tail = "\n" + (el.tail or "")
    for el in root.iter("td", "th"):
        el.tail = " " + (el.tail or "")
    return root.text_content()


def _plain_text(raw: bytes) -> str:
    s = raw.decode("utf-8", errors="replace")
    s = re.sub(r"<[^>\n]{1,40}>", " ", s)                 # EDGAR SGML tags (<PAGE>, <S>, <C>)
    return html.unescape(s)


def _primary_part(raw: bytes) -> bytes:
    """The first <DOCUMENT> of a full-submission text file (the form itself)."""
    if b"<DOCUMENT>" not in raw[:20000]:
        return raw
    m = re.search(rb"<DOCUMENT>(.*?)(?:</DOCUMENT>|\Z)", raw, re.S)
    body = m.group(1) if m else raw
    t = re.search(rb"<TEXT>(.*?)(?:</TEXT>|\Z)", body, re.S)
    return t.group(1) if t else body


def clean_text(raw: bytes) -> str:
    raw = _primary_part(raw)
    try:
        txt = _html_text(raw) if _HTMLISH.search(raw[:20000]) else _plain_text(raw)
    except Exception:  # noqa: BLE001 — malformed markup: fall back to tag stripping
        txt = _plain_text(re.sub(rb"<[^>]*>", b" ", raw))
    txt = txt.replace("\xa0", " ").replace("\r", "\n")
    keep = []
    for line in txt.split("\n"):
        line = line.strip()
        if not line or _TOC.match(line) or _digit_share(line) > 0.25:
            continue
        keep.append(line)
    return "\n".join(keep)


def tokens(text: str) -> Counter:
    return Counter(w for w in _WORD.findall(text.lower()) if w not in STOP and len(w) <= 30)


def risk_factors(text: str) -> str | None:
    """Item 1A of a 10-K: the longest span from a line-leading "Item 1A. Risk
    Factors" to the next line-leading Item 1B / Item 2 (the table of contents
    gives short spans)."""
    ends = [m.start() for m in _RF_END.finditer(text)]
    best = None
    for m in _RF_START.finditer(text):
        e = next((x for x in ends if x > m.end()), None)
        if e is None:
            continue
        if best is None or e - m.end() > best[1] - best[0]:
            best = (m.end(), e)
    if best is None:
        return None
    sec = text[best[0]:best[1]]
    return sec if len(sec.split()) >= 200 else None


def cosine(a: Counter, b: Counter) -> float:
    if not a or not b:
        return np.nan
    if len(a) > len(b):
        a, b = b, a
    dot = sum(v * b.get(k, 0) for k, v in a.items())
    na = np.sqrt(sum(v * v for v in a.values()))
    nb = np.sqrt(sum(v * v for v in b.values()))
    return float(dot / (na * nb))


def jaccard(a: Counter, b: Counter) -> float:
    if not a or not b:
        return np.nan
    sa, sb = set(a), set(b)
    return len(sa & sb) / len(sa | sb)


def doc_vectors(cik: int, accession: str, doc: str, group: str) -> dict | None:
    """Term counts of a filing's primary document (+ Item 1A for a 10-K).
    None when the download failed (retry later); empty counts when the
    document exists but holds no text."""
    acc = accession.replace("-", "")
    name = doc if isinstance(doc, str) and doc.strip() else f"{accession}.txt"
    r = _get(f"https://www.sec.gov/Archives/edgar/data/{cik}/{acc}/{name}")
    if r is None:
        return None
    if name.lower().endswith(".pdf") or r.headers.get("content-type", "").startswith("application/pdf"):
        return {"vec": Counter(), "rf": None}
    text = clean_text(r.content)
    rf = risk_factors(text) if group == "10-K" else None
    return {"vec": tokens(text), "rf": tokens(rf) if rf else None}


# ---------------------------------------------------------------------------
# stages 2-3: per-CIK chains of filings
# ---------------------------------------------------------------------------

def _cache_path(cik: int, group: str) -> str:
    return os.path.join(VEC, f"{cik}_{group}.pkl.z")


def _load_cache(cik: int, group: str) -> dict:
    path = _cache_path(cik, group)
    if not os.path.exists(path):
        return {}
    try:
        with open(path, "rb") as fh:
            return pickle.loads(zlib.decompress(fh.read()))
    except Exception:  # noqa: BLE001 — a torn write: rebuild from downloads
        return {}


def _save_cache(cik: int, group: str, cache: dict) -> None:
    path = _cache_path(cik, group)
    tmp = path + f".{os.getpid()}.tmp"
    with open(tmp, "wb") as fh:
        fh.write(zlib.compress(pickle.dumps(cache, protocol=5), 6))
    os.replace(tmp, path)


def _append_rows(rows: list[dict]) -> None:
    with open(OUT, "a", newline="") as fh:
        fcntl.flock(fh, fcntl.LOCK_EX)
        w = csv.DictWriter(fh, fieldnames=COLS)
        for r in rows:
            w.writerow(r)
        fh.flush()
        fcntl.flock(fh, fcntl.LOCK_UN)


def work(job: tuple) -> tuple[int, int, int]:
    """Process the pending filings of one CIK in filing order. Returns
    (cik, filings written, failures)."""
    cik, group, filings, done = job
    keep_last = 1 if group == "10-K" else 4
    cache = _load_cache(cik, group)
    accs = [f["accession"] for f in filings]
    n_ok = n_fail = 0
    for i, f in enumerate(filings):
        if f["accession"] in done:
            continue
        try:
            cur = cache.get(f["accession"]) or doc_vectors(cik, f["accession"], f["doc"], group)
            if cur is None:
                n_fail += 1
                continue
            ref = None
            if f["ref"]:
                ref = cache.get(f["ref"])
                if ref is None:
                    rf_row = next(g for g in filings if g["accession"] == f["ref"])
                    ref = doc_vectors(cik, rf_row["accession"], rf_row["doc"], group)
                    if ref is None:
                        n_fail += 1
                        continue
                    cache[f["ref"]] = ref
            cache[f["accession"]] = cur
            done.add(f["accession"])
            # keep what a pending filing still needs, plus the latest filings
            needed = {g["ref"] for g in filings if g["accession"] not in done and g["ref"]}
            needed |= set(accs[max(0, i + 1 - keep_last): i + 1])
            cache = {a: v for a, v in cache.items() if a in needed}
            _save_cache(cik, group, cache)
            v, rv = cur["vec"], (ref or {}).get("vec")
            base = {"cik": cik, "form": f["form"], "filed": f["filed"], "period": f["period"],
                    "accession": f["accession"], "ref_accession": f["ref"] or "",
                    "n_words": sum(v.values()),
                    "cos_sim": cosine(v, rv) if rv else np.nan,
                    "jac_sim": jaccard(v, rv) if rv else np.nan,
                    "n_words_rf": sum(cur["rf"].values()) if cur.get("rf") else np.nan}
            rr = (ref or {}).get("rf")
            base["cos_rf"] = cosine(cur["rf"], rr) if cur.get("rf") and rr else np.nan
            base["jac_rf"] = jaccard(cur["rf"], rr) if cur.get("rf") and rr else np.nan
            for c in ("cos_sim", "jac_sim", "cos_rf", "jac_rf"):
                if not np.isnan(base[c]):
                    base[c] = round(base[c], 6)
            _append_rows([{"ticker": t, **base} for t in f["tickers"].split("|")])
            n_ok += 1
        except Exception:  # noqa: BLE001 — log and leave it for a rerun
            n_fail += 1
            print(f"  ! {cik} {f['accession']}:\n{traceback.format_exc()}", flush=True)
    return cik, n_ok, n_fail


def run_group(ix: pd.DataFrame, group: str, workers: int, lock, nxt) -> None:
    f = with_refs(ix, group)
    f["ref"] = f["ref"].fillna("")
    f["doc"] = f["doc"].fillna("")
    f["period"] = f["period"].fillna("")
    done = set()
    if os.path.exists(OUT):
        done = set(pd.read_csv(OUT, usecols=["accession"])["accession"])
    jobs = []
    for cik, g in f.groupby("cik"):
        recs = g[["form", "filed", "period", "accession", "ref", "doc", "tickers"]].to_dict("records")
        d = {r["accession"] for r in recs if r["accession"] in done}
        if len(d) < len(recs):
            jobs.append((int(cik), group, recs, d))
    total = sum(len(j[2]) - len(j[3]) for j in jobs)
    print(f"=== {group}: {len(f)} filings, {len(f) - total} done, {total} to fetch "
          f"({len(jobs)} CIKs); {(f['ref'] != '').mean():.1%} have a comparable prior filing",
          flush=True)
    if not jobs:
        return
    jobs.sort(key=lambda j: -(len(j[2]) - len(j[3])))       # long chains first
    t0, n, fails = time.time(), 0, 0
    with mp.get_context("fork").Pool(workers, initializer=_init_worker, initargs=(lock, nxt)) as pool:
        for k, (cik, ok, bad) in enumerate(pool.imap_unordered(work, jobs), 1):
            n += ok
            fails += bad
            el = time.time() - t0
            if k % 10 == 0 or k == len(jobs):
                eta = (total - n) / max(n / el, 1e-9) / 3600
                print(f"  {group} {k}/{len(jobs)} CIKs, {n}/{total} filings, {fails} failed, "
                      f"{n / el:.2f}/s, ETA {eta:.1f}h [{el / 60:.0f} min]", flush=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--forms", nargs="*", default=["10-K", "10-Q"], choices=list(GROUPS))
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--refresh-index", action="store_true")
    args = ap.parse_args()
    os.makedirs(VEC, exist_ok=True)
    lock, nxt = mp.Lock(), mp.Value("d", 0.0)
    _init_worker(lock, nxt)
    if args.refresh_index or not os.path.exists(INDEX):
        with open(os.path.join(DIR, "us_ohlcv.pkl"), "rb") as fh:
            tickers = [t for t in pickle.load(fh)["close"].columns if t not in US_ETFS]
        build_index(tickers)
    ix = pd.read_csv(INDEX, dtype={"period": str, "doc": str})
    if not os.path.exists(OUT):
        with open(OUT, "w", newline="") as fh:
            csv.DictWriter(fh, fieldnames=COLS).writeheader()
    for group in args.forms:
        run_group(ix, group, args.workers, lock, nxt)
    print("done", flush=True)


if __name__ == "__main__":
    main()