"""LSTM inference — predict actual_temperature for every data point in a dataset.

Requires the same PyTorch environment used for training.
"""
from __future__ import annotations

import math
from pathlib import Path
from typing import Any

import numpy as np

from app.ml.lstm_trainer import (
    FEATURE_COLS,
    TARGET_IDX,
    _LSTMModel,
    _check_torch,
    _extract_features,
)

# Maximum number of chart points stored in DB (subsample if dataset is larger)
CHART_MAX_POINTS = 1500


def predict_lstm(
    records: list[dict[str, Any]],
    model_path: Path,
) -> dict[str, Any]:
    """
    Run inference on *records* using the saved model at *model_path*.

    Returns a dict::

        {
            'total_count': int,          # total prediction points
            'mae': float,
            'rmse': float,
            'chart_data': [              # sampled up to CHART_MAX_POINTS
                {'index': int, 'time': str, 'actual': float, 'predicted': float},
                ...
            ],
        }
    """
    _check_torch()

    import torch

    checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)

    seq_len: int = int(checkpoint.get('seq_len', 20))
    hidden_size: int = int(checkpoint.get('hidden_size', 64))
    num_layers: int = int(checkpoint.get('num_layers', 2))
    input_size: int = int(checkpoint.get('input_size', len(FEATURE_COLS)))
    data_min = np.array(checkpoint['data_min'], dtype=np.float32)
    data_max = np.array(checkpoint['data_max'], dtype=np.float32)

    data_range = data_max - data_min
    data_range[data_range == 0] = 1.0

    # ── prepare data ─────────────────────────────────────────────────────────
    raw = _extract_features(records)
    if len(raw) < seq_len + 1:
        raise ValueError(
            f'数据量不足：需至少 {seq_len + 1} 条有效记录，当前仅 {len(raw)} 条。'
        )

    data_norm = (raw - data_min) / data_range

    # Build all sequences at once → (N-seq_len, seq_len, features)
    n = len(data_norm)
    indices = np.arange(n - seq_len)
    X = np.stack([data_norm[i : i + seq_len] for i in indices], axis=0).astype(np.float32)

    # ── model inference ───────────────────────────────────────────────────────
    model = _LSTMModel(input_size, hidden_size, num_layers)
    model.load_state_dict(checkpoint['model_state'])
    model.eval()

    with torch.no_grad():
        X_t = torch.tensor(X)
        # Batch to avoid OOM on very large datasets
        batch_sz = 512
        preds_norm: list[float] = []
        for start in range(0, len(X_t), batch_sz):
            out = model(X_t[start : start + batch_sz])
            preds_norm.extend(out.numpy().tolist())

    preds_norm_arr = np.array(preds_norm, dtype=np.float32)

    # Denormalize
    t_min = float(data_min[TARGET_IDX])
    t_range = float(data_range[TARGET_IDX])
    preds_real = preds_norm_arr * t_range + t_min
    actuals_real = raw[seq_len:, TARGET_IDX]

    # ── metrics ───────────────────────────────────────────────────────────────
    errors = preds_real - actuals_real
    mae = float(np.mean(np.abs(errors)))
    mse = float(np.mean(errors ** 2))
    rmse = float(math.sqrt(mse))

    # MAPE — skip points where actual == 0 to avoid division by zero
    nonzero_mask = actuals_real != 0
    if nonzero_mask.any():
        mape: float | None = float(
            np.mean(np.abs(errors[nonzero_mask] / actuals_real[nonzero_mask])) * 100
        )
    else:
        mape = None

    # R² (coefficient of determination)
    ss_res = float(np.sum(errors ** 2))
    ss_tot = float(np.sum((actuals_real - float(np.mean(actuals_real))) ** 2))
    r2: float | None = (1.0 - ss_res / ss_tot) if ss_tot != 0 else None

    # ── build result points ───────────────────────────────────────────────────
    times = [str(r.get('time', '')) for r in records]
    total = len(preds_real)

    # Subsample for chart storage
    if total <= CHART_MAX_POINTS:
        sample_idx = list(range(total))
    else:
        step = total / CHART_MAX_POINTS
        sample_idx = [int(i * step) for i in range(CHART_MAX_POINTS)]

    chart_data = [
        {
            'index': seq_len + i,
            'time': times[seq_len + i] if (seq_len + i) < len(times) else str(seq_len + i),
            'actual': round(float(actuals_real[i]), 4),
            'predicted': round(float(preds_real[i]), 4),
            'current': round(float(raw[seq_len + i, 0]), 4),
            'voltage': round(float(raw[seq_len + i, 1]), 4),
        }
        for i in sample_idx
    ]

    return {
        'total_count': total,
        'mae': round(mae, 6),
        'mse': round(mse, 6),
        'rmse': round(rmse, 6),
        'mape': round(mape, 4) if mape is not None else None,
        'r2': round(r2, 6) if r2 is not None else None,
        'chart_data': chart_data,
    }
