"""Factor Library: maintains the growing collection of admitted alpha factors.

Implements the admission rules from the paper (Eq. 10, 11):
- Admission: IC(alpha) >= tau_IC AND max_{g in L} |rho(alpha, g)| < theta
- Replacement: IC(alpha) >= 0.10 AND IC(alpha) >= 1.3 * IC(g) AND only 1 correlated factor

内存优化策略：库中因子不再长期持有 (M,T) numpy signals。
相关性检查改为按需调用 LocalFactorEvaluator 重算。

本地引擎重算成本低，内存优先。
"""

from __future__ import annotations

import logging
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING

import numpy as np
from scipy.stats import spearmanr

if TYPE_CHECKING:
    from src.factorminer.evaluation.local_engine import LocalFactorEvaluator

logger = logging.getLogger(__name__)


@dataclass
class Factor:
    """A single admitted alpha factor."""

    id: int
    name: str
    formula: str
    category: str  # e.g., "VWAP", "Regime-switching", "Momentum", etc.
    ic_mean: float  # Mean IC across evaluation period
    icir: float  # IC Information Ratio
    ic_win_rate: float  # Fraction of periods with positive IC
    max_correlation: float  # Max |rho| with any other library factor at admission
    batch_number: int  # Which mining batch admitted this factor
    admission_date: str = ""
    signals: Optional[np.ndarray] = field(default=None, repr=False)  # (M, T)
    pool_metrics: Dict[str, dict] = field(default_factory=dict)
    research_metrics: dict = field(default_factory=dict)
    provenance: dict = field(default_factory=dict)
    metadata: dict = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.admission_date:
            self.admission_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    @staticmethod
    def _sanitize_pool_metrics(pool_metrics: Dict[str, dict]) -> Dict[str, dict]:
        sanitized: Dict[str, dict] = {}
        for pool_name, stats in pool_metrics.items():
            sanitized[pool_name] = {}
            for key, value in stats.items():
                if key == "ic_series":
                    continue
                if isinstance(value, np.ndarray):
                    sanitized[pool_name][key] = value.tolist()
                else:
                    sanitized[pool_name][key] = value
        return sanitized

    def to_dict(self) -> dict:
        """Serialize to a JSON-compatible dictionary (excludes signals)."""
        return {
            "id": self.id,
            "name": self.name,
            "formula": self.formula,
            "category": self.category,
            "ic_mean": self.ic_mean,
            "icir": self.icir,
            "ic_win_rate": self.ic_win_rate,
            "max_correlation": self.max_correlation,
            "batch_number": self.batch_number,
            "admission_date": self.admission_date,
            "pool_metrics": self._sanitize_pool_metrics(self.pool_metrics),
            "research_metrics": self.research_metrics,
            "provenance": self.provenance,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "Factor":
        """Reconstruct a Factor from a dictionary."""
        return cls(
            id=d["id"],
            name=d["name"],
            formula=d["formula"],
            category=d["category"],
            ic_mean=d["ic_mean"],
            icir=d["icir"],
            ic_win_rate=d["ic_win_rate"],
            max_correlation=d["max_correlation"],
            batch_number=d["batch_number"],
            admission_date=d.get("admission_date", ""),
            pool_metrics=d.get("pool_metrics", {}),
            research_metrics=d.get("research_metrics", {}),
            provenance=d.get("provenance", {}),
            metadata=d.get("metadata", {}),
        )


class FactorLibrary:
    """The factor library L that maintains admitted alpha factors.

    内存优化：因子不再持有 signals，按需通过 LocalFactorEvaluator 重算。

    Parameters
    ----------
    correlation_threshold : float
        Maximum allowed |rho| for admission (theta). Default 0.5 for A-shares.
    ic_threshold : float
        Minimum IC for admission (tau_IC). Default 0.04.
    """

    def __init__(
        self,
        correlation_threshold: float = 0.5,
        ic_threshold: float = 0.04,
    ) -> None:
        self.factors: Dict[int, Factor] = {}
        self._next_id: int = 1
        self.correlation_threshold = correlation_threshold
        self.ic_threshold = ic_threshold
        # 内部状态：仅用于诊断（按需计算，不持久化）
        self._correlation_cache: Optional[np.ndarray] = None
        self._id_to_index: Dict[int, int] = {}

    # ------------------------------------------------------------------
    # Correlation computation
    # ------------------------------------------------------------------

    def compute_correlation(
        self, signals_a: np.ndarray, signals_b: np.ndarray
    ) -> float:
        """Compute time-average cross-sectional Spearman correlation rho(alpha, beta).

        rho(alpha, beta) = (1/|T|) * sum_t Corr_rank(s_t^(alpha), s_t^(beta))

        Parameters
        ----------
        signals_a, signals_b : np.ndarray, shape (M, T)
            Cross-sectional signal matrices.

        Returns
        -------
        float
            Mean absolute Spearman rank correlation across time steps.
        """
        if signals_a.shape != signals_b.shape:
            raise ValueError(
                f"Signal shapes must match: {signals_a.shape} vs {signals_b.shape}"
            )
        M, T = signals_a.shape
        correlations = np.empty(T, dtype=np.float64)

        for t in range(T):
            col_a = signals_a[:, t]
            col_b = signals_b[:, t]
            # Mask NaNs from both columns
            valid = ~(np.isnan(col_a) | np.isnan(col_b))
            n_valid = valid.sum()
            if n_valid < 3:
                correlations[t] = np.nan
                continue
            rho, _ = spearmanr(col_a[valid], col_b[valid])
            correlations[t] = rho

        return float(np.nanmean(np.abs(correlations)))

    def _compute_correlation_vectorized(
        self, signals_a: np.ndarray, signals_b: np.ndarray
    ) -> float:
        """Faster vectorized Spearman correlation using rankdata.

        For large M and T this avoids per-timestep Python loops by ranking
        each column and computing Pearson on the ranks.
        """
        from scipy.stats import rankdata

        M, T = signals_a.shape
        # Mask invalid entries
        mask = ~(np.isnan(signals_a) | np.isnan(signals_b))

        corr_sum = 0.0
        n_valid_t = 0
        for t in range(T):
            valid = mask[:, t]
            n = valid.sum()
            if n < 3:
                continue
            ra = rankdata(signals_a[valid, t])
            rb = rankdata(signals_b[valid, t])
            # Pearson on ranks == Spearman
            ra_c = ra - ra.mean()
            rb_c = rb - rb.mean()
            denom = np.sqrt((ra_c**2).sum() * (rb_c**2).sum())
            if denom < 1e-12:
                continue
            corr_sum += abs((ra_c * rb_c).sum() / denom)
            n_valid_t += 1

        if n_valid_t == 0:
            return 0.0
        return corr_sum / n_valid_t

    # ------------------------------------------------------------------
    # Admission and replacement
    # ------------------------------------------------------------------

    def check_admission(
        self,
        candidate_ic: float,
        candidate_signals: np.ndarray,
        evaluator: Optional["LocalFactorEvaluator"] = None,
    ) -> Tuple[bool, str]:
        """Check if candidate passes admission criteria (Eq. 10).

        Admission rule:
            IC(alpha) >= tau_IC  AND  max_{g in L} |rho(alpha, g)| < theta

        内存优化：当提供 evaluator 时，按需重算库中因子信号进行比较。

        Parameters
        ----------
        candidate_ic : float
            The candidate factor's mean IC.
        candidate_signals : np.ndarray, shape (M, T)
            The candidate's realized signals.
        evaluator : LocalFactorEvaluator, optional
            本地因子评估器，用于按需重算库中因子信号。

        Returns
        -------
        (admitted, reason) : Tuple[bool, str]
        """
        if candidate_ic < self.ic_threshold:
            return False, (f"IC {candidate_ic:.4f} below threshold {self.ic_threshold}")

        if self.size == 0:
            return True, "First factor in library"

        max_corr = self._max_correlation_with_library(candidate_signals, evaluator)

        if max_corr >= self.correlation_threshold:
            return False, (
                f"Max correlation {max_corr:.4f} >= threshold "
                f"{self.correlation_threshold} with existing library factor"
            )

        return True, (f"Admitted: IC={candidate_ic:.4f}, max_corr={max_corr:.4f}")

    def check_replacement(
        self,
        candidate_ic: float,
        candidate_signals: np.ndarray,
        evaluator: Optional["LocalFactorEvaluator"] = None,
        ic_min: float = 0.10,
        ic_ratio: float = 1.3,
    ) -> Tuple[bool, Optional[int], str]:
        """Check replacement mechanism (Eq. 11).

        Replacement rule:
            IC(alpha) >= 0.10
            AND IC(alpha) >= 1.3 * IC(g)
            AND |{g in L : |rho(alpha, g)| > theta}| = 1

        If exactly one library factor g is correlated above theta AND the
        candidate's IC dominates g's IC by the required ratio, replace g.

        内存优化：当提供 evaluator 时，按需重算库中因子信号。

        Parameters
        ----------
        candidate_ic : float
            The candidate's mean IC.
        candidate_signals : np.ndarray, shape (M, T)
            The candidate's realized signals.
        evaluator : LocalFactorEvaluator, optional
            本地因子评估器，用于按需重算库中因子信号。
        ic_min : float
            Absolute IC floor for replacement (default 0.10).
        ic_ratio : float
            Required IC ratio over the correlated factor (default 1.3).

        Returns
        -------
        (should_replace, factor_to_replace_id, reason) : Tuple[bool, Optional[int], str]
        """
        if candidate_ic < ic_min:
            return (
                False,
                None,
                (f"IC {candidate_ic:.4f} below replacement floor {ic_min}"),
            )

        if self.size == 0:
            return False, None, "Library is empty, use admission instead"

        # Find all library factors correlated above theta
        correlated_factors = []
        for fid, factor in self.factors.items():
            # 按需重算因子信号
            factor_signals = self._get_factor_signals(factor, evaluator)
            if factor_signals is None:
                continue
            corr = self._compute_correlation_vectorized(
                candidate_signals, factor_signals
            )
            if corr >= self.correlation_threshold:
                correlated_factors.append((fid, corr, factor.ic_mean))

        if len(correlated_factors) != 1:
            return (
                False,
                None,
                (
                    f"Found {len(correlated_factors)} correlated factors "
                    f"(need exactly 1 for replacement)"
                ),
            )

        fid, corr, existing_ic = correlated_factors[0]
        if candidate_ic < ic_ratio * existing_ic:
            return (
                False,
                None,
                (
                    f"IC {candidate_ic:.4f} < {ic_ratio} * {existing_ic:.4f} = "
                    f"{ic_ratio * existing_ic:.4f}"
                ),
            )

        return (
            True,
            fid,
            (
                f"Replace factor {fid}: candidate IC {candidate_ic:.4f} > "
                f"{ic_ratio} * {existing_ic:.4f}, corr={corr:.4f}"
            ),
        )

    # ------------------------------------------------------------------
    # Library mutations
    # ------------------------------------------------------------------

    def admit_factor(self, factor: Factor) -> int:
        """Add a factor to the library.

        内存优化：不再保存 signals 到 Factor 对象，相关性检查改为按需重算。

        Parameters
        ----------
        factor : Factor
            The factor to add. Its ``id`` field is overwritten with the
            next available ID. signals 字段将被忽略。

        Returns
        -------
        int
            The assigned factor ID.
        """
        factor.id = self._next_id
        self._next_id += 1

        # 不再持有 signals，清理以节省内存
        factor.signals = None

        self.factors[factor.id] = factor

        # 清理相关性缓存
        self._correlation_cache = None

        logger.info(
            "Admitted factor %d '%s' (IC=%.4f, max_corr=%.4f, category=%s)",
            factor.id,
            factor.name,
            factor.ic_mean,
            factor.max_correlation,
            factor.category,
        )
        return factor.id

    def replace_factor(self, old_id: int, new_factor: Factor) -> None:
        """Replace an existing factor with a better one.

        内存优化：不再持有 signals，清理相关性缓存。

        Parameters
        ----------
        old_id : int
            ID of the factor being replaced.
        new_factor : Factor
            The replacement factor.
        """
        if old_id not in self.factors:
            raise KeyError(f"Factor {old_id} not in library")

        old_factor = self.factors[old_id]
        new_factor.id = self._next_id
        self._next_id += 1

        # 不再持有 signals，清理以节省内存
        new_factor.signals = None

        # Remove old factor
        del self.factors[old_id]

        # Insert new factor
        self.factors[new_factor.id] = new_factor

        # 清理相关性缓存
        self._correlation_cache = None

        logger.info(
            "Replaced factor %d with %d '%s' (IC=%.4f)",
            old_id,
            new_factor.id,
            new_factor.name,
            new_factor.ic_mean,
        )

    def remove_factor(self, factor_id: int) -> None:
        """Remove a factor from the library and rebuild correlation state."""
        if factor_id not in self.factors:
            raise KeyError(f"Factor {factor_id} not in library")

        removed = self.factors.pop(factor_id)
        self.update_correlation_matrix()

        logger.info(
            "Removed factor %d '%s' from library",
            factor_id,
            removed.name,
        )

    # ------------------------------------------------------------------
    # Correlation matrix management (按需重算)
    # ------------------------------------------------------------------

    def _get_factor_signals(
        self,
        factor: Factor,
        evaluator: Optional["LocalFactorEvaluator"],
    ) -> Optional[np.ndarray]:
        """获取因子信号。

        优先使用 evaluator 按需重算，否则返回 Factor 对象中缓存的 signals。

        Args:
            factor: 因子对象
            evaluator: 本地因子评估器

        Returns:
            (M, T) 信号矩阵或 None
        """
        # 如果 formula 以 # TODO 开头，说明不支持，跳过
        if factor.formula.startswith("# TODO"):
            return None

        if evaluator is not None:
            try:
                return evaluator.evaluate_single(factor.name, factor.formula)
            except Exception as e:
                logger.warning(
                    "计算因子 %s 信号失败: %s",
                    factor.name,
                    e,
                )
                return None

        # 回退到缓存的 signals（已废弃，仅兼容）
        return factor.signals

    def _max_correlation_with_library(
        self,
        candidate_signals: np.ndarray,
        evaluator: Optional["LocalFactorEvaluator"] = None,
    ) -> float:
        """Compute max |rho| between candidate and all library factors.

        内存优化：按需通过 evaluator 重算库中因子信号。

        Args:
            candidate_signals: 候选因子信号 (M, T)
            evaluator: 本地因子评估器

        Returns:
            最大绝对相关值
        """
        max_corr = 0.0
        for factor in self.factors.values():
            factor_signals = self._get_factor_signals(factor, evaluator)
            if factor_signals is None:
                continue
            corr = self._compute_correlation_vectorized(
                candidate_signals, factor_signals
            )
            max_corr = max(max_corr, corr)
        return max_corr

    def update_correlation_matrix(self) -> None:
        """废弃方法。

        内存优化后，相关性矩阵不再持久化。如需计算相关性，
        请使用 LocalFactorEvaluator 按需重算。

        此方法仅清理内部缓存。
        """
        self._correlation_cache = None
        self._id_to_index.clear()
        logger.warning(
            "[factor_library] update_correlation_matrix 已废弃，相关性计算改为按需重算"
        )

    # ------------------------------------------------------------------
    # Queries and diagnostics
    # ------------------------------------------------------------------

    @property
    def size(self) -> int:
        """Number of factors currently in the library."""
        return len(self.factors)

    def get_factor(self, factor_id: int) -> Factor:
        """Retrieve a factor by ID."""
        if factor_id not in self.factors:
            raise KeyError(f"Factor {factor_id} not in library")
        return self.factors[factor_id]

    def list_factors(self) -> List[Factor]:
        """Return all factors sorted by ID."""
        return [self.factors[k] for k in sorted(self.factors)]

    def get_factors_by_category(self, category: str) -> List[Factor]:
        """Return all factors matching a given category."""
        return [f for f in self.factors.values() if f.category == category]

    def get_diagnostics(self) -> dict:
        """Library diagnostics: per-category counts and IC statistics.

        注意：内存优化后，相关性矩阵不再持久化。
        如需相关性统计，请使用 LocalFactorEvaluator 按需计算。

        Returns
        -------
        dict with keys:
            - size: int
            - category_counts: dict[str, int]
            - category_avg_ic: dict[str, float]
        """
        diag: dict = {"size": self.size}

        # Category breakdown
        cat_counts: Dict[str, int] = defaultdict(int)
        cat_ic_sums: Dict[str, float] = defaultdict(float)
        for f in self.factors.values():
            cat_counts[f.category] += 1
            cat_ic_sums[f.category] += f.ic_mean

        diag["category_counts"] = dict(cat_counts)
        diag["category_avg_ic"] = {
            cat: cat_ic_sums[cat] / cat_counts[cat] for cat in cat_counts
        }

        return diag

    def get_state_summary(self) -> dict:
        """Summary for memory retrieval: size, categories, recent admissions.

        Returns a lightweight dictionary suitable for inclusion in LLM prompts
        or memory store entries.
        """
        factors_sorted = sorted(self.factors.values(), key=lambda f: f.id, reverse=True)
        recent = factors_sorted[:5]  # Last 5 admissions

        categories = defaultdict(int)
        for f in self.factors.values():
            categories[f.category] += 1

        return {
            "library_size": self.size,
            "categories": dict(categories),
            "recent_admissions": [
                {
                    "id": f.id,
                    "name": f.name,
                    "category": f.category,
                    "ic_mean": f.ic_mean,
                    "batch": f.batch_number,
                }
                for f in recent
            ],
            "correlation_threshold": self.correlation_threshold,
            "ic_threshold": self.ic_threshold,
        }
