import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
import gradio as gr
import numpy as np
import cv2
from PIL import Image

# ==============================================================================
# 1. CORE CONFIGURATION & MODEL REGISTRY
# ==============================================================================
IMG_SIZE = 128
LABELS = [
    "Adenocarcinoma (Lung Cancer Type)", 
    "Large Cell Carcinoma (Lung Cancer Type)", 
    "Normal Lung (No Cancer Detected)", 
    "Squamous Cell Carcinoma (Lung Cancer Type)"
]

class MedicalCLAHEEqualization(object):
    """
    Applies standard contrast adjustments to neutralize image lighting,
    ensuring the model looks only at scan shapes.
    """
    def __call__(self, img):
        img_np = np.array(img)
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        equalized = clahe.apply(gray)
        return Image.fromarray(cv2.cvtColor(equalized, cv2.COLOR_GRAY2RGB))

# Standard image sizing and color adjustments matching the model's training
eval_transforms = transforms.Compose([
    MedicalCLAHEEqualization(),
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

device = torch.device("cpu")

def load_bioset_cccm():
    """
    Loads the underlying architecture and fills it with our trained weights.
    """
    model = models.efficientnet_b0(weights=None)
    in_features = model.classifier[1].in_features
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.4, inplace=True),
        nn.Linear(in_features, len(LABELS))
    )
    
    # Safely load the saved training weights file
    if os.path.exists('chestnet_efficientnet_weights.pth'):
        model.load_state_dict(torch.load('chestnet_efficientnet_weights.pth', map_location=device))
    elif os.path.exists('chestnet_efficientnet_full.pth'):
        checkpoint = torch.load('chestnet_efficientnet_full.pth', map_location=device, weights_only=False)
        if hasattr(checkpoint, 'state_dict'):
            model.load_state_dict(checkpoint.state_dict())
        else:
            model = checkpoint
            
    model.eval()
    return model

# Initialize the model
model = load_bioset_cccm()

# ==============================================================================
# 2. IMAGE PREDICTION PIPELINE
# ==============================================================================
def predict_chest_slice(input_image):
    if input_image is None:
        return "### Please upload a scan to begin...", {}
        
    # Convert image array into a standard image format
    pil_img = Image.fromarray(input_image.astype('uint8'), 'RGB')
    tensor_input = eval_transforms(pil_img).unsqueeze(0).to(device)
    
    with torch.no_grad():
        logits = model(tensor_input)
        probabilities = F.softmax(logits, dim=1).squeeze(0).numpy()
        
    # Sort results from highest probability to lowest
    indexed_results = [
        {"label": LABELS[i], "prob": float(probabilities[i])} 
        for i in range(len(LABELS))
    ]
    indexed_results.sort(key=lambda x: x["prob"], reverse=True)
    
    # Create the top result header text
    top_winner = indexed_results[0]
    output_summary = f"### Main Finding: **{top_winner['label']}** ({top_winner['prob']*100:.1f}% Match)"
    
    # FIXED: Return all 4 conditions to show the full picture to the user
    gradio_label_output = {item["label"]: item["prob"] for item in indexed_results}
    
    return output_summary, gradio_label_output

# ==============================================================================
# 3. USER INTERFACE LAYOUT
# ==============================================================================
with gr.Blocks(title="Bioset CCCM Lung Analysis") as demo:
    
    # Title Block
    gr.Markdown(
        """
        # Bioset CCCM
        ### Part of the **Bioset Model Collection** by **Infinitode** • Artificial Intelligence + Biology
        ---
        """
    )
    
    # Simplified Medical Disclaimer
    gr.Markdown(
        """
        > ⚠️ **Important Safety Notice**
        > The Chest Cancer Classification Model (CCCM) is an experimental prototype built purely for educational and scientific research. It is **not a certified medical tool**, does not have FDA clearance, and must never be used to diagnose medical conditions or replace the expert advice of a real doctor or radiologist.
        """
    )
    
    # Main Interactive Split Screen
    with gr.Row():
        # Left Side: Image Upload Controls
        with gr.Column(scale=5):
            gr.Markdown("#### 📥 Step 1: Upload Lung Scan")
            
            input_image = gr.Image(
                label="Chest CT Scan Slice",
                sources=["upload", "clipboard", "webcam"],
                type="numpy"
            )
            
            with gr.Row():
                clear_btn = gr.Button("Clear Image", variant="secondary")
                submit_btn = gr.Button("Analyze Scan", variant="primary")
                
        # Right Side: Results Presentation
        with gr.Column(scale=7):
            gr.Markdown("#### 📊 Step 2: Analysis Results")
            
            output_text = gr.Markdown("### Awaiting image upload...")
            
            output_labels = gr.Label(
                num_top_classes=4, # FIXED: Set to 4 to ensure all conditions are shown every time
                label="Match Likelihood for Each Condition"
            )
            
            gr.Markdown(
                """
                * **How it works:** This app automatically smooths out lighting and shadows on your uploaded scan. This makes it easier for the AI to focus entirely on the physical shapes and textures in the lung tissue rather than the brightness settings of the scanner machine.
                """
            )

    # Performance Stats Accordion
    with gr.Accordion("📋 Model Profile & Performance Summary", open=True):
        gr.Markdown(
            """
            ### Accuracy & Performance Testing Data (CCCM v1.4)
            These scores show how well this model performed during final testing on our verified lung scan dataset:
            
            | Test Category | Score | What it means |
            | :--- | :--- | :--- |
            | **Overall Reliability** | **93.06%** | The average accuracy rate during final verification tests. |
            | **Highest Match Certainty** | **98.7%** | The maximum confidence score achieved on clear positive scans. |
            | **App File Size** | **16.6 MB** | A highly compressed, efficient size designed to load and run quickly. |
            """
        )
        
    # Connect the buttons to our prediction function
    submit_btn.click(
        fn=predict_chest_slice,
        inputs=[input_image],
        outputs=[output_text, output_labels]
    )
    
    clear_btn.click(
        fn=lambda: ("### Awaiting image upload...", {}),
        inputs=None,
        outputs=[output_text, output_labels]
    )

if __name__ == "__main__":
    demo.launch()