import os
import time
import pickle
import warnings
from typing import Dict, Optional, Tuple
from huggingface_hub import hf_hub_download

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image as PILImage
import timm
import albumentations as A
from albumentations.pytorch import ToTensorV2
from pytorch_grad_cam import GradCAMPlusPlus
from pytorch_grad_cam.utils.image import show_cam_on_image
import gradio as gr

warnings.filterwarnings("ignore")

# =============================================================================
# Constants & Config
# =============================================================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IMG_SIZE = 512
NUM_CLASSES = 5
BACKBONE = "tf_efficientnetv2_m"
FV_BACKBONE = "convnext_tiny.fb_in22k"

GRADE_MAP = {0: "No DR", 1: "Mild DR", 2: "Moderate DR", 3: "Severe DR", 4: "Proliferative DR"}
GRADE_COLORS = ["#2ecc71", "#f1c40f", "#e67e22", "#e74c3c", "#8e44ad"]
SEVERITY_ICONS = ["🟢", "🟡", "🟠", "🔴", "🟣"]
CLINICAL_ACTION = [
    "No DR detected. Routine annual screening recommended.",
    "Mild NPDR. Optimise glycaemic and blood-pressure control. Follow up in 12 months.",
    "Moderate NPDR. Ophthalmology referral within 3–6 months.",
    "Severe NPDR. Urgent ophthalmology referral. Consider anti-VEGF or laser assessment.",
    "Proliferative DR. URGENT referral — high blindness risk. Same-week appointment required.",
]

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

STRONG_HEUR = {"red_dominance": 0.55, "disc_coverage": 0.25, "edge_density": 0.008}
WEAK_HEUR   = {"red_dominance": 0.30, "disc_coverage": 0.12, "edge_density": 0.006}

# =============================================================================
# Model Definitions
# =============================================================================

class GeM(nn.Module):
    def __init__(self, p: float = 3.0, eps: float = 1e-6):
        super().__init__()
        self.p = nn.Parameter(torch.tensor(p))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.adaptive_avg_pool2d(x.clamp(min=self.eps).pow(self.p), 1).pow(1.0 / self.p)


class CoralHead(nn.Module):
    def __init__(self, in_features: int, num_classes: int):
        super().__init__()
        self.linear = nn.Linear(in_features, 1, bias=False)
        self.bias = nn.Parameter(torch.zeros(num_classes - 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x) + self.bias


class DRHead(nn.Module):
    def __init__(self, in_features: int, num_classes: int = 5, dropout: float = 0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, 512), nn.BatchNorm1d(512), nn.SiLU(inplace=True), nn.Dropout(dropout),
            nn.Linear(512, 256), nn.BatchNorm1d(256), nn.SiLU(inplace=True), nn.Dropout(dropout / 2),
            nn.Linear(256, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class DRModel(nn.Module):
    def __init__(self, backbone: str = BACKBONE, num_classes: int = 5, pretrained: bool = False):
        super().__init__()
        self.backbone = timm.create_model(backbone, pretrained=pretrained, num_classes=0, global_pool="")
        in_feat = self.backbone.num_features
        self.pool = GeM(p=3.0)
        self.head_cls = DRHead(in_feat, num_classes)
        self.head_coral = CoralHead(in_feat, num_classes)

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        feat = self.pool(self.backbone(x)).flatten(1)
        return {"logits": self.head_cls(feat), "coral": self.head_coral(feat)}


class FundusValidator(nn.Module):
    def __init__(self, backbone_name: str = FV_BACKBONE):
        super().__init__()
        self.backbone = timm.create_model(backbone_name, pretrained=False, num_classes=0, global_pool="avg")
        for p in self.backbone.parameters():
            p.requires_grad = False
        with torch.no_grad():
            feat_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[1]
        self.head = nn.Sequential(nn.Linear(feat_dim, 256), nn.GELU(), nn.Dropout(0.3), nn.Linear(256, 2))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            feats = self.backbone(x)
        return self.head(feats)


# =============================================================================
# Preprocessing
# =============================================================================
_clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))


def apply_clahe_lab(rgb: np.ndarray) -> np.ndarray:
    lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
    lab[:, :, 0] = _clahe.apply(lab[:, :, 0])
    return cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)


def retinal_mask(rgb: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    _, m = cv2.threshold(gray, 15, 255, cv2.THRESH_BINARY)
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
    return cv2.morphologyEx(cv2.morphologyEx(m, cv2.MORPH_CLOSE, k), cv2.MORPH_OPEN, k)


def crop_retinal_disc(rgb: np.ndarray, pad: int = 10) -> np.ndarray:
    contours, _ = cv2.findContours(retinal_mask(rgb), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return rgb
    x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
    x, y = max(0, x - pad), max(0, y - pad)
    x2, y2 = min(rgb.shape[1], x + w + 2 * pad), min(rgb.shape[0], y + h + 2 * pad)
    return rgb[y:y2, x:x2]


def fundus_heuristics(rgb: np.ndarray) -> Dict[str, float]:
    r, g, b = rgb[..., 0].astype(np.float32), rgb[..., 1].astype(np.float32), rgb[..., 2].astype(np.float32)
    bright = (r > 15) | (g > 15) | (b > 15)
    n_bright = int(bright.sum())
    red_dom = float(((r > g) & (r > b))[bright].mean()) if n_bright > 100 else float(((r > g) & (r > b)).mean())
    disc_cov = float(retinal_mask(rgb).mean() / 255.0)
    gx = cv2.Sobel(cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY), cv2.CV_32F, 1, 0, ksize=3)
    gy = cv2.Sobel(cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY), cv2.CV_32F, 0, 1, ksize=3)
    mag = np.sqrt(gx**2 + gy**2)
    edge_den = float(mag[bright].mean() / 255.0) if n_bright > 100 else float(mag.mean() / 255.0)
    return {"red_dominance": red_dom, "disc_coverage": disc_cov, "edge_density": edge_den}


# =============================================================================
# Model Loading — downloads weights from HF repo at runtime
# =============================================================================
SPACE_REPO = "Gokul-G1/Project2-Models"


def _get_weight(filename: str) -> str:
    """Return local path if available, otherwise download from HF Space repo."""
    if os.path.exists(filename):
        return filename
    print(f"  Downloading {filename} from {SPACE_REPO}…")
    return hf_hub_download(
        repo_id=SPACE_REPO,
        filename=filename,
        repo_type="model",
    )


def load_dr_model():
    m = DRModel(BACKBONE, NUM_CLASSES, pretrained=False).to(DEVICE)
    path = _get_weight("best_model.pt")
    ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
    state = ckpt.get("model_state", ckpt)
    m.load_state_dict({k: v for k, v in state.items() if k in m.state_dict()}, strict=False)
    return m.eval()


def load_fv_model():
    m = FundusValidator(FV_BACKBONE).to(DEVICE)
    path = _get_weight("fundus_validator.pt")
    ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
    head_state = ckpt.get("head_state_dict", ckpt)
    m.head.load_state_dict(head_state, strict=False)
    return m.eval()


def load_calib():
    path = _get_weight("calibration.pkl") if os.path.exists("calibration.pkl") or True else None
    try:
        path = _get_weight("calibration.pkl")
        with open(path, "rb") as f:
            return pickle.load(f)
    except Exception:
        return None


print("Loading models…")
model   = load_dr_model()
fv_model = load_fv_model()
calib   = load_calib()
print(f"  DR model loaded | FV model loaded | calib={'yes' if calib else 'no'} | device={DEVICE}")

# Transforms
inf_tfm = A.Compose([A.Resize(IMG_SIZE, IMG_SIZE), A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2()])
fv_tfm  = A.Compose([A.Resize(224, 224),          A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2()])

# Grad-CAM setup
class _SingleLogit(nn.Module):
    def __init__(self, inner): super().__init__(); self.inner = inner
    def forward(self, x): return self.inner(x)["logits"]

_cam_model = _SingleLogit(model).to(DEVICE).eval()
try:
    _target_layers = [model.backbone.blocks[-1]]
    cam_obj = GradCAMPlusPlus(model=_cam_model, target_layers=_target_layers)
    GRADCAM_OK = True
except Exception as e:
    print(f"Grad-CAM init failed: {e}")
    GRADCAM_OK = False


# =============================================================================
# Fundus Gating
# =============================================================================
def gate_fundus(pil_img: PILImage.Image) -> Tuple[bool, float, Dict, str]:
    rgb_raw = np.array(pil_img.convert("RGB"))
    try:
        rgb = apply_clahe_lab(crop_retinal_disc(rgb_raw))
    except Exception:
        rgb = rgb_raw

    h = fundus_heuristics(rgb)
    x = fv_tfm(image=rgb)["image"].unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        p = float(F.softmax(fv_model(x), dim=1).cpu().numpy()[0][1])

    strong_ok = all(h[k] >= v for k, v in STRONG_HEUR.items())
    weak_ok   = all(h[k] >= v for k, v in WEAK_HEUR.items())

    if strong_ok:                             return True, p, h, "Accepted (strong visual signatures)."
    if p >= 0.75:                             return True, p, h, f"Accepted (validator: {p*100:.1f}%)."
    if p >= 0.50 and weak_ok:                return True, p, h, "Accepted (validator + heuristics)."
    if h["disc_coverage"] >= 0.25 and p >= 0.35: return True, p, h, "Accepted (clear retinal disc)."
    if p >= 0.90:                             return True, p, h, "Accepted (high validator confidence)."
    return False, p, h, "Not a fundus image. Please upload a colour retinal fundus photograph."


# =============================================================================
# Grad-CAM Rendering & Lesion ROIs
# =============================================================================
def render_cam(vis: np.ndarray, tensor: torch.Tensor) -> Optional[np.ndarray]:
    if not GRADCAM_OK:
        return None
    try:
        gc = cam_obj(input_tensor=tensor, targets=None)[0, :]
        gc = cv2.resize(gc, (vis.shape[1], vis.shape[0]))
        overlay = show_cam_on_image(vis.astype(np.float32) / 255.0, gc, use_rgb=True)
        return (overlay * 255).astype(np.uint8)
    except Exception:
        return None

def render_rois(vis: np.ndarray, tensor: torch.Tensor) -> Optional[np.ndarray]:
    if not GRADCAM_OK:
        return None
    try:
        gc = cam_obj(input_tensor=tensor, targets=None)[0, :]
        gc = cv2.resize(gc, (vis.shape[1], vis.shape[0]))
        _, thresh = cv2.threshold((gc * 255).astype(np.uint8), 127, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        roi_img = vis.copy()
        cv2.drawContours(roi_img, contours, -1, (255, 0, 0), 2)  # Highlight in red
        return roi_img
    except Exception:
        return None


# =============================================================================
# Inference
# =============================================================================
def predict(uploaded):
    t0 = time.time()
    if uploaded is None:
        return None, None, None, "<p style='color:var(--body-text-color);padding:20px'>Upload a fundus image to begin.</p>"

    pil_img = PILImage.fromarray(uploaded.astype(np.uint8))
    accepted, p_fundus, heur, reason = gate_fundus(pil_img)

    if not accepted:
        rd  = heur.get("red_dominance", 0)
        dc  = heur.get("disc_coverage", 0)
        ed  = heur.get("edge_density", 0)
        html = f"""
        <div class="error-box">
          <h3 style="margin:0 0 8px">⚠️ Image Rejected</h3>
          <p style="margin:0">{reason}</p>
          <div style="margin-top:12px;font-size:12px;opacity:0.8">
            <strong>Validator score:</strong> {p_fundus*100:.1f}% (need ≥75%)<br>
            <strong>Red dominance:</strong> {rd:.3f} (need ≥{WEAK_HEUR['red_dominance']:.2f})
            {'✓' if rd >= WEAK_HEUR['red_dominance'] else '✗'}<br>
            <strong>Disc coverage:</strong> {dc:.3f} (need ≥{WEAK_HEUR['disc_coverage']:.2f})
            {'✓' if dc >= WEAK_HEUR['disc_coverage'] else '✗'}<br>
            <strong>Edge density:</strong>  {ed:.4f} (need ≥{WEAK_HEUR['edge_density']:.3f})
            {'✓' if ed >= WEAK_HEUR['edge_density'] else '✗'}
          </div>
        </div>"""
        return None, None, None, html

    # Pre-process - Enforce CLAHE
    rgb = np.array(pil_img.convert("RGB"))
    try:
        rgb = crop_retinal_disc(rgb)
        rgb = apply_clahe_lab(rgb)
    except Exception:
        pass
    vis    = cv2.resize(rgb, (320, 320))
    tensor = inf_tfm(image=cv2.resize(rgb, (IMG_SIZE, IMG_SIZE)))["image"].unsqueeze(0).to(DEVICE)

    # 4-view TTA
    with torch.no_grad():
        views = [tensor,
                 torch.flip(tensor, dims=[-1]),
                 torch.flip(tensor, dims=[-2]),
                 torch.rot90(tensor, 2, dims=(-2, -1))]
        logits_all, coral_all = [], []
        for v in views:
            out = model(v)
            logits_all.append(out["logits"].cpu().numpy())
            coral_all.append(torch.sigmoid(out["coral"]).cpu().numpy())

    logits_mean = np.mean(logits_all, axis=0)  # (1, 5)
    coral_mean  = np.mean(coral_all,  axis=0)  # (1, 4)

    # Probabilities
    if calib is not None:
        try:
            probs = calib.apply(logits_mean)[0]
        except Exception:
            ex = np.exp(logits_mean - logits_mean.max())
            probs = (ex / ex.sum())[0]
    else:
        ex = np.exp(logits_mean - logits_mean.max())
        probs = (ex / ex.sum())[0]

    # Fusion: CORAL ordinal + softmax
    argmax_pred = int(probs.argmax())
    coral_pred  = int((coral_mean[0] > 0.5).sum())
    coral_conf  = float(coral_mean[0, coral_pred - 1]) if coral_pred > 0 else 1.0
    pred = coral_pred if coral_conf > float(probs[argmax_pred]) else argmax_pred
    pred = max(0, min(4, pred))
    conf = float(probs[pred])

    # Grad-CAM and ROIs
    cam_img = render_cam(vis, tensor)
    roi_img = render_rois(vis, tensor)

    proc_time = time.time() - t0

    # Build probability bars
    bar_html = ""
    for i, (lbl, col) in enumerate(zip(GRADE_MAP.values(), GRADE_COLORS)):
        w = float(probs[i]) * 100
        weight = "700" if i == pred else "400"
        bar_html += f"""
        <div style="margin-bottom:6px">
          <div style="display:flex;justify-content:space-between;font-size:13px;font-weight:{weight}">
            <span>{SEVERITY_ICONS[i]} {lbl}</span><span>{w:.1f}%</span>
          </div>
          <div style="height:6px;background:var(--border-color-primary);border-radius:3px;overflow:hidden">
            <div style="width:{w}%;background:{col};height:100%"></div>
          </div>
        </div>"""

    html = f"""
    <div class="report-card">
      <!-- Grade header -->
      <div class="report-header">
        <span style="font-size:52px">{SEVERITY_ICONS[pred]}</span>
        <div>
          <div style="font-size:22px;font-weight:700;color:{GRADE_COLORS[pred]}">{GRADE_MAP[pred]}</div>
          <div style="font-size:13px;opacity:0.7">Grade {pred} / 4 &nbsp;·&nbsp; ICDR Scale</div>
        </div>
      </div>

      <!-- Confidence bars -->
      <div class="report-bars">
        <div style="font-size:13px;font-weight:600;margin-bottom:8px">Grade Probabilities</div>
        {bar_html}
      </div>

      <!-- Clinical action -->
      <div class="action-box" style="border-left-color: {GRADE_COLORS[pred]}">
        <div style="font-size:12px;font-weight:600;margin-bottom:4px;opacity:0.8">CLINICAL RECOMMENDATION</div>
        <div style="font-size:14px">{CLINICAL_ACTION[pred]}</div>
      </div>

      <!-- Metrics footer -->
      <div style="display:flex;justify-content:space-between;font-size:11px;opacity:0.6;
                  border-top:1px solid var(--border-color-primary);padding-top:10px">
        <span>⏱ Processing: <strong>{proc_time:.2f}s</strong></span>
        <span>🔍 Validator: <strong>{p_fundus*100:.1f}%</strong></span>
        <span>🎯 Confidence: <strong>{conf*100:.1f}%</strong></span>
      </div>
    </div>"""

    return vis, cam_img, roi_img, html


def safe_predict(uploaded):
    """Wrapper to catch network/processing errors and prevent crashes."""
    try:
        return predict(uploaded)
    except Exception as e:
        err_msg = f"<div class='error-box'><h3>⚠️ Processing Error</h3><p>{str(e)}</p><p>Please try uploading the image again.</p></div>"
        return None, None, None, err_msg


# =============================================================================
# Gradio UI — Google Premium Material Design 3 (Dark & Light)
# =============================================================================
premium_theme = gr.themes.Default(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
).set(
    body_background_fill="var(--background-fill-primary)",
    body_background_fill_dark="var(--background-fill-primary)",
    block_background_fill="var(--block-background-fill)",
    block_border_width="1px",
    block_shadow="0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06)",
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    button_primary_text_color="white",
)

css = """
.report-card { background: var(--block-background-fill); padding: 20px; border-radius: 14px; box-shadow: 0 4px 16px rgba(0,0,0,.08); font-family: 'Inter', sans-serif; line-height: 1.5; color: var(--body-text-color); }
.report-header { display: flex; align-items: center; gap: 14px; margin-bottom: 16px; }
.report-bars { background: var(--background-fill-secondary); border-radius: 8px; padding: 12px; margin-bottom: 14px; }
.action-box { padding: 12px; background: var(--background-fill-secondary); border-left: 4px solid var(--primary-500); border-radius: 6px; margin-bottom: 14px; }
.hero-banner { text-align: center; padding: 28px 20px 22px; background: linear-gradient(135deg, #1A73E8 0%, #0D47A1 100%); color: #fff; border-radius: 16px; margin-bottom: 20px; }
.error-box { padding: 20px; border-radius: 10px; background: rgba(220, 38, 38, 0.1); color: #ef4444; border: 1px solid rgba(220, 38, 38, 0.2); font-family: 'Inter', sans-serif; }
"""

with gr.Blocks(theme=premium_theme, css=css, title="DR Grading AI — Clinical Decision Support") as demo:

    gr.HTML("""
    <div class="hero-banner">
      <div style="font-size:2.4em;font-weight:800;letter-spacing:-0.5px">
        👁 Diabetic Retinopathy Grading
      </div>
      <div style="margin-top:8px;opacity:.88;font-size:1.05em">
        EfficientNetV2-M · CLAHE · CORAL Ordinal · Grad-CAM++ · Lesion ROIs
      </div>
      <div style="margin-top:6px;opacity:.7;font-size:.85em">
        For research & educational use only — not a medical device
      </div>
    </div>
    """)

    with gr.Row(equal_height=True):
        # ── LEFT PANEL ──────────────────────────────────────────────────────
        with gr.Column(scale=1, min_width=340):
            inp = gr.Image(
                label="Upload Retinal Fundus Photograph",
                type="numpy",
                height=380,
                sources=["upload", "clipboard"],
            )
            with gr.Row():
                btn_clear   = gr.Button("🗑 Clear",   variant="secondary", size="sm")
                btn_analyze = gr.Button("🔬 Analyze", variant="primary",   size="lg")

            gr.Markdown("""
            **How to use**
            1. Upload a colour fundus photograph (JPEG / PNG)
            2. Click **Analyze** — the AI validates it first
            3. Review the grade, heatmap, and recommendation

            > *Grades 0–4 follow the International Clinical DR (ICDR) severity scale.*
            """)

        # ── RIGHT PANEL ─────────────────────────────────────────────────────
        with gr.Column(scale=1, min_width=340):
            with gr.Row():
                out_vis = gr.Image(label="Enhanced Fundus (CLAHE)", interactive=False, height=200)
                out_cam = gr.Image(label="Pathology Heatmap (Grad-CAM++)", interactive=False, height=200)
            out_roi = gr.Image(label="Detected Lesion ROIs", interactive=False, height=200)
            out_report = gr.HTML()

    # ── Events ──────────────────────────────────────────────────────────────
    btn_analyze.click(
        fn=safe_predict,
        inputs=inp,
        outputs=[out_vis, out_cam, out_roi, out_report],
    )
    btn_clear.click(
        fn=lambda: (None, None, None, None, ""),
        inputs=[],
        outputs=[inp, out_vis, out_cam, out_roi, out_report],
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=5)
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)