"""LSTM temperature forecasting trainer.

"""
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(
    train_records: list[dict],
    val_records: list[dict],
    params: dict,
    save_path: Path,
    on_progress: Callable[[int, int, float, float | None], None],
    cancel_event: threading.Event,
) -> dict[str, float | None]:
    """
    Train an LSTM model on *train_records* validated against *val_records*.

    Args:
        train_records: list of dicts for the training set.
        val_records:   list of dicts for the validation set.
        params:        hyper-parameter dict (seq_len, hidden_size, num_layers,
                       epochs, batch_size, learning_rate,
                       lr_factor, lr_patience, lr_min).
        save_path:     destination .pt file.
        on_progress:   callback(pct, epoch, 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))

    # ── scheduler 超参数 ──────────────────────────────────────────────────────
    lr_factor = float(params.get('lr_factor', 0.5))
    lr_patience = max(1, int(params.get('lr_patience', 5)))
    lr_min = float(params.get('lr_min', 1e-6))

    # ── data preparation ────────────────────────────────────────────────────
    train_data = _extract_features(train_records)
    val_data = _extract_features(val_records)

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

    # min-max normalisation fitted on train set only
    data_min = train_data.min(axis=0)
    data_max = train_data.max(axis=0)
    data_range = data_max - data_min
    data_range[data_range == 0] = 1.0

    train_norm = (train_data - data_min) / data_range
    X_train, y_train = _make_sequences(train_norm, seq_len)

    has_val = len(val_data) >= seq_len + 1
    X_val_t = y_val_t = None
    if has_val:
        val_norm = (val_data - data_min) / data_range
        X_val, y_val = _make_sequences(val_norm, seq_len)
        if len(X_val) == 0:
            has_val = False
        else:
            X_val_t = torch.tensor(X_val)
            y_val_t = torch.tensor(y_val)

    device = torch.device('cpu')
    X_train_t = torch.tensor(X_train).to(device)
    y_train_t = torch.tensor(y_train).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()

    # ── scheduler ────────────────────────────────────────────────────────────
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode='min',
        factor=lr_factor,
        patience=lr_patience,
        min_lr=lr_min,
    )

    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()

            # 根据 val_loss 动态调整学习率
            scheduler.step(val_loss)

        pct = int((epoch + 1) / epochs * 100)
        on_progress(pct, epoch + 1, 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,
    }

