"""Training script for the LSTM price direction predictor.

Updated to produce proper (batch, T=100, F=8) shaped input instead of
the incorrect (batch, 1, N) that treated all features as a single timestep.

Per-candle features (8): close, rsi, macdHist, bbPercentB, atr, obv, funding, fearGreed
"""

import os
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from .data_loader import create_sequences, get_training_symbols, load_ohlcv

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

# Per-candle features: matches lstm_predictor.py PER_CANDLE_FEATURES
NUM_FEATURES = 8
SEQUENCE_LENGTH = 100


class PriceLSTM(nn.Module):
    """LSTM model for price direction classification.

    Input shape: (batch, T=100, F=8)
    Output: 3-class logits (down, sideways, up)
    """

    def __init__(self, input_size: int = NUM_FEATURES, hidden_size: int = 64, num_layers: int = 2):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 32),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(32, 3),  # 3 classes: down, sideways, up
        )

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        last_hidden = lstm_out[:, -1, :]
        return self.fc(last_hidden)


def create_enriched_sequences(
    df, window: int = SEQUENCE_LENGTH, horizon: int = 4
) -> tuple[np.ndarray, np.ndarray]:
    """Create training sequences with per-candle TA features.

    Instead of just OHLCV (5 features), each candle gets 8 features:
    close, rsi, macdHist, bbPercentB, atr, obv, funding, fearGreed.

    Falls back to create_sequences() if TA columns are missing (legacy data).
    """
    required_cols = ["close", "rsi", "macd_hist", "bb_percent_b", "atr", "obv"]
    has_ta = all(col in df.columns for col in required_cols)

    if not has_ta:
        # Legacy path: use raw OHLCV from data_loader
        return create_sequences(df, window, horizon)

    feature_cols = [
        "close", "rsi", "macd_hist", "bb_percent_b", "atr", "obv",
        "funding_rate", "fear_greed",
    ]
    # Fill missing optional columns with defaults
    if "funding_rate" not in df.columns:
        df = df.copy()
        df["funding_rate"] = 0.0
    if "fear_greed" not in df.columns:
        df = df.copy()
        df["fear_greed"] = 50.0

    data = df[feature_cols].values
    closes = df["close"].values

    X, y = [], []
    for i in range(window, len(data) - horizon):
        window_data = data[i - window: i].copy()

        # Normalize close prices by first close in window
        base_price = window_data[0, 0]
        if base_price == 0:
            continue
        window_data[:, 0] /= base_price

        X.append(window_data)

        # Label from future price change
        future_close = closes[i + horizon - 1]
        current_close = closes[i - 1]
        pct_change = (future_close - current_close) / current_close * 100

        if pct_change > 1.0:
            y.append(2)  # up
        elif pct_change < -1.0:
            y.append(0)  # down
        else:
            y.append(1)  # sideways

    return np.array(X, dtype=np.float32), np.array(y, dtype=np.int64)


def fit_scaler(X: np.ndarray):
    """Fit a StandardScaler on training data and return it.

    Reshapes (N, T, F) → (N*T, F) for fitting, then back.
    """
    from sklearn.preprocessing import StandardScaler

    n_samples, seq_len, n_features = X.shape
    flat = X.reshape(-1, n_features)
    scaler = StandardScaler()
    scaler.fit(flat)
    return scaler


def apply_scaler(X: np.ndarray, scaler) -> np.ndarray:
    """Apply a fitted StandardScaler to sequence data."""
    n_samples, seq_len, n_features = X.shape
    flat = X.reshape(-1, n_features)
    scaled = scaler.transform(flat)
    return scaled.reshape(n_samples, seq_len, n_features).astype(np.float32)


def train(days: int = 90, epochs: int = 50, batch_size: int = 32, lr: float = 1e-3):
    """Train LSTM on historical OHLCV data from the available symbol universe."""
    print(f"Loading data ({days} days)...")

    all_X, all_y = [], []
    for symbol in get_training_symbols():
        try:
            df = load_ohlcv(symbol, "4h", days)
            if len(df) < 200:
                print(f"  {symbol}: insufficient data ({len(df)} rows), skipping")
                continue
            X, y = create_enriched_sequences(df, window=SEQUENCE_LENGTH, horizon=4)
            all_X.append(X)
            all_y.append(y)
            print(f"  {symbol}: {len(X)} sequences, shape {X.shape}")
        except Exception as e:
            print(f"  {symbol}: error — {e}")

    if not all_X:
        print("No training data available. Ensure data collector has run.")
        return

    X = np.concatenate(all_X)
    y = np.concatenate(all_y)
    print(f"Total: {len(X)} sequences (shape {X.shape}), class distribution: {np.bincount(y)}")

    # Fit and save scaler for inference normalization
    scaler = fit_scaler(X)
    X = apply_scaler(X, scaler)

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    import joblib
    scaler_path = MODEL_DIR / "lstm_predictor_scaler.joblib"
    joblib.dump(scaler, scaler_path)
    print(f"Scaler saved to {scaler_path}")

    # Split 80/20
    split = int(len(X) * 0.8)
    X_train, X_val = X[:split], X[split:]
    y_train, y_val = y[:split], y[split:]

    train_ds = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))
    val_ds = TensorDataset(torch.tensor(X_val), torch.tensor(y_val))
    train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_dl = DataLoader(val_ds, batch_size=batch_size)

    model = PriceLSTM(input_size=NUM_FEATURES)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    best_val_acc = 0.0
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for X_batch, y_batch in train_dl:
            optimizer.zero_grad()
            output = model(X_batch)
            loss = criterion(output, y_batch)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        # Validation
        model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for X_batch, y_batch in val_dl:
                output = model(X_batch)
                preds = output.argmax(dim=-1)
                correct += (preds == y_batch).sum().item()
                total += len(y_batch)

        val_acc = correct / max(1, total)
        avg_loss = total_loss / len(train_dl)

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch + 1}/{epochs} — loss: {avg_loss:.4f}, val_acc: {val_acc:.3f}")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model, MODEL_DIR / "lstm_predictor.pt")

    print(f"Training complete. Best validation accuracy: {best_val_acc:.3f}")
    print(f"Model saved to {MODEL_DIR / 'lstm_predictor.pt'}")


if __name__ == "__main__":
    train()
