import cv2
import torch
import numpy as np
import torch.nn.functional as F
from torch import nn
from transformers import AutoImageProcessor, Swinv2ForImageClassification, SegformerForSemanticSegmentation
import streamlit as st
from PIL import Image
import io
import zipfile

# --- GlaucomaModel Class ---
class GlaucomaModel(object):
    def __init__(self, 
                 cls_model_path="pamixsun/swinv2_tiny_for_glaucoma_classification", 
                 seg_model_path='pamixsun/segformer_for_optic_disc_cup_segmentation',
                 device=torch.device('cpu')):
        self.device = device
        # Classification model for glaucoma
        self.cls_extractor = AutoImageProcessor.from_pretrained(cls_model_path)
        self.cls_model = Swinv2ForImageClassification.from_pretrained(cls_model_path).to(device).eval()
        # Segmentation model for optic disc and cup
        self.seg_extractor = AutoImageProcessor.from_pretrained(seg_model_path)
        self.seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_path).to(device).eval()
        # Mapping for class labels
        self.cls_id2label = self.cls_model.config.id2label

    def glaucoma_pred(self, image):
        inputs = self.cls_extractor(images=image.copy(), return_tensors="pt")
        with torch.no_grad():
            inputs.to(self.device)
            outputs = self.cls_model(**inputs).logits
            probs = F.softmax(outputs, dim=-1)
            disease_idx = probs.cpu()[0, :].numpy().argmax()
            confidence = probs.cpu()[0, disease_idx].item() * 100
        return disease_idx, confidence

    def optic_disc_cup_pred(self, image):
        inputs = self.seg_extractor(images=image.copy(), return_tensors="pt")
        with torch.no_grad():
            inputs.to(self.device)
            outputs = self.seg_model(**inputs)
        logits = outputs.logits.cpu()
        upsampled_logits = nn.functional.interpolate(
            logits, size=image.shape[:2], mode="bilinear", align_corners=False
        )
        seg_probs = F.softmax(upsampled_logits, dim=1)
        pred_disc_cup = upsampled_logits.argmax(dim=1)[0]
        
        # Calculate segmentation confidence based on probability distribution
        # For each pixel classified as cup/disc, check how confident the model is
        cup_mask = pred_disc_cup == 2
        disc_mask = pred_disc_cup == 1
        
        # Get confidence only for pixels predicted as cup/disc
        cup_confidence = seg_probs[0, 2, cup_mask].mean().item() * 100 if cup_mask.any() else 0
        disc_confidence = seg_probs[0, 1, disc_mask].mean().item() * 100 if disc_mask.any() else 0
        
        return pred_disc_cup.numpy().astype(np.uint8), cup_confidence, disc_confidence

    def process(self, image):
        disease_idx, cls_confidence = self.glaucoma_pred(image)
        disc_cup, cup_confidence, disc_confidence = self.optic_disc_cup_pred(image)

        try:
            vcdr = simple_vcdr(disc_cup)
        except:
            vcdr = np.nan

        mask = (disc_cup > 0).astype(np.uint8)
        x, y, w, h = cv2.boundingRect(mask)
        padding = max(50, int(0.2 * max(w, h)))
        x = max(x - padding, 0)
        y = max(y - padding, 0)
        w = min(w + 2 * padding, image.shape[1] - x)
        h = min(h + 2 * padding, image.shape[0] - y)

        cropped_image = image[y:y+h, x:x+w] if w >= 50 and h >= 50 else image.copy()
        _, disc_cup_image = add_mask(image, disc_cup, [1, 2], [[0, 255, 0], [255, 0, 0]], 0.2)

        return disease_idx, disc_cup_image, vcdr, cls_confidence, cup_confidence, disc_confidence, cropped_image

# --- Utility Functions ---
def simple_vcdr(mask):
    disc_area = np.sum(mask == 1)
    cup_area = np.sum(mask == 2)
    if disc_area == 0:
        return np.nan
    vcdr = cup_area / disc_area
    return vcdr

def add_mask(image, mask, classes, colors, alpha=0.5):
    overlay = image.copy()
    for class_id, color in zip(classes, colors):
        overlay[mask == class_id] = color
    output = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
    return output, overlay

def get_confidence_level(confidence):
    if confidence >= 90:
        return "Very High"
    elif confidence >= 75:
        return "High"
    elif confidence >= 60:
        return "Moderate"
    elif confidence >= 45:
        return "Low"
    else:
        return "Very Low"

# --- Streamlit Interface ---
def main():
    st.set_page_config(layout="wide", page_title="Glaucoma Screening Tool")
    
    # Header with better styling
    st.markdown("""
        <h1 style='text-align: center;'>Glaucoma Screening from Retinal Fundus Images</h1>
        <p style='text-align: center; color: gray;'>Upload retinal images for automated glaucoma detection and optic disc/cup segmentation</p>
        """, unsafe_allow_html=True)
    
    # Sidebar with better organization
    with st.sidebar:
        st.markdown("### Upload Settings")
        uploaded_files = st.file_uploader("Upload Retinal Images", 
            type=['png', 'jpeg', 'jpg'], 
            accept_multiple_files=True,
            help="Support multiple images in PNG, JPEG formats")
        
        st.markdown("### Analysis Settings")
        st.info("📊 Set confidence threshold to filter results")
        confidence_threshold = st.slider(
            "Classification Confidence Threshold (%)", 
            0, 100, 70,
            help="Images with confidence above this threshold will be marked as reliable predictions")
    
    if uploaded_files:
        for uploaded_file in uploaded_files:
            image = Image.open(uploaded_file).convert('RGB')
            image_np = np.array(image).astype(np.uint8)
            
            with st.spinner(f'🔄 Processing {uploaded_file.name}...'):
                model = GlaucomaModel(device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
                disease_idx, disc_cup_image, vcdr, cls_conf, cup_conf, disc_conf, cropped_image = model.process(image_np)
                
                # Create expandable section for each image
                with st.expander(f"📊 Analysis Results: {uploaded_file.name}", expanded=True):
                    # Image display section
                    cols = st.columns(3)
                    cols[0].image(image_np, caption="Original Image", use_column_width=True)
                    cols[1].image(disc_cup_image, caption="Segmentation Overlay", use_column_width=True)
                    cols[2].image(cropped_image, caption="Region of Interest", use_column_width=True)
                    
                    # Metrics section with clear separation
                    st.markdown("---")
                    metric_cols = st.columns(3)
                    
                    # Classification Results
                    with metric_cols[0]:
                        st.markdown("### 🔍 Classification")
                        diagnosis = model.cls_id2label[disease_idx]
                        is_confident = cls_conf >= confidence_threshold
                        
                        # Color-coded diagnosis
                        if diagnosis == "Glaucoma":
                            st.markdown(f"<div style='padding: 10px; background-color: #ffebee; border-radius: 5px;'>"
                                      f"<h4 style='color: #c62828;'>Diagnosis: {diagnosis}</h4></div>", 
                                      unsafe_allow_html=True)
                        else:
                            st.markdown(f"<div style='padding: 10px; background-color: #e8f5e9; border-radius: 5px;'>"
                                      f"<h4 style='color: #2e7d32;'>Diagnosis: {diagnosis}</h4></div>", 
                                      unsafe_allow_html=True)
                        
                        st.metric("Classification Confidence", f"{cls_conf:.1f}%")
                        if not is_confident:
                            st.warning("⚠️ Below confidence threshold")
                    
                    # Segmentation Results
                    with metric_cols[1]:
                        st.markdown("### 🎯 Segmentation Quality")
                        st.metric("Optic Cup Confidence", f"{cup_conf:.1f}%")
                        st.metric("Optic Disc Confidence", f"{disc_conf:.1f}%")
                        
                        # Confidence level explanation
                        cup_level = get_confidence_level(cup_conf)
                        disc_level = get_confidence_level(disc_conf)
                        st.info(f"Cup Detection: {cup_level}\nDisc Detection: {disc_level}")
                    
                    # Clinical Metrics
                    with metric_cols[2]:
                        st.markdown("### 📏 Clinical Metrics")
                        st.metric("Cup-to-Disc Ratio (CDR)", f"{vcdr:.3f}")
                        
                        # CDR interpretation
                        if vcdr > 0.7:
                            st.warning("⚠️ Elevated CDR (>0.7)")
                        elif vcdr > 0.5:
                            st.info("ℹ️ Borderline CDR (0.5-0.7)")
                        else:
                            st.success("✅ Normal CDR (<0.5)")
        
        # Download section
        if download_confident_images:
            st.sidebar.markdown("---")
            st.sidebar.markdown("### Download Results")
            with zipfile.ZipFile("confident_cropped_images.zip", "w") as zf:
                for cropped_image, name in download_confident_images:
                    img_buffer = io.BytesIO()
                    Image.fromarray(cropped_image).save(img_buffer, format="PNG")
                    zf.writestr(f"{name}_cropped.png", img_buffer.getvalue())
            
            st.sidebar.download_button(
                label="📥 Download Analysis Results",
                data=open("confident_cropped_images.zip", "rb"),
                file_name="glaucoma_analysis_results.zip",
                mime="application/zip",
                help="Download cropped images and analysis results"
            )
    else:
        # Welcome message when no files are uploaded
        st.markdown("""
            <div style='text-align: center; padding: 50px;'>
                <h3>👋 Welcome to the Glaucoma Screening Tool</h3>
                <p>Upload retinal fundus images using the sidebar to begin analysis</p>
            </div>
            """, unsafe_allow_html=True)

if __name__ == '__main__':
    main()