"""Point-in-time research context for industry, themes, and catalysts.

The module intentionally has no dependency on the M1 pipeline.  An
orchestrator can pass the ``[research_context]`` portion of a TOML file as a
plain mapping and decide where, or whether, to persist the returned tables.
All date filters are applied before rolling calculations so a historical run
cannot accidentally use observations that arrived in the future.
"""
from __future__ import annotations

from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Any, Mapping, Sequence
import json

import numpy as np
import pandas as pd


RESEARCH_CONTEXT_MODEL_VERSION = "0.1.0"

_EMPTY_LIFECYCLE_COLUMNS = [
    "context_key", "date", "hotspot_score", "smoothed_score", "score_change",
    "peak_to_date", "drawdown_from_peak", "hotspot_stage",
]
_EMPTY_MAPPING_COLUMNS = ["source_code", "industry_code", "theme_id", "theme_name", "weight", "role"]
_EMPTY_CATALYST_COLUMNS = [
    "catalyst_id", "source_code", "theme_id", "catalyst_date", "available_date",
    "status", "impact", "probability", "catalyst_score", "days_to_catalyst",
    "catalyst_proximity", "title",
]


def _empty(columns: Sequence[str]) -> pd.DataFrame:
    return pd.DataFrame(columns=list(columns))


def _normalise_stock_code(frame: pd.DataFrame) -> pd.DataFrame:
    result = frame.copy()
    if "source_code" not in result and "stock_code" in result:
        result = result.rename(columns={"stock_code": "source_code"})
    if "source_code" in result:
        result["source_code"] = result["source_code"].astype(str).str.strip().str.lower()
    return result


def _normalise_date(values: pd.Series) -> pd.Series:
    return pd.to_datetime(values, errors="coerce").dt.normalize()


def _column_or_default(frame: pd.DataFrame, column: str, default: object) -> pd.Series:
    """Return a Series even when an optional input column is absent."""
    return frame[column] if column in frame else pd.Series(default, index=frame.index)


def _read_table(value: object, base_dir: Path, preferred_json_key: str | None = None) -> pd.DataFrame:
    """Read CSV/JSON input without treating an omitted optional input as an error."""
    if value is None or (isinstance(value, str) and value == ""):
        return pd.DataFrame()
    if isinstance(value, pd.DataFrame):
        return value.copy()
    if isinstance(value, (list, tuple)):
        return pd.json_normalize(value)
    if isinstance(value, Mapping):
        raw: object = value
        if preferred_json_key and isinstance(raw.get(preferred_json_key), list):
            raw = raw[preferred_json_key]
        elif isinstance(raw.get("records"), list):
            raw = raw["records"]
        return pd.json_normalize(raw)
    path = Path(str(value))
    path = path if path.is_absolute() else base_dir / path
    if not path.exists():
        return pd.DataFrame()
    if path.suffix.lower() == ".csv":
        return pd.read_csv(path)
    if path.suffix.lower() == ".json":
        raw = json.loads(path.read_text(encoding="utf-8"))
        return _read_table(raw, base_dir, preferred_json_key)
    raise ValueError(f"research context only supports CSV/JSON: {path}")


def _context_key(frame: pd.DataFrame) -> pd.Series:
    for column in ("theme_id", "industry_code", "context_key", "theme_name", "industry_name"):
        if column in frame:
            return frame[column].astype(str).str.strip()
    return pd.Series("", index=frame.index, dtype=str)


def _infer_score(frame: pd.DataFrame, date_column: str) -> pd.Series:
    for column in ("hotspot_score", "signal_score", "score"):
        if column in frame:
            return pd.to_numeric(frame[column], errors="coerce").clip(0, 100)

    components = [column for column in ("return_5d", "return_20d", "activity", "attention", "breadth") if column in frame]
    if not components:
        return pd.Series(np.nan, index=frame.index, dtype=float)
    ranks = [pd.to_numeric(frame[column], errors="coerce").groupby(frame[date_column]).rank(pct=True) * 100 for column in components]
    return pd.concat(ranks, axis=1).mean(axis=1)


def classify_hotspot_lifecycle(
    industry_data: pd.DataFrame,
    as_of: str | date | pd.Timestamp | None = None,
    *,
    short_window: int = 5,
    group_column: str | None = None,
    date_column: str = "date",
) -> pd.DataFrame:
    """Classify each dated industry/theme observation using only data available then.

    Stages are ``萌芽``, ``升温``, ``高潮``, ``退潮`` and ``冷却``.  ``as_of`` first
    removes future rows (and rows whose ``available_date`` is not yet known),
    then all smoothing, peaks, and changes are calculated in chronological
    order.  Thus the output is safe to use in a historical backtest.
    """
    if industry_data is None or industry_data.empty:
        return _empty(_EMPTY_LIFECYCLE_COLUMNS)
    frame = industry_data.copy()
    if date_column not in frame:
        return _empty(_EMPTY_LIFECYCLE_COLUMNS)
    frame["date"] = _normalise_date(frame[date_column])
    frame = frame.loc[frame["date"].notna()].copy()
    cutoff = pd.Timestamp(as_of).normalize() if as_of is not None else None
    if cutoff is not None:
        frame = frame.loc[frame["date"].le(cutoff)].copy()
        for availability_column in ("available_date", "available_on", "published_date"):
            if availability_column in frame:
                available = _normalise_date(frame[availability_column])
                frame = frame.loc[available.isna() | available.le(cutoff)].copy()
                break
    if frame.empty:
        return _empty(_EMPTY_LIFECYCLE_COLUMNS)
    if group_column and group_column in frame:
        frame["context_key"] = frame[group_column].astype(str).str.strip()
    else:
        frame["context_key"] = _context_key(frame)
    frame = frame.loc[frame["context_key"].ne("")].copy()
    if frame.empty:
        return _empty(_EMPTY_LIFECYCLE_COLUMNS)
    frame["hotspot_score"] = _infer_score(frame, "date")
    frame = frame.sort_values(["context_key", "date"]).drop_duplicates(["context_key", "date"], keep="last")
    grouped = frame.groupby("context_key", sort=False)["hotspot_score"]
    frame["smoothed_score"] = grouped.transform(lambda item: item.rolling(short_window, min_periods=1).mean())
    frame["score_change"] = frame.groupby("context_key", sort=False)["smoothed_score"].diff(short_window)
    frame["peak_to_date"] = frame.groupby("context_key", sort=False)["smoothed_score"].cummax()
    frame["drawdown_from_peak"] = frame["peak_to_date"] - frame["smoothed_score"]

    score = frame["smoothed_score"]
    change = frame["score_change"].fillna(0.0)
    drawdown = frame["drawdown_from_peak"].fillna(0.0)
    near_peak = score.ge(frame["peak_to_date"] - 4)
    frame["hotspot_stage"] = np.select(
        [
            score.lt(30),
            drawdown.ge(12) & change.le(0),
            score.ge(80) & (change.gt(0) | near_peak),
            score.ge(60) & change.gt(0),
            score.ge(40) & change.ge(0),
        ],
        ["冷却", "退潮", "高潮", "升温", "萌芽"],
        default="冷却",
    )
    return frame.reset_index(drop=True)


def _expand_theme_json(raw: object) -> pd.DataFrame:
    """Expand convenient JSON theme definitions into normal membership rows."""
    if not isinstance(raw, Mapping) or not isinstance(raw.get("themes"), list):
        return pd.DataFrame()
    rows: list[dict[str, object]] = []
    for theme in raw["themes"]:
        if not isinstance(theme, Mapping):
            continue
        common = {key: theme.get(key) for key in ("theme_id", "theme_name", "weight", "role")}
        for member_key, output_key in (("stocks", "source_code"), ("members", "source_code"), ("industry_codes", "industry_code")):
            values = theme.get(member_key, [])
            if not isinstance(values, list):
                continue
            for value in values:
                row = dict(common)
                if isinstance(value, Mapping):
                    row.update(value)
                    if output_key not in row and "stock_code" in row:
                        row[output_key] = row["stock_code"]
                else:
                    row[output_key] = value
                rows.append(row)
    return pd.DataFrame(rows)


def load_theme_mapping(value: object = None, base_dir: Path | str = ".") -> pd.DataFrame:
    """Load custom stock/industry-to-theme membership from CSV or JSON."""
    raw_json: object | None = None
    if isinstance(value, Mapping):
        raw_json = value
    elif value is not None and not isinstance(value, pd.DataFrame) and value != "":
        path = Path(str(value))
        path = path if path.is_absolute() else Path(base_dir) / path
        if path.suffix.lower() == ".json" and path.exists():
            raw_json = json.loads(path.read_text(encoding="utf-8"))
    expanded = _expand_theme_json(raw_json)
    table = _read_table(value, Path(base_dir), "members")
    frame = pd.concat([table, expanded], ignore_index=True, sort=False) if not expanded.empty else table
    if frame.empty:
        return _empty(_EMPTY_MAPPING_COLUMNS)
    frame = _normalise_stock_code(frame)
    if "theme_id" not in frame:
        frame["theme_id"] = frame.get("theme_name", pd.Series("", index=frame.index)).astype(str)
    if "theme_name" not in frame:
        frame["theme_name"] = frame["theme_id"]
    for column in ("source_code", "industry_code", "theme_id", "theme_name", "role"):
        if column not in frame:
            frame[column] = ""
        frame[column] = frame[column].fillna("").astype(str).str.strip()
    frame["weight"] = pd.to_numeric(_column_or_default(frame, "weight", 1.0), errors="coerce").fillna(1.0).clip(lower=0)
    frame = frame.loc[frame["theme_id"].ne("") & (frame["source_code"].ne("") | frame["industry_code"].ne(""))]
    return frame[_EMPTY_MAPPING_COLUMNS].drop_duplicates().reset_index(drop=True)


def map_stock_themes(stocks: pd.DataFrame | Sequence[str], theme_mapping: pd.DataFrame) -> pd.DataFrame:
    """Map stock codes (and optional industry codes) to custom themes."""
    if isinstance(stocks, pd.DataFrame):
        universe = _normalise_stock_code(stocks)
    else:
        universe = pd.DataFrame({"source_code": list(stocks)})
        universe = _normalise_stock_code(universe)
    if "source_code" not in universe:
        raise ValueError("stocks needs a source_code or stock_code column")
    universe = universe.drop_duplicates("source_code")
    if theme_mapping is None or theme_mapping.empty:
        return _empty(_EMPTY_MAPPING_COLUMNS)
    mapping = load_theme_mapping(theme_mapping)
    direct = universe[["source_code"]].merge(mapping.loc[mapping["source_code"].ne("")], on="source_code", how="inner")
    by_industry = _empty(_EMPTY_MAPPING_COLUMNS)
    if "industry_code" in universe and mapping["industry_code"].ne("").any():
        by_industry = universe[["source_code", "industry_code"]].merge(
            mapping.loc[mapping["industry_code"].ne("")].drop(columns="source_code"), on="industry_code", how="inner"
        )
    result = pd.concat([direct, by_industry], ignore_index=True, sort=False)
    if result.empty:
        return _empty(_EMPTY_MAPPING_COLUMNS)
    result["weight"] = pd.to_numeric(result["weight"], errors="coerce").fillna(1.0)
    return result[_EMPTY_MAPPING_COLUMNS].drop_duplicates(["source_code", "theme_id"]).reset_index(drop=True)


def evaluate_catalysts(
    catalysts: pd.DataFrame,
    as_of: str | date | pd.Timestamp | None = None,
    *,
    horizon_days: int = 45,
    grace_days: int = 7,
) -> pd.DataFrame:
    """Return catalysts known at ``as_of`` with status and timing-aware scores."""
    if catalysts is None or catalysts.empty:
        return _empty(_EMPTY_CATALYST_COLUMNS)
    frame = _normalise_stock_code(catalysts)
    if "catalyst_date" not in frame:
        for alternative in ("expected_date", "event_date", "date"):
            if alternative in frame:
                frame = frame.rename(columns={alternative: "catalyst_date"})
                break
    if "catalyst_date" not in frame:
        return _empty(_EMPTY_CATALYST_COLUMNS)
    frame["catalyst_date"] = _normalise_date(frame["catalyst_date"])
    cutoff = pd.Timestamp(as_of).normalize() if as_of is not None else pd.Timestamp.today().normalize()
    available_column = next((item for item in ("available_date", "available_on", "published_date") if item in frame), None)
    if available_column:
        frame["available_date"] = _normalise_date(frame[available_column])
        frame = frame.loc[frame["available_date"].isna() | frame["available_date"].le(cutoff)].copy()
    else:
        frame["available_date"] = pd.NaT
    if frame.empty:
        return _empty(_EMPTY_CATALYST_COLUMNS)
    frame["catalyst_id"] = _column_or_default(frame, "catalyst_id", _column_or_default(frame, "id", range(len(frame)))).astype(str)
    frame["title"] = _column_or_default(frame, "title", _column_or_default(frame, "name", "")).fillna("").astype(str)
    frame["theme_id"] = _column_or_default(frame, "theme_id", "").fillna("").astype(str).str.strip()
    if "source_code" not in frame:
        frame["source_code"] = ""
    frame["status"] = _column_or_default(frame, "status", "计划").fillna("计划").astype(str).str.strip().str.lower()
    frame["impact"] = pd.to_numeric(_column_or_default(frame, "impact", 1.0), errors="coerce").fillna(1.0)
    frame["probability"] = pd.to_numeric(_column_or_default(frame, "probability", 1.0), errors="coerce").fillna(1.0)
    frame["impact"] = np.where(frame["impact"].gt(1), frame["impact"] / 100.0, frame["impact"]).clip(0, 1)
    frame["probability"] = np.where(frame["probability"].gt(1), frame["probability"] / 100.0, frame["probability"]).clip(0, 1)
    frame["days_to_catalyst"] = (frame["catalyst_date"] - cutoff).dt.days
    days = frame["days_to_catalyst"]
    frame["catalyst_proximity"] = np.select(
        [days.between(0, horizon_days), days.between(-grace_days, -1)],
        [1 - days.clip(lower=0) / max(horizon_days, 1), 0.25],
        default=0.0,
    )
    cancelled = frame["status"].isin({"取消", "cancelled", "canceled", "expired", "失效"})
    confirmed = frame["status"].isin({"确认", "confirmed", "已确认"})
    confidence = np.where(confirmed, 1.0, np.where(frame["status"].isin({"传闻", "rumor"}), 0.6, 0.8))
    frame["catalyst_score"] = (100 * frame["impact"] * frame["probability"] * frame["catalyst_proximity"] * confidence).where(~cancelled, 0.0)
    return frame[[column for column in _EMPTY_CATALYST_COLUMNS if column in frame]].reset_index(drop=True)


def _weighted_average(frame: pd.DataFrame, value: str) -> pd.DataFrame:
    valid = frame.loc[pd.to_numeric(frame[value], errors="coerce").notna()].copy()
    if valid.empty:
        return pd.DataFrame(columns=["source_code", value])
    valid["_weight"] = pd.to_numeric(valid.get("weight", 1.0), errors="coerce").fillna(1.0).clip(lower=0)
    valid["_value"] = pd.to_numeric(valid[value], errors="coerce")
    return valid.groupby("source_code", as_index=False).apply(
        lambda group: pd.Series({value: np.average(group["_value"], weights=group["_weight"])})
    ).reset_index(drop=True)


@dataclass(frozen=True)
class ResearchContext:
    """Loaded point-in-time research inputs and a stock-level scoring facade."""

    as_of: pd.Timestamp | None
    industry_data: pd.DataFrame
    lifecycle: pd.DataFrame
    theme_mapping: pd.DataFrame
    catalysts: pd.DataFrame
    weights: Mapping[str, float]

    def score_stocks(self, stocks: pd.DataFrame | Sequence[str], as_of: str | date | pd.Timestamp | None = None) -> pd.DataFrame:
        """Score a stock universe; absent optional data contributes no artificial score."""
        if isinstance(stocks, pd.DataFrame):
            universe = _normalise_stock_code(stocks)
        else:
            universe = _normalise_stock_code(pd.DataFrame({"source_code": list(stocks)}))
        if "source_code" not in universe:
            raise ValueError("stocks needs a source_code or stock_code column")
        universe = universe.drop_duplicates("source_code").copy()
        membership = map_stock_themes(universe, self.theme_mapping)
        lifecycle = self.lifecycle.copy()
        cutoff = pd.Timestamp(as_of).normalize() if as_of is not None else self.as_of
        if cutoff is not None and not lifecycle.empty:
            lifecycle = lifecycle.loc[lifecycle["date"].le(cutoff)]
        latest = lifecycle.sort_values("date").drop_duplicates("context_key", keep="last") if not lifecycle.empty else _empty(_EMPTY_LIFECYCLE_COLUMNS)

        themed = membership.merge(latest[["context_key", "hotspot_score", "hotspot_stage"]], left_on="theme_id", right_on="context_key", how="left") if not membership.empty else pd.DataFrame()
        hotspot = _weighted_average(themed, "hotspot_score") if not themed.empty else pd.DataFrame(columns=["source_code", "hotspot_score"])
        stage = (themed.sort_values("weight", ascending=False).drop_duplicates("source_code")[["source_code", "hotspot_stage"]]
                 if not themed.empty else pd.DataFrame(columns=["source_code", "hotspot_stage"]))

        industry = pd.DataFrame(columns=["source_code", "industry_score"])
        if "industry_code" in universe and not latest.empty and "industry_code" in self.industry_data:
            industry = universe[["source_code", "industry_code"]].merge(
                latest[["context_key", "hotspot_score"]], left_on="industry_code", right_on="context_key", how="left"
            )[["source_code", "hotspot_score"]].rename(columns={"hotspot_score": "industry_score"})

        catalyst = self.catalysts.copy()
        if cutoff is not None:
            catalyst = evaluate_catalysts(catalyst, cutoff)
        direct_cat = catalyst.loc[catalyst["source_code"].ne("")][["source_code", "catalyst_score"]] if not catalyst.empty else pd.DataFrame()
        theme_cat = membership.merge(catalyst.loc[catalyst["theme_id"].ne("")][["theme_id", "catalyst_score"]], on="theme_id", how="inner") if not membership.empty and not catalyst.empty else pd.DataFrame()
        theme_cat = _weighted_average(theme_cat, "catalyst_score") if not theme_cat.empty else pd.DataFrame(columns=["source_code", "catalyst_score"])
        cat_parts = pd.concat([direct_cat, theme_cat], ignore_index=True) if not direct_cat.empty or not theme_cat.empty else pd.DataFrame(columns=["source_code", "catalyst_score"])
        cat_score = cat_parts.groupby("source_code", as_index=False)["catalyst_score"].max() if not cat_parts.empty else cat_parts

        result = universe.merge(industry, on="source_code", how="left").merge(hotspot, on="source_code", how="left").merge(stage, on="source_code", how="left").merge(cat_score, on="source_code", how="left")
        result["theme_ids"] = membership.groupby("source_code")["theme_id"].agg(lambda values: ";".join(dict.fromkeys(values))).reindex(result["source_code"]).to_numpy() if not membership.empty else ""
        result["theme_names"] = membership.groupby("source_code")["theme_name"].agg(lambda values: "；".join(dict.fromkeys(values))).reindex(result["source_code"]).to_numpy() if not membership.empty else ""
        component_columns = ["industry_score", "hotspot_score", "catalyst_score"]
        configured = {key: float(self.weights.get(key, 0.0)) for key in ("industry", "hotspot", "catalyst")}
        total_weight = sum(configured.values()) or 1.0
        result["research_context_score"] = (
            result["industry_score"].fillna(0) * configured["industry"]
            + result["hotspot_score"].fillna(0) * configured["hotspot"]
            + result["catalyst_score"].fillna(0) * configured["catalyst"]
        ) / total_weight
        result["research_context_coverage"] = result[component_columns].notna().mean(axis=1)
        result["model_version"] = RESEARCH_CONTEXT_MODEL_VERSION
        return result


def load_research_context(config: Mapping[str, Any] | None = None, as_of: str | date | pd.Timestamp | None = None) -> ResearchContext:
    """Load optional CSV/JSON research inputs specified by a plain TOML mapping.

    Accepted keys are ``industry_data_path``, ``theme_mapping_path`` and
    ``catalyst_path`` (the shorter ``industry_data``, ``theme_mapping`` and
    ``catalysts`` aliases also work), plus ``base_dir`` and ``weights``.
    """
    raw = dict(config or {})
    if isinstance(raw.get("research_context"), Mapping):
        raw = dict(raw["research_context"])
    base_dir = Path(raw.get("base_dir", "."))
    industry = _read_table(raw.get("industry_data_path", raw.get("industry_data")), base_dir, "industry_data")
    mapping = load_theme_mapping(raw.get("theme_mapping_path", raw.get("theme_mapping")), base_dir)
    catalyst_input = _read_table(raw.get("catalyst_path", raw.get("catalysts")), base_dir, "catalysts")
    cutoff = pd.Timestamp(as_of).normalize() if as_of is not None else None
    lifecycle = classify_hotspot_lifecycle(
        industry, cutoff,
        short_window=int(raw.get("short_window", 5)),
        group_column=raw.get("group_column"),
        date_column=str(raw.get("date_column", "date")),
    )
    catalysts = evaluate_catalysts(
        catalyst_input, cutoff,
        horizon_days=int(raw.get("catalyst_horizon_days", 45)),
        grace_days=int(raw.get("catalyst_grace_days", 7)),
    )
    weights = {"industry": 0.35, "hotspot": 0.35, "catalyst": 0.30}
    weights.update({key: float(value) for key, value in dict(raw.get("weights", {})).items() if key in weights})
    return ResearchContext(cutoff, industry, lifecycle, mapping, catalysts, weights)