strata/trainer.py

"""
STRATA-TRAINER: Walk-Forward Weight Optimizer
==============================================
Optimizes STRATA model weights from historical OHLCV data using
coordinate-wise perturbation search (gradient-free).

This is the "training loop" equivalent for STRATA — analogous to
model.fit() in Keras or trainer.train() in HuggingFace Transformers,
but designed for the structured state-machine architecture of STRATA.

No PyTorch or TensorFlow required — pure Python, no external dependencies.

Usage:
    from strata.trainer import StrataTrainer

    trainer = StrataTrainer(asset="AAPL")
    windows = StrataTrainer.prepare_windows(candles_list, window_size=31)
    model   = trainer.train(windows, n_trials=100)
    model.save("aapl_model.json")
"""

import copy
import random
from typing import Dict, List, Optional, Tuple

from .core   import DEFAULT_WEIGHTS, initial_state, update_state, compute_confidence, classify_regime
from .memory import StrataMEMORY
from .decide import decide
from .guard  import StrataGUARD
from .sense  import sense


# ---------------------------------------------------------------------------
# Which weights are trainable (others are structural/fixed)
# ---------------------------------------------------------------------------
TRAINABLE_KEYS = [
    "w_trend_bias",
    "w_vol_uncertainty",
    "w_break_momentum",
    "w_liq_trap",
    "w_fake_break_bias",
    "w_trend_str_bias",
    "c_bias",
    "c_momentum",
    "c_trap",
    "decay_bias",
    "decay_momentum",
    "decay_trap_risk",
    "decay_uncertainty",
]

# Hard bounds per parameter to prevent pathological values
WEIGHT_BOUNDS: Dict[str, Tuple[float, float]] = {
    "w_trend_bias":       (0.05, 0.50),
    "w_vol_uncertainty":  (0.05, 0.50),
    "w_break_momentum":   (0.05, 0.60),
    "w_liq_trap":         (0.05, 0.40),
    "w_fake_break_bias":  (0.05, 0.60),
    "w_trend_str_bias":   (0.05, 0.50),
    "c_bias":             (0.50, 2.50),
    "c_momentum":         (0.30, 2.00),
    "c_trap":             (0.50, 2.50),
    "decay_bias":         (0.80, 0.99),
    "decay_momentum":     (0.70, 0.99),
    "decay_trap_risk":    (0.75, 0.99),
    "decay_uncertainty":  (0.60, 0.99),
}


def _clip_weights(weights: Dict[str, float]) -> Dict[str, float]:
    """Clip all trainable weights to their allowed bounds."""
    w = dict(weights)
    for k, (lo, hi) in WEIGHT_BOUNDS.items():
        if k in w:
            w[k] = max(lo, min(hi, w[k]))
    return w


def _run_episode(
    windows:  List[List[Dict]],
    weights:  Dict[str, float],
    asset:    Optional[str] = None,
) -> Dict[str, float]:
    """
    Run a full episode over all windows with given weights.
    Returns performance metrics used as the optimization objective.
    """
    state  = initial_state()
    memory = StrataMEMORY()
    guard  = StrataGUARD(asset=asset)

    n_long = n_short = n_hold = n_blocked = 0
    confidence_sum = 0.0
    trap_sum       = 0.0
    bias_abs_sum   = 0.0
    n_trending     = 0

    for window in windows:
        inp        = sense(window)
        mem_signal = memory.snapshot()
        state      = update_state(state, inp, mem_signal, weights=weights)
        decision   = decide(state, weights=weights)
        confidence = decision["confidence"]

        approved, _ = guard.evaluate(state, decision, confidence)

        action = decision["action"] if approved else "HOLD"
        if action == "LONG":
            n_long  += 1
        elif action == "SHORT":
            n_short += 1
        else:
            n_hold  += 1
        if not approved:
            n_blocked += 1

        confidence_sum += confidence
        trap_sum       += state["trap_risk"]
        bias_abs_sum   += abs(state["bias"])
        if decision["regime"] in ("TRENDING", "TRANSITIONING"):
            n_trending += 1

    total = len(windows)
    if total == 0:
        return {"score": -999.0}

    block_rate      = n_blocked / total
    action_rate     = (n_long + n_short) / total
    mean_confidence = confidence_sum / total
    mean_trap       = trap_sum / total
    mean_bias_abs   = bias_abs_sum / total
    trending_rate   = n_trending / total

    # Objective: maximize action quality
    # We want:
    #   - meaningful action rate (not too low = dormant, not too high = reckless)
    #   - high mean confidence on active bars
    #   - low mean trap risk
    #   - good bias signal strength
    #   - reasonable trend detection
    #
    # Penalties:
    #   - block_rate > 0.70: model too defensive
    #   - block_rate < 0.20: model not defensive enough (no risk management)
    #   - action_rate < 0.10: model mostly dormant

    score = (
        mean_confidence * 2.0
        + mean_bias_abs * 1.5
        + action_rate   * 1.0
        + trending_rate * 0.5
        - mean_trap     * 1.5
    )

    # Penalty for extreme block rates
    if block_rate > 0.75:
        score -= (block_rate - 0.75) * 3.0
    if block_rate < 0.15:
        score -= (0.15 - block_rate) * 2.0

    return {
        "score":       score,
        "block_rate":  block_rate,
        "action_rate": action_rate,
        "mean_conf":   mean_confidence,
        "mean_trap":   mean_trap,
        "trending":    trending_rate,
    }


class StrataTrainer:
    """
    Walk-forward coordinate optimizer for STRATA model weights.

    Finds weights that maximize signal quality over the training window
    using gradient-free coordinate perturbation search.

    Analogous to model.compile() + model.fit() in Keras.
    """

    def __init__(
        self,
        asset:   Optional[str] = None,
        verbose: bool          = True,
        seed:    int           = 42,
    ):
        self.asset   = asset
        self.verbose = verbose
        self._rng    = random.Random(seed)

    @staticmethod
    def prepare_windows(
        candles:     List[Dict],
        window_size: int = 31,
    ) -> List[List[Dict]]:
        """
        Convert a flat list of OHLCV candles into sliding windows
        suitable for model.fit() and predict().

        Args:
            candles:     flat list of OHLCV dicts, oldest → newest
            window_size: number of candles per window (minimum 3, recommend 31)

        Returns:
            list of windows, each window is a list of `window_size` candles
        """
        if len(candles) < window_size:
            raise ValueError(
                f"Need at least {window_size} candles, got {len(candles)}"
            )
        return [
            candles[i : i + window_size]
            for i in range(len(candles) - window_size + 1)
        ]

    def optimize(
        self,
        windows:   List[List[Dict]],
        seed:      Optional[Dict] = None,
        n_trials:  int            = 50,
        step_size: float          = 0.02,
    ) -> Dict[str, float]:
        """
        Coordinate-wise perturbation search over TRAINABLE_KEYS.

        Each trial: pick a random trainable weight, perturb +/- step_size,
        keep change if score improves. Repeat n_trials times.

        Args:
            windows:   training windows from prepare_windows()
            seed:      starting weights (defaults to DEFAULT_WEIGHTS)
            n_trials:  number of perturbation attempts
            step_size: perturbation magnitude per step

        Returns:
            optimized weights dict
        """
        weights = _clip_weights(copy.deepcopy(seed or DEFAULT_WEIGHTS))
        weights.pop("_meta", None)

        best_metrics = _run_episode(windows, weights, self.asset)
        best_score   = best_metrics["score"]

        if self.verbose:
            print(f"[StrataTrainer] Starting optimization | asset={self.asset} | "
                  f"windows={len(windows)} | trials={n_trials}")
            print(f"  Baseline score: {best_score:.4f}  "
                  f"block={best_metrics['block_rate']:.1%}  "
                  f"conf={best_metrics['mean_conf']:.3f}")

        improvements = 0
        for trial in range(n_trials):
            # Pick random trainable key
            key = self._rng.choice(TRAINABLE_KEYS)
            lo, hi = WEIGHT_BOUNDS[key]

            # Try both directions
            for direction in (+1, -1):
                candidate = dict(weights)
                candidate[key] = max(lo, min(hi, candidate[key] + direction * step_size))

                metrics = _run_episode(windows, candidate, self.asset)
                if metrics["score"] > best_score:
                    best_score   = metrics["score"]
                    weights      = candidate
                    improvements += 1
                    if self.verbose and improvements % 5 == 0:
                        print(f"  [{trial+1:>4}/{n_trials}] improved → score={best_score:.4f}  "
                              f"block={metrics['block_rate']:.1%}  "
                              f"conf={metrics['mean_conf']:.3f}  "
                              f"key={key}{'+' if direction > 0 else '-'}")
                    break   # take first improvement

            # Cool step size over time
            step_size = max(0.005, step_size * 0.998)

        if self.verbose:
            final = _run_episode(windows, weights, self.asset)
            print(f"\n[StrataTrainer] Done | improvements={improvements}/{n_trials}")
            print(f"  Final score:    {final['score']:.4f}")
            print(f"  block_rate:     {final['block_rate']:.1%}")
            print(f"  action_rate:    {final['action_rate']:.1%}")
            print(f"  mean_conf:      {final['mean_conf']:.3f}")
            print(f"  mean_trap:      {final['mean_trap']:.3f}")
            print(f"  trending_rate:  {final['trending']:.1%}")

        return weights

    def train(
        self,
        windows:   List[List[Dict]],
        n_trials:  int   = 50,
        step_size: float = 0.02,
    ) -> "StrataModel":
        """
        Full training pipeline: optimize weights and return a fitted StrataModel.

        Args:
            windows:   training windows from prepare_windows()
            n_trials:  optimization passes
            step_size: perturbation magnitude

        Returns:
            fitted StrataModel ready for predict() and save()

        Example:
            trainer = StrataTrainer(asset="AAPL")
            windows = StrataTrainer.prepare_windows(my_candles)
            model   = trainer.train(windows, n_trials=100)
            model.save("aapl_model.json")
        """
        from .model import StrataModel

        best_weights = self.optimize(
            windows   = windows,
            n_trials  = n_trials,
            step_size = step_size,
        )

        model = StrataModel(asset=self.asset, weights=best_weights)
        model._trained = True
        model._meta = {
            "asset":       self.asset or "UNKNOWN",
            "version":     StrataModel.VERSION,
            "trained":     True,
            "n_windows":   len(windows),
            "n_trials":    n_trials,
            "description": f"Trained on {len(windows)} windows",
        }
        return model