"""
state_transitions.py

[Phase 4 -- new] Deterministic next-state probability estimation, sibling
to prediction_engine.py's existing _next_move_distribution() -- same
philosophy applied one layer up: an inspectable, config-driven heuristic
table (NOT ML, NOT statistically calibrated -- same explicit discipline as
prediction_engine.py's own docstring), operating on regime_engine.py's
market-state taxonomy instead of prediction_engine's five next-move states.

Reuses two readings that already exist rather than inventing a third bias
figure that would mean almost the same thing:
  - prediction.bullish_probability[horizon] (prediction_engine.py)
  - compression.breakout_pressure_score (compression_engine.py)

Pure function, no state, no I/O. Not consumed anywhere in the live bot yet.
"""

from dataclasses import dataclass, field

from config_loader import get_config
import regime_engine as re

DEFAULT_CONFIG = {
    "nudge_bullish_probability_weight": 0.4,
    "nudge_breakout_pressure_weight": 0.3,
    "prediction_horizon": "3m",
    # state: {candidate_next_state: base_weight, ...} -- weights need not
    # sum to 1.0, they're renormalized after nudging. Every regime_engine
    # state should appear as a key here; a state missing from this table
    # falls back to "UNKNOWN: 1.0" (see estimate_next_state()) rather than
    # raising, since a config typo or a not-yet-added state should fail
    # safe to "no opinion," not crash the caller.
    "base_transitions": {
        re.CHOP: {re.RANGE: 0.4, re.CHOP: 0.4, re.STRUCTURE_BREAKDOWN: 0.2},
        re.RANGE: {re.ACCUMULATION: 0.4, re.CHOP: 0.3, re.RANGE: 0.3},
        re.ACCUMULATION: {re.PRESSURE_BUILDING: 0.45, re.ACCUMULATION: 0.35, re.RANGE: 0.2},
        re.PRESSURE_BUILDING: {re.BREAKOUT_ATTEMPT: 0.5, re.ACCUMULATION: 0.35, re.STRUCTURE_BREAKDOWN: 0.15},
        re.BREAKOUT_ATTEMPT: {re.BREAKOUT_CONFIRMED: 0.5, re.FAILED_BREAKOUT: 0.35, re.BREAKOUT_ATTEMPT: 0.15},
        re.BREAKOUT_CONFIRMED: {re.EXPANSION: 0.55, re.PULLBACK: 0.30, re.FAILED_BREAKOUT: 0.15},
        re.EXPANSION: {re.CONTINUATION: 0.5, re.PULLBACK: 0.35, "EXTENDED": 0.15},
        re.CONTINUATION: {re.CONTINUATION: 0.45, re.PULLBACK: 0.30, "EXTENDED": 0.25},
        "EXTENDED": {re.PULLBACK: 0.55, re.CONTINUATION: 0.20, re.STRUCTURE_BREAKDOWN: 0.25},
        re.PULLBACK: {re.CONTINUATION: 0.45, re.VWAP_RECLAIM: 0.30, re.STRUCTURE_BREAKDOWN: 0.25},
        re.VWAP_RECLAIM: {re.CONTINUATION: 0.45, re.PULLBACK: 0.30, re.FAILED_BREAKOUT: 0.25},
        re.RECOVERY: {re.ACCUMULATION: 0.4, re.RECOVERY: 0.35, re.STRUCTURE_BREAKDOWN: 0.25},
        re.STRONG_UPTREND: {re.CONTINUATION: 0.4, re.STRONG_UPTREND: 0.35, re.PULLBACK: 0.25},
        re.WEAK_UPTREND: {re.STRONG_UPTREND: 0.4, re.PULLBACK: 0.35, re.CHOP: 0.25},
        re.DOWNTREND: {"DECLINING": 0.5, re.DOWNTREND: 0.3, re.RECOVERY: 0.2},
        "DECLINING": {re.STRUCTURE_BREAKDOWN: 0.4, "DECLINING": 0.35, re.RECOVERY: 0.25},
        "STALE": {"WEAK": 0.35, re.ACCUMULATION: 0.35, "STALE": 0.3},
        "WEAK": {"STALE": 0.4, "DECLINING": 0.35, "WEAK": 0.25},
        re.FAILED_BREAKOUT: {re.RECOVERY: 0.35, "DECLINING": 0.4, re.FAILED_BREAKOUT: 0.25},
        re.STRUCTURE_BREAKDOWN: {"DECLINING": 0.5, re.STRUCTURE_BREAKDOWN: 0.35, re.RECOVERY: 0.15},
        re.OPENING_VOLATILITY: {re.PRESSURE_BUILDING: 0.35, re.OPENING_VOLATILITY: 0.35, re.CHOP: 0.3},
        re.DISTRIBUTION: {"DECLINING": 0.45, re.DISTRIBUTION: 0.35, re.STRUCTURE_BREAKDOWN: 0.2},
    },
}


@dataclass
class TransitionEstimate:
    current_state: str
    probabilities: dict = field(default_factory=dict)   # {candidate_next_state: probability}
    insufficient_data: bool = False
    missing: list = field(default_factory=list)


def estimate_next_state(current_state: str, prediction=None, compression=None, cfg: dict = None) -> TransitionEstimate:
    """
    current_state: a regime_engine.py regime label (e.g. regime.regime).
    prediction: optional prediction_engine.PredictionReading -- used for the
        bullish-probability nudge. Omit to skip that nudge (base table only).
    compression: optional compression_engine.CompressionReading -- used for
        the breakout-pressure nudge. Omit to skip that nudge.
    """
    cfg = {**DEFAULT_CONFIG, **(cfg or get_config().get("state_transitions", {}))}
    table = cfg["base_transitions"]

    base = table.get(current_state, {"UNKNOWN": 1.0})

    bullish_prob = None
    if prediction is not None and not getattr(prediction, "insufficient_data", False):
        bullish_prob = prediction.bullish_probability.get(cfg["prediction_horizon"])

    pressure = None
    if compression is not None and not getattr(compression, "insufficient_data", False):
        pressure = compression.breakout_pressure_score / 100.0

    nudged = {}
    for candidate, weight in base.items():
        w = weight
        bullish_leaning = candidate in (re.CONTINUATION, re.EXPANSION, re.BREAKOUT_ATTEMPT,
                                         re.BREAKOUT_CONFIRMED, re.PRESSURE_BUILDING, re.RECOVERY)
        bearish_leaning = candidate in (re.STRUCTURE_BREAKDOWN, re.FAILED_BREAKOUT, "DECLINING", "WEAK")
        if bullish_prob is not None and bullish_leaning:
            w *= (1 + cfg["nudge_bullish_probability_weight"] * (bullish_prob - 0.5) * 2)
        if bullish_prob is not None and bearish_leaning:
            w *= (1 - cfg["nudge_bullish_probability_weight"] * (bullish_prob - 0.5) * 2)
        if pressure is not None and candidate in (re.BREAKOUT_ATTEMPT, re.PRESSURE_BUILDING):
            w *= (1 + cfg["nudge_breakout_pressure_weight"] * pressure)
        nudged[candidate] = max(0.0, w)

    total = sum(nudged.values()) or 1.0
    probabilities = {k: round(v / total, 3) for k, v in nudged.items()}

    return TransitionEstimate(current_state=current_state, probabilities=probabilities)