import gradio as gr
import torch
import torch.nn as nn
import numpy as np
import joblib

# 1. MODEL ARCHITECTURE (BHI_Net)
class BHI_Net(nn.Module):
    def __init__(self, num_classes=6, channels=19, samples=2500):
        super(BHI_Net, self).__init__()
        self.conv_block = nn.Sequential(
            nn.Conv1d(channels, 32, kernel_size=15, stride=1, padding=7),
            nn.BatchNorm1d(32), nn.ReLU(), nn.MaxPool1d(2),
            nn.Conv1d(32, 64, kernel_size=15, stride=1, padding=7),
            nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(2),
            nn.Conv1d(64, 128, kernel_size=15, stride=1, padding=7),
            nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(4)
        )
        self.avgpool = nn.AdaptiveAvgPool1d(10)
        self.classifier = nn.Sequential(
            nn.Linear(128 * 10, 256), nn.ReLU(), nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        x = self.conv_block(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)

# 2. LOAD ASSETS
device = torch.device("cpu")
label_names = ["Healthy", "Alzheimers", "Parkinsons", "Epilepsy", "Depression", "Anxiety"]

# Load Models
model_n = BHI_Net(num_classes=4)
model_n.load_state_dict(torch.load('neuro_best_model.pth', map_location=device))
model_n.eval()

model_m = BHI_Net(num_classes=3)
model_m.load_state_dict(torch.load('mental_best_model.pth', map_location=device))
model_m.eval()

# Load Data and Scalers
X_data = np.load('X_master.npy', mmap_mode='r')
Y_data = np.load('Y_master.npy')
n_test_idx = np.load('neuro_unseen_indices.npy')
m_test_idx = np.load('mental_unseen_indices.npy')
n_scaler = joblib.load('neuro_platt_scaler.pkl')
m_scaler = joblib.load('mental_platt_scaler.pkl')

# Recreate Test Indices (Logic from your notebook)
mental_mask = (Y_data == 0) | (Y_data == 4) | (Y_data == 5)
global_mental_indices = np.where(mental_mask)[0][m_test_idx]
all_test_indices = np.unique(np.concatenate([n_test_idx, global_mental_indices]))

# 3. DIAGNOSTIC FUNCTION
def get_preloaded_diagnosis(patient_id):
    idx = int(patient_id)
    input_tensor = torch.from_numpy(X_data[idx]).float().unsqueeze(0)
    actual_label = label_names[Y_data[idx]]

    with torch.no_grad():
        n_logits = model_n(input_tensor).numpy()
        n_probs = n_scaler.predict_proba(n_logits)[0]
        m_logits = model_m(input_tensor).numpy()
        m_probs = m_scaler.predict_proba(m_logits)[0]

    # Consensus Logic (Matching your notebook)
    STUCK_THRESHOLD = 0.305
    DISEASE_BOOST = 1.4

    if n_probs[3] > 0.85:
        diag, conf = "Epilepsy", n_probs[3] * 100
    elif n_probs[1] < STUCK_THRESHOLD and n_probs[0] < 0.5:
        m_idx = np.argmax(m_probs)
        diag = ["Healthy", "Depression", "Anxiety"][m_idx]
        conf = m_probs[m_idx] * 100
    else:
        n_probs[1:3] *= DISEASE_BOOST
        n_idx = np.argmax(n_probs)
        diag = ["Healthy", "Alzheimers", "Parkinsons", "Epilepsy"][n_idx]
        conf = n_probs[n_idx] * 100

    result_emoji = "✅ CORRECT" if diag == actual_label else "❌ INCORRECT"

    report = f"🏥 CLINICAL EEG DIAGNOSTIC SUMMARY\n"
    report += "="*45 + "\n"
    report += f"PATIENT ID      : {idx}\n"
    report += f"ACTUAL DIAGNOSIS: {actual_label}\n"
    report += f"AI PREDICTION   : {diag} ({conf:.1f}%)\n"
    report += f"RESULT          : {result_emoji}\n"
    report += "="*45
    return report

# 4. INTERFACE
demo = gr.Interface(
    fn=get_preloaded_diagnosis,
    inputs=gr.Dropdown(choices=[str(i) for i in all_test_indices], label="Select Patient from Study"),
    outputs=gr.Textbox(label="Clinical Report", lines=10),
    title="Clinical EEG Specialist AI",
    description="Select a Patient ID from the unseen test set to view the AI's diagnostic summary."
)

demo.launch()