import json
import os
import numpy as np
import pandas as pd
import gradio as gr
import matplotlib.pyplot as plt

DATA_DIR = "data"
VARIANTS = ["WT", "var1", "var2", "var3", "var4"]

# SWI weights from the paper context [cite: 52]
swi_weights = {'A': 0.835, 'C': 0.111, 'D': 0.921, 'E': 0.963, 'F': 0.354, 
               'G': 0.732, 'H': 0.654, 'I': 0.443, 'K': 0.901, 'L': 0.512, 
               'M': 0.456, 'N': 0.812, 'P': 0.678, 'Q': 0.854, 'R': 0.887, 
               'S': 0.776, 'T': 0.798, 'V': 0.567, 'W': 0.312, 'Y': 0.412}

sequences_dict = {
    "WT": "EVKLVESGGDLVKPGGSLKLSCAASGFTFSSYGMSWVRQTPDKRLEWVATISRGGSYTYYPDSVKGRFTISRDNAKNTLYLQMSSLKSEDTAMYYCARRETYDEKGFAYWGQGTTVTVSSGGGGSGGGGSGGGGSDIELTQSPSSLTVTAGEKVTMSCKSSQSLLNSGNQKNYLTWYQQKPGQPPKLLIYWASTRESGVPDRFTGSGSGRDFTLTISSVQAEDLAVYYCQNDNSHPLTFGAGTKLELKR",
    "var1": "EVKLTESGGKTVKPGGSVTLSCQVSGIDFSSYGMSWVRQTPEDKLEWVATISRGGSYTYIPERYKGRVTISKDFSKSTLYLNMKNLKPSDTATYYCAIRETYDEKGFAHFGQGTKLTVTDDGPDPGSPDPSDPSNDIKLTQSPSEITVNEGEKVTITCNSSQSLLNSGNQKNYLTWYKKKPGEPPQLLIYWASTRESGVSSRFVGSGSGTSFTLTISSFKKEDVATYYCQNDNSHPLTFGQGTKLSLKP",
    "var2": "EVKLVESGGATVAPGGSVTLSCKASGFDFSSYGMSWVAQTPEGQLRWVATISRGGSYTYVPEAYAGRVTISRDFAASTAYLHMRDLQPSDTATYYCAVRETYDEKGFAHFGQGTQLTVADPGAAPGAPEPALPSDDIQLTQSPSEVTANLGETVTFTCTSSKSLLNSGNQKNYLTWYKQAPGQPPQLLIYWASTREPGVSSRFVGSGSGTNFTLTISSLQESDVATYYCQNDNSHPLTFGSGTKLSLKK",
    "var3": "MVQLVESGGATVAPGGSVTLSCKVSGIDFSSYGMSWVAQTPDNQLRWVATISRGGSYTYIPKSYAGRFTISVDHSKSTAYLNMSNLQPSDTATYYCAIRETYDEKGFAHFGQGTQLTVADPGAAPGAPLEDDPSNDIQLTQSPSEVTVNLGEKVTLTCNSSKSLLNSGNQKNYLTWYQQKPGEPPKLLIYWASTREEGVSSRFVGSGSGTSFTLTISSFQEEDVATYYCQNDNSHPLTFGQGTKLTLKP",
    "var4": "SVSLTESGGATVAPGGSVTLSCKASGFDFSSYGMSWVAQTPEGRLLWVATISRGGSYTYVPSRYAGRFTISRDDSKSTLYLNMSNLQPSDTATYYCAVRETYDEKGFAHFGEGTQLTVSDAGEAPGGAEPADPSPDIVLTQSPSSVTVNLGEKVTLTCNSSKSLLNSGNQKNYLTWYKQAPGQPPQLLIYWASTRESGVSSRFVGSGSGTSFTLTITSFQAEDVATYYCQNDNSHPLTFGSGTKLSLKP"
}

# Pre-load WT PAE for difference calculations
def get_wt_pae():
    path = os.path.join(DATA_DIR, "WT", "pae.json")
    with open(path) as f:
        data = json.load(f)
        return np.array(data.get("pae", data.get("predicted_aligned_error")))

WT_PAE_DATA = get_wt_pae()

def calculate_swi(seq):
    return round(sum(swi_weights.get(aa, 0.5) for aa in seq) / len(seq), 3)

def calculate_net_charge(seq):
    return seq.count('R') + seq.count('K') - seq.count('D') - seq.count('E')

def calculate_interface_pae(pae, h_range=(0, 115), l_range=(140, 250)):
    block1 = pae[h_range[0]:h_range[1], l_range[0]:l_range[1]]
    block2 = pae[l_range[0]:l_range[1], h_range[0]:h_range[1]]
    return round((np.mean(block1) + np.mean(block2)) / 2, 3)

def load_variant_data(variant):
    base = os.path.join(DATA_DIR, variant)
    with open(os.path.join(base, "scores.json")) as f:
        scores = json.load(f)
    with open(os.path.join(base, "pae.json")) as f:
        pae_data = json.load(f)
        pae = np.array(pae_data.get("pae", pae_data.get("predicted_aligned_error")))
    pdb_path = os.path.join(base, "structure.pdb")
    avg_plddt = np.mean(scores["plddt"]) if isinstance(scores["plddt"], list) else scores["plddt"]
    ptm = np.mean(scores.get("ptm", 0)) if isinstance(scores.get("ptm"), list) else scores.get("ptm", 0)
    return avg_plddt, ptm, pae, pdb_path

def view_structure(pdb_path):
    with open(pdb_path) as f:
        pdb_data = f.read()
    viewer_id = "mol_view_" + str(hash(pdb_path))
    html = f"""
    <script src="https://3dmol.org/build/3Dmol-min.js"></script>
    <div id="{viewer_id}" style="height: 500px; width: 100%;"></div>
    <script>
        var viewer = $3Dmol.createViewer("{viewer_id}", {{}});
        viewer.addModel(`{pdb_data}`, "pdb");
        viewer.setStyle({{'cartoon': {{'colorscheme': {{'prop':'b','gradient': 'roygb','min':50,'max':90}}}}}});
        viewer.zoomTo();
        viewer.render();
    </script>
    """
    return html

def update_variant(variant):
    avg_plddt, ptm, pae, pdb_path = load_variant_data(variant)
    i_pae = calculate_interface_pae(pae)
    seq = sequences_dict.get(variant, "")
    swi_val = calculate_swi(seq)
    charge_val = calculate_net_charge(seq)
    pos82 = seq[81] if len(seq) > 81 else "N/A"

    metrics = pd.DataFrame({
        "Metric": ["Avg pLDDT", "pTM", "Interface PAE", "SWI", "Net Charge", "Pos 82 Residue"],
        "Value": [round(float(avg_plddt), 2), round(float(ptm), 3), i_pae, swi_val, charge_val, pos82]
    })

    # Generate Difference PAE plot
    fig, ax = plt.subplots(figsize=(5, 5))
    if variant == "WT":
        im = ax.imshow(pae, cmap="viridis", vmin=0, vmax=30)
        ax.set_title("WT: Absolute PAE Map")
        cbar_label = "Expected Error (Å)"
    else:
        # Subtract WT PAE to show relative changes 
        diff_pae = pae - WT_PAE_DATA
        im = ax.imshow(diff_pae, cmap="bwr", vmin=-5, vmax=5)
        ax.set_title(f"{variant}: ΔPAE (Relative to WT)")
        cbar_label = "Error Difference (Å)"
    
    fig.colorbar(im, ax=ax, label=cbar_label)
    ax.set_xlabel("Residue Index")
    ax.set_ylabel("Residue Index")
    
    structure_html = view_structure(pdb_path)
    return metrics, fig, structure_html

with gr.Blocks() as demo:
    gr.Markdown(
        """
        <center>
        <h1>AI-Assisted Intrabody Design Explorer</h1>
        <p>This demo visualizes precomputed AlphaFold2 results for WT and ProteinMPNN-designed scFv variants, following the methodology  described in the original paper (Galindo et al. (2026))</p>
        </center>
        """
        )
    variant = gr.Dropdown(VARIANTS, value="WT", label="Select Design Variant")
    
    with gr.Row():
        metrics_table = gr.Dataframe(label="Structural & Solubility Metrics")
        with gr.Column():
            pae_plot = gr.Plot(label="PAE Analysis")
            gr.Markdown("""
            ### How to read the ΔPAE Map:
            * **🟦 Blue Regions**: **Improved stability.** Error is lower than WT. 
            * **⬜ White Regions**: **No significant change.** Structural confidence is identical to WT. 
            * **🟥 Red Regions**: **Reduced stability.** Error is higher than WT (potential distortion). 
            """)
    
    structure_view = gr.HTML(label="3D Structure (Confidence Colored)")

    variant.change(update_variant, inputs=variant, outputs=[metrics_table, pae_plot, structure_view])
    demo.load(update_variant, inputs=variant, outputs=[metrics_table, pae_plot, structure_view])

demo.launch(
    auth=(os.environ.get("app_user"), os.environ.get("app_pass")),
    ssr_mode=False
)