"""Standalone inference for the UA-ANN colorimetric biosensor model.

Loads the trained 5-seed MLP ensemble and predicts uric-acid concentration (μM)
from smartphone RGB readings.

Two input modes are supported:

1. **With per-phone blank** (recommended, matches training conditions).
   Caller supplies (R_blank, G_blank, B_blank) — the values measured on the
   same phone for the same buffer at zero analyte. ΔRGB is computed exactly
   as in training.

2. **Without blank** (less accurate, single-shot inference).
   The script falls back to the per-buffer mean blanks computed across the 6
   training phones.

Usage from the command line
---------------------------

    # Single sample with blank reference
    python inference.py --R 50 --G 70 --B 90 --buffer DI \\
                        --R0 60 --G0 78 --B0 104

    # Batch from CSV (columns: R, G, B, Buffer, [R0, G0, B0])
    python inference.py --csv samples.csv --out predictions.csv

Usage from Python
-----------------

    from inference import load_model, predict
    bundle = load_model("model.pt")
    pred_uM = predict(bundle, R=50, G=70, B=90, buffer="DI",
                      R0=60, G0=78, B0=104)
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
import torch
import torch.nn as nn


# ------------------------------------------------------------------ #
# Per-buffer mean blanks across the 6 training phones (used as fallback
# when the caller does not supply a per-device blank).
# ------------------------------------------------------------------ #
MEAN_BLANKS = {
    "DI":   {"R": 59.864,  "G": 77.204,  "B": 103.901},
    "pH4":  {"R": 80.980,  "G": 125.474, "B": 157.914},
    "pH11": {"R": 96.932,  "G": 125.474, "B": 145.695},
    "SBF":  {"R": 80.434,  "G": 127.047, "B": 153.688},
}


class MLP(nn.Module):
    def __init__(self, in_dim, hidden, dropout=0.10):
        super().__init__()
        layers, prev = [], in_dim
        for h in hidden:
            layers += [nn.Linear(prev, h), nn.ReLU(), nn.Dropout(dropout)]
            prev = h
        layers.append(nn.Linear(prev, 1))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x).squeeze(-1)


def load_model(path: str | Path = "model.pt") -> dict:
    """Load checkpoint and rebuild the 5-seed ensemble.

    Returns a dict with keys: ensemble, scaler_mean, scaler_scale, feature_cols,
    feature_set, buffer_order, y_mean, y_std.
    """
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    in_dim = ckpt["in_dim"]
    hidden = ckpt["architecture"]
    ensemble = []
    for sd in ckpt["ensemble_state_dicts"]:
        m = MLP(in_dim=in_dim, hidden=hidden)
        m.load_state_dict(sd)
        m.eval()
        ensemble.append(m)
    return {
        "ensemble":      ensemble,
        "scaler_mean":   np.asarray(ckpt["scaler_mean"], dtype=np.float32),
        "scaler_scale":  np.asarray(ckpt["scaler_scale"], dtype=np.float32),
        "feature_cols":  ckpt["feature_cols"],
        "feature_set":   ckpt["feature_set"],
        "buffer_order":  ckpt["buffer_order"],
        # The released checkpoint was trained directly on raw μM targets
        # (no y-normalization), so the identity transform is the correct default.
        "y_mean":        float(ckpt.get("y_mean", 0.0)),
        "y_std":         float(ckpt.get("y_std",  1.0)),
    }


def _build_features(R: float, G: float, B: float, buffer: str,
                    R0: Optional[float], G0: Optional[float], B0: Optional[float],
                    bundle: dict) -> np.ndarray:
    """Construct a 1-D feature vector matching the model's input layout."""
    if buffer not in bundle["buffer_order"]:
        raise ValueError(f"Unknown buffer {buffer!r}; expected one of {bundle['buffer_order']}")

    if R0 is None or G0 is None or B0 is None:
        mb = MEAN_BLANKS[buffer]
        R0 = R0 if R0 is not None else mb["R"]
        G0 = G0 if G0 is not None else mb["G"]
        B0 = B0 if B0 is not None else mb["B"]

    dR, dG, dB = R0 - R, G0 - G, B0 - B
    ohe = [1.0 if b == buffer else 0.0 for b in bundle["buffer_order"]]
    # The released model was trained on feature_set="all":
    # [R, G, B, dR, dG, dB] + 4 buffer one-hot dims = 10 features.
    feats = [R, G, B, dR, dG, dB] + ohe
    if len(feats) != len(bundle["feature_cols"]):
        raise RuntimeError(f"Feature length mismatch: built {len(feats)}, "
                           f"expected {len(bundle['feature_cols'])}")
    return np.asarray(feats, dtype=np.float32)


def predict(bundle: dict, R: float, G: float, B: float, buffer: str,
            R0: Optional[float] = None, G0: Optional[float] = None,
            B0: Optional[float] = None) -> float:
    """Predict UA concentration (μM) for a single RGB sample."""
    x = _build_features(R, G, B, buffer, R0, G0, B0, bundle)
    x_n = (x - bundle["scaler_mean"]) / bundle["scaler_scale"]
    xt = torch.tensor(x_n, dtype=torch.float32).unsqueeze(0)
    with torch.no_grad():
        preds = [m(xt).item() for m in bundle["ensemble"]]
    yhat_n = float(np.mean(preds))
    return yhat_n * bundle["y_std"] + bundle["y_mean"]


def predict_batch(bundle: dict, df: pd.DataFrame) -> np.ndarray:
    """Predict for a DataFrame with columns R, G, B, Buffer (and optionally R0, G0, B0)."""
    has_blank = all(c in df.columns for c in ("R0", "G0", "B0"))
    out = np.empty(len(df), dtype=np.float32)
    for i, row in enumerate(df.itertuples(index=False)):
        out[i] = predict(
            bundle,
            R=row.R, G=row.G, B=row.B, buffer=row.Buffer,
            R0=row.R0 if has_blank else None,
            G0=row.G0 if has_blank else None,
            B0=row.B0 if has_blank else None,
        )
    return out


def _cli() -> int:
    p = argparse.ArgumentParser(description="UA-ANN biosensor inference")
    p.add_argument("--model", default="model.pt", help="Path to model.pt")
    p.add_argument("--csv",   help="Batch CSV with R, G, B, Buffer columns")
    p.add_argument("--out",   default="predictions.csv",
                   help="Output path for batch predictions")
    p.add_argument("--R", type=float); p.add_argument("--G", type=float); p.add_argument("--B", type=float)
    p.add_argument("--R0", type=float); p.add_argument("--G0", type=float); p.add_argument("--B0", type=float)
    p.add_argument("--buffer", choices=["DI", "pH4", "pH11", "SBF"])
    args = p.parse_args()

    if not Path(args.model).exists():
        print(f"Model file not found: {args.model}", file=sys.stderr); return 1

    bundle = load_model(args.model)

    if args.csv:
        df = pd.read_csv(args.csv)
        df["predicted_UA_uM"] = predict_batch(bundle, df)
        df.to_csv(args.out, index=False)
        print(f"Wrote {args.out} ({len(df)} predictions)")
        print(df.head())
    else:
        for k in ("R", "G", "B", "buffer"):
            if getattr(args, k) is None:
                print(f"--{k} required for single-sample mode", file=sys.stderr); return 1
        y = predict(bundle, R=args.R, G=args.G, B=args.B, buffer=args.buffer,
                    R0=args.R0, G0=args.G0, B0=args.B0)
        used_blank = "supplied" if (args.R0 is not None) else "mean-blank fallback"
        print(json.dumps({
            "predicted_UA_uM": round(y, 2),
            "buffer":          args.buffer,
            "blank_source":    used_blank,
        }, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(_cli())