"""
app.py — HistoPath DX  (single-image inference server)

Matches the Kaggle inference script exactly:
  - Backbone : facebook/mask2former-swin-small-coco-instance
  - SEG_CKPT : best_m2f.pth           (Mask2Former backbone weights)
  - CLS_CKPT : best_m2f_classifier.pth (full StrongerMask2FormerClassifier)
  - Vahadane : fit once on REFERENCE_IMAGE_PATH at startup
  - /predict : POST image → Vahadane → processor → model → JSON + overlay PNG
"""

import io, os, base64
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.cm as cm_lib

from PIL import Image as PILImage
from flask import Flask, request, jsonify, render_template
from flask_cors import CORS
from transformers import (Mask2FormerForUniversalSegmentation,
                          Mask2FormerImageProcessor)

# =============================================================================
# ── CONFIGURATION  (edit these paths before running) ─────────────────────────
# =============================================================================
BASE_DIR             = os.path.dirname(os.path.abspath(__file__))
REFERENCE_IMAGE_PATH = os.path.join(BASE_DIR, "reference.jpg")
SEG_CKPT             = os.path.join(BASE_DIR, "best_m2f.pth")
CLS_CKPT             = os.path.join(BASE_DIR, "best_m2f_classifier.pth")
_BACKBONE_HUB        = "facebook/mask2former-swin-small-coco-instance"

SEG_THR   = 0.5    # instance segmentation confidence threshold
CLS_THR   = 0.5    # classification decision threshold
DEVICE    = torch.device("cuda" if torch.cuda.is_available() else "cpu")

id2label  = {0: "ganglion_cell"}
label2id  = {"ganglion_cell": 0}
CLS_NAMES = {0: "non_diseased", 1: "diseased"}

# =============================================================================
# ── VAHADANE NORMALISER ───────────────────────────────────────────────────────
# =============================================================================
try:
    import spams
    SPAMS_AVAILABLE = True
except ImportError:
    from sklearn.decomposition import NMF
    SPAMS_AVAILABLE = False


def _rgb_to_od(img):
    return -np.log(np.clip(img.astype(np.float64), 1, 254) / 255.0)

def _od_to_rgb(od):
    return np.clip(np.exp(-od) * 255.0, 0, 255).astype(np.uint8)

def _tissue_mask(img, thresh=0.15):
    return _rgb_to_od(img).sum(axis=2) > thresh


class VahadaneNormalizer:
    def __init__(self, lambda1=0.1, max_iter=3):
        self.lambda1  = lambda1
        self.max_iter = max_iter
        self.target_stain_matrix       = None
        self.target_concentrations_max = None

    def _stain_matrix(self, img):
        mask = _tissue_mask(img)
        OD   = _rgb_to_od(img)
        OD_t = OD[mask].T
        if OD_t.shape[1] < 10:
            return np.array([[0.5626, 0.2159],
                             [0.7201, 0.8012],
                             [0.4062, 0.5581]])
        if SPAMS_AVAILABLE:
            D = spams.trainDL(
                np.asfortranarray(OD_t.astype(np.float64)),
                K=2, lambda1=self.lambda1, iter=self.max_iter,
                mode=2, modeD=0, posAlpha=True, posD=True, verbose=False)
        else:
            model = NMF(n_components=2, init="nndsvda",
                        max_iter=500, random_state=42)
            model.fit(np.maximum(OD_t.T, 0))
            D = model.components_.T
        D = D / (np.linalg.norm(D, axis=0, keepdims=True) + 1e-8)
        if D[2, 0] > D[2, 1]:
            D = D[:, [1, 0]]
        return D

    def fit(self, ref_img):
        self.target_stain_matrix = self._stain_matrix(ref_img)
        OD   = _rgb_to_od(ref_img)
        mask = _tissue_mask(ref_img)
        C, _, _, _ = np.linalg.lstsq(self.target_stain_matrix,
                                      OD[mask].T, rcond=None)
        self.target_concentrations_max = np.percentile(C, 99, axis=1)
        return self

    def normalize(self, img):
        h, w = img.shape[:2]
        W    = self._stain_matrix(img)
        OD   = _rgb_to_od(img).reshape(-1, 3)
        C, _, _, _ = np.linalg.lstsq(W, OD.T, rcond=None)
        maxC = np.percentile(C, 99, axis=1, keepdims=True)
        maxC = np.where(maxC < 1e-6, 1e-6, maxC)
        C    = C / maxC * self.target_concentrations_max[:, None]
        return _od_to_rgb((self.target_stain_matrix @ C).T.reshape(h, w, 3))


# =============================================================================
# ── MODEL ARCHITECTURE  (must match training exactly) ─────────────────────────
# =============================================================================
class SpatialAttentionPool(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.attn = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 1, kernel_size=1, bias=False),
        )
    def forward(self, x):
        w = self.attn(x).flatten(2).softmax(dim=-1)
        return (x.flatten(2) * w).sum(dim=-1)


class MaskGuidedPool(nn.Module):
    def forward(self, features, mask_labels):
        B, C, Hp, Wp = features.shape
        pooled = []
        for b in range(B):
            masks = mask_labels[b]
            if masks.numel() == 0 or masks.shape[0] == 0:
                pooled.append(features[b].mean(dim=(-2, -1)))
                continue
            union   = masks.max(dim=0).values
            union_r = F.interpolate(
                union.unsqueeze(0).unsqueeze(0).float(),
                size=(Hp, Wp), mode="bilinear", align_corners=False,
            ).squeeze(0).squeeze(0).to(features.device)
            w_sum = union_r.sum().clamp(min=1e-6)
            pooled.append(
                (features[b] * union_r.unsqueeze(0)).sum(dim=(-2, -1)) / w_sum)
        return torch.stack(pooled, dim=0)


class StrongerMask2FormerClassifier(nn.Module):
    def __init__(self, m2f_backbone, in_channels=256, num_classes=2):
        super().__init__()
        self.m2f          = m2f_backbone
        self.spatial_pool = SpatialAttentionPool(in_channels)
        self.mask_pool    = MaskGuidedPool()
        self.cls_head     = nn.Sequential(
            nn.Linear(in_channels * 2, 256),
            nn.LayerNorm(256),
            nn.GELU(),
            nn.Dropout(0.4),
            nn.Linear(256, 64),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(64, num_classes),
        )

    def forward(self, pixel_values, pixel_mask=None,
                mask_labels=None, class_labels=None):
        outputs = self.m2f(
            pixel_values=pixel_values,
            pixel_mask=pixel_mask,
            mask_labels=mask_labels,
            class_labels=class_labels,
            output_hidden_states=True,
        )
        feat      = outputs.pixel_decoder_last_hidden_state
        v_spatial = self.spatial_pool(feat)
        # mask_labels=None at inference → spatial attention fallback
        v_mask    = (self.mask_pool(feat, mask_labels)
                     if mask_labels is not None
                     else self.spatial_pool(feat))
        cls_logits = self.cls_head(torch.cat([v_spatial, v_mask], dim=1))
        seg_loss   = outputs.loss if mask_labels is not None else None
        return cls_logits, seg_loss


# =============================================================================
# ── STARTUP: Vahadane → processor → model ────────────────────────────────────
# =============================================================================
print(f"Device: {DEVICE}")

# --- Vahadane ---
vahadane = None
if os.path.exists(REFERENCE_IMAGE_PATH):
    ref_np   = np.array(PILImage.open(REFERENCE_IMAGE_PATH).convert("RGB"))
    vahadane = VahadaneNormalizer(lambda1=0.1, max_iter=3)
    vahadane.fit(ref_np)
else:
    pass

# --- Processor (same config as training) ---
processor = Mask2FormerImageProcessor.from_pretrained(
    _BACKBONE_HUB,
    do_resize=True,
    size={"shortest_edge": 512, "longest_edge": 1024},
    do_normalize=True,
)
print("✅ Processor ready.")

# --- Model ---
model = None
try:
    assert os.path.exists(SEG_CKPT), f"SEG_CKPT not found: {SEG_CKPT}"
    assert os.path.exists(CLS_CKPT), f"CLS_CKPT not found: {CLS_CKPT}"

    backbone = Mask2FormerForUniversalSegmentation.from_pretrained(
        _BACKBONE_HUB,
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True,
    )
    backbone.load_state_dict(
        torch.load(SEG_CKPT, map_location="cpu"), strict=False)
    print(f"✅ Segmentation backbone loaded: {SEG_CKPT}")

    model = StrongerMask2FormerClassifier(backbone, num_classes=2).to(DEVICE)
    model.load_state_dict(
        torch.load(CLS_CKPT, map_location=DEVICE), strict=False)
    model.eval()
    print(f"✅ Classifier loaded: {CLS_CKPT}")

except Exception as e:
    print(f"❌ Model loading failed: {e}")


# =============================================================================
# ── INFERENCE HELPER ──────────────────────────────────────────────────────────
# =============================================================================
CMAP = cm_lib.get_cmap("tab10")


def run_inference(raw_np: np.ndarray) -> dict:
    """
    Mirrors save_inference_viz() from the Kaggle script exactly.

    raw_np : H×W×3 uint8 numpy array
    Returns dict:
        prediction   : 0 | 1
        probability  : float  P(diseased)
        label        : "non_diseased" | "diseased"
        n_segments   : int
        overlay_b64  : base64-encoded PNG (two-panel: original | normalised+seg)
    """
    # 1. Vahadane normalisation
    if vahadane is not None:
        try:
            norm_np = vahadane.normalize(raw_np)
        except Exception as e:
            print(f"⚠️  Vahadane failed ({e}) — using raw image.")
            norm_np = raw_np
    else:
        norm_np = raw_np

    norm_pil = PILImage.fromarray(norm_np)
    raw_pil  = PILImage.fromarray(raw_np)

    # 2. Processor — same settings as TestDataset in Kaggle script
    inputs   = processor(images=[norm_pil], return_tensors="pt")
    pv       = inputs["pixel_values"].to(DEVICE)   # (1, 3, H', W')
    pm       = inputs["pixel_mask"].to(DEVICE)     # (1, H', W')
    _, _, proc_h, proc_w = pv.shape

    # 3. Forward pass
    with torch.no_grad():
        # Raw M2F output needed for post_process_instance_segmentation
        raw_out = model.m2f(pixel_values=pv, pixel_mask=pm)
        # Full classifier forward (mask_labels=None → spatial fallback)
        cls_logits, _ = model(pv, pm, mask_labels=None)

    probs      = torch.softmax(cls_logits, dim=1)[0].cpu()
    prob_dis   = float(probs[1].item())
    pred_class = int(prob_dis >= CLS_THR)
    label      = CLS_NAMES[pred_class]

    # 4. Instance segmentation post-processing (mirrors Kaggle script)
    res      = processor.post_process_instance_segmentation(
        raw_out,
        target_sizes=[(proc_h, proc_w)],
        threshold=SEG_THR,
    )[0]
    pred_seg = res["segmentation"].cpu().numpy()   # (proc_h, proc_w) int
    segments = res["segments_info"]

    # 5. Resize segmentation map to display (original) image size
    display_w, display_h = norm_pil.size
    seg_pil  = PILImage.fromarray(pred_seg.astype(np.int32)).resize(
        (display_w, display_h), resample=PILImage.NEAREST)
    seg_disp = np.array(seg_pil)

    # 6. Two-panel figure  (original | Vahadane-normalised + seg overlay)
    #    Matches save_inference_viz() layout exactly
    fig, axes = plt.subplots(1, 2, figsize=(14, 6), facecolor="#0a0f1e")
    for ax in axes:
        ax.set_facecolor("#0a0f1e")

    # Left — original image (no overlay)
    axes[0].imshow(raw_pil)
    axes[0].set_title("Input Image",
                    color="white", fontsize=12, pad=10)
    axes[0].axis("off")

    # Right — Vahadane-normalised + instance mask overlays
    axes[1].imshow(raw_pil)
    for si, seg in enumerate(segments):
        overlay     = np.zeros((*seg_disp.shape, 4))
        mask_region = seg_disp == seg["id"]
        colour      = CMAP(si % 10)[:3]
        overlay[mask_region] = (*colour, 0.45)
        axes[1].imshow(overlay)
    axes[1].set_title(
        f"Segmentation  |  {len(segments)} segment(s) detected",
        color="white", fontsize=12, pad=10)
    axes[1].axis("off")

    cls_colour = "#ff4444" if pred_class == 1 else "#44ff88"
    verdict    = ("⚠  DISEASED — Ganglion cells detected"
              if pred_class == 1
              else "✓  NON-DISEASED — No ganglion cells detected")
    fig.suptitle(
        f"{verdict}  |  Confidence: {prob_dis:.1%}",
        color=cls_colour, fontsize=14, fontweight="bold", y=1.01)

    plt.tight_layout()

    buf = io.BytesIO()
    plt.savefig(buf, format="png", dpi=130,
                bbox_inches="tight", facecolor="#0a0f1e")
    plt.close(fig)
    buf.seek(0)
    overlay_b64 = base64.b64encode(buf.read()).decode("utf-8")

    return {
        "prediction":  pred_class,
        "probability": round(prob_dis, 4),
        "label":       label,
        "n_segments":  len(segments),
        "overlay_b64": overlay_b64,
    }


# =============================================================================
# ── FLASK ROUTES ──────────────────────────────────────────────────────────────
# =============================================================================
app = Flask(__name__)
CORS(app)


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/predict", methods=["POST"])
def predict():
    # Validate upload
    if "image" not in request.files:
        return jsonify({"error": "No image file in request."}), 400
    file = request.files["image"]
    if not file or file.filename == "":
        return jsonify({"error": "Empty filename."}), 400
    if not file.filename.lower().endswith((".jpg", ".jpeg", ".png")):
        return jsonify({"error": "Only JPG and PNG are accepted."}), 400

    # Load raw image
    try:
        raw_np = np.array(
            PILImage.open(io.BytesIO(file.read())).convert("RGB"))
    except Exception as e:
        return jsonify({"error": f"Could not read image: {e}"}), 400

    if model is None:
        return jsonify({"error": "Model not loaded — check server logs."}), 503

    try:
        result = run_inference(raw_np)
    except Exception as e:
        import traceback; traceback.print_exc()
        return jsonify({"error": f"Inference failed: {e}"}), 500

    return jsonify(result)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)