"""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).
        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))

    # ── 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'\u8bad\u7ec3\u96c6\u6709\u6548\u6570\u636e\u91cf\u4e0d\u8db3\uff1a\u9700\u81f3\u5c11 {min_required} \u6761\uff0c\u5f53\u524d\u4ec5 {len(train_data)} \u6761\u3002'
            '\u8bf7\u68c0\u67e5\u6570\u636e\u5305\u5185\u5bb9\u6216\u51cf\u5c0f\u5e8f\u5217\u957f\u5ea6\u3002'
        )

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

    train_loss = 0.0
    val_loss: float | None = None

    for epoch in range(epochs):
        if cancel_event.is_set():
            raise InterruptedError('\u8bad\u7ec3\u5df2\u53d6\u6d88')

        # ── 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, 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,
    }

    """
    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 is ignored – fixed 70/15/15 split is used.
        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, 'test_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))

    # ── 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)

    # Fixed 70 / 15 / 15 split
    n_total = len(X)
    n_train = max(1, int(n_total * _TRAIN_RATIO))
    n_val   = max(1, int(n_total * _VAL_RATIO))
    # test gets the remainder so the three parts always sum to n_total
    X_train, y_train = X[:n_train], y[:n_train]
    X_val,   y_val   = X[n_train : n_train + n_val], y[n_train : n_train + n_val]
    X_test,  y_test  = X[n_train + n_val :],          y[n_train + n_val :]

    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
    has_test = len(X_test) > 0
    if has_val:
        X_val_t = torch.tensor(X_val).to(device)
        y_val_t = torch.tensor(y_val).to(device)
    if has_test:
        X_test_t = torch.tensor(X_test).to(device)
        y_test_t = torch.tensor(y_test).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)

    # ── test-set evaluation (completely hidden during training) ──────────────
    test_loss: float | None = None
    if has_test:
        model.eval()
        with torch.no_grad():
            test_pred = model(X_test_t)
            test_loss = criterion(test_pred, y_test_t).item()

    # ── 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,
        'test_loss':  round(float(test_loss),  6) if test_loss  is not None else None,
    }
