"""
STRATA-CORE: Cognitive Engine
State transition function F(State(t-1), Input(t), Memory) -> State(t)
"""

import math
from typing import Dict


# ---------------------------------------------------------------------------
# V1 weights (Phase 1 seed — archived for comparison)
# ---------------------------------------------------------------------------
WEIGHTS_V1 = {
    "w_trend_bias":       0.20,
    "w_vol_uncertainty":  0.10,
    "w_break_momentum":   0.25,
    "w_liq_trap":         0.15,
    "w_fake_break_bias":  0.30,
    "w_trend_str_bias":   0.20,
    "alpha":              0.40,
    "beta":               0.35,
    "gamma":              0.25,
    "c_bias":             1.20,
    "c_momentum":         0.80,
    "c_trap":             1.00,
    "decay_bias":         0.95,
    "decay_momentum":     0.90,
    "decay_trap_risk":    0.92,
    "decay_uncertainty":  0.93,
    "regime_momentum_mix": 0.60,
}

# ---------------------------------------------------------------------------
# V2 weights (Phase 2 seed — calibrated from 5-year AAPL replay)
#
# Changes vs V1 and rationale:
#   w_vol_uncertainty  0.10 → 0.25  : vol signal was too weak to influence state
#   alpha              0.40 → 0.30  : reduce raw signal dominance
#   gamma              0.25 → 0.35  : raise pattern memory influence to fix LONG bias
#   beta               0.35 → 0.35  : unchanged — structure weight balanced
#   c_momentum         0.80 → 1.00  : momentum was underweighted in confidence
#   c_trap             1.00 → 1.20  : raise trap penalty in confidence (more cautious)
#   w_break_momentum   0.25 → 0.30  : structure breaks should drive momentum harder
#   w_liq_trap         0.15 → 0.20  : liquidity spikes should register more trap risk
# ---------------------------------------------------------------------------
DEFAULT_WEIGHTS = {
    # Δ_signal
    "w_trend_bias":        0.20,
    "w_vol_uncertainty":   0.25,   # v1: 0.10
    # Δ_structure
    "w_break_momentum":    0.30,   # v1: 0.25
    "w_liq_trap":          0.15,   # kept at v1; 0.20 caused trap saturation in early steps
    # Δ_pattern
    "w_fake_break_bias":   0.30,
    "w_trend_str_bias":    0.20,
    # Conflict resolution
    "alpha":               0.30,   # v1: 0.40  signal weight
    "beta":                0.35,   # unchanged  structure weight
    "gamma":               0.35,   # v1: 0.25  pattern weight
    # Confidence
    "c_bias":              1.20,
    "c_momentum":          1.00,   # v1: 0.80
    "c_trap":              1.20,   # v1: 1.00
    # Decay
    "decay_bias":          0.95,
    "decay_momentum":      0.90,
    "decay_trap_risk":     0.92,
    "decay_uncertainty":   0.80,   # v2: faster decay 0.93→0.80; needed because w_vol_uncertainty raised to 0.25
    "regime_momentum_mix": 0.60,
}

STATE_KEYS = ["bias", "momentum", "trap_risk", "uncertainty"]
DERIVED_KEYS = ["regime_score"]  # computed, not delta-driven


def initial_state() -> Dict[str, float]:
    """Return a zeroed initial state."""
    return {k: 0.0 for k in STATE_KEYS}


def _clamp(value: float, lo: float = -1.0, hi: float = 1.0) -> float:
    """Hard clamp — prevents state explosion."""
    return max(lo, min(hi, value))


def _sigmoid(x: float) -> float:
    return 1.0 / (1.0 + math.exp(-x))


# ---------------------------------------------------------------------------
# Delta components
# ---------------------------------------------------------------------------

def _delta_signal(inp: Dict[str, float], w: Dict[str, float]) -> Dict[str, float]:
    """
    STRATA-SENSE output → state deltas.
    inp keys: trend (-1..1), vol (0..1), liquidity (-1..1)
    """
    return {
        "bias":        inp.get("trend", 0.0) * w["w_trend_bias"],
        "uncertainty": inp.get("vol", 0.0)   * w["w_vol_uncertainty"],
    }


def _delta_structure(inp: Dict[str, float], w: Dict[str, float]) -> Dict[str, float]:
    """
    Market structure signals → state deltas.
    inp keys: break_structure (bool/0-1), liquidity_above (bool/0-1)
    """
    # Downside pressure: large negative trend + high vol = trap risk rising.
    # Represents the danger of being long during a sharp directional move down.
    # Scaled by 0.3 to contribute signal without saturating trap_risk.
    trend    = inp.get("trend", 0.0)
    vol      = inp.get("vol",   0.0)
    downside_pressure = max(0.0, -trend) * vol * w["w_liq_trap"] * 0.15
    return {
        "momentum":  inp.get("break_structure", 0.0) * w["w_break_momentum"],
        "trap_risk": inp.get("liquidity_above", 0.0) * w["w_liq_trap"] + downside_pressure,
    }


def _delta_pattern(memory: Dict[str, float], w: Dict[str, float]) -> Dict[str, float]:
    """
    Memory pattern scores → state deltas.
    memory keys: fake_breakout (0..1), trend_strength (0..1)
    """
    fake_bo = memory.get("fake_breakout", 0.0)
    bias_delta = (
        - fake_bo * w["w_fake_break_bias"]
        + memory.get("trend_strength", 0.0) * w["w_trend_str_bias"]
    )
    # fake_breakout is a trap signal — push trap_risk up directly
    trap_delta = fake_bo * w["w_fake_break_bias"] * 0.25
    return {"bias": bias_delta, "trap_risk": trap_delta}


def _apply_decay(state: Dict[str, float], w: Dict[str, float]) -> Dict[str, float]:
    """Exponential decay — prevents state accumulating stale history."""
    return {
        "bias":        state["bias"]        * w["decay_bias"],
        "momentum":    state["momentum"]     * w["decay_momentum"],
        "trap_risk":   state["trap_risk"]    * w["decay_trap_risk"],
        "uncertainty": state["uncertainty"]  * w["decay_uncertainty"],
    }


def _compute_regime_score(state: Dict[str, float], w: Dict[str, float]) -> float:
    """
    regime_score is a derived scalar: how strongly the system is in a directional regime.
    Combines bias magnitude and momentum, range [-1, 1].
    Positive = trending long, Negative = trending short, ~0 = ranging/choppy.
    """
    mix = w["regime_momentum_mix"]
    raw = state["bias"] * (1.0 - mix) + state["momentum"] * mix
    return _clamp(raw)


def _resolve_bias(
    signal_bias: float,
    structure_bias: float,
    pattern_bias: float,
    w: Dict[str, float],
) -> float:
    """
    Weighted conflict resolution across layers.
    No single layer wins — all contribute proportionally.
    """
    return (
        signal_bias    * w["alpha"] +
        structure_bias * w["beta"]  +
        pattern_bias   * w["gamma"]
    )


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def update_state(
    state: Dict[str, float],
    inp: Dict[str, float],
    memory: Dict[str, float],
    weights: Dict[str, float] = None,
) -> Dict[str, float]:
    """
    F(State(t-1), Input(t), Memory) -> State(t)

    State(t) = State(t-1) + Δ_signal + Δ_structure + Δ_pattern - Δ_decay

    All state values are clamped to [-1.0, 1.0] after update.
    """
    w = weights if weights is not None else DEFAULT_WEIGHTS

    d_sig  = _delta_signal(inp, w)
    d_str  = _delta_structure(inp, w)
    d_pat  = _delta_pattern(memory, w)

    # Accumulate deltas into a working copy
    next_state = dict(state)

    for k in STATE_KEYS:
        next_state[k] += (
            d_sig.get(k, 0.0) +
            d_str.get(k, 0.0) +
            d_pat.get(k, 0.0)
        )

    # Conflict resolution for bias specifically
    next_state["bias"] = _resolve_bias(
        signal_bias    = d_sig.get("bias", 0.0),
        structure_bias = d_str.get("bias", 0.0),
        pattern_bias   = d_pat.get("bias", 0.0),
        w              = w,
    ) + state["bias"]

    # Decay
    next_state = _apply_decay(next_state, w)

    # Clamp core state values
    for k in STATE_KEYS:
        next_state[k] = _clamp(next_state[k])

    # Derived: regime_score computed from updated state (not delta-driven)
    next_state["regime_score"] = _compute_regime_score(next_state, w)

    return next_state


def compute_confidence(state: Dict[str, float], weights: Dict[str, float] = None) -> float:
    """
    confidence = sigmoid(|bias|*c1 + momentum*c2 - trap_risk*c3)
    Returns value in (0, 1).
    """
    w = weights if weights is not None else DEFAULT_WEIGHTS
    score = (
        abs(state["bias"])    * w["c_bias"]     +
        state["momentum"]     * w["c_momentum"] -
        state["trap_risk"]    * w["c_trap"]
    )
    return _sigmoid(score)


def classify_regime(state: Dict[str, float]) -> str:
    """Derive market regime label from current state."""
    bias      = state["bias"]
    momentum  = state["momentum"]
    trap_risk = state["trap_risk"]

    if trap_risk > 0.6:
        return "CHOPPY"
    if abs(bias) > 0.5 and momentum > 0.3:
        return "TRENDING"
    if abs(bias) < 0.2 and momentum < 0.2:
        return "RANGING"
    return "TRANSITIONING"