import numpy as np
import cv2
import tensorflow as tf
import streamlit as st
import matplotlib.pyplot as plt
from lime import lime_image
from skimage.segmentation import mark_boundaries
from keras.layers import BatchNormalization, DepthwiseConv2D, TFSMLayer
import os
from io import BytesIO
import base64

# FIXED CSS - Removed animations and stabilized background
st.markdown(
    """
    <style>
    /* Main App Styling - FIXED: Stable background */
    .stApp {
        background: #f8fafc !important;
        /* Removed gradient and animations */
    }
    
    /* Header Styling - FIXED: No animations */
    .main-header {
        background: #1e40af;
        color: white;
        padding: 1.5rem;
        border-radius: 12px;
        margin-bottom: 2rem;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        /* Removed gradient animations */
    }
    
    /* FIXED: Stable flex container */
    .flex-row {
        display: flex;
        gap: 2rem;
        align-items: stretch;
        margin-top: 1rem;
    }
    .flex-row > div {
        flex: 1;
        display: flex;
        flex-direction: column;
    }
    
    /* FIXED: Stable medical cards */
    .medical-card {
        background: white;
        padding: 1.5rem;
        border-radius: 12px;
        border-left: 4px solid #3b82f6;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        flex-grow: 1;
        border: 1px solid #e2e8f0;
        /* Removed gradient and animations */
    }
    
    .medical-card h3 {
        margin-top: 0;
        border-bottom: 2px solid #e2e8f0;
        padding-bottom: 0.5rem;
    }
    
    /* FIXED: Removed conflicting prediction styles */
    .prediction-card {
        background: white;
        padding: 2rem;
        border-radius: 16px;
        margin: 2rem 0;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        border: 1px solid #e2e8f0;
        /* Removed all animations and gradients */
    }
    
    /* FIXED: Stable processing container */
    .processing-container {
        background: white;
        border-radius: 16px;
        padding: 2rem;
        margin: 2rem 0;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        border: 1px solid #e2e8f0;
        /* Removed animations */
    }
    
    /* FIXED: Stable LIME container */
    .lime-container {
        background: white;
        border-radius: 16px;
        padding: 2rem;
        margin: 2rem 0;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        border: 1px solid #e2e8f0;
        /* Removed animations */
    }
    
    /* FIXED: Stable button styling */
    .stDownloadButton > button {
        background: #3b82f6;
        color: white;
        border: none;
        border-radius: 12px;
        padding: 0.75rem 1.5rem;
        font-weight: 600;
        font-size: 1rem;
        width: 100%;
        margin-top: 1rem;
        /* Removed all hover animations and transitions */
    }
    
    /* FIXED: Stable upload instructions */
    .upload-instructions {
        background: #f0f9ff;
        border: 2px solid #3b82f6;
        border-radius: 12px;
        padding: 3rem;
        text-align: center;
        margin: 2rem 0;
        /* Removed gradient */
    }
    
    .upload-instructions h3 {
        color: #1e40af;
        margin-bottom: 1rem;
        font-size: 1.5rem;
    }
    
    .upload-instructions p {
        color: #64748b;
        margin-bottom: 1rem;
    }
    
    /* FIXED: Stable feature grid */
    .feature-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(280px, 1fr));
        gap: 2rem;
        margin: 2rem 0;
    }
    
    .feature-card {
        background: white;
        border-radius: 12px;
        padding: 1.5rem;
        text-align: center;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        border: 1px solid #e2e8f0;
        /* Removed hover animations */
    }
    
    .feature-icon {
        font-size: 2.5rem;
        margin-bottom: 1rem;
    }
    
    .feature-title {
        font-size: 1.1rem;
        font-weight: 600;
        color: #1e40af;
        margin-bottom: 0.5rem;
    }
    
    .feature-description {
        color: #6b7280;
        font-size: 0.9rem;
    }
    
    /* FIXED: Stable confidence bar */
    .confidence-bar {
        background: #e2e8f0;
        border-radius: 10px;
        overflow: hidden;
        margin: 1rem 0;
        height: 12px;
        position: relative;
    }
    
    .confidence-fill {
        height: 100%;
        border-radius: 10px;
        position: relative;
        /* Removed transitions */
    }
    
    .confidence-fill.high {
        background: #16a34a;
    }
    
    .confidence-fill.medium {
        background: #f59e0b;
    }
    
    .confidence-fill.low {
        background: #ef4444;
    }
    
    /* FIXED: Stable sidebar */
    .sidebar-content {
        background: white;
        border-radius: 12px;
        padding: 1rem;
        margin: 1rem 0;
        border: 1px solid #e2e8f0;
    }
    
    /* FIXED: Stable image container */
    .image-container {
        background: white;
        border-radius: 12px;
        padding: 1rem;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
        border: 1px solid #e2e8f0;
    }
    
    /* FIXED: Stable metrics */
    .metrics-row {
        display: flex;
        justify-content: space-around;
        margin: 1.5rem 0;
        padding: 1rem;
        background: #f8fafc;
        border-radius: 8px;
    }
    
    .metric-item {
        text-align: center;
        flex: 1;
    }
    
    .metric-value {
        font-size: 1.5rem;
        font-weight: 700;
        color: #1e40af;
        margin-bottom: 0.25rem;
    }
    
    .metric-label {
        font-size: 0.875rem;
        color: #6b7280;
        text-transform: uppercase;
        letter-spacing: 0.1em;
    }
    
    /* FIXED: Stable typography */
    .prediction-title {
        font-size: 1.75rem;
        font-weight: 700;
        color: #1e40af;
        margin-bottom: 1rem;
        text-align: center;
    }
    
    .confidence-text {
        font-size: 1.2rem;
        font-weight: 600;
        color: #374151;
        text-align: center;
        margin-top: 0.5rem;
    }
    
    /* FIXED: Stable processing steps */
    .processing-step {
        background: white;
        border-radius: 8px;
        padding: 1rem;
        margin: 0.5rem 0;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
        border-left: 3px solid #3b82f6;
    }
    
    /* FIXED: Stable message styling */
    .stSuccess {
        background: #f0fdf4;
        border-left: 4px solid #22c55e;
        border-radius: 8px;
    }
    
    .stWarning {
        background: #fffbeb;
        border-left: 4px solid #f59e0b;
        border-radius: 8px;
    }
    
    .stError {
        background: #fef2f2;
        border-left: 4px solid #ef4444;
        border-radius: 8px;
    }
    
    /* FIXED: Remove any potential animation triggers */
    * {
        transition: none !important;
        animation: none !important;
        transform: none !important;
    }
    
    /* FIXED: Ensure stable viewport */
    .block-container {
        padding-top: 1rem;
        padding-bottom: 1rem;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# --- Fix deserialization issues ---
original_bn = BatchNormalization.from_config
BatchNormalization.from_config = classmethod(
    lambda cls, config, *a, **k: original_bn(
        config if not isinstance(config.get("axis"), list) else {**config, "axis": config["axis"][0]}, *a, **k
    )
)
original_dw = DepthwiseConv2D.from_config
DepthwiseConv2D.from_config = classmethod(
    lambda cls, config, *a, **k: original_dw({k: v for k, v in config.items() if k != "groups"}, *a, **k)
)

# --- FIXED: Simplified background function (no dynamic changes) ---
def set_background():
    """Set a stable, consistent background"""
    st.markdown("""
        <style>
        .stApp {
            background: #f8fafc !important;
        }
        [data-testid="stSidebar"] > div:first-child {
            background: #e0f7fa !important;  /* Light cyan */
            border-radius: 0 15px 15px 0;
            padding: 1rem;
        }
        </style>
    """, unsafe_allow_html=True)


# Apply stable background
set_background()

# --- Constants ---
IMG_SIZE = (224, 224)
CLASS_NAMES = [
    'Normal', 'Diabetic Retinopathy', 'Glaucoma', 'Cataract',
    'Age-related Macular Degeneration (AMD)', 'Hypertension', 'Myopia', 'Others'
]
LIME_EXPLAINER = lime_image.LimeImageExplainer()

# --- Load Model ---
@st.cache_resource
def load_model():
    model_path = "Model"
    if not os.path.exists(model_path):
        st.error(f"🚨 Model folder '{model_path}' not found.")
        st.stop()
    try:
        model = tf.keras.Sequential([TFSMLayer(model_path, call_endpoint="serving_default")])
        return model
    except Exception as e:
        st.error(f"🚨 Error loading model: {e}")
        st.stop()

# --- Prediction ---
def predict(images, model):
    images = np.array(images)
    preds = model.predict(images, verbose=0)
    if isinstance(preds, dict):
        for v in preds.values():
            if isinstance(v, (np.ndarray, list)):
                return np.array(v)
        return np.array(list(preds.values())[0])
    else:
        return preds

# --- FIXED: Stable preprocessing with consistent styling ---
def preprocess_with_steps(img):
    h, w = img.shape[:2]
    center, radius = (w // 2, h // 2), min(w, h) // 2
    Y, X = np.ogrid[:h, :w]
    dist = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)
    mask = (dist <= radius).astype(np.uint8)

    circ = img.copy()
    white_bg = np.ones_like(circ, dtype=np.uint8) * 255
    circ = np.where(mask[:, :, np.newaxis] == 1, circ, white_bg)

    lab = cv2.cvtColor(circ, cv2.COLOR_RGB2LAB)
    cl = cv2.createCLAHE(clipLimit=2.0).apply(lab[:, :, 0])
    merged = cv2.merge((cl, lab[:, :, 1], lab[:, :, 2]))
    clahe_img = cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)

    sharp = cv2.addWeighted(clahe_img, 4, cv2.GaussianBlur(clahe_img, (0, 0), 10), -4, 128)
    resized = cv2.resize(sharp, IMG_SIZE) / 255.0

    # FIXED: Stable visualization with consistent styling
    fig, axs = plt.subplots(1, 4, figsize=(16, 4))
    fig.patch.set_facecolor('white')  # Fixed to white background
    
    for ax, image, title in zip(
        axs, [img, circ, clahe_img, resized],
        ["Original", "Circular Crop", "CLAHE", "Sharpen + Resize"]
    ):
        ax.imshow(image)
        ax.set_title(title, fontsize=14, fontweight='bold', color='#1e40af')
        ax.axis("off")
    
    plt.tight_layout()
    st.pyplot(fig)
    plt.close(fig)
    return resized

# FIXED: Stable explanation text (no dynamic styling)
explanation_text = {
    'Normal': """
    <div class="medical-card">
        <h3 style="color:#059669; font-weight:bold;">✅ Normal Retina</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>🟢 <strong>Clear retinal structure</strong> - No pathological lesions detected</li>
            <li>🩺 <strong>Healthy blood vessels</strong> - Normal caliber and branching pattern</li>
            <li>👁 <strong>Intact optic disc & macula</strong> - Proper anatomical structure</li>
            <li>✅ <strong>No disease indicators</strong> - Excellent retinal health</li>
        </ul>
    </div>
    """,

    'Diabetic Retinopathy': """
    <div class="medical-card">
        <h3 style="color:#dc2626; font-weight:bold;">⚠️ Diabetic Retinopathy</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>🔴 <strong>Microhemorrhages</strong> - Red spots indicating vessel damage</li>
            <li>🩸 <strong>Vascular leakage</strong> - Fluid accumulation in retinal tissue</li>
            <li>👁 <strong>Macular involvement</strong> - Possible diabetic macular edema</li>
            <li>🔬 <strong>Requires monitoring</strong> - Regular ophthalmologic follow-up needed</li>
        </ul>
    </div>
    """,

    'Glaucoma': """
    <div class="medical-card">
        <h3 style="color:#7c3aed; font-weight:bold;">👁 Glaucoma</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>🔴 <strong>Optic nerve damage</strong> - Thinning of nerve fiber layer</li>
            <li>⚪ <strong>Increased cup-to-disc ratio</strong> - Optic disc cupping</li>
            <li>📉 <strong>Visual field risk</strong> - Potential peripheral vision loss</li>
            <li>💊 <strong>Pressure management</strong> - IOP control essential</li>
        </ul>
    </div>
    """,

    'Cataract': """
    <div class="medical-card">
        <h3 style="color:#f59e0b; font-weight:bold;">🌫️ Cataract</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>☁️ <strong>Lens opacity</strong> - Clouding affecting image clarity</li>
            <li>🔍 <strong>Reduced contrast</strong> - Decreased retinal detail visibility</li>
            <li>👁 <strong>Fundus visualization</strong> - Limited view of posterior structures</li>
            <li>🏥 <strong>Surgical consideration</strong> - May benefit from cataract extraction</li>
        </ul>
    </div>
    """,

    'Age-related Macular Degeneration (AMD)': """
    <div class="medical-card">
        <h3 style="color:#be185d; font-weight:bold;">🧓 Age-related Macular Degeneration</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>🟡 <strong>Drusen deposits</strong> - Yellow spots near macular region</li>
            <li>👁 <strong>Central vision impact</strong> - Macula-specific changes</li>
            <li>📈 <strong>Progressive condition</strong> - Age-related degenerative process</li>
            <li>🔬 <strong>Monitoring required</strong> - Regular assessment for progression</li>
        </ul>
    </div>
    """,

    'Hypertension': """
    <div class="medical-card">
        <h3 style="color:#dc2626; font-weight:bold;">🩸 Hypertensive Retinopathy</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>⭐ <strong>Cotton wool spots</strong> - Nerve fiber layer infarcts</li>
            <li>🔴 <strong>Flame hemorrhages</strong> - Superficial retinal bleeding</li>
            <li>🩸 <strong>Arteriovenous nicking</strong> - Vessel caliber changes</li>
            <li>💊 <strong>BP management</strong> - Systemic hypertension control needed</li>
        </ul>
    </div>
    """,

    'Myopia': """
    <div class="medical-card">
        <h3 style="color:#2563eb; font-weight:bold;">👓 Myopic Changes</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>🔵 <strong>Axial elongation signs</strong> - Elongated eyeball morphology</li>
            <li>⚪ <strong>Peripapillary atrophy</strong> - Tissue thinning around optic disc</li>
            <li>📐 <strong>Disc tilting</strong> - Oblique optic disc orientation</li>
            <li>👁 <strong>Refractive changes</strong> - Associated with high myopia</li>
        </ul>
    </div>
    """,

    'Others': """
    <div class="medical-card">
        <h3 style="color:#6b7280; font-weight:bold;">🔍 Unclassified Findings</h3>
        <ul style="font-size:16px; line-height:1.8; color:#374151; margin:0;">
            <li>❓ <strong>Atypical presentation</strong> - Unusual retinal patterns</li>
            <li>🔬 <strong>Further evaluation</strong> - Additional testing recommended</li>
            <li>🩺 <strong>Specialist referral</strong> - Ophthalmologist consultation advised</li>
            <li>📋 <strong>Comprehensive exam</strong> - Complete ocular assessment needed</li>
        </ul>
    </div>
    """
}

# --- FIXED: Stable LIME Display ---
def show_lime(img, model, pred_idx, pred_label, all_probs):
    with st.spinner("🔬 Generating LIME explanation..."):
        explanation = LIME_EXPLAINER.explain_instance(
            image=img,
            classifier_fn=lambda imgs: predict(imgs, model),
            top_labels=1,
            hide_color=0,
            num_samples=200,
        )
        temp, mask = explanation.get_image_and_mask(
            label=pred_idx, positive_only=True, num_features=10, hide_rest=False
        )
        lime_img = mark_boundaries(temp, mask)

        buf = BytesIO()
        plt.imsave(buf, lime_img, format="png")
        buf.seek(0)
        lime_data = buf.getvalue()

        # FIXED: Stable layout
        col1, col2 = st.columns(2)
        with col1:
            st.markdown("""
                <div class="image-container">
                    <h3 style="color:#1e40af; margin-bottom:1rem;">🔬 LIME Explanation</h3>
                </div>
            """, unsafe_allow_html=True)
            st.image(lime_data, width=280, output_format="PNG")
            st.download_button(
                "📥 Download LIME Analysis",
                lime_data,
                file_name=f"{pred_label}_LIME_Analysis.png",
                mime="image/png"
            )
        
        with col2:
            st.markdown(explanation_text.get(pred_label, "<p>No explanation available.</p>"), unsafe_allow_html=True)

# --- FIXED: Stable confidence display ---
def show_confidence(confidence, pred_label):
    # FIXED: Determine confidence level without dynamic styling
    if confidence >= 80:
        icon = "🎯"
        level = "high"
    elif confidence >= 60:
        icon = "⚠️"
        level = "medium"
    else:
        icon = "🔍"
        level = "low"
    
    st.markdown(f"""
        <div class="prediction-card">
            <h2 style="margin:0; color:#1e40af;">{icon} Diagnosis: <strong>{pred_label}</strong></h2>
            <div class="confidence-bar">
                <div class="confidence-fill {level}" style="width:{confidence}%"></div>
            </div>
            <p style="margin:0.5rem 0 0 0; font-size:18px; font-weight:bold;">
                Confidence: {confidence:.1f}%
            </p>
        </div>
    """, unsafe_allow_html=True)

# --- FIXED: Stable Streamlit App UI ---
st.set_page_config(
    page_title="👁️ Retina AI Classifier", 
    layout="wide",
    initial_sidebar_state="expanded"
)

# FIXED: Stable main header
st.markdown("""
    <div class="main-header">
        <h1 style="margin:0; font-size:2.5rem;">👁️ Retina Disease Classifier</h1>
        <p style="margin:0.5rem 0 0 0; font-size:1.2rem;">
            AI-Powered Retinal Analysis with LIME Explainability
        </p>
    </div>
""", unsafe_allow_html=True)

model = load_model()

# FIXED: Stable sidebar
with st.sidebar:
    st.markdown("""
        <div class="sidebar-content">
            <h3 style="color:#1e40af; margin-top:0;">📂 Upload Images</h3>
            <p style="color:#6b7280; margin-bottom:1rem;">
                Upload retinal fundus images for AI analysis
            </p>
        </div>
    """, unsafe_allow_html=True)
    
    uploaded_files = st.file_uploader(
        "Choose retinal images",
        type=["jpg", "jpeg", "png"],
        accept_multiple_files=True,
        help="Upload high-quality fundus photographs"
    )
    
    selected_filename = None
    if uploaded_files:
        st.markdown("""
            <div class="sidebar-content">
                <h4 style="color:#1e40af; margin-top:0;">🎯 Select Image</h4>
            </div>
        """, unsafe_allow_html=True)
        filenames = [f.name for f in uploaded_files]
        selected_filename = st.selectbox(
            "Choose image for analysis",
            filenames,
            help="Select which image to analyze with LIME"
        )

# FIXED: Stable main content area
if uploaded_files and selected_filename:
    file = next(f for f in uploaded_files if f.name == selected_filename)
    file.seek(0)
    bgr = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)

    # FIXED: Stable processing steps section
    st.markdown("""
        <div class="processing-container">
            <h3 style="color:#1e40af; margin-top:0; font-size:1.5rem;">🔬 Image Preprocessing Pipeline</h3>
            <p style="color:#6b7280; margin-bottom:1rem; font-size:1.1rem;">
                Standardized preprocessing steps for optimal AI analysis
            </p>
        </div>
    """, unsafe_allow_html=True)
    
    preprocessed = preprocess_with_steps(rgb)
    input_tensor = np.expand_dims(preprocessed, axis=0)

    # Prediction
    preds = predict(input_tensor, model)
    pred_idx = np.argmax(preds)
    pred_label = CLASS_NAMES[pred_idx]
    confidence = np.max(preds) * 100

    # FIXED: Stable prediction display
    show_confidence(confidence, pred_label)
    
    # FIXED: Stable LIME explanation section
    st.markdown("""
        <div class="lime-container">
            <h3 style="color:#1e40af; margin-top:0; font-size:1.5rem;">🧠 AI Explanation & Clinical Insights</h3>
            <p style="color:#6b7280; margin-bottom:1rem; font-size:1.1rem;">
                Understanding how AI identified the diagnosis with medical context
            </p>
        </div>
    """, unsafe_allow_html=True)
    
    # LIME explanation
    show_lime(preprocessed, model, pred_idx, pred_label, preds)
    
else:
    # FIXED: Stable welcome screen
    st.markdown("""
        <div class="upload-instructions">
            <h3>Welcome to the Retina AI Classifier</h3>
            <p>Upload retinal fundus images to begin AI-powered analysis</p>
            <p style="font-size:0.9rem;">Drag and drop your images or use the sidebar to get started</p>
        </div>
    """, unsafe_allow_html=True)
    
    # FIXED: Stable feature grid
    st.markdown("""
        <div class="feature-grid">
            <div class="feature-card">
                <div class="feature-icon">🔬</div>
                <div class="feature-title">AI-Powered Analysis</div>
                <div class="feature-description">Advanced deep learning models trained on thousands of retinal images</div>
            </div>
            <div class="feature-card">
                <div class="feature-icon">👁️</div>
                <div class="feature-title">8 Conditions Detected</div>
                <div class="feature-description">Normal, Diabetic Retinopathy, Glaucoma, Cataract, AMD, Hypertension, Myopia, Others</div>
            </div>
            <div class="feature-card">
                <div class="feature-icon">🔍</div>
                <div class="feature-title">LIME Explanations</div>
                <div class="feature-description">Visual explanations showing which areas influenced the AI's decision</div>
            </div>
            <div class="feature-card">
                <div class="feature-icon">🏥</div>
                <div class="feature-title">Clinical Grade</div>
                <div class="feature-description">Designed for healthcare professionals with detailed medical insights</div>
            </div>
        </div>
    """, unsafe_allow_html=True)