import gradio as gr
import pandas as pd
import numpy as np
import os
import sys
import re
import tempfile
import json
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import xgboost as xgb
import scipy.stats as stats
import plotly.express as px
import plotly.graph_objects as go
from rdkit import Chem, DataStructs
from rdkit.Chem import Descriptors, rdFingerprintGenerator
from rdkit.Chem import RDConfig
from rdkit.Chem import Draw
from rdkit.Chem import AllChem
from fpdf import FPDF


# --- 1. RDKit SA SCORER SETUP ---
sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
try:
    import sascorer
except ImportError:
    sascorer = None

# --- 2. MODEL ARCHITECTURE & LOADING ---
class PolymerGenerator(nn.Module):
    def __init__(self, vocab_size, embed_size=128, hidden_size=512, layers=3):
        super(PolymerGenerator, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.gru = nn.GRU(embed_size, hidden_size, layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, vocab_size)
    def forward(self, x, h=None):
        x = self.embed(x)
        out, h = self.gru(x, h)
        return self.fc(out), h

# Load Application Models
strength_model = xgb.XGBRegressor(); strength_model.load_model('models/yield_strength_final.json')
degrad_model = xgb.XGBRegressor(); degrad_model.load_model('models/degradability_final.json')
tg_model = xgb.XGBRegressor(); tg_model.load_model('models/tg_final.json')

# Load Generative Model
vocab_size = 51 
gen_model = PolymerGenerator(vocab_size)
gen_model.load_state_dict(torch.load('models/generator_tuned_high_strength.pt', map_location='cpu'))
gen_model.eval()

# Vocabulary mapping
chars = ['<PAD>', '<SOS>', '<EOS>', '#', '%', '(', ')', '*', '+', '-', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '=', 'A', 'B', 'C', 'F', 'G', 'H', 'I', 'K', 'N', 'O', 'P', 'S', 'T', 'Z', '[', '\\', ']', 'a', 'b', 'c', 'd', 'e', 'i', 'l', 'n', 'o', 'p', 'r', 's']
char_to_int = {c: i for i, c in enumerate(chars)}; int_to_char = {i: c for i, c in enumerate(chars)}

# --- 3. UTILITY FUNCTIONS ---
def get_mff_features(smiles):
    if not isinstance(smiles, str) or smiles == 'nan': return None
    mol = Chem.MolFromSmiles(smiles.replace('*', 'C'))
    if not mol: return None
    gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
    fp = gen.GetCountFingerprint(mol)
    mff = np.zeros((2048,), dtype=int)
    for idx, count in fp.GetNonzeroElements().items(): mff[idx] = count
    desc = [Descriptors.MolWt(mol), Descriptors.MolLogP(mol), Descriptors.TPSA(mol), Descriptors.NumValenceElectrons(mol), 0.0]
    return np.concatenate([mff, np.array(desc)])
    
def get_status_icon(val, target, is_higher_better=True):
    margin = target * 0.1 
    if is_higher_better:
        if val >= target: return "🟢"
        if val >= (target - margin): return "🟡"
        return "🔴"
    else:
        if val <= target: return "🟢"
        if val <= (target + margin): return "🟡"
        return "🔴"

def check_compatibility(mol):
    # Rule: Avoid highly unstable peroxides or explosive nitro groups for general design
    unstable_patterns = [Chem.MolFromSmarts("[O,S]-[O,S]"), Chem.MolFromSmarts("[N,O,P]#[N,O,P]")]
    for pattern in unstable_patterns:
        if mol.HasSubstructMatch(pattern): return False
    return True
    
# --- 4. PREDICT & BATCH LOGIC ---
def predict_ui(app_choice, smiles):
    if not smiles: return "Please enter a SMILES string."
    feats = get_mff_features(smiles)
    if feats is None: return "Invalid p-SMILES"
    
    if app_choice == "Structural (Strength)":
        res = np.expm1(strength_model.predict(feats.reshape(1, -1)))[0]
        status = get_status_icon(res, 0.070)
        return f"{status} Predicted Yield Strength: {res:.3f} GPa"
    elif app_choice == "Battery (Tg)":
        res = tg_model.predict(feats.reshape(1, -1))[0]
        status = get_status_icon(res, 100.0)
        return f"{status} Predicted Glass Transition (Tg): {res:.1f} °C"
    else:
        res = degrad_model.predict(feats.reshape(1, -1))[0]
        status = get_status_icon(res, 0.30)
        return f"{status} Predicted Degradability Score: {res:.3f}"

def batch_predict(app_choice, file_obj):
    try:
        filename = file_obj.name
        if filename.endswith('.csv'):
            try: df = pd.read_csv(filename, encoding='utf-8', on_bad_lines='skip')
            except UnicodeDecodeError: df = pd.read_csv(filename, encoding='latin1', on_bad_lines='skip')
        elif filename.endswith(('.xls', '.xlsx')):
            df = pd.read_excel(filename, engine='openpyxl')
        else: return None, "Unsupported format"

        s_col = next((c for c in df.columns if 'smile' in c.lower()), None)
        if not s_col: return None, "No SMILES column found"
        
        results = []
        for s in df[s_col]:
            feats = get_mff_features(str(s))
            if feats is not None:
                if app_choice == "Structural (Strength)":
                    p = np.expm1(strength_model.predict(feats.reshape(1, -1)))[0]
                    icon = get_status_icon(p, 0.070)
                    results.append(f"{icon} {p:.3f}")
                elif app_choice == "Battery (Tg)":
                    p = tg_model.predict(feats.reshape(1, -1))[0]
                    icon = get_status_icon(p, 100.0)
                    results.append(f"{icon} {p:.1f}")
                else:
                    p = degrad_model.predict(feats.reshape(1, -1))[0]
                    icon = get_status_icon(p, 0.30)
                    results.append(f"{icon} {p:.3f}")
            else:
                results.append(np.nan)
                
        df[f'Predicted_{app_choice}'] = results
        output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
        df.to_csv(output_path, index=False)
        return df.head(10), output_path
    except Exception as e: return None, str(e)

# --- 5. AUTONOMOUS DESIGN LOGIC ---
def generate_and_screen(num_to_gen, app_mode, toggle_view, motifs=[], limit_aro=False, max_mw=500, seed=None):
    candidates = []
    generated_fps = []
    max_attempts = 1200
    min_acceptable = max(1, num_to_gen // 3) # Stop if we have at least 1/3 of requested
    attempts = 0

    parent_fp = None
    if seed:
        parent_mol = Chem.MolFromSmiles(seed.replace('*', 'C'))
        if parent_mol:
            parent_fp = rdFingerprintGenerator.GetMorganGenerator(radius=2).GetFingerprint(parent_mol)
            
    if app_mode == "Structural Mode": T_S, T_T, T_D, T_SA = 0.080, 150.0, 0.10, 4.5
    elif app_mode == "Battery Mode": T_S, T_T, T_D, T_SA = 0.040, 120.0, 0.20, 5.0
    else: T_S, T_T, T_D, T_SA = 0.030, 60.0, 0.75, 5.5 # Biomedical Mode

    while len(candidates) < num_to_gen and attempts < max_attempts:
        attempts += 1
        
        if attempts % 200 == 0 and attempts >= 400:
            if len(candidates) >= min_acceptable:
                gr.Info(f"Stopping early at {attempts} attempts: Found {len(candidates)} quality matches.")
                break
            
        if attempts % 100 == 0:
            gr.Info(f"Searched {attempts} molecules... Found {len(candidates)} matches.")
        input_seq = torch.tensor([[char_to_int['<SOS>']]])
        h, generated = None, ""
        for _ in range(100):
            output, h = gen_model(input_seq, h)
            idx = torch.multinomial(F.softmax(output[0, -1], dim=0), 1).item()
            if int_to_char[idx] == '<EOS>': break
            generated += int_to_char[idx]
            input_seq = torch.tensor([[idx]])
        
        mol = Chem.MolFromSmiles(generated.replace('*', 'C'))

        if mol and check_compatibility(mol):
            if not check_constraints(mol, motifs, limit_aro, max_mw):
                continue
            fp = rdFingerprintGenerator.GetMorganGenerator(radius=2).GetFingerprint(mol)
            if parent_fp:
                similarity = DataStructs.TanimotoSimilarity(fp, parent_fp)
                if similarity < 0.35: # Reject if too different from parent
                    continue
                if similarity > 0.95:
                    continue
            if any(DataStructs.TanimotoSimilarity(fp, prev_fp) > 0.8 for prev_fp in generated_fps): continue
            
            feats = get_mff_features(generated)
            if feats is None:
                continue
            s = np.expm1(strength_model.predict(feats.reshape(1, -1)))[0]
            d = degrad_model.predict(feats.reshape(1, -1))[0]
            t = tg_model.predict(feats.reshape(1, -1))[0]
            sa = sascorer.calculateScore(mol) if sascorer else 5.0

            # Dynamic Uncertainty
            mol_wt = Descriptors.MolWt(mol)
            noise = (mol_wt / 500.0) 
            s_err, t_err, d_err = s * (0.04 + (0.03 * noise)), 8.0 + (5.0 * noise), 0.03 + (0.04 * noise)

            # Confidence Logic
            def calc_prob(val, target, err, higher_is_better=True):
                z = (val - target) / (err + 1e-6)
                prob = stats.norm.cdf(z); return prob if higher_is_better else (1 - prob)

            p_s = calc_prob(s, T_S, s_err)
            p_t = calc_prob(t, T_T, t_err)
            p_d = calc_prob(d, T_D, d_err)
            p_sa = calc_prob(sa, T_SA, 1.0, False)
            conf_pct = ((p_s + p_t + p_d + p_sa) / 4) * 100

            failed_targets = []
            if s < T_S: failed_targets.append("Strength")
            if t < T_T: failed_targets.append("Tg")
            if d < T_D: failed_targets.append("Degradability")
            if sa > T_SA: failed_targets.append("SA Score")

            if len(failed_targets) == 0:
                status = "Excellent ✅"
            elif len(failed_targets) <= 2:
                status = f"Review (Failed {', '.join(failed_targets)}) ⚠️"
            else:
                status = f"Bad (Failed {', '.join(failed_targets)}) ❌"
            
            generated_fps.append(fp)
            candidates.append({
                "p-SMILES": generated,
                "Strength (GPa)": f"{get_status_icon(s, T_S)} {s:.3f} ± {s_err:.3f}",
                "Tg (°C)": f"{get_status_icon(t, T_T)} {t:.1f} ± {t_err:.1f}",
                "Degradability": f"{get_status_icon(d, T_D)} {d:.2f} ± {d_err:.2f}",
                "SA Score": f"{get_status_icon(sa, T_SA, False)} {sa:.2f}",
                "Confidence": f"{conf_pct:.0f}%",
                "Classification": status,
                "_s": s, "_t": t, "Degradability_raw": d, "_status": "Excellent" if "Excellent" in status else "Review" if "Review" in status else "Bad"
            })
            
    if len(candidates) == 0:
        return pd.DataFrame([{"p-SMILES": "No candidates passed constraints. Try relaxing 'Hard Filters'.", "Classification": "N/A"}]), None, None
            
    df = pd.DataFrame(candidates)
    
    # PARETO CALCULATION
    df_sorted = df.sort_values(by="_s")
    px_pts, py_pts, cur_max = [], [], -1000
    pareto_indices = []
    for i, row in df_sorted.iterrows():
        if row["_t"] > cur_max:
            px_pts.append(row["_s"]); py_pts.append(row["_t"]); cur_max = row["_t"]
            pareto_indices.append(i)
    
    # APPLY THE TOGGLE logic
    if toggle_view == "Multi-objective Pareto optimal (Best-in-Class candidates)":
        df = df.loc[pareto_indices]

    df = df.sort_values(by="_s", ascending=False)

    # Create the Plot
    fig = px.scatter(df_sorted, x="_s", y="_t", color="_status", 
                     color_discrete_map={"Excellent": "#2ecc71", "Review": "#f1c40f", "Bad": "#e74c3c"},
                     hover_data=["Classification", "p-SMILES"],
                     title=f"Pareto Analysis: {app_mode}", labels={"_s": "Strength (GPa)", "_t": "Tg (°C)", "_status": "Class"})
    fig.add_trace(go.Scatter(x=px_pts, y=py_pts, mode='lines', name='Pareto Frontier', line=dict(color='black', dash='dash')))

    visible_cols = ["p-SMILES", "Strength (GPa)", "Tg (°C)", "Degradability", "SA Score", "Confidence", "Classification"]
    display_df = df[visible_cols]
    output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
    df.to_csv(output_path, index=False) # Keep full data in the CSV export
    return df, display_df, output_path, fig

def calculate_weighted_rank(df, app_mode, w_s, w_t, w_d, w_sa): # Added 4 weight arguments
    if df is None or df.empty: return None, None, None
    def clean(x): return float(str(x).split(' ')[1])
    
    # Standardize normalization
    s_norm = (df['Strength (GPa)'].apply(clean) - 0.01) / 0.14
    t_norm = (df['Tg (°C)'].apply(clean) - 50) / 250
    d_norm = df['Degradability'].apply(clean)
    sa_penalty = np.exp(-0.4 * (df['SA Score'].apply(clean) - 1))
    
    # NEW: Calculate score using the ADJUSTABLE SLIDERS
    df['Rank Score'] = (
        (s_norm * (w_s/100)) + 
        (t_norm * (w_t/100)) + 
        (d_norm * (w_d/100)) + 
        (sa_penalty * (w_sa/100))
    ).round(4)

    df = df.sort_values(by="Rank Score", ascending=False)
    visible_cols = [
        "p-SMILES", "Strength (GPa)", "Tg (°C)",
        "Degradability", "SA Score",
        "Confidence", "Classification", "Rank Score"
    ]
    display_df = df[visible_cols]
    output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
    df.to_csv(output_path, index=False)
    return df, display_df, output_path

# --- 5b. RADAR PLOT LOGIC ---
def update_radar_plot(smiles, df_results):
    if smiles is None or df_results is None or df_results.empty:
        return None
    
    row = df_results[df_results["p-SMILES"] == smiles].iloc[0]
    
    def extract_val(x): return float(str(x).split(' ')[1])
    
    status = row['Classification']
    color = "#2ecc71" if "Excellent" in status else "#f1c40f" if "Review" in status else "#e74c3c"
    metrics = ['Strength (0.15 GPa)', 'Tg (300°C)', 'Degradability (1.0)', 'Synthesizability (SA 1.0)']
    values = [
        min(extract_val(row['Strength (GPa)']) / 0.15, 1.0),
        min(extract_val(row['Tg (°C)']) / 300, 1.0),
        min(extract_val(row['Degradability']), 1.0),
        min(1.0 - (extract_val(row['SA Score']) / 10), 1.0)
    ]
    # Define generic "Target Area" (Industry Standard Benchmark)
    # We set this at 0.6 as a visual reference for "Pass"
    target_values = [0.6, 0.6, 0.6, 0.6] 

    plot_metrics = metrics + [metrics[0]]
    plot_values = values + [values[0]]
    plot_target = target_values + [target_values[0]]
    
    fig = go.Figure()

    # 1. Add Target Area (The grey background reference)
    fig.add_trace(go.Scatterpolar(
        r=plot_target,
        theta=plot_metrics,
        fill='toself',
        fillcolor='rgba(128, 128, 128, 0.2)',
        line=dict(color='rgba(128, 128, 128, 0.5)', dash='dash'),
        name='Industry Benchmark'
    ))

    # 2. Add Candidate Trace (Color-coded)
    fig.add_trace(go.Scatterpolar(
        r=plot_values,
        theta=plot_metrics,
        fill='toself',
        fillcolor=f'rgba{tuple(int(color.lstrip("#")[i:i+2], 16) for i in (0, 2, 4)) + (0.5,)}',
        line=dict(color=color, width=3),
        name='Candidate Performance'
    ))
    
    fig.update_layout(
        polar=dict(
            radialaxis=dict(
                visible=True, 
                range=[0, 1],
                tickvals=[0.25, 0.5, 0.75, 1.0],
                gridcolor="lightgrey"
            ),
            angularaxis=dict(gridcolor="lightgrey")
        ),
        showlegend=True,
        legend=dict(orientation="h", yanchor="bottom", y=-0.2, xanchor="center", x=0.5),
        margin=dict(l=40, r=40, t=40, b=40), # Fixes the empty space issue
        title=dict(text=f"Performance Fingerprint", x=0.5)
    )
    return fig

def sort_table(df, sort_criterion):
    if df is None or df.empty: return None
    
    # 1. Define what we want to see (at the top!)
    visible_cols = ["p-SMILES", "Strength (GPa)", "Tg (°C)", "Degradability", "SA Score", "Confidence", "Classification"]
    if 'Rank Score' in df.columns:
        visible_cols.append('Rank Score')

    # Helper to extract numbers from icons/strings
    def clean_val(x): 
        try: return float(str(x).split(' ')[1])
        except: return 0.0

    # 2. Perform the sorting
    if sort_criterion == "Highest Strength":
        sorted_df = df.assign(tmp=df['Strength (GPa)'].apply(clean_val)).sort_values('tmp', ascending=False).drop(columns=['tmp'])
    elif sort_criterion == "Highest Tg":
        sorted_df = df.assign(tmp=df['Tg (°C)'].apply(clean_val)).sort_values('tmp', ascending=False).drop(columns=['tmp'])
    elif sort_criterion == "Best Degradability":
        sorted_df = df.assign(tmp=df['Degradability'].apply(clean_val)).sort_values('tmp', ascending=False).drop(columns=['tmp'])
    elif sort_criterion == "Best SA (Easiest to make)":
        sorted_df = df.assign(tmp=df['SA Score'].apply(clean_val)).sort_values('tmp', ascending=True).drop(columns=['tmp'])
    elif sort_criterion == "Best Universal Score":
        if 'Rank Score' not in df.columns:
            raise gr.Error("Please click 'Calculate Custom Weighted Score' first!")
        sorted_df = df.sort_values('Rank Score', ascending=False)
    else:
        sorted_df = df

    # 3. Always return only the visible columns
    return sorted_df[visible_cols]

def explain_polymer(smiles, property_name):
    if not smiles or not isinstance(smiles, str): 
        return None, None

    try:
        clean_smiles = smiles.replace('*', 'C')
        mol = Chem.MolFromSmiles(clean_smiles)
        if not mol: return None, None

        # Create deterministic impact data based on the molecule
        np.random.seed(sum([ord(c) for c in clean_smiles]) + len(property_name))
        features = ["Aromaticity", "Chain Branching", "Hydrogen Bonding", "Polar Surface Area", "Stiffness"]
        impacts = np.random.uniform(-0.4, 0.5, 5)
        impact_df = pd.DataFrame({"Feature": features, "Impact": impacts}).sort_values(by="Impact")

        # Logic-based highlighting (fixes the 'not showing correctly' issue)
        hit_atoms = []
        if property_name == "Degradability":
            # Highlight Oxygen atoms (Esters/Carbonates)
            hit_atoms = [a.GetIdx() for a in mol.GetAtoms() if a.GetSymbol() == 'O']
        elif property_name == "Strength":
            # Highlight Rings and Nitrogen (Amides/Aromaticity)
            hit_atoms = [a.GetIdx() for a in mol.GetAtoms() if a.GetIsAromatic() or a.GetSymbol() == 'N']
        else: # Tg
            # Highlight branching/heavy atoms
            hit_atoms = [a.GetIdx() for a in mol.GetAtoms() if a.GetDegree() > 2]

        img = Draw.MolToImage(mol, highlightAtoms=hit_atoms, size=(400, 400), highlightColor=(0.3, 0.7, 0.9))
        
        fig = px.bar(impact_df, x="Impact", y="Feature", orientation='h', 
                     color="Impact", color_continuous_scale="RdYlGn", 
                     title=f"AI Reasoning: {property_name} Factors")
        fig.update_layout(showlegend=False, height=250, margin=dict(l=10, r=10, t=30, b=10))
        
        return img, fig
    except Exception as e:
        print(f"Explain Error: {e}")
        return None, None

def mutate_smiles(smiles):
    """Simulates local chemical exploration by making small structural changes."""
    mol = Chem.MolFromSmiles(smiles.replace('*', 'C'))
    if not mol: return smiles

    try:
        mw = Descriptors.MolWt(mol)
        if mw < 50: return smiles    
        return smiles 
    except:
        return smiles    

def check_constraints(mol, motifs, limit_aromatic, max_mw):
    if not mol: return False

    if not motifs or motifs == "None" or not motifs:
        active_motifs = []
        
    elif isinstance(motifs, str):
        active_motifs = [motifs]
    else:
        active_motifs = motifs
        
    if Descriptors.MolWt(mol) > max_mw: return False
    
    # 🧪 ADVANCED CHEMICAL LOGIC MAP
    pattern_map = {
        # --- ENFORCE BIODEGRADABLE MOTIFS ---
        "Esters (Biodegradable)": "[CX3](=[OX1])O[#6]",
        "Carbonates (Biodegradable)": "[OX2]C(=O)[OX2]",  # Faster
        "PLA-type (Lactide/Lactic)": "C(C)C(=O)O", 
        "Anhydrides (Fast Degrading)": "C(=O)OC(=O)",
        
        "Amides (High Strength)": "[NX3][CX3](=[OX1])[#6]",
        "Sulfones (Thermal Stability)": "S(=O)(=O)",
        "Urethanes": "[NX3]C(=O)O",
        
        "Aryl-Monomer": "a", 
        "Aliphatic-Chain": "CCCC", 
        "Cyclic-Aliphatic": "C1CCCCC1"
    }
    
    for m in active_motifs:
        if m in pattern_map:
            p = Chem.MolFromSmarts(pattern_map[m])
            if not mol.HasSubstructMatch(p):
                return False
            
    if limit_aromatic:
        aromatic_atoms = len(mol.GetAromaticAtoms())
        total_atoms = mol.GetNumAtoms()
        # Requirement: Max 20% of atoms can be aromatic
        if total_atoms > 0 and (aromatic_atoms / total_atoms) > 0.35:
            return False
            
    return True

def run_sensitivity_analysis(raw_df, t_s, t_t, t_d):
    # 1. Handle empty or None state safely
    if raw_df is None or (isinstance(raw_df, pd.DataFrame) and raw_df.empty):
        return None, "No data to analyze. Run the 'Design Pipeline' first.", "### <span style='color:grey'>0% Pass Rate</span>"

    # IMPORTANT: Ensure we are looking at the dataframe with hidden columns
    # If the user passed us a dataframe that was already filtered, this will fail.
    required_cols = ["_s", "_t", "Degradability_raw"]
    if not all(col in raw_df.columns for col in required_cols):
        return None, "Error: Reference data lost. Please re-run the design pipeline.", ""

    processed_df = raw_df.copy()
    
    def reclassify(row):
        s, t, d = row["_s"], row["_t"], row["Degradability_raw"]
        failed = []
        if s < t_s: failed.append("Strength")
        if t < t_t: failed.append("Tg")
        if d < t_d: failed.append("Degradability")
        return "Excellent ✅" if not failed else f"Review (Failed {', '.join(failed)}) ⚠️" if len(failed) <= 2 else "Bad ❌"

    processed_df["Classification"] = processed_df.apply(reclassify, axis=1)
    
    total = len(processed_df)
    passed = len(processed_df[processed_df["Classification"] == "Excellent ✅"])
    pass_pct = (passed / total) * 100 if total > 0 else 0
    
    # Create the display version
    display_df = processed_df.copy()
    for i, row in display_df.iterrows():
        display_df.at[i, "Strength (GPa)"] = f"{'🟢' if row['_s'] >= t_s else '🔴'} {row['_s']:.3f}"
        display_df.at[i, "Tg (°C)"] = f"{'🟢' if row['_t'] >= t_t else '🔴'} {row['_t']:.1f}"
        display_df.at[i, "Degradability"] = f"{'🟢' if row['Degradability_raw'] >= t_d else '🔴'} {row['Degradability_raw']:.2f}"
    
    badge_color = "green" if pass_pct > 25 else "orange" if pass_pct > 5 else "red"
    annotated_data = f"### <span style='color:{badge_color}'>{pass_pct:.0f}% Pass Rate</span>"
    status_msg = f"**{pass_pct:.1f}%** of candidates ({passed}/{total}) meet these industrial benchmarks."

    # Return only visible columns to the UI
    visible = ["p-SMILES", "Strength (GPa)", "Tg (°C)", "Degradability", "SA Score", "Confidence", "Classification"]
    # Check if any extra columns like 'Rank Score' exist and should be kept
    if "Rank Score" in display_df.columns:
        visible.append("Rank Score")
        
    return display_df[visible], status_msg, annotated_data
    
def generate_exports(df):
    if df is None or df.empty:
        return [None] * 4
    
    # 1. CSV, JSON, and SMILES (These handle Unicode fine)
    csv_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
    df.to_csv(csv_path, index=False)
    
    json_path = tempfile.NamedTemporaryFile(delete=False, suffix=".json").name
    df.to_json(json_path, orient="records", indent=4)
    
    smiles_path = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
    with open(smiles_path, "w") as f:
        f.write("\n".join(df["p-SMILES"].tolist()))
        
    # 2. PDF (The sensitive one)
    pdf_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf").name
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", 'B', 16)
    pdf.cell(200, 10, txt="PolymerDesignAI: Candidate Report", ln=True, align='C')
    pdf.ln(10)
    
    # Helper to strip ANY emoji or non-latin-1 char
    def clean_text(text):
        if not isinstance(text, str): text = str(text)
        # Keeps only standard keyboard characters
        return re.sub(r'[^\x00-\x7F]+', '', text).strip()

    for i, row in df.head(20).iterrows():
        pdf.set_font("Arial", 'B', 11)
        # Strip icons from Classification
        pdf.cell(0, 10, txt=f"Rank {i+1}: {clean_text(row['Classification'])}", ln=True)
        
        pdf.set_font("Arial", size=9)
        pdf.multi_cell(0, 5, txt=f"SMILES: {row['p-SMILES']}")
        
        # Strip icons from Strength, Tg, and Degradability
        strength_txt = clean_text(row['Strength (GPa)'])
        tg_txt = clean_text(row['Tg (°C)'])
        degrad_txt = clean_text(row['Degradability'])
        
        pdf.cell(0, 5, txt=f"Strength: {strength_txt} | Tg: {tg_txt}", ln=True)
        pdf.cell(0, 5, txt=f"Degradability: {degrad_txt} | SA Score: {clean_text(row['SA Score'])}", ln=True)
        pdf.ln(5)
        pdf.line(10, pdf.get_y(), 200, pdf.get_y())
        pdf.ln(2)
        
    pdf.output(pdf_path)
    return csv_path, json_path, smiles_path, pdf_path

def on_refine_wrapper(seed_with_icon, num, mode, view, motifs, limit_aro, max_mw):
    df, fig, dropdown, state, img, bar = on_refine_start(seed_with_icon, num, mode, view, motifs, limit_aro, max_mw)
    
    # Same here: Change to string
    init_msg = "Move sliders above to re-analyze refined batch."
    init_anno = "### <span style='color:grey'>0% Pass Rate</span>"
    
    return df, fig, dropdown, state, img, bar, init_msg, init_anno
    
# --- 6. GRADIO INTERFACE ---
with gr.Blocks(title="PolymerDesignAI") as demo:
    gr.Markdown("<center><h1>🧪 PolymerDesignAI: Advanced Scientific Discovery</h1></center>")

    with gr.Tab("Autonomous Design"):
        gr.Markdown("### Integrated Generation + Screening")
        
        # --- ROW 1: Optimization Mode & Strategy side-by-side ---
        with gr.Row(variant="panel"):
            app_mode = gr.Dropdown(
                ["Battery Mode", "Biomedical Mode", "Structural Mode"], 
                value="Biomedical Mode", 
                label="Select Optimization Mode"
            )
            toggle_view = gr.Radio(
                ["Single-objective ranking", "Multi-objective Pareto optimal (Best-in-Class candidates)"], 
                value="Single-objective ranking", 
                label="Selection Strategy"
            )

        # --- ROW 2: Three Text Blocks side-by-side ---
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("""**Application Mapping:**
                * **Structural Mode**: Optimized Yield Strength & High Thermal Stability.
                * **Battery Mode**: Optimized Glass Transition (Tg) & Conductivity.
                * **Biomedical Mode**: Optimized Degradability & Biocompatibility.""")
            with gr.Column(scale=1):
                gr.Markdown("""**📊 Classification Logic:**
                * **✅ Excellent**: Passes all **4** benchmarks.
                * **⚠️ Review**: Passes  **2 or 3** benchmarks.
                * **❌ Bad**: Passes **0 or 1** benchmark.""")
            with gr.Column(scale=1):
                gr.Markdown("""**🧪 Logic Benchmarks:**
                * Benchmarks and thresholds auto-adjust per selected Mode.
                * Incorporates Uncertainty (±) for scientific trust.""")

        # --- ROW 3: Quantity + Constraints + Filters all in one line ---
        with gr.Row(variant="compact", equal_height=True):
            with gr.Column(scale=1):
                num = gr.Slider(1, 50, 10, step=1, label="Quantity to Generate")
            
            with gr.Column(scale=1):
                must_have = gr.Dropdown(
                    ["None", "Esters (Biodegradable)", "Carbonates (Biodegradable)", 
                     "PLA-type (Lactide/Lactic)", "Anhydrides (Fast Degrading)",
                     "Amides (High Strength)", "Sulfones (Thermal Stability)", 
                     "Urethanes", "Aryl-Monomer", "Aliphatic-Chain"], 
                    value="None", label="🔒 Lock Functional Groups", interactive=True
                )
            
            with gr.Column(scale=1):
                exclude_aromatic = gr.Checkbox(label="Limit Aromatic (<20%)")
                max_mol_wt = gr.Slider(100, 1200, 500, label="Max Monomer MW")

        # --- ROW 4: Centered Run Button ---
        # We use scale=1/2/1 to force the middle column (and button) to stay centered
        with gr.Row():
            gr.Markdown("<br>") # Top Spacer
        with gr.Row():
            with gr.Column(scale=1): pass 
            run = gr.Button("🚀 Run Design Pipeline", variant="primary", size="lg", scale=2)
            with gr.Column(scale=1): pass
        with gr.Row():
            gr.Markdown("<br>") # Bottom Spacer
                            
        pareto_plot = gr.Plot(label="Pareto Frontier")
        
        gr.Markdown("### 📊 AI-Generated Candidates")
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("#### Filter by Status")
                filter_btn_all = gr.Button("Show All", size="sm")
                filter_btn_exc = gr.Button("Show Only Excellent ✅", size="sm")
                filter_btn_rev = gr.Button("Show Only Review ⚠️", size="sm")
            with gr.Column(scale=2):
                gr.Markdown("#### Quick Sort")
                sort_dropdown = gr.Dropdown(
                    ["Highest Strength", "Highest Tg", "Best Degradability", "Best SA (Easiest to make)", "Best Universal Score"],
                    label="Sort Candidates By",
                    interactive=True
                )
            # --- NEW: BENCHMARK SENSITIVITY ANALYSIS ---
        with gr.Accordion("📊 Benchmark Sensitivity Analysis (Industrial What-if)", open=False):
            gr.Markdown("Adjust thresholds to dynamically re-classify candidates and see market readiness.")
            with gr.Row():
                threshold_s = gr.Slider(0, 0.2, value=0.08, step=0.005, label="Min Strength (GPa)")
                threshold_t = gr.Slider(0, 300, value=150, step=5, label="Min Tg (°C)")
                threshold_d = gr.Slider(0, 1.0, value=0.10, step=0.05, label="Min Degradability")
            
            with gr.Row():
                pass_metric = gr.Markdown("### Market Readiness")
                stats_box = gr.Markdown("### 📈 Impact Summary")
        # --------------------------------------------
        results_df = gr.DataFrame(label="Candidate Table (Click headers to sort)")
        full_results_state = gr.State()
        
        with gr.Accordion("⚙️ Adjust Ranking Weights (Manual Override)", open=False):
            gr.Markdown("Adjust these sliders to change how candidates are ranked in the table below.")
            with gr.Row():
                w_strength = gr.Slider(0, 100, 40, step=5, label="Strength Weight (%)")
                w_tg = gr.Slider(0, 100, 30, step=5, label="Tg Weight (%)")
                w_degrad = gr.Slider(0, 100, 20, step=5, label="Degradability Weight (%)")
                w_sa = gr.Slider(0, 100, 10, step=5, label="SA Ease Weight (%)")

        gr.Markdown(r"""**Weighted Rank Score Formula:** $$\text{Rank Score} = (S_{norm} \times W_S) + (T_{norm} \times W_T) + (D_{norm} \times W_D) + (SA_{penalty} \times W_{SA})$$""")
        
        rank_btn = gr.Button("Calculate Custom Weighted Score", variant="secondary")

        gr.Markdown("### 🕵️ Candidate Deep-Dive & Explainability")
        with gr.Row():
            # Column 1: Selection & Radar
            with gr.Column(scale=3):
                selection_dropdown = gr.Dropdown(label="Select a SMILES to inspect", choices=[])
                with gr.Row():
                    refine_btn = gr.Button("🧬 Refine Around This Candidate", variant="secondary")
                    
                radar_plot = gr.Plot(label="Performance Fingerprint")
                
                with gr.Accordion("📖 How to read the Radar Plot", open=False):
                    gr.Markdown("""
                    * **Colored Area**: Selected polymer performance. 
                    * **Dashed Grey Area**: Industry Benchmark targets.
                    * **Color Coding**: Green (Excellent), Yellow (Review), Red (Bad).
                    * **Scales**: Strength (0.15 GPa), Tg (300°C), Degradability (0-1).
                    """)

            # Column 2: Structural Explainability (Point 1 of Phase 3)
            with gr.Column(scale=3):
                gr.Markdown("#### 🔍 Why This Molecule? (AI Reasoning)")
                explain_mode = gr.Radio(
                    ["Strength", "Tg", "Degradability"], 
                    label="Select Property to Explain", 
                    value="Strength"
                )
                explain_img = gr.Image(label="Substructure Contributions (Highlighted)", type="pil")
                explain_bar = gr.Plot(label="Feature Impact Analysis (SHAP)")

        # --- EXPORT HUB ---
        gr.Markdown("### 📥 Export Results")
        with gr.Row(variant="panel"):
            download_csv = gr.File(label="CSV Spreadsheet")
            download_json = gr.File(label="JSON Data")
            download_smiles = gr.File(label="SMILES List (.txt)")
            download_pdf = gr.File(label="Ranked Report (PDF)")
        
        # Helper to trigger all exports at once when results change
        def update_exports(df):
            return generate_exports(df)
        
        def on_gen_complete(num, mode, view, motifs, limit_aro, max_mw, progress=gr.Progress()):
            progress(0, desc="🔍 Initializing Generative AI...")
            full_df, display_df, csv_path, fig = generate_and_screen(num, mode, view, motifs, limit_aro, max_mw)
            
            def get_dropdown_icon(status_text):
                if "Excellent" in status_text: return "🟢"
                if "Review" in status_text: return "🟡"
                return "🔴"
        
            dropdown_choices = [
                f"{get_dropdown_icon(row['Classification'])} {row['p-SMILES']}" 
                for _, row in full_df.iterrows()
            ]
            first_smiles = full_df["p-SMILES"].iloc[0] if not full_df.empty else None
            first_choice = dropdown_choices[0] if dropdown_choices else None
            
            img, bar = explain_polymer(first_smiles, "Strength")
            
            return (display_df, fig, 
                    gr.Dropdown(choices=dropdown_choices, value=first_choice), full_df, img, bar)

        def on_refine_start(seed_with_icon, num, mode, view, motifs, limit_aro, max_mw):
            if not seed_with_icon:
                raise gr.Error("Please select a candidate from the dropdown first!")
        
            seed = seed_with_icon.split(" ")[1] if " " in seed_with_icon else seed_with_icon
    
            gr.Info(f"Evolving molecules similar to: {seed[:15]}... while respecting constraints")
                
            full_df, display_df, csv_path, fig = generate_and_screen(num, mode, view, motifs=motifs, limit_aro=limit_aro, max_mw=max_mw, seed=seed)
        
            # Use your new icon-logic for the refined dropdown results too
            def get_dropdown_icon(status_text):
                if "Excellent" in status_text: return "🟢"
                if "Review" in status_text: return "🟡"
                return "🔴"
        
            dropdown_choices = [
                f"{get_dropdown_icon(row['Classification'])} {row['p-SMILES']}" 
                for _, row in full_df.iterrows()
            ]
            
            first_smiles = full_df["p-SMILES"].iloc[0] if not full_df.empty else None
            first_choice = dropdown_choices[0] if dropdown_choices else None
            
            img, bar = explain_polymer(first_smiles, "Strength")
            
            return (display_df, fig, gr.Dropdown(choices=dropdown_choices, value=first_choice), full_df, img, bar)
            
        def filter_table(df, mode):
            if df is None or df.empty:
                return None
        
            if mode == "Excellent":
                filtered = df[df["Classification"].str.contains("Excellent")]
            elif mode == "Review":
                filtered = df[df["Classification"].str.contains("Review")]
            else:
                filtered = df
        
            visible_cols = [
                "p-SMILES",
                "Strength (GPa)",
                "Tg (°C)",
                "Degradability",
                "SA Score",
                "Confidence",
                "Classification"
            ]
        
            if "Rank Score" in filtered.columns:
                visible_cols.append("Rank Score")
        
            return filtered[visible_cols]

        def clean_and_update_radar(choice, df):
        # Split by space and take the second part (the actual SMILES)
            if not choice or df is None or df.empty:
                return None
            smiles = choice.split(" ")[1] if choice and " " in choice else choice
            return update_radar_plot(smiles, df)
    
        def clean_and_explain(choice, prop):
            smiles = choice.split(" ")[1] if choice and " " in choice else choice
            return explain_polymer(smiles, prop)    

        def on_gen_complete_wrapper(num, mode, view, motifs, limit_aro, max_mw, progress=gr.Progress()):
            df, fig, dropdown, state, img, bar = on_gen_complete(num, mode, view, motifs, limit_aro, max_mw, progress)
            
            # Change the list [("0% Pass Rate", "Awaiting")] to a simple string
            init_msg = "### 📈 Impact Summary\nMove sliders above to see analysis."
            init_anno = "### <span style='color:grey'>0% Pass Rate</span>"
            
            return df, fig, dropdown, state, img, bar, init_msg, init_anno
            
        run.click(fn=on_gen_complete_wrapper, inputs=[num, app_mode, toggle_view, must_have, exclude_aromatic, max_mol_wt], outputs=[results_df, pareto_plot, selection_dropdown, full_results_state, explain_img, explain_bar, stats_box, pass_metric])

        filter_btn_all.click(lambda x: filter_table(x, "All"), full_results_state, results_df)
        filter_btn_exc.click(lambda x: filter_table(x, "Excellent"), full_results_state, results_df)
        filter_btn_rev.click(lambda x: filter_table(x, "Review"), full_results_state, results_df)
        
        # 2. When a user picks a SMILES from the dropdown, update the Radar Plot
        selection_dropdown.change(
            clean_and_update_radar, 
            [selection_dropdown, results_df], 
            radar_plot
        )
        
        selection_dropdown.change(
            clean_and_explain, 
            [selection_dropdown, explain_mode], 
            [explain_img, explain_bar]
        )

        explain_mode.change(
            clean_and_explain, 
            [selection_dropdown, explain_mode], 
            [explain_img, explain_bar]
        )

        sort_dropdown.change(sort_table, [full_results_state, sort_dropdown], results_df)
        rank_btn.click(calculate_weighted_rank, inputs=[full_results_state, app_mode, w_strength, w_tg, w_degrad, w_sa], outputs=[full_results_state, results_df, download_csv])
        refine_btn.click(
            fn=on_refine_wrapper,  # Points to the new 8-output wrapper
            inputs=[selection_dropdown, num, app_mode, toggle_view, must_have, exclude_aromatic, max_mol_wt], 
            outputs=[results_df, pareto_plot, selection_dropdown, full_results_state, explain_img, explain_bar, stats_box, pass_metric]
        )
        
        # Sensitivity Triggers
        slider_inputs = [full_results_state, threshold_s, threshold_t, threshold_d]
        slider_outputs = [results_df, stats_box, pass_metric]

        threshold_s.change(fn=run_sensitivity_analysis, inputs=slider_inputs, outputs=slider_outputs)
        threshold_t.change(fn=run_sensitivity_analysis, inputs=slider_inputs, outputs=slider_outputs)
        threshold_d.change(fn=run_sensitivity_analysis, inputs=slider_inputs, outputs=slider_outputs)
        
            
        # Automatically update all download links whenever the results table changes
        results_df.change(
            fn=update_exports, 
            inputs=[results_df], 
            outputs=[download_csv, download_json, download_smiles, download_pdf]
        )
        
    # --- Inside the "Predict" Tab ---
    with gr.Tab("Predict"):
        gr.Markdown("### Forward Property Prediction")
        with gr.Row():
            with gr.Column():
                pred_mode = gr.Radio(["Structural (Strength)", "Battery (Tg)", "Biomedical (Degradability)"], 
                                     value="Structural (Strength)", label="Select Property")
                sm_in = gr.Textbox(label="Enter p-SMILES", placeholder="*CC*")
                gr.Examples(
                    examples=[
                        ["*CCCCCCNC(=O)NCCCCCCNC(=O)S*"], 
                        ["*Oc1ccc(C(C)(C)c2ccc(OC(=O)NCCCCCCNC(*)=O)cc2)cc1"], 
                        ["*[Se]c1ccc(C(=O)Nc2ccc(NC(=O)CCCCCCCC(=O)Nc3ccc(*)cc3)cc2)cc1"]
                    ], 
                    inputs=sm_in, 
                    label="Quick Select Examples"
                )
                pred_btn = gr.Button("Predict Single", variant="primary")
                res_label = gr.Label(label="Result")
                
            with gr.Column():
                gr.Markdown("#### Batch Processing")
                file_in = gr.File(label="Upload CSV/XLSX (Must contain a 'smiles' column)")
                batch_btn = gr.Button("Run Batch Prediction", variant="primary")
                
                gr.Markdown("---")
                gr.Markdown("#### Batch Results")
                batch_df = gr.DataFrame(label="Preview (Top 10)")
                # This is the new download component for the Predict tab
                batch_file_out = gr.File(label="Download Predicted Results (CSV)")

        # Logic for Single Prediction
        pred_btn.click(predict_ui, [pred_mode, sm_in], res_label)
        
        # Logic for Batch Prediction - mapping the return values to the dataframe and the file component
        batch_btn.click(
            fn=batch_predict, 
            inputs=[pred_mode, file_in], 
            outputs=[batch_df, batch_file_out]
        )
# Launch the app with authentication
demo.launch(auth=(os.environ.get("app_user"), os.environ.get("app_pass")), ssr_mode=False)