"""
HMOC binary classifier -- 1D CNN architecture + 5-model ensemble inference.

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

Input convention
----------------
Each raw spectrum is a 1-D array of intensities (variable length). Before the
network it is (1) truncated/zero-padded to ``SEQ_LEN`` (1000) bins and (2)
min-max normalised per spectrum to [0, 1], matching training.

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

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

SEQ_LEN = 1000
SEEDS = [42, 100, 2026, 777, 888]


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, target_len=SEQ_LEN):
    """Pad/truncate to target_len, then per-spectrum min-max to [0, 1]."""
    x = np.asarray(spectrum, dtype=np.float32)
    x = pad_or_truncate(x, target_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)
        x = x.squeeze(-1)  # (batch, 128)
        x = self.fc(x)
        return 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, device="cpu"):
    """
    Score raw 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]
        Raw spectra (each a 1-D array, any length).

    Returns
    -------
    np.ndarray of shape (N,) with one HMOC probability in [0, 1] per spectrum.
    """
    models = [load_model(w, device=device) for w in weights_paths]
    X = torch.tensor(np.stack([preprocess(s) 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 from this folder and score a dummy batch.
    d = default_weights_dir()
    paths = [os.path.join(d, f"hmoc_model_seed_{s}.pth") for s in SEEDS]
    for p in paths:
        assert os.path.exists(p), f"missing weight file: {p}"
    dummy = [np.random.rand(900 + 37 * i).astype(np.float32) for i in range(4)]
    out = predict_ensemble(paths, dummy)
    print("Loaded 5 seed models and scored", len(dummy), "dummy spectra:")
    print("  HMOC probabilities =", np.round(out, 4))