"""
EMR (Extended Mass Range) binary classifier -- channels 700-1017 (318 bins).

Self-contained: no external data paths. Load a single seed model or the full
5-model ensemble and score CDA TOF mass spectra for EMR membership.

This is a shallow 1D CNN (same architecture as the HMOC classifier), trained with
only 34 known EMR positives, so it is best used as a *ranked triage list*: sort
unknown spectra by the ensemble probability and review the top candidates.

Input convention
----------------
The network was trained on the spectrum **cropped to channels 700..1017** (318
bins), then min-max normalised per spectrum to [0, 1]. Pass either a full raw
spectrum (``preprocess`` crops it for you) or an already-cropped 318-length array.

Files in this repo
------------------
    emr_model_seed_{42,100,2026,777,888}.pth   -> model.state_dict() per seed
    metrics.json                                -> per-fold validation metrics
"""
import os

import numpy as np
import torch
import torch.nn as nn

SEQ_LEN = 318
CHANNEL_START = 700   # inclusive
CHANNEL_END = 1018    # exclusive  -> spectrum[700:1018] is 318 bins
SEEDS = [42, 100, 2026, 777, 888]


def crop(spectrum):
    """Extract the model region (channels 700..1017) from a full spectrum."""
    return np.asarray(spectrum)[CHANNEL_START:CHANNEL_END]


def _pad_or_truncate(x, target_len=SEQ_LEN):
    if len(x) > target_len:
        return x[:target_len]
    if len(x) < target_len:
        return np.pad(x, (0, target_len - len(x)))
    return x


def preprocess(spectrum, already_cropped=False):
    """
    -> 318-bin, min-max [0,1] float32 array ready for the network.

    Pass a full raw spectrum (default), or set ``already_cropped=True`` if you
    already have the 700-1017 region (e.g. a ``spectrum_cropped`` column).
    """
    x = np.asarray(spectrum, dtype=np.float32)
    if not already_cropped:
        x = crop(x)
    x = _pad_or_truncate(x, SEQ_LEN)
    mn, mx = x.min(), x.max()
    return (x - mn) / (mx - mn + 1e-8)


class Spectra1DCNN(nn.Module):
    """3x Conv1D (32->64->128) + BatchNorm + ReLU + GAP + FC(128->64) + Dropout + FC(64->1)."""

    def __init__(self, seq_len=SEQ_LEN):
        super(Spectra1DCNN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv1d(1, 32, kernel_size=7, padding=3),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.Conv1d(32, 64, kernel_size=5, padding=2),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Conv1d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(1),
        )
        self.fc = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(64, 1),
        )

    def forward(self, x):
        # x shape: (batch, 1, seq_len)
        x = self.conv(x).squeeze(-1)  # (batch, 128)
        return self.fc(x)


def load_model(weights_path, seq_len=SEQ_LEN, device="cpu"):
    model = Spectra1DCNN(seq_len).to(device)
    model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
    model.eval()
    return model


def predict_ensemble(weights_paths, spectra, already_cropped=False, device="cpu"):
    """
    Score spectra by averaging the sigmoid output of every seed model.

    Parameters
    ----------
    weights_paths : list[str]
        Paths to the per-seed ``.pth`` files (usually the 5 in this repo).
    spectra : iterable[ndarray]
        Full raw spectra, or already-cropped 700-1017 arrays.
    already_cropped : bool
        Set True if each spectrum is already the 318-bin 700-1017 region.

    Returns
    -------
    np.ndarray of shape (N,) with one EMR probability in [0, 1] per spectrum.
    """
    models = [load_model(w, device=device) for w in weights_paths]
    X = torch.tensor(
        np.stack([preprocess(s, already_cropped=already_cropped) for s in spectra]),
        dtype=torch.float32).unsqueeze(1).to(device)  # (N, 1, SEQ_LEN)
    probs = torch.zeros(X.size(0), device=device)
    with torch.no_grad():
        for m in models:
            m.eval()
            probs += torch.sigmoid(m(X)).squeeze(1)
    probs /= len(models)
    return probs.cpu().numpy()


def default_weights_dir():
    return os.path.dirname(os.path.abspath(__file__))


if __name__ == "__main__":
    # Quick self-test: load all 5 seeds and score a dummy batch of full spectra.
    d = default_weights_dir()
    paths = [os.path.join(d, f"emr_model_seed_{s}.pth") for s in SEEDS]
    for p in paths:
        assert os.path.exists(p), f"missing weight file: {p}"
    # dummy FULL spectra (1018 channels) so the 700-1017 crop is exercised
    rng = np.random.default_rng(0)
    dummy = [rng.random(1018).astype(np.float32) for _ in range(4)]
    out = predict_ensemble(paths, dummy)
    print("Loaded 5 seed models and scored", len(dummy), "dummy spectra:")
    print("  EMR probabilities =", np.round(out, 4))