"""LSTM time-series predictor for price direction.

Input: 100-candle windows with 8 per-candle features (close, rsi, macdHist,
       bbPercentB, atr, obv, funding, fearGreed) → shape (1, 100, 8)
Output: price direction probability over 1h/4h/1d horizons

Trained on historical klines from PostgreSQL via training/train_lstm.py.
"""

import os
from pathlib import Path

import numpy as np

MODEL_DIR = Path(os.getenv("MODEL_DIR", "models"))

# Per-candle features expected by the LSTM: 8 features per timestep.
# The training pipeline (train_lstm.py) creates sequences with these columns.
PER_CANDLE_FEATURES = ["close", "rsi", "macdHist", "bbPercentB", "atr", "obv", "funding", "fearGreed"]
NUM_FEATURES = len(PER_CANDLE_FEATURES)
SEQUENCE_LENGTH = 100


class LSTMPredictor:
    def __init__(self):
        self.version = "0.2.0"
        self.is_loaded = False
        self.last_trained: str | None = None
        self.accuracy: float | None = None
        self.model = None
        self._scaler = None

    def load(self):
        """Load trained LSTM model from disk, or initialize with heuristic fallback."""
        model_path = MODEL_DIR / "lstm_predictor.pt"
        scaler_path = MODEL_DIR / "lstm_predictor_scaler.joblib"
        if model_path.exists():
            try:
                import torch

                self.model = torch.load(model_path, weights_only=True)
                self.is_loaded = True
                self.last_trained = str(model_path.stat().st_mtime)
                # Load fitted scaler if available
                if scaler_path.exists():
                    try:
                        import joblib
                        self._scaler = joblib.load(scaler_path)
                    except Exception:
                        self._scaler = None
            except Exception:
                self._init_heuristic()
        else:
            self._init_heuristic()

    def _init_heuristic(self):
        """Fallback: use a simple heuristic until trained model is available."""
        self.is_loaded = True
        self.version = "0.2.0-heuristic"

    def predict(self, ohlcv_window: list[dict], indicators: dict) -> dict:
        """Predict price direction from OHLCV + indicators.

        Returns:
            dict with keys: direction, probability, model
        """
        if self.model is not None:
            return self._predict_model(ohlcv_window, indicators)
        return self._predict_heuristic(indicators)

    def _predict_model(self, ohlcv_window: list[dict], indicators: dict) -> dict:
        """Run inference through trained LSTM model.

        Restructures input to proper (1, T, F) shape where T=100 timesteps
        and F=8 per-candle features, instead of the incorrect (1, 1, N) shape
        that treated all features as a single timestep.
        """
        import torch

        candles = ohlcv_window[-SEQUENCE_LENGTH:]

        # Build per-candle feature matrix: shape (T, F)
        # Each candle gets its own feature vector from available OHLCV + indicators
        feature_matrix = []
        for candle in candles:
            row = [
                float(candle.get("close", 0)),
                float(candle.get("rsi", indicators.get("rsi", 50))),
                float(candle.get("macdHist", indicators.get("macdHistogram", 0))),
                float(candle.get("bbPercentB", indicators.get("bollingerPercentB", 0.5))),
                float(candle.get("atr", indicators.get("atr", 0))),
                float(candle.get("obv", indicators.get("obv", 0))),
                float(candle.get("funding", indicators.get("fundingRate", 0))),
                float(candle.get("fearGreed", indicators.get("fearGreed", 50))),
            ]
            feature_matrix.append(row)

        # Pad if fewer than SEQUENCE_LENGTH candles (repeat first candle)
        while len(feature_matrix) < SEQUENCE_LENGTH:
            feature_matrix.insert(0, feature_matrix[0])

        arr = np.array(feature_matrix, dtype=np.float32)  # (100, 8)

        # Apply fitted scaler if available (normalize per-feature)
        if self._scaler is not None:
            arr = self._scaler.transform(arr)

        # Shape: (1, T=100, F=8) — proper LSTM input
        tensor = torch.tensor(arr, dtype=torch.float32).unsqueeze(0)

        with torch.no_grad():
            output = self.model(tensor)
            probs = torch.softmax(output, dim=-1).squeeze().numpy()

        directions = ["up", "sideways", "down"]
        idx = int(np.argmax(probs))

        return {
            "direction": directions[idx],
            "probability": float(probs[idx]),
            "model": f"lstm-predictor-{self.version}",
        }

    def _predict_heuristic(self, indicators: dict) -> dict:
        """Simple heuristic based on RSI + MACD until model is trained."""
        rsi = indicators.get("rsi", 50)
        macd_hist = indicators.get("macdHistogram", 0)

        score = 0.0
        if rsi < 30:
            score += 0.3
        elif rsi > 70:
            score -= 0.3
        if macd_hist > 0:
            score += 0.2
        elif macd_hist < 0:
            score -= 0.2

        if score > 0.1:
            return {"direction": "up", "probability": 0.5 + score, "model": "lstm-heuristic"}
        elif score < -0.1:
            return {"direction": "down", "probability": 0.5 + abs(score), "model": "lstm-heuristic"}
        return {"direction": "sideways", "probability": 0.5, "model": "lstm-heuristic"}
