from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch

try:
    from .config import CORE_FEATURE_NAMES, make_experiment_config
    from .dataset import build_record_from_file, normalization_from_payload
    from .model import build_model
except ImportError:
    from config import CORE_FEATURE_NAMES, make_experiment_config
    from dataset import build_record_from_file, normalization_from_payload
    from model import build_model


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Predict task1 RMS from a single waveform CSV.")
    parser.add_argument("--file", type=str, required=True)
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--checkpoint", type=str, default=None)
    return parser.parse_args()


def resolve_device(device_name: str | None, default_device: str) -> torch.device:
    requested = device_name or default_device
    if requested.startswith("cuda") and not torch.cuda.is_available():
        return torch.device("cpu")
    return torch.device(requested)


def resolve_checkpoint_path(project_root: Path, checkpoint_dir: str, best_model_name: str, checkpoint_arg: str | None) -> Path:
    if checkpoint_arg:
        return Path(checkpoint_arg).resolve()
    return project_root / checkpoint_dir / "task1_feature_mlp" / best_model_name


def save_prediction_figure(
    file_path: Path,
    record,
    pred_rms: float,
    true_rms: float | None,
    output_dir: Path,
) -> Path:
    time_middle = record.time_middle.detach().cpu().numpy().reshape(-1)
    x_middle = record.x_middle.detach().cpu().numpy().reshape(-1)
    y_middle = record.y_middle.detach().cpu().numpy().reshape(-1)
    sampling_rate = record.sampling_rate
    fft_values = np.fft.rfft(x_middle)
    freqs = np.fft.rfftfreq(x_middle.size, d=1.0 / sampling_rate)
    amplitudes = np.abs(fft_values)

    fig, axes = plt.subplots(3, 1, figsize=(12, 10))
    fig.suptitle(f"Task1 Single-File Prediction | {file_path.name}")

    axes[0].plot(time_middle, x_middle, color="tab:blue")
    axes[0].set_title("Input Base Excitation (Middle Segment)")
    axes[0].set_xlabel("Time")
    axes[0].set_ylabel("Acceleration")
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(freqs, amplitudes, color="tab:purple")
    axes[1].axvline(record.frequency_hz, color="tab:red", linestyle="--", label=f"Dominant freq = {record.frequency_hz:.4f} Hz")
    axes[1].set_xlim(0.0, 5.0)
    axes[1].set_title("Input Spectrum")
    axes[1].set_xlabel("Frequency (Hz)")
    axes[1].set_ylabel("Amplitude")
    axes[1].grid(True, alpha=0.3)
    axes[1].legend()

    labels = ["Pred RMS"] if true_rms is None else ["True RMS", "Pred RMS"]
    values = [pred_rms] if true_rms is None else [true_rms, pred_rms]
    colors = ["tab:orange"] if true_rms is None else ["tab:green", "tab:orange"]
    axes[2].bar(labels, values, color=colors)
    title = f"Predicted RMS = {pred_rms:.6f}"
    if true_rms is not None:
        error_percent = abs(pred_rms - true_rms) / max(abs(true_rms), 1e-12) * 100.0
        title = f"True RMS = {true_rms:.6f} | Pred RMS = {pred_rms:.6f} | Error = {error_percent:.2f}%"
    axes[2].set_title(title)
    axes[2].set_ylabel("RMS")
    axes[2].grid(True, axis="y", alpha=0.3)

    plt.tight_layout()
    output_dir.mkdir(parents=True, exist_ok=True)
    figure_path = output_dir / f"{file_path.stem}_prediction.png"
    plt.savefig(figure_path, dpi=180, bbox_inches="tight")
    plt.close(fig)
    return figure_path


def main() -> None:
    args = parse_args()
    config = make_experiment_config()
    device = resolve_device(args.device, config.train.device)
    checkpoint_path = resolve_checkpoint_path(
        config.data.project_root,
        config.train.checkpoint_dir,
        config.train.best_model_name,
        args.checkpoint,
    )
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    normalization = normalization_from_payload(checkpoint["normalization"])
    model = build_model(config).to(device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    file_path = Path(args.file).resolve()
    record = build_record_from_file(file_path, split="predict", config=config.data)
    feature_norm = ((record.features - normalization.feature_mean) / normalization.feature_std).unsqueeze(0).to(device)
    x_rms_raw = record.x_rms.unsqueeze(0).to(device)
    target_mean = normalization.target_mean.to(device).view(1, 1)
    target_std = normalization.target_std.to(device).view(1, 1)

    with torch.no_grad():
        pred_norm = model(feature_norm)
        pred_tr = torch.clamp(pred_norm * target_std + target_mean, min=1e-6)
        pred_rms = float((pred_tr * x_rms_raw).detach().cpu().numpy().reshape(-1)[0])

    true_rms = float(record.y_rms.item()) if record.y_rms.numel() > 0 else None
    output_dir = config.data.project_root / "evaluation_outputs" / "task1_feature_mlp"
    figure_path = save_prediction_figure(file_path, record, pred_rms, true_rms, output_dir)

    print(f"Checkpoint: {checkpoint_path}")
    print(f"Input file: {file_path}")
    print(f"Extracted frequency (Hz): {record.frequency_hz:.6f}")
    for feature_name, feature_value in zip(CORE_FEATURE_NAMES, record.features.tolist()):
        print(f"{feature_name}: {feature_value:.6f}")
    print(f"Predicted RMS: {pred_rms:.6f}")
    if true_rms is not None:
        error_percent = abs(pred_rms - true_rms) / max(abs(true_rms), 1e-12) * 100.0
        print(f"True RMS: {true_rms:.6f}")
        print(f"Relative RMS Error (%): {error_percent:.4f}")
    print(f"Figure saved to: {figure_path}")


if __name__ == "__main__":
    main()