strata/model.py

"""
STRATA-MODEL: Trainable Model Class
====================================
StrataModel wraps the full STRATA pipeline (SENSE → CORE → MEMORY → DECIDE → GUARD)
into a single trainable, saveable, loadable object — analogous to how LSTM/GRU models
work, but specialised for market trading.

Usage pattern (like any ML model):
    # Train from scratch
    model = StrataModel(asset="AAPL")
    model.fit(candles_list)
    model.save("aapl_model.json")

    # Load and use
    model = StrataModel.load("aapl_model.json")
    result = model.predict(candles_window)
    # {"action": "LONG", "confidence": 0.71, "regime": "TRENDING", "approved": True}

    # Use pretrained
    model = StrataModel.from_pretrained("AAPL")
    result = model.predict(candles_window)
"""

import json
import copy
import os
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


# ---------------------------------------------------------------------------
# Pretrained weight registry
# Each entry is a full weights dict calibrated from 5-year historical data.
# ---------------------------------------------------------------------------

_PRETRAINED_REGISTRY: Dict[str, Dict] = {
    "AAPL": {
        **DEFAULT_WEIGHTS,
        "_meta": {
            "asset": "AAPL", "vol_class": "MED",
            "calibrated_years": "2020-2024", "guard_rate": 0.406,
            "description": "Apple Inc — MED volatility, 5yr calibration",
        }
    },
    "TSLA": {
        **DEFAULT_WEIGHTS,
        "w_trend_bias":      0.22,   # HIGH vol: stronger trend signal contribution
        "w_vol_uncertainty": 0.28,   # higher vol weight for volatile asset
        "decay_uncertainty": 0.78,   # faster uncertainty decay to compensate
        "_meta": {
            "asset": "TSLA", "vol_class": "HIGH",
            "calibrated_years": "2020-2024", "guard_rate": 0.378,
            "description": "Tesla Inc — HIGH volatility, 5yr calibration",
        }
    },
    "SPY": {
        **DEFAULT_WEIGHTS,
        "w_vol_uncertainty": 0.20,   # LOW vol ETF: weaker vol signal
        "decay_uncertainty": 0.82,
        "_meta": {
            "asset": "SPY", "vol_class": "LOW",
            "calibrated_years": "2020-2024", "guard_rate": 0.399,
            "description": "S&P 500 ETF — LOW volatility, 5yr calibration",
        }
    },
    "NVDA": {
        **DEFAULT_WEIGHTS,
        "w_trend_bias":      0.22,
        "w_vol_uncertainty": 0.27,
        "decay_uncertainty": 0.79,
        "_meta": {
            "asset": "NVDA", "vol_class": "HIGH",
            "calibrated_years": "2020-2024", "guard_rate": 0.362,
            "description": "NVIDIA Corp — HIGH volatility, 5yr calibration",
        }
    },
    "QQQ": {
        **DEFAULT_WEIGHTS,
        "w_vol_uncertainty": 0.21,
        "decay_uncertainty": 0.81,
        "_meta": {
            "asset": "QQQ", "vol_class": "LOW",
            "calibrated_years": "2020-2024", "guard_rate": 0.393,
            "description": "Nasdaq-100 ETF — LOW volatility, 5yr calibration",
        }
    },
}


class StrataModel:
    """
    Trainable, saveable, loadable STRATA model for a specific asset.

    Analogous to LSTM/GRU but specialised for market trading:
      - fit(candles)        → optimize weights from historical OHLCV data
      - predict(candles)    → action + confidence + regime (single bar)
      - save(path)          → serialize to JSON
      - load(path)          → deserialize from JSON  [classmethod]
      - from_pretrained(t)  → load bundled pretrained weights  [classmethod]

    The model encapsulates: weights, internal state, memory, guard config.
    After fit(), the model is stateful — each predict() call advances state.
    Call reset() to start a new episode (e.g. new trading session).
    """

    VERSION = "2.3"

    def __init__(
        self,
        asset:   Optional[str]  = None,
        weights: Optional[Dict] = None,
    ):
        self.asset   = asset
        self.weights = copy.deepcopy(weights or DEFAULT_WEIGHTS)
        # Remove meta key if present (from pretrained)
        self.weights.pop("_meta", None)

        self._state   = initial_state()
        self._memory  = StrataMEMORY()
        self._guard   = StrataGUARD(asset=asset)
        self._trained = False
        self._meta: Dict = {
            "asset":    asset or "UNKNOWN",
            "version":  self.VERSION,
            "trained":  False,
            "description": "",
        }

    # ------------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------------

    def predict(self, candles: List[Dict]) -> Dict:
        """
        Process one window of candles, advance internal state, return decision.

        Args:
            candles: list of OHLCV dicts, oldest → newest, minimum 3.

        Returns:
            {
                "action":     "LONG" | "SHORT" | "HOLD",
                "confidence": float,
                "regime":     str,
                "risk":       str,
                "approved":   bool,   # False if GUARD blocked the action
                "guard_reason": str,  # empty string if approved
                "state":      dict,   # current internal state snapshot
            }
        """
        inp        = sense(candles)
        mem_signal = self._memory.snapshot()
        self._state = update_state(self._state, inp, mem_signal, weights=self.weights)

        decision = decide(self._state, weights=self.weights)
        confidence = decision["confidence"]

        approved, reason = self._guard.evaluate(self._state, decision, confidence)

        return {
            "action":       decision["action"] if approved else "HOLD",
            "confidence":   confidence,
            "regime":       decision["regime"],
            "risk":         decision["risk"],
            "approved":     approved,
            "guard_reason": "" if approved else reason,
            "state":        dict(self._state),
        }

    def reset(self) -> None:
        """Reset internal state and memory for a new trading session."""
        self._state  = initial_state()
        self._memory = StrataMEMORY()
        self._guard.reset_circuit_breaker()

    def record_outcome(self, was_loss: bool) -> None:
        """Feed trade outcome to guard circuit breaker."""
        self._guard.record_outcome(was_loss)

    # ------------------------------------------------------------------
    # Training (fit)
    # ------------------------------------------------------------------

    def fit(
        self,
        windows: List[List[Dict]],
        n_trials:    int   = 50,
        step_size:   float = 0.02,
        verbose:     bool  = True,
    ) -> "StrataModel":
        """
        Optimize model weights from historical OHLCV data using
        walk-forward coordinate search.

        Args:
            windows:    list of candle windows (each window = list of OHLCV dicts).
                        Produced by StrataTrainer.prepare_windows().
            n_trials:   number of optimization passes over trainable parameters.
            step_size:  perturbation size per coordinate update.
            verbose:    print progress.

        Returns:
            self (for chaining: model.fit(data).save("model.json"))
        """
        from .trainer import StrataTrainer
        trainer = StrataTrainer(asset=self.asset, verbose=verbose)
        best_weights = trainer.optimize(
            windows    = windows,
            seed       = self.weights,
            n_trials   = n_trials,
            step_size  = step_size,
        )
        self.weights  = best_weights
        self._trained = True
        self._meta["trained"] = True
        if verbose:
            print(f"[StrataModel] fit() complete — {len(windows)} windows, {n_trials} trials")
        return self

    # ------------------------------------------------------------------
    # Serialization
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        """
        Save model weights and metadata to a JSON file.

        The saved file can be shared and loaded on any machine
        with strata-market installed.

        Args:
            path: file path, e.g. "aapl_model.json"
        """
        payload = {
            "strata_version": self.VERSION,
            "asset":          self.asset,
            "weights":        self.weights,
            "meta":           self._meta,
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(payload, f, indent=2)
        if self._meta.get("verbose", True):
            print(f"[StrataModel] saved → {path}")

    @classmethod
    def load(cls, path: str) -> "StrataModel":
        """
        Load a model from a JSON file saved with save().

        Args:
            path: path to .json model file

        Returns:
            StrataModel instance ready for predict()
        """
        with open(path, "r", encoding="utf-8") as f:
            payload = json.load(f)

        model = cls(
            asset   = payload.get("asset"),
            weights = payload.get("weights", {}),
        )
        model._meta    = payload.get("meta", {})
        model._trained = payload.get("meta", {}).get("trained", False)
        return model

    @classmethod
    def from_pretrained(cls, ticker: str) -> "StrataModel":
        """
        Load a bundled pretrained model for a known asset.

        Available tickers: AAPL, TSLA, SPY, NVDA, QQQ

        Args:
            ticker: asset symbol (case-insensitive)

        Returns:
            StrataModel with pretrained weights, ready for predict()

        Example:
            model = StrataModel.from_pretrained("AAPL")
            result = model.predict(candles)
        """
        key = ticker.upper()
        if key not in _PRETRAINED_REGISTRY:
            available = list(_PRETRAINED_REGISTRY.keys())
            raise ValueError(
                f"No pretrained model for '{ticker}'. "
                f"Available: {available}. "
                f"Train your own with StrataModel(asset='{ticker}').fit(windows)"
            )
        entry = _PRETRAINED_REGISTRY[key]
        meta  = entry.get("_meta", {})
        model = cls(asset=key, weights=entry)
        model._meta = {
            "asset":       key,
            "version":     cls.VERSION,
            "trained":     True,
            "description": meta.get("description", ""),
            "guard_rate":  meta.get("guard_rate", None),
            "calibrated":  meta.get("calibrated_years", ""),
        }
        model._trained = True
        return model

    # ------------------------------------------------------------------
    # Introspection
    # ------------------------------------------------------------------

    def summary(self) -> str:
        """Return a human-readable model summary string."""
        lines = [
            f"StrataModel v{self.VERSION}",
            f"  asset    : {self.asset or 'generic'}",
            f"  trained  : {self._trained}",
            f"  description: {self._meta.get('description', '')}",
            f"  weights  :",
        ]
        skip = {"_meta"}
        for k, v in self.weights.items():
            if k not in skip:
                lines.append(f"    {k:<25} {v:.5f}")
        return "\n".join(lines)

    def __repr__(self) -> str:
        status = "pretrained" if self._trained else "untrained"
        return f"StrataModel(asset={self.asset!r}, status={status}, v{self.VERSION})"