"""Markov Chain Market State Analysis — markov.py

Classifies the current market into BULL / BEAR / RANGING using 3 objective
criteria (EMA alignment, RSI zone, volume ratio) on CLOSED candles only.
No lookahead bias: all computations use iloc[:-1] or iloc[-2] as 'current'.

Functions:
  classify_state(df)            → (state_str, confidence_0_to_1, reasons)
  build_transition_matrix(df)   → {from_state: {to_state: prob, ...}, ...}
  forecast_states(state, matrix, days) → {day: {state: prob, ...}, ...}
  persistence_score(matrix, state)     → float 0.0–1.0
  analyze_symbol(src, symbol)          → full analysis dict

States: "BULL", "BEAR", "RANGING"
"""
from __future__ import annotations
import math, time
from typing import Literal

# ─── State type ──────────────────────────────────────────────────────────────
State = Literal["BULL", "BEAR", "RANGING"]
ALL_STATES: list[State] = ["BULL", "BEAR", "RANGING"]

# ─── Cache: 30-min TTL per symbol ────────────────────────────────────────────
_mk_cache: dict = {}
MK_TTL = 1800   # 30 minutes — matrix is stable; no need to re-compute every scan


# ─────────────────────────────────────────────────────────────────────────────
# CORE: Classify a single candle/row into a state
# Uses only CLOSED data — no lookahead. Call on df.iloc[:-1] for live candle.
# ─────────────────────────────────────────────────────────────────────────────

def _classify_row(close_series, vol_series, idx: int, lookback: int = 50) -> State:
    """Classify state at index `idx` using data UP TO idx (inclusive).
    Requires at least `lookback` rows before idx."""
    start = max(0, idx - lookback + 1)
    closes = close_series.iloc[start:idx + 1]
    vols   = vol_series.iloc[start:idx + 1]

    if len(closes) < 20:
        return "RANGING"

    c = closes.values

    # ── EMA alignment ────────────────────────────────────────────────────────
    # Fast EMA(20) vs Slow EMA(50)
    def _ema(arr, span):
        k = 2.0 / (span + 1)
        result = [arr[0]]
        for v in arr[1:]:
            result.append(result[-1] * (1 - k) + v * k)
        return result

    if len(c) >= 50:
        e20 = _ema(c, 20)[-1]
        e50 = _ema(c, 50)[-1]
        ema_bull = e20 > e50 * 1.001   # 0.1% buffer to avoid noise
        ema_bear = e20 < e50 * 0.999
    elif len(c) >= 20:
        e20 = _ema(c, 20)[-1]
        e10 = _ema(c, 10)[-1]
        ema_bull = e10 > e20 * 1.001
        ema_bear = e10 < e20 * 0.999
    else:
        ema_bull = ema_bear = False

    # ── RSI(14) zone ─────────────────────────────────────────────────────────
    rsi_len = min(14, len(c) - 1)
    if rsi_len >= 2:
        deltas = [c[i] - c[i-1] for i in range(1, len(c))]
        gains  = [max(d, 0) for d in deltas[-rsi_len:]]
        losses = [max(-d, 0) for d in deltas[-rsi_len:]]
        avg_g  = sum(gains) / rsi_len
        avg_l  = sum(losses) / rsi_len
        if avg_l == 0:
            rsi = 100.0
        elif avg_g == 0:
            rsi = 0.0
        else:
            rs  = avg_g / avg_l
            rsi = 100.0 - (100.0 / (1.0 + rs))
    else:
        rsi = 50.0

    rsi_bull = rsi >= 55
    rsi_bear = rsi <= 45

    # ── Volume ratio (recent 5 vs 20-bar avg) ─────────────────────────────────
    if len(vols) >= 20:
        vol_avg = float(vols.iloc[-20:].mean())
        vol_recent = float(vols.iloc[-5:].mean()) if len(vols) >= 5 else vol_avg
        vol_ratio = vol_recent / vol_avg if vol_avg > 0 else 1.0
    else:
        vol_ratio = 1.0

    # Volume above average strengthens the directional signal
    vol_confirms = vol_ratio >= 1.10   # 10% above avg = confirming

    # ── Vote: majority rule across 3 criteria ─────────────────────────────────
    bull_votes = sum([ema_bull, rsi_bull, (vol_confirms and ema_bull)])
    bear_votes = sum([ema_bear, rsi_bear, (vol_confirms and ema_bear)])

    if bull_votes >= 2:
        return "BULL"
    elif bear_votes >= 2:
        return "BEAR"
    else:
        return "RANGING"


# ─────────────────────────────────────────────────────────────────────────────
# PUBLIC: classify_state — current market state + confidence
# ─────────────────────────────────────────────────────────────────────────────

def classify_state(df) -> tuple[State, float, list[str]]:
    """Classify current market state from a klines DataFrame.

    Uses only closed candles (iloc[:-1] = last confirmed candle).
    Returns (state, confidence_0_to_1, reasons).
    """
    reasons = []

    if df is None or len(df) < 22:
        return "RANGING", 0.40, ["Insufficient data"]

    # Use last CLOSED candle = iloc[-2] for live; for building matrix use all
    closes = df["close"]
    vols   = df["volume"] if "volume" in df.columns else df.get("vol", df["close"] * 0)

    # ── EMA alignment on full history (closed candles only = iloc[:-1]) ──────
    idx = len(closes) - 2   # last confirmed closed candle
    c   = closes.values[:idx+1]

    def _ema_arr(arr, span):
        k = 2.0 / (span + 1)
        result = [arr[0]]
        for v in arr[1:]:
            result.append(result[-1] * (1 - k) + v * k)
        return result

    e20 = _ema_arr(c, 20)[-1] if len(c) >= 20 else c[-1]
    e50 = _ema_arr(c, 50)[-1] if len(c) >= 50 else e20

    ema_bull = e20 > e50 * 1.001
    ema_bear = e20 < e50 * 0.999

    if ema_bull:
        reasons.append(f"EMA20 ({e20:.4g}) > EMA50 ({e50:.4g}) — bullish alignment")
    elif ema_bear:
        reasons.append(f"EMA20 ({e20:.4g}) < EMA50 ({e50:.4g}) — bearish alignment")
    else:
        reasons.append(f"EMA20 ≈ EMA50 — no trend")

    # ── RSI(14) ──────────────────────────────────────────────────────────────
    rsi_len = 14
    deltas = [float(c[i]) - float(c[i-1]) for i in range(1, len(c))]
    gains  = [max(d, 0) for d in deltas[-rsi_len:]]
    losses = [max(-d, 0) for d in deltas[-rsi_len:]]
    avg_g  = sum(gains) / rsi_len
    avg_l  = sum(losses) / rsi_len
    if avg_l == 0:
        rsi = 100.0
    elif avg_g == 0:
        rsi = 0.0
    else:
        rsi = 100.0 - 100.0 / (1.0 + avg_g / avg_l)

    rsi_bull = rsi >= 55
    rsi_bear = rsi <= 45
    reasons.append(f"RSI14 = {rsi:.1f} ({'bullish' if rsi_bull else 'bearish' if rsi_bear else 'neutral'})")

    # ── Volume ratio ─────────────────────────────────────────────────────────
    vol_vals = vols.values[:idx+1]
    if len(vol_vals) >= 20:
        vol_avg    = sum(vol_vals[-20:]) / 20
        vol_recent = sum(vol_vals[-5:]) / 5 if len(vol_vals) >= 5 else vol_avg
        vol_ratio  = vol_recent / vol_avg if vol_avg > 0 else 1.0
    else:
        vol_ratio = 1.0
    vol_up = vol_ratio >= 1.10
    reasons.append(f"Volume ratio {vol_ratio:.2f}× 20-bar avg ({'elevated' if vol_up else 'normal/low'})")

    # ── State + confidence ────────────────────────────────────────────────────
    bull_score = sum([ema_bull, rsi_bull, (vol_up and ema_bull)])
    bear_score = sum([ema_bear, rsi_bear, (vol_up and ema_bear)])

    if bull_score >= 2:
        state = "BULL"
        # Confidence: how strongly all signals agree
        confidence = 0.50 + 0.15 * bull_score   # 2→0.80, 3→0.95
    elif bear_score >= 2:
        state = "BEAR"
        confidence = 0.50 + 0.15 * bear_score
    else:
        state = "RANGING"
        # Lower confidence when signals are mixed, higher when cleanly flat
        mix = abs(bull_score - bear_score)
        confidence = 0.55 if mix == 0 else 0.48

    confidence = round(min(0.95, max(0.35, confidence)), 2)
    return state, confidence, reasons


# ─────────────────────────────────────────────────────────────────────────────
# PUBLIC: build_transition_matrix — empirical from 500 candles
# ─────────────────────────────────────────────────────────────────────────────

def build_transition_matrix(df, lookback: int = 500) -> dict[str, dict[str, float]]:
    """Walk through last `lookback` closed candles, classify each, count transitions.

    Returns:
      {
        "BULL":    {"BULL": 0.78, "BEAR": 0.08, "RANGING": 0.14},
        "BEAR":    {"BULL": 0.09, "BEAR": 0.76, "RANGING": 0.15},
        "RANGING": {"BULL": 0.28, "BEAR": 0.27, "RANGING": 0.45},
      }
    """
    closes = df["close"]
    vols   = df["volume"] if "volume" in df.columns else df.get("vol", df["close"] * 0)

    n = len(closes)
    # We need at least 50 rows to meaningfully classify; cap lookback
    start = max(50, n - lookback)
    end   = n - 1   # exclude live (open) candle

    # Classify each candle
    labels: list[State] = []
    for i in range(start, end):
        labels.append(_classify_row(closes, vols, i))

    # Count transitions
    counts: dict[str, dict[str, int]] = {
        s: {t: 0 for t in ALL_STATES} for s in ALL_STATES
    }
    for i in range(len(labels) - 1):
        counts[labels[i]][labels[i+1]] += 1

    # Normalize rows
    matrix: dict[str, dict[str, float]] = {}
    for state in ALL_STATES:
        total = sum(counts[state].values())
        if total == 0:
            # No observations: uniform prior
            matrix[state] = {s: 1/3 for s in ALL_STATES}
        else:
            matrix[state] = {s: round(counts[state][s] / total, 4) for s in ALL_STATES}

    return matrix


# ─────────────────────────────────────────────────────────────────────────────
# PUBLIC: forecast_states — matrix exponentiation for multi-day forecast
# ─────────────────────────────────────────────────────────────────────────────

def _mat_multiply(A: dict, B: dict) -> dict:
    """Multiply two 3×3 dicts-of-dicts."""
    result = {s: {t: 0.0 for t in ALL_STATES} for s in ALL_STATES}
    for i in ALL_STATES:
        for k in ALL_STATES:
            for j in ALL_STATES:
                result[i][j] += A[i][k] * B[k][j]
    return result


def _mat_pow(M: dict, n: int) -> dict:
    """Raise transition matrix M to the nth power."""
    if n <= 0:
        # Identity
        return {s: {t: (1.0 if s == t else 0.0) for t in ALL_STATES} for s in ALL_STATES}
    if n == 1:
        return {s: dict(M[s]) for s in ALL_STATES}
    half = _mat_pow(M, n // 2)
    result = _mat_multiply(half, half)
    if n % 2 == 1:
        result = _mat_multiply(result, M)
    return result


def forecast_states(
    current_state: State,
    matrix: dict[str, dict[str, float]],
    days: list[int] = None,
    candles_per_day: int = 24,   # 1h candles → 24 per day
) -> dict[int, dict[str, float]]:
    """Forecast probability distribution over states for each horizon.

    Args:
      current_state:    Current state ("BULL", "BEAR", "RANGING")
      matrix:           Transition matrix from build_transition_matrix()
      days:             List of forecast horizons in days [1, 3, 7]
      candles_per_day:  Number of 1h candles per day (24 for 1h bars)

    Returns:
      {1: {"BULL": 0.74, "BEAR": 0.14, "RANGING": 0.12},
       3: {"BULL": 0.58, ...},
       7: {"BULL": 0.45, ...}}
    """
    if days is None:
        days = [1, 3, 7]

    result = {}
    for d in days:
        steps = d * candles_per_day
        Mn = _mat_pow(matrix, steps)
        row = Mn[current_state]
        result[d] = {s: round(row[s], 3) for s in ALL_STATES}

    return result


# ─────────────────────────────────────────────────────────────────────────────
# PUBLIC: persistence_score — how sticky is the current state?
# ─────────────────────────────────────────────────────────────────────────────

def persistence_score(matrix: dict[str, dict[str, float]], state: State) -> float:
    """Self-transition probability = P(state → same state).

    Returns 0.0–1.0. High = state tends to persist.
    """
    return round(matrix.get(state, {}).get(state, 0.5), 4)


# ─────────────────────────────────────────────────────────────────────────────
# PUBLIC: analyze_symbol — full Markov analysis for one symbol
# ─────────────────────────────────────────────────────────────────────────────

def analyze_symbol(src, symbol: str) -> dict:
    """Fetch 1h klines and run full Markov analysis.

    Cached for MK_TTL seconds per symbol.
    Returns:
      {
        symbol, state, confidence, reasons,
        matrix, forecast, persistence,
        candles_used, ts
      }
    """
    cached = _mk_cache.get(symbol)
    if cached and time.time() - cached["ts"] < MK_TTL:
        return cached["result"]

    try:
        df = src.klines(symbol, "1h")

        state, conf, reasons = classify_state(df)
        matrix = build_transition_matrix(df, lookback=500)
        forecast = forecast_states(state, matrix, days=[1, 3, 7])
        persist  = persistence_score(matrix, state)

        result = {
            "symbol":      symbol,
            "state":       state,
            "confidence":  conf,
            "reasons":     reasons,
            "matrix":      matrix,
            "forecast": {
                "1d": forecast[1],
                "3d": forecast[3],
                "7d": forecast[7],
            },
            "persistence":    persist,
            "candles_used":   min(len(df), 500),
            "lookahead_safe": True,   # badge for frontend: closed candles only
            "ts":             time.time(),
        }

        _mk_cache[symbol] = {"result": result, "ts": time.time()}
        return result

    except Exception as e:
        err = {
            "symbol":     symbol,
            "state":      "RANGING",
            "confidence": 0.40,
            "reasons":    [f"Error: {str(e)[:80]}"],
            "matrix":     {s: {t: 1/3 for t in ALL_STATES} for s in ALL_STATES},
            "forecast":   {
                "1d": {s: 1/3 for s in ALL_STATES},
                "3d": {s: 1/3 for s in ALL_STATES},
                "7d": {s: 1/3 for s in ALL_STATES},
            },
            "persistence":    0.50,
            "candles_used":   0,
            "lookahead_safe": True,
            "error":          str(e)[:80],
            "ts":             time.time(),
        }
        _mk_cache[symbol] = {"result": err, "ts": time.time()}
        return err