"""Point-in-time, stratified instrument universe with an immutable manifest (docs/DATA_CARD.md §universe).

Selection information set = instrument master (first/last date, delisted status) + daily bars up to the selection
date only (average dollar volume, daily volatility). Strata: asset class (stock / ETF), dollar-volume tercile,
volatility tercile, SIC division (stocks, from SEC submissions when available; "unknown" otherwise) or ETF category
(keyword rule on the vendor name). Delisted instruments with 1-minute history in the sample window are eligible; the
only survival requirement is presence for >= `min_months_in_sample` months after the sample start (documented bias).
"""
from __future__ import annotations

import hashlib
import json
import math
import random
from dataclasses import dataclass, asdict
from datetime import date, datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

from .client import FinDataClient

ETF_RULES = [
    ("leveraged_inverse", ["2X", "3X", "ULTRA", "INVERSE", "SHORT ", "BEAR", "BULL", "-1X", "-2X", "-3X", "PROSHARES"]),
    ("fixed_income", ["BOND", "TREASURY", "CREDIT", "MUNI", "AGGREGATE", "FIXED INCOME", "TIPS", "YIELD", "DEBT"]),
    ("commodity", ["GOLD", "SILVER", "OIL", "COMMODIT", "NATURAL GAS", "METAL", "ENERGY FUND", "CRUDE"]),
    ("international", ["EMERGING", "INTERNATIONAL", "EAFE", "CHINA", "JAPAN", "EUROPE", "GLOBAL", "ASIA", "BRAZIL",
                       "INDIA", "KOREA", "MEXICO", "CANADA", "WORLD", "FRONTIER", "LATIN"]),
    ("sector", ["SECTOR", "TECHNOLOGY", "FINANCIAL", "HEALTH", "UTILITIES", "ENERGY", "MATERIALS", "INDUSTRIAL",
                "CONSUMER", "REAL ESTATE", "SEMICONDUCTOR", "BIOTECH", "BANK", "RETAIL", "HOMEBUILDER", "TRANSPORT"]),
    ("volatility_alt", ["VIX", "VOLATILITY", "CURRENCY", "DOLLAR", "YEN", "EURO ", "BITCOIN", "CRYPTO"]),
    ("broad_equity", ["S&P", "500", "RUSSELL", "NASDAQ", "DOW", "TOTAL", "MARKET", "VALUE", "GROWTH", "DIVIDEND",
                      "MID-CAP", "SMALL-CAP", "LARGE-CAP", "EQUITY", "INDEX", "CORE", "QUALITY", "MOMENTUM"]),
]

SIC_DIVISIONS = [(100, 999, "agriculture"), (1000, 1499, "mining"), (1500, 1799, "construction"),
                 (2000, 3999, "manufacturing"), (4000, 4999, "transport_utilities"), (5000, 5199, "wholesale"),
                 (5200, 5999, "retail"), (6000, 6799, "finance_insurance_realestate"), (7000, 8999, "services"),
                 (9100, 9999, "public_admin")]


def sic_division(sic) -> str:
    try:
        s = int(sic)
    except (TypeError, ValueError):
        return "unknown"
    for lo, hi, name in SIC_DIVISIONS:
        if lo <= s <= hi:
            return name
    return "unknown"


def etf_category(name: str | None) -> str:
    n = (name if isinstance(name, str) else "").upper()
    for cat, kws in ETF_RULES:
        if any(k in n for k in kws):
            return cat
    return "other"


@dataclass
class UniverseConfig:
    selection_date: str = "2017-12-29"
    adv_window_start: str = "2017-12-01"
    sample_start: str = "2018-01-02"
    sample_end: str = "2025-12-31"
    min_history_months: int = 6         # first_date <= selection_date - 6 months
    min_months_in_sample: int = 6       # last_date >= sample_start + 6 months (survival requirement, documented)
    min_adv_usd: float = 5e6            # average daily dollar volume floor at selection (minute bars need activity)
    n_target: int = 300
    etf_share: float = 0.25
    holdout_share: float = 0.20
    seed: int = 41
    min_coverage: float = 0.0           # regular-hours minute coverage over the sample window (P1 coverage audit gate)
    coverage_file: str | None = None    # results/metrics/coverage_audit_candidates.csv


def daily_cross_sections(client: FinDataClient, start: str, end: str, asset_class: str) -> pd.DataFrame:
    """Daily bars for all symbols on each session in [start, end] (one scoped SQL per session date)."""
    from .calendar import sessions
    view = f"market_bars_{asset_class}_1d_raw"
    frames = []
    for s in sessions(start, end):
        d = s.date.date().isoformat()
        cols, rows, meta = client.sql(f"SELECT symbol_as_reported AS symbol, open, high, low, close, volume "
                                      f"FROM {view} WHERE session_date = DATE '{d}'", max_rows=100_000)
        df = pd.DataFrame(rows, columns=cols)
        df["session_date"] = d
        frames.append(df)
    out = pd.concat(frames, ignore_index=True)
    out["asset_class"] = asset_class
    return out


def selection_stats(cross: pd.DataFrame) -> pd.DataFrame:
    """ADV (USD), daily log-return volatility and number of sessions per symbol from the cross-sections."""
    cross = cross.sort_values(["symbol", "session_date"]).copy()
    cross["dollar_volume"] = cross["close"] * cross["volume"]
    cross["lr"] = np.log(cross["close"]).groupby(cross["symbol"]).diff()
    g = cross.groupby("symbol")
    out = pd.DataFrame({"adv_usd": g["dollar_volume"].mean(), "daily_vol": g["lr"].std(), "n_days": g.size(),
                        "last_close": g["close"].last(), "asset_class": g["asset_class"].first()})
    return out.reset_index()


def load_sic_map(sec_tickers_json: Path | None, sec_submissions_dir: Path | None) -> dict[str, str]:
    """ticker -> SIC division from the SEC company_tickers file + cached submissions (public data, optional)."""
    if not sec_tickers_json or not Path(sec_tickers_json).exists():
        return {}
    data = json.loads(Path(sec_tickers_json).read_text())
    if isinstance(data, dict) and "data" in data and "cik_str" not in data.get("0", {}):
        data = data["data"]
    tick2cik = {}
    for v in (data.values() if isinstance(data, dict) else data):
        tick2cik[str(v["ticker"]).upper()] = int(v["cik_str"])
    out = {}
    if sec_submissions_dir and Path(sec_submissions_dir).exists():
        for t, cik in tick2cik.items():
            p = Path(sec_submissions_dir) / f"CIK{cik:010d}.json"
            if p.exists():
                try:
                    out[t] = sic_division(json.loads(p.read_text()).get("sic"))
                except Exception:
                    out[t] = "unknown"
    return out


def tercile(x: pd.Series) -> pd.Series:
    r = x.rank(pct=True, method="first")
    return pd.cut(r, [0, 1 / 3, 2 / 3, 1.0001], labels=["low", "mid", "high"], include_lowest=True).astype(str)


def build_candidates(client: FinDataClient, cfg: UniverseConfig, sic_map: dict[str, str]) -> pd.DataFrame:
    sel = pd.Timestamp(cfg.selection_date)
    first_max = sel - pd.DateOffset(months=cfg.min_history_months)
    last_min = pd.Timestamp(cfg.sample_start) + pd.DateOffset(months=cfg.min_months_in_sample)
    inst = []
    for ac in ("stock", "etf"):
        for r in client.instruments(asset_class=ac):
            inst.append(r)
    master = pd.DataFrame(inst)
    master = master[master["asset_class"].isin(["stock", "etf"])].copy()
    master["first_date"] = pd.to_datetime(master["first_date"])
    master["last_date"] = pd.to_datetime(master["last_date"])
    elig = master[(master["first_date"] <= first_max) & (master["last_date"] >= last_min)].copy()
    stats = pd.concat([selection_stats(daily_cross_sections(client, cfg.adv_window_start, cfg.selection_date, ac))
                       for ac in ("stock", "etf")], ignore_index=True)
    stats = stats.drop(columns=["asset_class"])
    cand = elig.merge(stats, on="symbol", how="inner")
    cand = cand[(cand["adv_usd"] >= cfg.min_adv_usd) & (cand["n_days"] >= 15) & cand["daily_vol"].notna()].copy()
    if cfg.coverage_file:
        cov = pd.read_csv(cfg.coverage_file)[["symbol", "coverage", "sessions_present"]]
        cand = cand.merge(cov, on="symbol", how="left")
        cand["coverage"] = cand["coverage"].fillna(0.0)
        cand = cand[cand["coverage"] >= cfg.min_coverage].copy()
    cand["base_symbol"] = cand["symbol"].str.replace("-DELISTED", "", regex=False)
    cand["delisted"] = cand["status"].eq("delisted")
    cand["group"] = np.where(cand["asset_class"].eq("etf"), cand["name"].map(etf_category),
                             cand["base_symbol"].map(lambda t: sic_map.get(t, "unknown")))
    cand["adv_tercile"] = cand.groupby("asset_class")["adv_usd"].transform(tercile)
    cand["vol_tercile"] = cand.groupby("asset_class")["daily_vol"].transform(tercile)
    cand["stratum"] = cand["asset_class"] + "|" + cand["adv_tercile"] + "|" + cand["vol_tercile"]
    return cand.sort_values("symbol").reset_index(drop=True)


def stratified_sample(cand: pd.DataFrame, cfg: UniverseConfig) -> pd.DataFrame:
    """Proportional allocation across strata (asset class × ADV tercile × vol tercile), delisted instruments kept with
    their natural weight, sector/ETF category balanced inside each stratum by round-robin."""
    rng = random.Random(cfg.seed)
    n_etf = int(round(cfg.n_target * cfg.etf_share))
    n_stock = cfg.n_target - n_etf
    chosen = []
    for ac, n_ac in (("stock", n_stock), ("etf", n_etf)):
        sub = cand[cand["asset_class"] == ac]
        strata = sorted(sub["stratum"].unique())
        if not strata:
            continue
        quota = {s: n_ac // len(strata) for s in strata}
        for s in rng.sample(strata, n_ac - sum(quota.values())):
            quota[s] += 1
        for s in strata:
            pool = sub[sub["stratum"] == s]
            groups = {}
            for sym, grp in zip(pool["symbol"], pool["group"]):
                groups.setdefault(grp, []).append(sym)
            for g in groups.values():
                rng.shuffle(g)
            order = sorted(groups.keys(), key=lambda k: rng.random())
            picked = []
            while len(picked) < min(quota[s], pool.shape[0]):
                for g in order:
                    if groups[g] and len(picked) < quota[s]:
                        picked.append(groups[g].pop())
            chosen.extend(picked)
    out = cand[cand["symbol"].isin(chosen)].copy().sort_values("symbol").reset_index(drop=True)
    syms = out["symbol"].tolist()
    rng2 = random.Random(cfg.seed + 1)
    rng2.shuffle(syms)
    n_hold = int(round(len(syms) * cfg.holdout_share))
    hold = set(syms[:n_hold])
    out["role"] = np.where(out["symbol"].isin(hold), "heldout_instrument", "pooled")
    return out


def manifest(sample: pd.DataFrame, cand: pd.DataFrame, cfg: UniverseConfig, dataset_version: str | None,
             extra: dict | None = None) -> dict:
    rows = sample.to_dict(orient="records")
    for r in rows:
        for k, v in list(r.items()):
            if isinstance(v, (pd.Timestamp, datetime, date)):
                r[k] = str(v)[:10]
            elif isinstance(v, float) and math.isnan(v):
                r[k] = None
            elif isinstance(v, (np.integer,)):
                r[k] = int(v)
            elif isinstance(v, (np.floating,)):
                r[k] = float(v)
            elif isinstance(v, (np.bool_,)):
                r[k] = bool(v)
    body = {"config": asdict(cfg), "created_at": datetime.now(timezone.utc).isoformat(),
            "n_candidates": int(cand.shape[0]), "n_selected": len(rows),
            "n_heldout": int((sample["role"] == "heldout_instrument").sum()),
            "n_delisted": int(sample["delisted"].sum()), "instrument_master_version": dataset_version,
            "selection_rules": [
                f"first_date <= {cfg.selection_date} - {cfg.min_history_months} months",
                f"last_date >= {cfg.sample_start} + {cfg.min_months_in_sample} months (documented survival requirement)",
                f"average daily dollar volume over [{cfg.adv_window_start}, {cfg.selection_date}] >= {cfg.min_adv_usd:,.0f} USD",
                "proportional allocation over asset class x ADV tercile x volatility tercile; round-robin over SIC division / ETF category inside strata",
                f"{cfg.holdout_share:.0%} of selected instruments held out entirely (zero-shot), seed {cfg.seed}",
                f"regular-hours 1-minute coverage over the sample window >= {cfg.min_coverage:.0%} (coverage audit, availability filter)",
            ], "instruments": rows}
    if extra:
        body.update(extra)
    body["sha256"] = hashlib.sha256(json.dumps(body, sort_keys=True, default=str).encode()).hexdigest()
    return body