"""Inference for DeepECG-SSL-CLSA-CVD (Model A / Model A+).
Requires: torch, safetensors, numpy, fairseq-signals (https://github.com/Jwoo5/fairseq-signals) and the public
DeepECG-SSL v2 (amplitude-preserved) pretrained encoder checkpoint, which defines the architecture.
ECG input: (N, 12, 2500) float32 in millivolts, 250 Hz, 10 s, standard lead order, no per-lead normalisation.

    from inference import DeepECGSSLCVD
    m = DeepECGSSLCVD("weights/", "path/to/deepecg-ssl-v2-wcr-amp-preserved-encoder.pt", device="cuda")
    p_ecg = m.ecg_probability(ecg_mv)                       # Model A (calibrated ECG-only probability)
    risk = m.model_a_plus(ecg_mv, age, female, hypertension)  # Model A+ (no blood test)
"""
import json, os, numpy as np, torch, torch.nn as nn
from safetensors.torch import load_file
from fairseq_signals.utils import checkpoint_utils

HERE = os.path.dirname(os.path.abspath(__file__))
lg = lambda p: np.log(np.clip(p, 1e-6, 1 - 1e-6) / (1 - np.clip(p, 1e-6, 1 - 1e-6)))


class Net(nn.Module):
    def __init__(self, base_ckpt):
        super().__init__()
        models, _, _ = checkpoint_utils.load_model_and_task(base_ckpt)
        self.enc = models[0] if isinstance(models, list) else models
        for n in ["quantizer", "project_q", "final_proj", "target_glu"]:
            if hasattr(self.enc, n): setattr(self.enc, n, None)
        self.head = nn.Sequential(nn.Dropout(0.3), nn.Linear(768, 1))

    def forward(self, x):
        return self.head(self.enc(source=x, mask=False, features_only=True)["x"].mean(1)).squeeze(1)


class DeepECGSSLCVD:
    def __init__(self, weights_dir, base_ckpt, device="cpu"):
        self.p = json.load(open(os.path.join(HERE, "model_params.json"))); self.device = device; self.models = []
        for k in range(5):
            m = Net(base_ckpt); m.load_state_dict(load_file(os.path.join(weights_dir, f"deepecg_ssl_clsa_cvd_fold{k}.safetensors")), strict=False)
            self.models.append(m.to(device).eval())

    @torch.no_grad()
    def ecg_probability(self, ecg_mv, batch=64):
        x = torch.as_tensor(np.asarray(ecg_mv, dtype=np.float32)); out = []
        for k, m in enumerate(self.models):
            pk = []
            for i in range(0, len(x), batch):
                pk.append(torch.sigmoid(m(x[i:i + batch].to(self.device))).float().cpu().numpy())
            a, b = self.p["per_fold_platt"][str(k)]["a"], self.p["per_fold_platt"][str(k)]["b"]
            out.append(1 / (1 + np.exp(-(a * lg(np.concatenate(pk)) + b))))
        return np.mean(out, 0)

    def model_a_plus(self, ecg_mv, age, female, hypertension):
        c = self.p["model_A_plus"]; e = lg(self.ecg_probability(ecg_mv))
        z = c["intercept"] + c["b_ecg"] * e + c["b_age"] * np.asarray(age, float) + c["b_female"] * np.asarray(female, float) + c["b_htn"] * np.asarray(hypertension, float)
        return 1 / (1 + np.exp(-z))