#!/usr/bin/env python3
"""Evaluate unconditional waveform samples with EQT-feature FID/PRDC metrics.

This is the portable Train-bundle equivalent of the older SWaG evaluator.  It
does not read labels or conditions: it compares only the ``data`` datasets in
one real HDF5 and one generated HDF5.  Features are extracted from the hidden
representation of SeisBench's EQTransformer, then FID, Density, Coverage,
PRDC precision, and PRDC recall are computed.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
from pathlib import Path
from typing import Iterable

import h5py
import numpy as np
import torch
from scipy.linalg import LinAlgWarning, sqrtm
from sklearn.neighbors import NearestNeighbors
from tqdm import tqdm

DEFAULT_EQT_CHECKPOINT = Path(__file__).resolve().parent / "eqtransformer_original_state_dict.pt"


def resolve(path: str | Path) -> Path:
    value = Path(path).expanduser()
    return value.resolve()


def load_eqt(checkpoint: str, device: torch.device) -> torch.nn.Module:
    try:
        import seisbench.models as sbm
    except ImportError as exc:
        raise RuntimeError(
            "This evaluator requires SeisBench. Install it with: pip install seisbench"
        ) from exc

    key = str(checkpoint).strip().lower()
    if key in {"original", "stead"}:
        model = sbm.EQTransformer.from_pretrained("original")
    elif key == "random_init":
        model = sbm.EQTransformer(in_samples=6000)
    else:
        state = torch.load(resolve(checkpoint), map_location="cpu", weights_only=False)
        state_dict = state.get("state_dict", state.get("model", state))
        keys = set(state_dict.keys())
        pretrained_lstm = any(".lstm.cell_f." in key or ".lstm.cell_b." in key for key in keys)
        model = (
            sbm.EQTransformer.from_pretrained("original")
            if pretrained_lstm
            else sbm.EQTransformer(in_samples=6000)
        )
        model.load_state_dict(state_dict, strict=True)
    return model.eval().to(device)


class EQTHidden(torch.nn.Module):
    """Return the same hidden feature used by the archived SWaG evaluator."""

    def __init__(self, model: torch.nn.Module):
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.model.encoder(x)
        x = self.model.res_cnn_stack(x)
        x = self.model.bi_lstm_stack(x)
        x, _ = self.model.transformer_d0(x)
        x, _ = self.model.transformer_d(x)
        return x


def prepare_batch(batch: np.ndarray, target_npts: int = 6000) -> torch.Tensor:
    """Normalize [N,T,3] or [N,3,T] waveforms like the archived evaluator."""
    from scipy.signal import resample

    prepared: list[np.ndarray] = []
    for waveform in np.asarray(batch, dtype=np.float32):
        if waveform.ndim != 2:
            raise ValueError(f"Expected one waveform with 2 dimensions, got {waveform.shape}")
        if waveform.shape[-1] == 3:
            channels_last = waveform
        elif waveform.shape[0] == 3:
            channels_last = waveform.T
        else:
            raise ValueError(f"Expected waveform shape [T,3] or [3,T], got {waveform.shape}")
        if channels_last.shape[0] != target_npts:
            channels_last = resample(channels_last, target_npts, axis=0)
        mean = channels_last.mean(axis=0, keepdims=True)
        std = channels_last.std(axis=0, keepdims=True)
        std[std < 1e-6] = 1.0
        prepared.append(((channels_last - mean) / std).T.astype(np.float32, copy=False))
    if not prepared:
        return torch.zeros((0, 3, target_npts), dtype=torch.float32)
    return torch.from_numpy(np.stack(prepared, axis=0))


def read_indices(path: Path, indices: np.ndarray, batch_start: int, batch_end: int) -> np.ndarray:
    with h5py.File(path, "r") as handle:
        data = handle["data"]
        selected = indices[batch_start:batch_end]
        if selected.size == 0:
            return np.zeros((0,) + tuple(data.shape[1:]), dtype=np.float32)
        order = np.argsort(selected, kind="stable")
        sorted_indices = selected[order]
        sorted_batch = np.asarray(data[sorted_indices], dtype=np.float32)
        inverse = np.empty_like(order)
        inverse[order] = np.arange(order.size, dtype=order.dtype)
        return sorted_batch[inverse]


def row_count(path: Path) -> int:
    with h5py.File(path, "r") as handle:
        if "data" not in handle or len(handle["data"].shape) != 3:
            raise ValueError(f"Expected a 3D 'data' dataset in {path}")
        return int(handle["data"].shape[0])


def select_indices(path: Path, sample_size: int, seed: int) -> np.ndarray:
    count = row_count(path)
    n = count if sample_size <= 0 else min(int(sample_size), count)
    rng = np.random.default_rng(seed)
    return np.sort(rng.choice(count, size=n, replace=False)) if n < count else np.arange(count, dtype=np.int64)


@torch.no_grad()
def extract_features(
    path: Path,
    indices: np.ndarray,
    feature_model: torch.nn.Module,
    device: torch.device,
    batch_size: int,
    num_gpus: int,
) -> np.ndarray:
    features: list[np.ndarray] = []
    for start in tqdm(range(0, len(indices), batch_size), desc=f"EQT features: {path.name}"):
        end = min(start + batch_size, len(indices))
        waves = read_indices(path, indices, start, end)
        model_input = prepare_batch(waves).to(device, non_blocking=True)
        hidden = feature_model(model_input)
        hidden = hidden.detach().float().cpu().numpy()
        features.append(np.concatenate([hidden.mean(axis=-1), hidden.std(axis=-1)], axis=1).astype(np.float32))
    if not features:
        raise ValueError(f"No samples selected from {path}")
    return np.concatenate(features, axis=0)


def frechet_distance(real: np.ndarray, fake: np.ndarray, eps: float = 1e-6) -> float:
    real = np.asarray(real, dtype=np.float64)
    fake = np.asarray(fake, dtype=np.float64)
    mu_real, mu_fake = real.mean(axis=0), fake.mean(axis=0)
    cov_real = np.atleast_2d(np.cov(real, rowvar=False))
    cov_fake = np.atleast_2d(np.cov(fake, rowvar=False))
    cov_real += np.eye(cov_real.shape[0]) * eps
    cov_fake += np.eye(cov_fake.shape[0]) * eps
    with np.errstate(all="ignore"):
        import warnings
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", LinAlgWarning)
            covmean = sqrtm(cov_real @ cov_fake)
    if np.iscomplexobj(covmean):
        covmean = covmean.real
    diff = mu_real - mu_fake
    return float(np.real(diff @ diff + np.trace(cov_real + cov_fake - 2.0 * covmean)))


def prdc(real: np.ndarray, fake: np.ndarray, nearest_k: int) -> dict[str, float]:
    real = np.asarray(real, dtype=np.float64)
    fake = np.asarray(fake, dtype=np.float64)
    if len(real) <= nearest_k or len(fake) == 0:
        raise ValueError(f"Need more than nearest_k={nearest_k} real samples")
    real_neighbors = NearestNeighbors(n_neighbors=nearest_k + 1).fit(real)
    real_distances, _ = real_neighbors.kneighbors(real)
    real_radii = real_distances[:, -1]

    fake_neighbors = NearestNeighbors(n_neighbors=nearest_k + 1).fit(fake)
    fake_distances, _ = fake_neighbors.kneighbors(fake)
    fake_radii = fake_distances[:, -1]

    precision_hits = 0
    density_sum = 0.0
    coverage = np.zeros(len(real), dtype=bool)
    recall_hits = 0
    for start in range(0, len(fake), 1024):
        fake_chunk = fake[start : start + 1024]
        distances = np.linalg.norm(fake_chunk[:, None, :] - real[None, :, :], axis=2)
        precision_hits += int(np.any(distances <= real_radii[None, :], axis=1).sum())
        density_sum += float(np.sum(distances <= real_radii[None, :]) / nearest_k)
        coverage |= np.any(distances <= real_radii[None, :], axis=0)
    for start in range(0, len(real), 1024):
        real_chunk = real[start : start + 1024]
        distances = np.linalg.norm(real_chunk[:, None, :] - fake[None, :, :], axis=2)
        recall_hits += int(np.any(distances <= fake_radii[None, :], axis=1).sum())

    return {
        "precision": precision_hits / len(fake),
        "recall": recall_hits / len(real),
        "density": density_sum / len(fake),
        "coverage": float(coverage.mean()),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--real-h5", required=True, type=Path)
    parser.add_argument("--fake-h5", required=True, type=Path)
    parser.add_argument("--output-dir", default="Train/results/distribution_eval", type=Path)
    parser.add_argument(
        "--eqt-checkpoint",
        default=str(DEFAULT_EQT_CHECKPOINT),
        help=(
            "Feature extractor checkpoint. The default is the bundled exact "
            "SeisBench-original state dict; do not replace it with the Iquique "
            "transfer-picker checkpoint unless you intentionally want a "
            "different metric space."
        ),
    )
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--sample-size", type=int, default=10000)
    parser.add_argument("--nearest-k", type=int, default=5)
    parser.add_argument("--seed", type=int, default=2026)
    parser.add_argument("--num-gpus", type=int, default=1)
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()
    if args.batch_size < 1 or args.sample_size == 1 or args.nearest_k < 1 or args.num_gpus < 1:
        raise ValueError("batch-size, nearest-k, and num-gpus must be positive; sample-size cannot be 1")

    real_h5, fake_h5 = resolve(args.real_h5), resolve(args.fake_h5)
    output_dir = resolve(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_csv = output_dir / "distribution_metrics_eqt.csv"
    if output_csv.exists() and not args.overwrite:
        print(output_csv)
        return

    if torch.cuda.is_available():
        if args.num_gpus > torch.cuda.device_count():
            raise ValueError(f"Requested {args.num_gpus} GPUs, only {torch.cuda.device_count()} are visible")
        device = torch.device("cuda:0")
    else:
        if args.num_gpus > 1:
            raise RuntimeError("--num-gpus > 1 requires CUDA")
        device = torch.device("cpu")

    real_indices = select_indices(real_h5, args.sample_size, args.seed)
    fake_indices = select_indices(fake_h5, args.sample_size, args.seed + 1)
    model = load_eqt(args.eqt_checkpoint, device)
    feature_model: torch.nn.Module = EQTHidden(model)
    if device.type == "cuda" and args.num_gpus > 1:
        feature_model = torch.nn.DataParallel(feature_model, device_ids=list(range(args.num_gpus)))
    feature_model.eval()

    real_features = extract_features(real_h5, real_indices, feature_model, device, args.batch_size, args.num_gpus)
    fake_features = extract_features(fake_h5, fake_indices, feature_model, device, args.batch_size, args.num_gpus)
    metrics = {
        "fid": frechet_distance(real_features, fake_features),
        **prdc(real_features, fake_features, args.nearest_k),
        "feature_model": "eqt",
        "eqt_checkpoint": str(args.eqt_checkpoint),
        "n_real": int(len(real_features)),
        "n_fake": int(len(fake_features)),
        "nearest_k": int(args.nearest_k),
        "seed": int(args.seed),
        "unconditional": True,
    }
    with output_csv.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(metrics))
        writer.writeheader()
        writer.writerow(metrics)
    np.save(output_dir / "real_features_eqt.npy", real_features)
    np.save(output_dir / "fake_features_eqt.npy", fake_features)
    (output_dir / "eval_config.json").write_text(json.dumps(vars(args), default=str, indent=2), encoding="utf-8")
    print(output_csv)
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()