"""
Point-in-time index membership reconstruction — fixes survivorship bias.

Approach: Wikipedia's "Selected changes to the list of S&P 500 components"
table lists every add/remove event (394 rows back to 1976, as of 2026). We
start from today's membership and walk the change log *backward*:
  - An 'Added' ticker on date D was NOT a member before D.
  - A 'Removed' ticker on date D WAS a member before D.
Applied iteratively, this yields the set of members on any historical date.

The membership info is cached in data/sp500_history.json so Wikipedia is hit
at most once per day. The cache stores per-ticker membership intervals:
    { "ticker": [[start, end_or_null], ...] }
where dates are YYYY-MM-DD strings.
"""

import io
import json
import os
import urllib.request
from datetime import date, datetime

import numpy as np
import pandas as pd

CACHE_DIR = "data"
_HEADERS = {"User-Agent": "Mozilla/5.0 (quant-backtest)"}


# ---------------------------------------------------------------------------
# Fetch + parse Wikipedia
# ---------------------------------------------------------------------------

SP500_LIST_URL = "https://en.wikipedia.org/wiki/List_of_S%26P_500_companies"
# The change log moved off the list page in 2026; the list page is still tried
# first in case it moves back.
SP500_CHANGES_URLS = [
    SP500_LIST_URL,
    "https://en.wikipedia.org/wiki/Historical_components_of_the_S%26P_500",
]


def _read_tables(url: str) -> list[pd.DataFrame]:
    req = urllib.request.Request(url, headers=_HEADERS)
    with urllib.request.urlopen(req) as resp:
        html = resp.read().decode("utf-8")
    return pd.read_html(io.StringIO(html))


def _find_changes_table(tables: list[pd.DataFrame]) -> pd.DataFrame | None:
    for t in tables:
        if any("Added" in "_".join(map(str, c if isinstance(c, tuple) else (c,))) for c in t.columns):
            changes = t.copy()
            changes.columns = [
                "_".join(c).strip() if isinstance(c, tuple) else c
                for c in changes.columns
            ]
            changes.columns = [
                c.replace("Effective Date_Effective Date", "Date") for c in changes.columns
            ]
            return changes
    return None


def _fetch_sp500_tables() -> tuple[pd.DataFrame, pd.DataFrame]:
    """Return (current_list, changes_log) from Wikipedia."""
    current = _read_tables(SP500_LIST_URL)[0]
    for url in SP500_CHANGES_URLS:
        changes = _find_changes_table(_read_tables(url))
        if changes is not None:
            return current, changes
    raise ValueError("S&P 500 change log table not found on Wikipedia")


def _normalize_ticker(t: str) -> str:
    """Yahoo Finance ticker format: BRK.B → BRK-B."""
    # Some change-log cells carry a trailing " |" separator.
    return str(t).split("|")[0].replace(".", "-").strip()


# ---------------------------------------------------------------------------
# Membership reconstruction
# ---------------------------------------------------------------------------

def build_sp500_history() -> dict[str, list[list[str | None]]]:
    """
    Reconstruct per-ticker membership intervals.

    Returns
    -------
    dict: ticker -> list of [start_date, end_date_or_None] pairs.
        end_date=None means the ticker is still a member as of today.
        Dates are YYYY-MM-DD strings.

    Algorithm: start from today's set of members, walk the change log from
    newest to oldest. For each event on date D:
      - The 'Added' ticker: its current (open) interval starts on D.
        Close it: [..., D] — it was NOT a member before D.
      - The 'Removed' ticker: it was a member up to D (exclusive).
        Open a new interval ending on D (start unknown for now; will be
        closed by an earlier event or left open-start).

    After the walk, any ticker still "open" (never closed backward) has an
    interval reaching back before the earliest logged change.
    """
    current, changes = _fetch_sp500_tables()

    current_tickers = {_normalize_ticker(s) for s in current["Symbol"].tolist()}

    # Parse change log
    changes["dt"] = pd.to_datetime(changes["Date"], errors="coerce")
    changes = changes.dropna(subset=["dt"]).sort_values("dt", ascending=False)

    # For each ticker, collect intervals [start, end].
    # We track a "current open interval" per ticker during the backward walk.
    # intervals[ticker] = list of [start, end] completed intervals (oldest-first).
    # open_start[ticker] = start date of the currently open (most-recent) interval.
    intervals: dict[str, list[list[str | None]]] = {}
    open_end: dict[str, str | None] = {}  # end of currently-open interval

    # Initialize: today's members have an open interval ending = None (still in)
    for t in current_tickers:
        open_end[t] = None  # still a member today
        intervals[t] = []

    # Track the start date of each open interval as we walk backward.
    # For a member today, the interval started at the last "Added" event in the
    # changes log, OR before the log begins if never added.
    # We'll close the interval when we hit the "Added" event going backward.
    open_start: dict[str, str | None] = {t: None for t in current_tickers}

    for _, row in changes.iterrows():
        d = row["dt"].strftime("%Y-%m-%d")
        added = row.get("Added_Ticker")
        removed = row.get("Removed_Ticker")

        if pd.notna(added):
            a = _normalize_ticker(added)
            # This ticker was added on d → its open interval starts on d.
            if a in open_end:
                open_start[a] = d
                # Finalize the current open interval
                intervals[a].append([d, open_end[a]])
                # Pop: no further open interval backward in time for this ticker
                # (unless 'Removed' opens a new older one below)
                del open_end[a]

        if pd.notna(removed):
            r = _normalize_ticker(removed)
            # This ticker was removed on d → it WAS a member before d.
            # Open a new interval ending on d (start unknown yet).
            if r not in open_end:
                intervals.setdefault(r, [])
                open_end[r] = d  # end of the new older interval

    # Any ticker still with an open interval → start predates the log.
    # Use the oldest logged date as a conservative "unknown earlier" marker: None.
    for t, end in open_end.items():
        intervals.setdefault(t, []).append([None, end])

    _close_unlogged_additions(intervals, current)

    # Sort intervals per ticker oldest→newest
    for t, ivs in intervals.items():
        ivs.sort(key=lambda iv: (iv[0] or "0000-00-00"))

    return intervals


# The change log is sparse before 2011 (0-2 logged additions a year in
# 2001-2006, 6-12 in 2007-2010, against 20-30 real changes), so a current
# member whose addition it missed would count as a member since before the
# log: 80 of the 94 current members added 2001-2010 did, e.g. ISRG, MA, BKNG.
# The list page's "Date added" starts those intervals. Later dates are left to
# the log, which is complete there, because "Date added" also records renames
# and re-listings (META 2013, CPAY 2018) as additions.
UNLOGGED_ADDITIONS_BEFORE = "2011-01-01"


def _close_unlogged_additions(intervals: dict, current: pd.DataFrame) -> None:
    col = next((c for c in current.columns if "added" in str(c).lower()), None)
    if col is None:
        print("--- WARNING: no 'Date added' column; pre-2011 additions stay open-start ---")
        return
    added = pd.to_datetime(current[col], errors="coerce")
    for symbol, day in zip(current["Symbol"], added):
        if pd.isna(day) or day >= pd.Timestamp(UNLOGGED_ADDITIONS_BEFORE):
            continue
        for iv in intervals.get(_normalize_ticker(symbol), []):
            if iv[0] is None and iv[1] is None:
                iv[0] = day.strftime("%Y-%m-%d")


# ---------------------------------------------------------------------------
# Cache I/O
# ---------------------------------------------------------------------------

def _cache_path() -> str:
    return os.path.join(CACHE_DIR, "sp500_history.json")


def load_sp500_history(force_refresh: bool = False, *, refresh: bool = True) -> dict[str, list[list[str | None]]]:
    """Load cached membership history, or rebuild if stale (>1 day old)."""
    path = _cache_path()
    if not force_refresh and os.path.exists(path):
        try:
            with open(path) as f:
                data = json.load(f)
            if not refresh or data.get("date") == str(date.today()):
                return data["intervals"]
        except Exception:
            pass
    print("--- Rebuilding S&P 500 membership history from Wikipedia ---")
    try:
        intervals = build_sp500_history()
    except Exception as exc:  # noqa: BLE001 — page layout changes must not break backtests
        if not os.path.exists(path):
            raise
        with open(path) as f:
            cached = json.load(f)
        print(f"--- WARNING: rebuild failed ({exc}); using membership history "
              f"cached on {cached.get('date')} ---")
        return cached["intervals"]
    os.makedirs(CACHE_DIR, exist_ok=True)
    with open(path, "w") as f:
        json.dump({"date": str(date.today()), "intervals": intervals}, f)
    print(f"--- Cached {len(intervals)} tickers' membership intervals ---")
    return intervals


# ---------------------------------------------------------------------------
# Convert intervals → aligned mask DataFrame
# ---------------------------------------------------------------------------

def membership_mask(dates: pd.DatetimeIndex,
                    intervals: dict[str, list[list[str | None]]] | None = None,
                    tickers: list[str] | None = None) -> pd.DataFrame:
    """
    Boolean DataFrame: rows = dates, columns = tickers.
    True where the ticker was an S&P 500 member on that date.

    If `tickers` is given, restrict columns to that list (useful for aligning
    with a price DataFrame). Otherwise, include all tickers ever a member.
    """
    if intervals is None:
        intervals = load_sp500_history()
    cols = tickers if tickers is not None else sorted(intervals.keys())
    # Tickers not in `intervals` (e.g. SPY, benchmarks, ETFs) are treated as
    # always-members so callers can pass the full price matrix through
    # mask_prices without zeroing out benchmark series.
    mask = pd.DataFrame(False, index=dates, columns=cols)
    for t in cols:
        if t not in intervals:
            mask[t] = True
            continue
        for start, end in intervals[t]:
            s = pd.Timestamp(start) if start else dates[0]
            e = pd.Timestamp(end) if end else dates[-1] + pd.Timedelta(days=1)
            # Interval semantics: member on [start, end). A ticker removed on
            # date D was no longer a member on D.
            mask.loc[(mask.index >= s) & (mask.index < e), t] = True
    return mask


def all_tickers_ever(intervals: dict | None = None) -> list[str]:
    """All tickers that were ever S&P 500 members (for price data fetching)."""
    if intervals is None:
        intervals = load_sp500_history()
    return sorted(intervals.keys())


def mask_prices(prices: pd.DataFrame,
                intervals: dict | None = None) -> pd.DataFrame:
    """
    Return a copy of `prices` with NaN set for (date, ticker) pairs where
    the ticker was not an S&P 500 member on that date.

    This is the key survivorship-bias fix: strategies compute signals from
    the masked price data, so they naturally cannot select stocks outside
    the point-in-time index membership.

    Warm-up note: a newly-added member needs sufficient non-NaN history for
    its rolling windows to produce a valid signal. For this codebase's
    ~252-day lookbacks, a stock becomes "selectable" roughly 1 year after
    joining. This is conservative but correct: before that, we have no
    legitimate signal anyway.
    """
    mask = membership_mask(prices.index, intervals, tickers=list(prices.columns))
    return prices.where(mask)


# ---------------------------------------------------------------------------
# CSI 300 (China A-shares)
# ---------------------------------------------------------------------------
#
# data/csi300_history.json holds monthly constituent snapshots
#     {"snapshots": {"YYYY-MM-DD": [yahoo tickers], ...}}
# seeded from baostock's query_hs300_stocks (first trading day of each month,
# 2008 onward). A snapshot applies from its date until the next one, so an
# index change is picked up within a month. load_csi300_history() appends a
# snapshot whenever the current constituent list (universe.get_csi300) differs
# from the latest one, which keeps the file current without baostock.

def _csi300_path() -> str:
    return os.path.join(CACHE_DIR, "csi300_history.json")


def baostock_to_yahoo(code: str) -> str:
    """'sh.600000' -> '600000.SS', 'sz.000001' -> '000001.SZ'."""
    ex, num = code.split(".")
    return f"{num}.{'SS' if ex == 'sh' else 'SZ'}"


def load_csi300_history(current: list[str] | None = None, *, refresh: bool = True) -> dict[str, list[str]]:
    """Monthly CSI 300 snapshots; appends ``current`` (or today's list) if it changed."""
    path = _csi300_path()
    with open(path) as f:
        data = json.load(f)
    snaps = data["snapshots"]
    if current is None and refresh:
        try:
            from universe import get_csi300
            current = get_csi300()
        except Exception as exc:  # noqa: BLE001 — offline runs use the stored history
            print(f"--- WARNING: CSI 300 list unavailable ({exc}); using stored history ---")
            current = None
    if current:
        last = snaps[max(snaps)]
        today = str(date.today())
        if sorted(current) != sorted(last) and today > max(snaps):
            snaps[today] = sorted(current)
            with open(path, "w") as f:
                json.dump(data, f, indent=0)
            print(f"--- CSI 300 constituents changed; stored snapshot {today} ---")
    return snaps


def pit_membership(market: str, dates: pd.DatetimeIndex, *, refresh: bool = False) -> pd.DataFrame:
    if market == "us":
        intervals = load_sp500_history(refresh=refresh)
        return membership_mask(dates, intervals, sorted(intervals))
    if market == "cn":
        snapshots = load_csi300_history(refresh=refresh)
        return csi300_membership_mask(dates, all_csi300_tickers_ever(snapshots), snapshots)
    raise ValueError(f"No PIT membership for {market}")


def csi300_membership_mask(dates: pd.DatetimeIndex, tickers: list[str],
                           snapshots: dict[str, list[str]] | None = None) -> pd.DataFrame:
    """Boolean mask (dates x tickers): True while the ticker was in the CSI 300."""
    snaps = snapshots if snapshots is not None else load_csi300_history()
    keys = sorted(snaps)
    col = {t: i for i, t in enumerate(tickers)}
    arr = np.zeros((len(dates), len(tickers)), dtype=bool)
    for i, k in enumerate(keys):
        start = pd.Timestamp(k)
        end = pd.Timestamp(keys[i + 1]) if i + 1 < len(keys) else dates[-1] + pd.Timedelta(days=1)
        rows = (dates >= start) & (dates < end)
        cols = [col[t] for t in snaps[k] if t in col]
        if rows.any() and cols:
            arr[np.ix_(rows, cols)] = True
    mask = pd.DataFrame(arr, index=dates, columns=tickers)
    # Tickers never in the index (benchmarks, ETFs) pass through unmasked.
    ever = {t for lst in snaps.values() for t in lst}
    for t in tickers:
        if t not in ever:
            mask[t] = True
    return mask


def all_csi300_tickers_ever(snapshots: dict[str, list[str]] | None = None,
                            since: str | None = None) -> list[str]:
    snaps = snapshots if snapshots is not None else load_csi300_history()
    return sorted({t for k, lst in snaps.items() if since is None or k >= since for t in lst})


# ---------------------------------------------------------------------------
# Market-generic helpers
# ---------------------------------------------------------------------------

def pit_tickers_ever(market: str, since: str | None = None) -> list[str] | None:
    """Every ticker that was an index member (since ``since``), or None when the
    market has no membership history."""
    if market == "us":
        return all_tickers_ever()
    if market == "cn":
        return all_csi300_tickers_ever(since=since)
    return None


def mask_prices_pit(prices: pd.DataFrame, market: str) -> pd.DataFrame:
    """NaN out (date, ticker) pairs outside point-in-time index membership."""
    if market == "us":
        return mask_prices(prices)
    if market == "cn":
        return prices.where(csi300_membership_mask(prices.index, list(prices.columns)))
    return prices