"""Causal, low-churn routing utilities for the non-promoting MLP head probe.

The router is deliberately a small state machine rather than a per-bar MoE
gate.  It uses only contemporaneous or historical columns and cannot change
state more often than ``min_dwell_bars``.  The heads are residual output heads
on top of the promoted dense model's frozen 55→16→8 shared trunk.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd


BEAR, NEUTRAL, BULL = -1, 0, 1


@dataclass(frozen=True)
class SlowRouterConfig:
    smoothing_span: int = 72
    trend_window: int = 28
    enter_threshold: float = 0.20
    exit_threshold: float = 0.05
    min_dwell_bars: int = 28


def _column(frame: pd.DataFrame, name: str) -> pd.Series:
    if name not in frame:
        raise ValueError(f"Slow regime router requires '{name}' in the chart export")
    return pd.to_numeric(frame[name], errors="coerce").fillna(0.0)


def router_signal(frame: pd.DataFrame, config: SlowRouterConfig = SlowRouterConfig()) -> np.ndarray:
    """Return the causal state signal from on-chain, trend, vol and liquidity.

    Price trend is a trailing return; all smoothers use ``adjust=False`` so a
    value at t is a function of rows through t only.
    """
    if config.smoothing_span < 1 or config.trend_window < 1:
        raise ValueError("smoothing_span and trend_window must be positive")
    if config.min_dwell_bars < 1:
        raise ValueError("min_dwell_bars must be positive")
    mvrv = _column(frame, "mvrv_zscore_cont")
    rvol = _column(frame, "rvol_norm")
    liq = _column(frame, "fed_net_liq_sign")
    close = _column(frame, "close")
    trend = (close / close.shift(config.trend_window) - 1.0).fillna(0.0).clip(-0.05, 0.05) * 20.0
    smooth = lambda values: values.ewm(span=config.smoothing_span, adjust=False).mean()
    # Volatility reduces confidence but cannot directly cause a directional
    # switch.  This prevents a volatility spike from becoming an implicit gate.
    return (0.55 * smooth(mvrv) + 0.25 * smooth(trend) + 0.20 * smooth(liq)
            - 0.10 * smooth(rvol).abs()).to_numpy(dtype=np.float64)


def slow_regime_states(frame: pd.DataFrame, config: SlowRouterConfig = SlowRouterConfig()) -> np.ndarray:
    """Map the causal signal to {-1, 0, +1} with hysteresis and minimum dwell."""
    signal = router_signal(frame, config)
    states = np.zeros(len(signal), dtype=np.int8)
    state, last_change = NEUTRAL, -config.min_dwell_bars
    for index, value in enumerate(signal):
        if index - last_change < config.min_dwell_bars:
            states[index] = state
            continue
        candidate = state
        if state == BULL:
            if value <= -config.enter_threshold:
                candidate = BEAR
            elif value < config.exit_threshold:
                candidate = NEUTRAL
        elif state == BEAR:
            if value >= config.enter_threshold:
                candidate = BULL
            elif value > -config.exit_threshold:
                candidate = NEUTRAL
        elif value >= config.enter_threshold:
            candidate = BULL
        elif value <= -config.enter_threshold:
            candidate = BEAR
        if candidate != state:
            state, last_change = candidate, index
        states[index] = state
    return states


def shared_trunk_forward(X: np.ndarray, layers: list[tuple[np.ndarray, np.ndarray]]) -> np.ndarray:
    """Forward through the frozen dense trunk, excluding its final output head."""
    if len(layers) != 3:
        raise ValueError("Tiny-head experiment requires a 3-layer dense MLP artifact")
    value = np.asarray(X, dtype=np.float64)
    for weights, bias in layers[:-1]:
        with np.errstate(over="ignore", divide="ignore", invalid="ignore"):
            value = np.tanh(value @ np.asarray(weights, dtype=np.float64).T + np.asarray(bias, dtype=np.float64))
    return value


def regime_head_scores(trunk: np.ndarray, states: np.ndarray,
                       base_head: tuple[np.ndarray, np.ndarray],
                       residual_heads: tuple[np.ndarray, np.ndarray]) -> np.ndarray:
    """Score a frozen trunk with one small residual output head per slow state."""
    state_index = np.asarray(states, dtype=np.int8) + 1
    if not np.isin(state_index, (0, 1, 2)).all():
        raise ValueError("states must contain only -1, 0, or 1")
    base_weights, base_bias = (np.asarray(base_head[0], dtype=np.float64),
                               np.asarray(base_head[1], dtype=np.float64))
    head_weights, head_bias = (np.asarray(residual_heads[0], dtype=np.float64),
                               np.asarray(residual_heads[1], dtype=np.float64))
    if head_weights.shape != (3, trunk.shape[1]) or head_bias.shape != (3,):
        raise ValueError("residual heads must have shapes (3, trunk_width) and (3,)")
    with np.errstate(over="ignore", divide="ignore", invalid="ignore"):
        logits = (np.asarray(trunk, dtype=np.float64) @ base_weights.T + base_bias).reshape(-1)
    logits += (np.asarray(trunk, dtype=np.float64) * head_weights[state_index]).sum(axis=1) + head_bias[state_index]
    return np.tanh(logits) * 1000.0