# src/core/market_regime.py
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Sequence

import numpy as np
import pandas as pd


REGIME_BULL_STRONG = "bull_strong"
REGIME_BULL_MODERATE = "bull_moderate"
REGIME_SIDEWAYS = "sideways"
REGIME_BEAR_MODERATE = "bear_moderate"
REGIME_BEAR_STRONG = "bear_strong"

REGIME_DETAIL_LABELS = (
    REGIME_BULL_STRONG,
    REGIME_BULL_MODERATE,
    REGIME_SIDEWAYS,
    REGIME_BEAR_MODERATE,
    REGIME_BEAR_STRONG,
)

REGIME_BULL = "bull"
REGIME_BEAR = "bear"
REGIME_GROUP_LABELS = (REGIME_BULL, REGIME_SIDEWAYS, REGIME_BEAR)


@dataclass(frozen=True)
class TrendScoreConfig:
    price_col: str = "close"
    ema_periods: Sequence[int] = (20, 50, 100, 200)

    bull_threshold: int = 2
    bull_strong_threshold: int = 4
    bear_threshold: int = -2
    bear_strong_threshold: int = -4

    bull_persistence_bars: int = 5
    sideways_persistence_bars: int = 3
    bear_persistence_bars: int = 2


def _classify_window_avg_score(avg_score: float) -> str:
    if avg_score >= 3.0:
        return REGIME_BULL_STRONG
    if avg_score >= 1.0:
        return REGIME_BULL_MODERATE
    if avg_score > -1.0:
        return REGIME_SIDEWAYS
    if avg_score > -3.0:
        return REGIME_BEAR_MODERATE
    return REGIME_BEAR_STRONG


def _classify_score_detail(
    score: float,
    *,
    bull_threshold: int,
    bull_strong_threshold: int,
    bear_threshold: int,
    bear_strong_threshold: int,
) -> str:
    if score >= bull_strong_threshold:
        return REGIME_BULL_STRONG
    if score >= bull_threshold:
        return REGIME_BULL_MODERATE
    if score <= bear_strong_threshold:
        return REGIME_BEAR_STRONG
    if score <= bear_threshold:
        return REGIME_BEAR_MODERATE
    return REGIME_SIDEWAYS


def _detail_to_group(label: str) -> str:
    if label in (REGIME_BULL_STRONG, REGIME_BULL_MODERATE):
        return REGIME_BULL
    if label in (REGIME_BEAR_STRONG, REGIME_BEAR_MODERATE):
        return REGIME_BEAR
    return REGIME_SIDEWAYS


def _apply_asymmetric_persistence_filter(
    labels: Sequence[str],
    *,
    bull_persistence_bars: int,
    sideways_persistence_bars: int,
    bear_persistence_bars: int,
) -> List[str]:
    if not labels:
        return []

    def family(label: str) -> str:
        if label in (REGIME_BULL_STRONG, REGIME_BULL_MODERATE):
            return REGIME_BULL
        if label in (REGIME_BEAR_STRONG, REGIME_BEAR_MODERATE):
            return REGIME_BEAR
        return REGIME_SIDEWAYS

    def required_bars(label: str) -> int:
        fam = family(label)
        if fam == REGIME_BULL:
            return max(int(bull_persistence_bars), 1)
        if fam == REGIME_BEAR:
            return max(int(bear_persistence_bars), 1)
        return max(int(sideways_persistence_bars), 1)

    confirmed = labels[0]
    candidate = labels[0]
    candidate_run = 0
    out: List[str] = []

    for label in labels:
        if family(label) == family(confirmed):
            confirmed = label
            candidate = label
            candidate_run = 0
            out.append(confirmed)
            continue

        if family(label) == family(candidate):
            candidate_run += 1
        else:
            candidate = label
            candidate_run = 1

        if candidate_run >= required_bars(candidate):
            confirmed = candidate
            candidate_run = 0

        out.append(confirmed)

    return out


def compute_regime_series(
    df: pd.DataFrame,
    config: TrendScoreConfig | None = None,
) -> pd.DataFrame:
    cfg = config or TrendScoreConfig()

    if df is None or df.empty:
        return pd.DataFrame(
            columns=[
                cfg.price_col,
                "trend_score",
                "raw_regime_detail",
                "stable_regime_detail",
                "stable_regime_group",
            ]
        )

    if cfg.price_col not in df.columns:
        raise ValueError(f"price_col '{cfg.price_col}' not found in DataFrame")

    out = pd.DataFrame(index=df.index.copy())
    out[cfg.price_col] = pd.to_numeric(df[cfg.price_col], errors="coerce").astype(float)

    score = pd.Series(0, index=out.index, dtype=int)

    for period in cfg.ema_periods:
        ema_col = f"ema_{int(period)}"
        out[ema_col] = out[cfg.price_col].ewm(
            span=int(period),
            adjust=False,
            min_periods=1,
        ).mean()
        score += np.where(out[cfg.price_col] >= out[ema_col], 1, -1)

    out["trend_score"] = score.astype(int)

    raw_detail = [
        _classify_score_detail(
            v,
            bull_threshold=cfg.bull_threshold,
            bull_strong_threshold=cfg.bull_strong_threshold,
            bear_threshold=cfg.bear_threshold,
            bear_strong_threshold=cfg.bear_strong_threshold,
        )
        for v in out["trend_score"].tolist()
    ]

    stable_detail = _apply_asymmetric_persistence_filter(
        raw_detail,
        bull_persistence_bars=cfg.bull_persistence_bars,
        sideways_persistence_bars=cfg.sideways_persistence_bars,
        bear_persistence_bars=cfg.bear_persistence_bars,
    )

    out["raw_regime_detail"] = raw_detail
    out["stable_regime_detail"] = stable_detail
    out["stable_regime_group"] = [_detail_to_group(x) for x in stable_detail]

    return out


def summarize_regimes_for_windows(
    regime_df: pd.DataFrame,
    windows: Iterable[Dict[str, Any]],
    *,
    detail_col: str = "stable_regime_detail",
    group_col: str = "stable_regime_group",
) -> List[Dict[str, Any]]:
    if regime_df is None or regime_df.empty:
        return []

    if detail_col not in regime_df.columns:
        raise ValueError(f"detail_col '{detail_col}' not found in regime_df")
    if group_col not in regime_df.columns:
        raise ValueError(f"group_col '{group_col}' not found in regime_df")

    summaries: List[Dict[str, Any]] = []

    for w in windows:
        test_start = pd.Timestamp(w["test_start"])
        test_end = pd.Timestamp(w["test_end"])
        window_id = int(w["window_id"])

        mask = (regime_df.index >= test_start) & (regime_df.index <= test_end)

        seg_detail = regime_df.loc[mask, detail_col]
        seg_group = regime_df.loc[mask, group_col]
        seg_score = regime_df.loc[mask, "trend_score"]

        total = int(len(seg_detail))

        if total == 0:
            summaries.append({
                "window": window_id,
                "regime": REGIME_SIDEWAYS,
                "regime_detail": REGIME_SIDEWAYS,
                "bull_pct": 0.0,
                "sideways_pct": 0.0,
                "bear_pct": 0.0,
                "bull_strong_pct": 0.0,
                "bull_moderate_pct": 0.0,
                "sideways_detail_pct": 0.0,
                "bear_moderate_pct": 0.0,
                "bear_strong_pct": 0.0,
                "n_bars": 0,
            })
            continue

        group_counts = seg_group.value_counts()
        detail_counts = seg_detail.value_counts()

        bull_pct = float(group_counts.get(REGIME_BULL, 0) / total)
        sideways_pct = float(group_counts.get(REGIME_SIDEWAYS, 0) / total)
        bear_pct = float(group_counts.get(REGIME_BEAR, 0) / total)

        bull_strong_pct = float(detail_counts.get(REGIME_BULL_STRONG, 0) / total)
        bull_moderate_pct = float(detail_counts.get(REGIME_BULL_MODERATE, 0) / total)
        sideways_detail_pct = float(detail_counts.get(REGIME_SIDEWAYS, 0) / total)
        bear_moderate_pct = float(detail_counts.get(REGIME_BEAR_MODERATE, 0) / total)
        bear_strong_pct = float(detail_counts.get(REGIME_BEAR_STRONG, 0) / total)

        majority_group = max(
            (
                (REGIME_BULL, bull_pct),
                (REGIME_SIDEWAYS, sideways_pct),
                (REGIME_BEAR, bear_pct),
            ),
            key=lambda x: x[1],
        )[0]

        avg_score = float(seg_score.mean()) if len(seg_score) else 0.0
        majority_detail = _classify_window_avg_score(avg_score)

        summaries.append({
            "window": window_id,
            "regime": majority_group,
            "regime_detail": majority_detail,
            "avg_score": avg_score,
            "bull_pct": bull_pct,
            "sideways_pct": sideways_pct,
            "bear_pct": bear_pct,
            "bull_strong_pct": bull_strong_pct,
            "bull_moderate_pct": bull_moderate_pct,
            "sideways_detail_pct": sideways_detail_pct,
            "bear_moderate_pct": bear_moderate_pct,
            "bear_strong_pct": bear_strong_pct,
            "n_bars": total,
        })

    return summaries


def compute_regimes_for_windows(
    benchmark_df: pd.DataFrame,
    windows: Iterable[Dict[str, Any]],
    config: TrendScoreConfig | None = None,
) -> Dict[str, Any]:
    cfg = config or TrendScoreConfig()
    regime_df = compute_regime_series(benchmark_df, config=cfg)
    by_window = summarize_regimes_for_windows(regime_df, windows)

    return {
        "config": {
            "price_col": cfg.price_col,
            "ema_periods": list(cfg.ema_periods),
            "bull_threshold": int(cfg.bull_threshold),
            "bull_strong_threshold": int(cfg.bull_strong_threshold),
            "bear_threshold": int(cfg.bear_threshold),
            "bear_strong_threshold": int(cfg.bear_strong_threshold),
            "bull_persistence_bars": int(cfg.bull_persistence_bars),
            "sideways_persistence_bars": int(cfg.sideways_persistence_bars),
            "bear_persistence_bars": int(cfg.bear_persistence_bars),
        },
        "by_bar": regime_df,
        "by_window": by_window,
    }