import gradio as gr
import torch
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification

# 1. Setup Device
device = "cuda" if torch.cuda.is_available() else "cpu"

# 2. Load Fine-Tuned Model and Processor
model_path = "actorcritic/twak"  # <-- Change this back to just the repo name
# model_path = "/home/shohog/Documents/twok"

try:
    # Add subfolder="model" here
    ft_processor = AutoImageProcessor.from_pretrained(model_path, subfolder="model")
    
    # Add subfolder="model" here too
    ft_model = AutoModelForImageClassification.from_pretrained(
        model_path,    
        subfolder="model",
        problem_type="multi_label_classification",
        use_safetensors=True
    )
    ft_model.to(device)
    ft_model.eval()
    print("✓ Fine-tuned model loaded successfully")
except Exception as e:
    print(f"Error loading fine-tuned model: {e}")
    ft_processor = None
    ft_model = None

# 3. Dictionary for Disease Information (Bengali)
disease_info = {
    "Acne": "ব্রণ (Acne) অবহেলা করলে বা নখ দিয়ে খুঁটলে ত্বকে স্থায়ী গর্ত বা দাগ (Scar) সৃষ্টি হতে পারে। তীব্র মাত্রার ব্রণে গভীর সংক্রমণ (Cystic infection) হতে পারে, যা মানসিক অবসাদ ও হীনমন্যতার অন্যতম কারণ।",
    
    "Arsenic": "আর্সেনিক (Arsenic) দূষণের দীর্ঘমেয়াদী প্রভাবে ত্বকের ক্যান্সার (Skin Cancer) হতে পারে। এছাড়া এটি ফুসফুস, কিডনি ও মূত্রথলির ক্যান্সার এবং মারাত্মক হৃদরোগের ঝুঁকি বহুগুণ বাড়িয়ে দেয়, যা প্রাণঘাতী।",
    
    "Atopic Dermatitis": "এটোপিক ডার্মাটাইটিস (Atopic Dermatitis) এর কারণে তীব্র চুলকানি থেকে ত্বকে ক্ষত সৃষ্টি হয়ে মারাত্মক ব্যাকটেরিয়াল ইনফেকশন (যেমন- Staph infection) হতে পারে। এটি রোগীর ঘুম ও স্বাভাবিক জীবনযাত্রা চরমভাবে ব্যাহত করে।",
    
    "Candidal Intertrigo": "ক্যানডিডাল ইন্টারট্রিগো (Candidal Intertrigo) সঠিক সময়ে চিকিৎসা না করলে ত্বকে গভীর ক্ষত বা আলসার তৈরি হতে পারে। ডায়াবেটিস বা রোগ প্রতিরোধ ক্ষমতা কম থাকলে এই ছত্রাক রক্তে ছড়িয়ে মারাত্মক আকার ধারণ করতে পারে।",
    
    "Contact Dermatitis": "কন্টাক্ট ডার্মাটাইটিস (Contact Dermatitis) এর কারণে ত্বকে বড় বড় ফোসকা পড়তে পারে এবং দীর্ঘমেয়াদী সংস্পর্শে ত্বক স্থায়ীভাবে পুরু, কালচে ও অকেজো হয়ে যেতে পারে, যা কর্মক্ষমতা কমিয়ে দেয়।",
    
    "Eczema": "একজিমা (Eczema) অবহেলা করলে ত্বকে 'একজিমা হারপেটিকাম' এর মতো মারাত্মক ভাইরাল বা ব্যাকটেরিয়াল ইনফেকশন হতে পারে, যা দ্রুত সারা শরীরে ছড়িয়ে পড়ে জীবনের জন্য হুমকিস্বরূপ হতে পারে।",
    
    "Psoriasis": "সোরিয়াসিস (Psoriasis) কেবল ত্বকের রোগ নয়; এটি জয়েন্ট বা অস্থিসন্ধি স্থায়ীভাবে বিকল করে দিতে পারে (সোরিয়াটিক আর্থ্রাইটিস)। এছাড়া এটি হৃদরোগ, ডায়াবেটিস এবং স্ট্রোকের ঝুঁকি মারাত্মকভাবে বাড়িয়ে দেয়।",
    
    "Scabies": "স্ক্যাবিস (Scabies) অত্যন্ত ছোঁয়াচে হওয়ায় দ্রুত পুরো পরিবার ও সমাজে ছড়িয়ে পড়ে। অতিরিক্ত চুলকানির ফলে সৃষ্ট ক্ষত থেকে ব্যাকটেরিয়াল ইনফেকশন হয়ে শিশুদের কিডনি বিকল (Glomerulonephritis) বা বাতজ্বর হতে পারে।",
    
    "Seborrheic Dermatitis": "সিবোরিক ডার্মাটাইটিস (Seborrheic Dermatitis) তীব্র আকার ধারণ করলে মাথার ত্বকে ঘা হতে পারে এবং স্থায়ীভাবে চুল পড়ে যেতে পারে (Hair loss), যা তীব্র মানসিক চাপের সৃষ্টি করে।",
    
    "Steroid Modified Tinea": "স্টেরয়েড-মডিফাইড টিনিয়া (Steroid Modified Tinea) এর কারণে ছত্রাক অত্যন্ত শক্তিশালী ও ওষুধ-প্রতিরোধী হয়ে ওঠে। এর ফলে ত্বক স্থায়ীভাবে পাতলা হয়ে যায়, ফেটে যায় এবং এই ইনফেকশন সারানো অত্যন্ত কঠিন ও ব্যয়বহুল হয়ে পড়ে।",
    
    "Tinea Corporis": "টিনিয়া কর্পোরিস (Tinea Corporis) বা দাদ দ্রুত শরীরের অন্যান্য অংশে ছড়িয়ে পড়তে পারে। দীর্ঘমেয়াদী সংক্রমণে ত্বকে স্থায়ী কালো দাগ ও চুলকানি থেকে ব্যাকটেরিয়াল ইনফেকশন হতে পারে।",
    
    "Tinea Cruris": "টিনিয়া ক্রুরিস (Tinea Cruris) অবহেলা করলে এটি যৌনাঙ্গ ও এর আশেপাশের সংবেদনশীল অংশে ছড়িয়ে তীব্র যন্ত্রণাদায়ক ঘা সৃষ্টি করতে পারে, যা দৈনন্দিন হাঁটাচলা ও স্বাভাবিক জীবনযাপন অসম্ভব করে তোলে।",
    
    "Tinea Faciei": "টিনিয়া ফেসিয়েই (Tinea Faciei) মুখের ত্বকে স্থায়ী দাগ বা ক্ষত সৃষ্টি করতে পারে। ভুল চিকিৎসায় (বিশেষ করে স্টেরয়েড ব্যবহারে) এটি ভয়াবহ রূপ নিতে পারে, যা রোগীর চেহারা বিকৃত করে মানসিক বিপর্যয় ডেকে আনে।",
    
    "Vitiligo": "শ্বেতী (Vitiligo) আক্রান্ত স্থানে মেলানিন না থাকায় তা সূর্যের অতিবেগুনি রশ্মির প্রতি অত্যন্ত সংবেদনশীল হয়ে পড়ে, যার ফলে মারাত্মক সানবার্ন এবং ত্বকের ক্যান্সারের (Skin Cancer) ঝুঁকি বহুগুণ বেড়ে যায়। এছাড়া এটি থাইরয়েড বা ডায়াবেটিসের মতো অন্যান্য রোগের পূর্বাভাস হতে পারে।"
}


# 15 Original Folders (Index 6 is the 'Others/Healthy' folder)
ORIGINAL_CLASSES =[
    "Acne", "Arsenic", "Atopic Dermatitis", "Candidal Intertrigo", 
    "Contact Dermatitis", "Eczema", "Healthy/Others", "Psoriasis", 
    "Scabies", "Seborrheic Dermatitis", "Steroid Modified Tinea", 
    "Tinea Corporis", "Tinea Cruris", "Tinea Faciei", "Vitiligo"
]

# Optimal Thresholds per Output Neuron (Targeting 80% Recall)
CLASS_THRESHOLDS = {
    0: 0.89, 1: 0.88, 2: 0.34, 3: 0.04, 4: 0.66, 5: 0.48, 
    6: 0.28, 7: 0.73, 8: 0.54, 9: 0.44, 10: 0.68, 11: 0.83, 
    12: 0.85, 13: 0.95
}

def get_original_class_idx(neuron_idx):
    """Maps the 14 output neurons back to the original 15 folder indices."""
    if neuron_idx < 6:
        return neuron_idx
    else:
        return neuron_idx + 1  # Shift back to account for the skipped Class 6

# 4. Prediction Function
def predict(image):
    if image is None:
        return "<p style='text-align:center; color:#ef4444; padding: 20px;'>Please upload an image.</p>"

    if ft_processor is None or ft_model is None:
        return "<p style='text-align:center; color:#ef4444; padding: 20px;'>Fine-tuned model not loaded. Please check the model path.</p>"

    # Convert image to RGB and run inference
    image_rgb = image.convert("RGB")
    inputs = ft_processor(images=image_rgb, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = ft_model(**inputs)

    logits = outputs.logits
    scores = torch.sigmoid(logits)[0]
    print(logits, scores)

    detected_diseases =[]

    # Check each of the 14 neurons against its specific threshold
    for neuron_idx in range(14):
        prob = scores[neuron_idx].item()
        if prob >= CLASS_THRESHOLDS[neuron_idx]:
            orig_idx = get_original_class_idx(neuron_idx)
            detected_diseases.append((orig_idx, prob))

    html_output = '<div class="result-container">'

    # SCENARIO A: No thresholds were crossed -> It's Class 6 (Healthy/Others)
    if len(detected_diseases) == 0:
        predicted_class = "Healthy / Others"
        info_text = "কোনো চর্মরোগ শনাক্ত হয়নি। ত্বক সুস্থ অথবা এটি অন্য কোনো অবস্থা হতে পারে।"
        
        # Added flex-direction: column here to stack them vertically
        html_output += f"""
        <div style="display: flex; flex-direction: column; align-items: center; justify-content: center; gap: 8px;">
            <span class="disease-name" style="color: #10b981;">{predicted_class}</span>
            <div class="tooltip">
                <span class="info-icon">i</span>
                <div class="tooltiptext">
                    <strong>{predicted_class}</strong><br><br>
                    {info_text}
                </div>
            </div>
        </div>
        """
        
    # SCENARIO B: One or more diseases detected
    else:
        # Sort by probability (highest confidence first)
        detected_diseases.sort(key=lambda x: x[1], reverse=True)
        
        for orig_idx, prob in detected_diseases:
            predicted_class = ORIGINAL_CLASSES[orig_idx]
            info_text = disease_info.get(predicted_class, "Fundamental details for this condition are not available.")
            confidence = prob * 100
            print(confidence)
            
            # Added flex-direction: column here to stack them vertically
            html_output += f"""
            <div style="display: flex; flex-direction: column; align-items: center; justify-content: center; gap: 8px; margin-bottom: 15px; width: 100%;">
                <span class="disease-name">{predicted_class} <span style="font-size: 16px; color: #6b7280;">({confidence:.1f}%)</span></span>
                <div class="tooltip">
                    <span class="info-icon">i</span>
                    <div class="tooltiptext">
                        <strong>{predicted_class}</strong><br><br>
                        {info_text}
                    </div>
                </div>
            </div>
            """
            
    html_output += '</div>'
    return html_output

# 5. Custom CSS
custom_css = """
.result-container {
    display: flex;
    flex-direction: column; /* This moves the button below the text */
    align-items: center;
    justify-content: center;
    gap: 12px;
    padding: 30px 20px;
    background: #ffffff;
    border-radius: 8px;
    border: 1px solid #e5e7eb;
    margin-top: 5px;
    min-height: 100px;
    box-shadow: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
    text-align: center;
}
.disease-name {
    font-size: 24px;
    font-weight: 600;
    color: #111827;
}
.disease-name.unknown {
    color: #f59e0b;
}
.confidence {
    font-size: 18px;
    font-weight: 400;
    color: #6b7280;
}
.tooltip {
    position: relative;
    display: inline-block;
    margin-top: 5px; /* Adds a little breathing room above the button */
}

/* --- Intuitive Info Icon --- */
.info-icon {
    cursor: pointer;
    background: #6366f1;
    color: white;
    border-radius: 16px;
    width: auto;
    height: 28px;
    padding: 0 12px;
    display: inline-flex;
    align-items: center;
    justify-content: center;
    font-size: 14px;
    font-weight: bold;
    font-family: sans-serif;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    transition: all 0.2s ease;
    animation: pulse 2s infinite;
}

.info-icon:hover { 
    background: #4f46e5; 
    box-shadow: 0 4px 6px rgba(0,0,0,0.15);
    animation: none;
}

@keyframes pulse {
    0% { box-shadow: 0 0 0 0 rgba(99, 102, 241, 0.4); }
    70% { box-shadow: 0 0 0 6px rgba(99, 102, 241, 0); }
    100% { box-shadow: 0 0 0 0 rgba(99, 102, 241, 0); }
}

/* --- Center-Aligned Tooltip Card --- */
.tooltip .tooltiptext {
    visibility: hidden;
    width: 280px; 
    max-width: 90vw; 
    background-color: #ffffff;
    color: #374151;
    text-align: left;
    border-radius: 8px;
    padding: 15px;
    position: absolute;
    z-index: 10;
    bottom: 130%;
    
    /* Center alignment logic */
    left: 50%;
    right: auto;
    margin-left: 0;
    
    opacity: 0;
    transition: opacity 0.3s, transform 0.3s;
    transform: translate(-50%, 10px); 
    
    font-size: 14px;
    font-weight: normal;
    box-shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.2), 0 8px 10px -6px rgba(0, 0, 0, 0.1);
    border: 1px solid #e5e7eb;
    line-height: 1.5;
    white-space: normal;
}

/* Center the little arrow */
.tooltip .tooltiptext::after {
    content: "";
    position: absolute;
    top: 100%;
    left: 50%;
    right: auto;
    margin-left: 0;
    transform: translateX(-50%);
    border-width: 8px;
    border-style: solid;
    border-color: #ffffff transparent transparent transparent;
}

.tooltip:hover .tooltiptext {
    visibility: visible;
    opacity: 1;
    transform: translate(-50%, 0); 
}
"""


# Tiny JavaScript to force light mode by removing the 'dark' class from the body
force_light_mode_js = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'light') {
        url.searchParams.set('__theme', 'light');
        window.location.href = url.href;
    }
}
"""

# Add the js parameter to your Blocks
with gr.Blocks(js=force_light_mode_js) as demo:
    # Changed gr.Markdown to gr.HTML
    gr.HTML(
        "<h1 style='text-align: center; color: darkgreen; font-weight: 900; font-family: Impact, sans-serif;'>Twok Test</h1>"
    )

    # Changed gr.Markdown to gr.HTML
    gr.HTML(
        "<p style='text-align: center; font-size: 16px; color: #4b5563;'>"
        "নির্ভুল চর্মরোগ নির্ণয় এবং জনস্বাস্থ্য সুরক্ষায় একটি বিশেষায়িত কৃত্রিম বুদ্ধিমত্তা ভিত্তিক সেবা।"
        "</p>"
    )

    image_input = gr.Image(type="pil", label="Upload Skin Image")
    predict_btn = gr.Button("Test")
    output_html = gr.HTML(label="Prediction Result")

    with gr.Accordion("Sample Image", open=False):
        gr.Examples(
            examples=["twok_sample.jpg", "twok_sample2.jpg"], 
            inputs=image_input,
            outputs=output_html,
            fn=predict,
            cache_examples=False  # <--- ADD THIS LINE

        )

    with gr.Accordion("Disclaimer", open=False):
        gr.Markdown("*(Its a prediagnosis ai, not a substitute for professional medical advice)*\n\n*(এটি একটি প্রাথমিক রোগ নির্ণয়কারী এআই বা কৃত্রিম বুদ্ধিমত্তা, এটি পেশাদার ডাক্তারের পরামর্শের বিকল্প নয়)*")

    predict_btn.click(fn=predict, inputs=image_input, outputs=output_html)

demo.launch(css=custom_css, theme=gr.themes.Default())