"""
Fetch point-in-time CSI 300 / CSI 500 membership and daily bars from baostock.

The live CN panel (data/cn.csv) holds today's CSI 300 members only, so every
CN backtest on it carries survivorship bias. baostock serves historical index
membership, delisted stocks, suspensions (tradestatus), ST flags, turnover and
valuation ratios, which lets research backtests trade the index as it was.

Usage:
    uv run --with baostock python -m research.lab.fetch_cn [--start 2007-01-01] [--workers 4]

Writes data/research/cn_ohlcv.pkl:
    fields: {"open","high","low","close","preclose","volume","amount","turn",
             "pctchg","tradestatus","isst","pe","pb","ps","pcf"} DataFrames
             (dates x codes; prices back-adjusted, 后复权)
    members: {"hs300": {date: [codes]}, "zz500": {date: [codes]}}
    index: DataFrame of CSI 300 / CSI 500 index closes
"""

import argparse
import os
import pickle
from datetime import date
from multiprocessing import Pool

import pandas as pd

OUT = os.path.join("data", "research", "cn_ohlcv.pkl")
K_FIELDS = ("date,code,open,high,low,close,preclose,volume,amount,turn,"
            "tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST")
NUMERIC = ["open", "high", "low", "close", "preclose", "volume", "amount", "turn",
           "pctChg", "peTTM", "pbMRQ", "psTTM", "pcfNcfTTM", "tradestatus", "isST"]
RENAME = {"pctChg": "pctchg", "peTTM": "pe", "pbMRQ": "pb", "psTTM": "ps",
          "pcfNcfTTM": "pcf", "isST": "isst"}


def _rows(rs) -> list[list[str]]:
    out = []
    while rs.error_code == "0" and rs.next():
        out.append(rs.get_row_data())
    return out


def _login() -> None:
    import baostock as bs
    bs.login()


def fetch_members(start: str) -> dict[str, dict[str, list[str]]]:
    path = os.path.join(BAR_DIR, f"members_{start}_{date.today()}.pkl")
    if os.path.exists(path):
        return pd.read_pickle(path)
    members = _fetch_members(start)
    pd.to_pickle(members, path)
    return members


def _fetch_members(start: str) -> dict[str, dict[str, list[str]]]:
    import baostock as bs
    bs.login()
    months = pd.date_range(start, date.today(), freq="MS")
    members: dict[str, dict[str, list[str]]] = {"hs300": {}, "zz500": {}}
    for d in months:
        ds = d.strftime("%Y-%m-%d")
        members["hs300"][ds] = [r[1] for r in _rows(bs.query_hs300_stocks(date=ds))]
        members["zz500"][ds] = [r[1] for r in _rows(bs.query_zz500_stocks(date=ds))]
    bs.logout()
    return members


BAR_DIR = os.path.join("data", "research", "cn_bars")


def fetch_one(args: tuple[str, str, str]) -> pd.DataFrame | None:
    """Fetch one stock's bars, cached per code so an interrupted run resumes."""
    code, start, end = args
    path = os.path.join(BAR_DIR, f"{code}_{end}.pkl")
    if os.path.exists(path):
        return pd.read_pickle(path)
    df = _fetch_one(code, start, end)
    pd.to_pickle(df, path)
    return df


def _fetch_one(code: str, start: str, end: str) -> pd.DataFrame | None:
    import baostock as bs
    for _ in range(3):
        rs = bs.query_history_k_data_plus(code, K_FIELDS, start_date=start, end_date=end,
                                          frequency="d", adjustflag="1")
        if rs.error_code == "0":
            rows = _rows(rs)
            if not rows:
                return None
            df = pd.DataFrame(rows, columns=rs.fields)
            df["date"] = pd.to_datetime(df["date"])
            for c in NUMERIC:
                df[c] = pd.to_numeric(df[c], errors="coerce")
            return df
        bs.logout()
        bs.login()
    print(f"  failed {code}: {rs.error_msg}", flush=True)
    return None


def fetch_index(start: str, end: str) -> pd.DataFrame:
    import baostock as bs
    bs.login()
    out = {}
    for code in ("sh.000300", "sh.000905"):
        rs = bs.query_history_k_data_plus(code, "date,close", start_date=start, end_date=end,
                                          frequency="d")
        df = pd.DataFrame(_rows(rs), columns=rs.fields)
        out[code] = pd.Series(pd.to_numeric(df["close"]).values, index=pd.to_datetime(df["date"]))
    bs.logout()
    return pd.DataFrame(out)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--start", default="2007-01-01")
    ap.add_argument("--member-start", default="2008-01-01")
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--hs300-only", action="store_true")
    args = ap.parse_args()
    end = date.today().strftime("%Y-%m-%d")

    os.makedirs(BAR_DIR, exist_ok=True)
    print("Fetching index membership...", flush=True)
    members = fetch_members(args.member_start)
    hs300 = sorted({c for lst in members["hs300"].values() for c in lst})
    codes = list(hs300)
    if not args.hs300_only:
        codes += sorted({c for lst in members["zz500"].values() for c in lst} - set(hs300))
    print(f"  {len(codes)} codes ever in the index set", flush=True)

    frames = []
    with Pool(args.workers, initializer=_login) as pool:
        for i, df in enumerate(pool.imap(fetch_one, [(c, args.start, end) for c in codes])):
            if df is not None:
                frames.append(df)
            if (i + 1) % 100 == 0:
                print(f"  bars {i + 1}/{len(codes)}", flush=True)
    bars = pd.concat(frames, ignore_index=True)
    fields = {}
    for c in NUMERIC:
        fields[RENAME.get(c, c)] = bars.pivot(index="date", columns="code", values=c).sort_index()

    panel = {"fields": fields, "members": members, "index": fetch_index(args.start, end)}
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "wb") as fh:
        pickle.dump(panel, fh)
    c = fields["close"]
    print(f"Saved {OUT}: {c.shape[1]} codes, {c.index[0].date()} → {c.index[-1].date()}")


if __name__ == "__main__":
    main()