"""LSTM temperature forecasting trainer.

Uses PyTorch if available. If not installed, raises a descriptive RuntimeError
so the train service can mark the task as failed with a helpful message.
"""
from __future__ import annotations

import threading
from pathlib import Path
from typing import Callable

import numpy as np

# ── optional torch import ────────────────────────────────────────────────────
_TORCH_AVAILABLE = False
try:
    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader, TensorDataset

    _TORCH_AVAILABLE = True

    class _LSTMModel(nn.Module):
        def __init__(self, input_size: int, hidden_size: int, num_layers: int) -> None:
            super().__init__()
            dropout = 0.2 if num_layers > 1 else 0.0
            self.lstm = nn.LSTM(
                input_size,
                hidden_size,
                num_layers,
                batch_first=True,
                dropout=dropout,
            )
            self.fc = nn.Linear(hidden_size, 1)

        def forward(self, x: 'torch.Tensor') -> 'torch.Tensor':
            out, _ = self.lstm(x)
            return self.fc(out[:, -1, :]).squeeze(-1)

except ImportError:
    pass  # _TORCH_AVAILABLE stays False


# ── constants ────────────────────────────────────────────────────────────────
FEATURE_COLS = ['current', 'voltage', 'set_temperature', 'actual_temperature']
TARGET_COL = 'actual_temperature'
TARGET_IDX = FEATURE_COLS.index(TARGET_COL)


def _check_torch() -> None:
    if not _TORCH_AVAILABLE:
        raise RuntimeError(
            'PyTorch 未安装，请执行：\n'
            '  pip install torch --index-url https://download.pytorch.org/whl/cpu\n'
            '安装后重启后端服务。'
        )


# ── data helpers ─────────────────────────────────────────────────────────────

def _extract_features(records: list[dict]) -> np.ndarray:
    """Convert records list → (N, 4) float32 array; rows with parse errors skipped."""
    rows: list[list[float]] = []
    for r in records:
        try:
            row = [float(r.get(col) or 0) for col in FEATURE_COLS]
            if any(v != v for v in row):  # NaN guard
                continue
            rows.append(row)
        except (TypeError, ValueError):
            continue
    return np.array(rows, dtype=np.float32)


def _make_sequences(data: np.ndarray, seq_len: int) -> tuple[np.ndarray, np.ndarray]:
    """Return (X, y): X.shape=(N, seq_len, features), y.shape=(N,)."""
    xs: list[np.ndarray] = []
    ys: list[float] = []
    for i in range(len(data) - seq_len):
        xs.append(data[i : i + seq_len])
        ys.append(data[i + seq_len, TARGET_IDX])
    return np.array(xs, dtype=np.float32), np.array(ys, dtype=np.float32)


# ── public training entry point ───────────────────────────────────────────────

def train_lstm(
    records: list[dict],
    params: dict,
    save_path: Path,
    on_progress: Callable[[int, float, float | None], None],
    cancel_event: threading.Event,
) -> dict[str, float | None]:
    """
    Train an LSTM model on *records* and persist it to *save_path*.

    Args:
        records:      list of dicts with keys in FEATURE_COLS.
        params:       hyper-parameter dict (seq_len, hidden_size, num_layers,
                      epochs, batch_size, learning_rate, train_ratio).
        save_path:    destination .pt file.
        on_progress:  callback(pct, train_loss, val_loss) called after each epoch.
        cancel_event: when set, training stops with InterruptedError.

    Returns:
        {'train_loss': float, 'val_loss': float|None}
    """
    _check_torch()

    seq_len = max(1, int(params.get('seq_len', 20)))
    hidden_size = max(1, int(params.get('hidden_size', 64)))
    num_layers = max(1, int(params.get('num_layers', 2)))
    epochs = max(1, int(params.get('epochs', 50)))
    batch_size = max(1, int(params.get('batch_size', 32)))
    lr = float(params.get('learning_rate', 0.001))
    train_ratio = min(0.99, max(0.5, float(params.get('train_ratio', 0.8))))

    # ── data preparation ────────────────────────────────────────────────────
    data = _extract_features(records)
    min_required = seq_len + 10
    if len(data) < min_required:
        raise ValueError(
            f'有效数据量不足：需至少 {min_required} 条，当前仅 {len(data)} 条。'
            '请检查数据包内容或减小序列长度。'
        )

    # min-max normalisation per feature
    data_min = data.min(axis=0)
    data_max = data.max(axis=0)
    data_range = data_max - data_min
    data_range[data_range == 0] = 1.0
    data_norm = (data - data_min) / data_range

    X, y = _make_sequences(data_norm, seq_len)
    n_train = max(1, int(len(X) * train_ratio))
    X_train, y_train = X[:n_train], y[:n_train]
    X_val, y_val = X[n_train:], y[n_train:]

    device = torch.device('cpu')
    X_train_t = torch.tensor(X_train).to(device)
    y_train_t = torch.tensor(y_train).to(device)
    has_val = len(X_val) > 0
    if has_val:
        X_val_t = torch.tensor(X_val).to(device)
        y_val_t = torch.tensor(y_val).to(device)

    train_loader = DataLoader(
        TensorDataset(X_train_t, y_train_t),
        batch_size=batch_size,
        shuffle=True,
    )

    # ── model ────────────────────────────────────────────────────────────────
    input_size = len(FEATURE_COLS)
    model = _LSTMModel(input_size, hidden_size, num_layers).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()

    train_loss = 0.0
    val_loss: float | None = None

    for epoch in range(epochs):
        if cancel_event.is_set():
            raise InterruptedError('训练已取消')

        # ── train step ───────────────────────────────────────────────────────
        model.train()
        epoch_loss = 0.0
        for xb, yb in train_loader:
            optimizer.zero_grad()
            pred = model(xb)
            loss = criterion(pred, yb)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item() * len(xb)

        train_loss = epoch_loss / len(X_train)

        # ── val step ─────────────────────────────────────────────────────────
        if has_val:
            model.eval()
            with torch.no_grad():
                val_pred = model(X_val_t)
                val_loss = criterion(val_pred, y_val_t).item()

        pct = int((epoch + 1) / epochs * 100)
        on_progress(pct, train_loss, val_loss)

    # ── persist ──────────────────────────────────────────────────────────────
    save_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            'model_state': model.state_dict(),
            'params': params,
            'data_min': data_min.tolist(),
            'data_max': data_max.tolist(),
            'feature_cols': FEATURE_COLS,
            'target_col': TARGET_COL,
            'input_size': input_size,
            'hidden_size': hidden_size,
            'num_layers': num_layers,
            'seq_len': seq_len,
        },
        save_path,
    )

    return {
        'train_loss': round(float(train_loss), 6),
        'val_loss': round(float(val_loss), 6) if val_loss is not None else None,
    }
