import os, sys, subprocess, shutil, tempfile, json, secrets, copy
from pathlib import Path

os.environ["DGLBACKEND"] = "pytorch"

# Install protenix + its deps, holding torch/numpy/DGL pins in place
def _install_protenix():
    try:
        import protenix  # noqa: F401
        # Need 2.0.0: it's the only version that has the 'protenix-v2' checkpoint name.
        # Also verify the runner top-level package is present.
        assert getattr(protenix, "__version__", "") == "2.0.0", \
            f"need 2.0.0, have {getattr(protenix, '__version__', 'unknown')}"
        from runner.batch_inference import protenix_cli  # noqa: F401
        from ml_collections.config_dict import ConfigDict  # noqa: F401
        from biotite.structure import AtomArray  # noqa: F401
        import biotite.structure.io.pdbx as _pdbx
        assert hasattr(_pdbx.convert, "PDBX_BOND_TYPE_ID_TO_TYPE"), "biotite<1.0 installed, need >=1.0"
        return  # all required modules present
    except (ImportError, AssertionError) as e:
        print(f"[protenix-install] guard failed ({e}), installing...")

    print("Installing protenix==2.0.0 from PyPI...")
    # 2.0.0 declares Requires-Python>=3.11 but is pure-Python and runs fine on 3.10.
    # --ignore-requires-python bypasses that metadata check.
    # --no-deps avoids pulling torch==2.7.1 which conflicts with the DGL pin.
    r = subprocess.run(
        [sys.executable, "-m", "pip", "install", "protenix==2.0.0",
         "--no-deps", "--ignore-requires-python", "-q"],
        check=False, capture_output=True, text=True,
    )
    if r.returncode != 0:
        print(f"[protenix-install] install failed: {r.stderr[-500:]}")
    else:
        print("[protenix-install] protenix==2.0.0 OK")
    # Install runtime deps one-by-one; skip any that fail so a single
    # bad package can't abort the whole setup.
    runtime_deps = [
        "ml-collections==1.1.0",
        "biotite>=1.0,<2",   # 1.x has PDBX_BOND_TYPE_ID_TO_TYPE; requires numpy>=1.25 only
        "gemmi==0.6.7",
        "modelcif==1.4",
        "optree==0.17.0",
        "fair-esm==2.0.0",
        "wandb==0.21.1",
        "protobuf==6.31.1",
        "tqdm",
        "ipdb",
        "pdbeccdutils==1.0.0",
        "scikit-learn-extra==0.3.0",
        "rdkit",
        "ninja",   # required by torch cpp_extension.load for FusedLayerNorm JIT compile
    ]
    for dep in runtime_deps:
        r = subprocess.run(
            [sys.executable, "-m", "pip", "install", dep, "-q"],
            check=False, capture_output=True, text=True,
        )
        if r.returncode != 0:
            print(f"[protenix-install] FAILED {dep}: {r.stderr[-300:]}")
        else:
            print(f"[protenix-install] OK {dep}")
    print("protenix install complete.")

_install_protenix()


if not os.path.exists('./SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'):
    print('Downloading model weights 1')
    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt')
    print('Successfully Downloaded')

if not os.path.exists('./SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'):
    print('Downloading model weights 2')
    os.system('wget http://files.ipd.uw.edu/pub/sequence_diffusion/checkpoints/SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt')
    print('Successfully Downloaded')

import numpy as np
import gradio as gr
import py3Dmol
from io import StringIO
import matplotlib.pyplot as plt
from utils.sampler import HuggingFace_sampler
from utils.parsers_inference import parse_pdb
from model.util import writepdb
from utils.inpainting_util import *


plt.rcParams.update({'font.size': 13})

with open('./tmp/args.json','r') as f:
    args = json.load(f)

args['checkpoint'] = None
args['dump_trb'] = False
args['dump_args'] = True
args['save_best_plddt'] = True
args['T'] = 25
args['strand_bias'] = 0.0
args['loop_bias'] = 0.0
args['helix_bias'] = 0.0


# ── shared helpers ────────────────────────────────────────────────────────────

def _get_file_path(file_obj):
    if isinstance(file_obj, str):
        return file_obj
    return file_obj.name

def _iframe(html_body):
    x = f'<!DOCTYPE html><html></center> {html_body} </center></html>'
    return (
        f"""<iframe height="500px" width="100%" name="result" """
        f"""allow="midi; geolocation; microphone; camera; display-capture; encrypted-media;" """
        f"""sandbox="allow-modals allow-forms allow-scripts allow-same-origin allow-popups """
        f"""allow-top-navigation-by-user-activation allow-downloads" """
        f"""allowfullscreen="" allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
    )


# ── sequence diffusion ────────────────────────────────────────────────────────

def protein_diffusion_model(sequence, seq_len, helix_bias, strand_bias, loop_bias,
                    secondary_structure, aa_bias, aa_bias_potential,
                    num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
                    contigs, pssm, seq_mask, str_mask, rewrite_pdb):

    dssp_checkpoint = './SEQDIFF_230205_dssp_hotspots_25mask_EQtasks_mod30.pt'
    og_checkpoint = './SEQDIFF_221219_equalTASKS_nostrSELFCOND_mod30.pt'

    model_args = copy.deepcopy(args)

    S = HuggingFace_sampler(args=model_args)

    S.out_prefix = './tmp/'+secrets.token_hex(nbytes=10).upper()

    S.args['checkpoint'] = None
    S.args['dump_trb'] = False
    S.args['dump_args'] = True
    S.args['save_best_plddt'] = True
    S.args['T'] = 20
    S.args['strand_bias'] = 0.0
    S.args['loop_bias'] = 0.0
    S.args['helix_bias'] = 0.0
    S.args['potentials'] = None
    S.args['potential_scale'] = None
    S.args['aa_composition'] = None

    alt_aa_dict = {'B':['D','N'],'J':['I','L'],'U':['C'],'Z':['E','Q'],'O':['K']}
    if sequence not in ['',None]:
        L = len(sequence)
        aa_seq = []
        for aa in sequence.upper():
            if aa in alt_aa_dict.keys():
                aa_seq.append(np.random.choice(alt_aa_dict[aa]))
            else:
                aa_seq.append(aa)

        S.args['sequence'] = aa_seq
    elif contigs not in ['',None]:
        S.args['contigs'] = [contigs]
    else:
        S.args['contigs'] = [f'{seq_len}']
        L = int(seq_len)

    if rewrite_pdb not in ['',None]:
        S.args['pdb'] = _get_file_path(rewrite_pdb)

    if seq_mask not in ['',None]:
        S.args['inpaint_seq'] = [seq_mask]
    if str_mask not in ['',None]:
        S.args['inpaint_str'] = [str_mask]

    if secondary_structure in ['',None]:
        secondary_structure = None
    else:
        secondary_structure = ''.join(['E' if x == 'S' else x for x in secondary_structure])
        if L < len(secondary_structure):
            secondary_structure = secondary_structure[:len(sequence)]
        elif L == len(secondary_structure):
            pass
        else:
            dseq = L - len(secondary_structure)
            secondary_structure += secondary_structure[-1]*dseq

    potential_list = []
    potential_bias_list = []

    if aa_bias not in ['',None]:
        potential_list.append('aa_bias')
        S.args['aa_composition'] = aa_bias
        if aa_bias_potential in ['',None]:
            aa_bias_potential = 3
        potential_bias_list.append(str(aa_bias_potential))

    if hydrophobic_target_score not in ['',None]:
        potential_list.append('hydrophobic')
        S.args['hydrophobic_score'] = float(hydrophobic_target_score)
        if hydrophobic_potential in ['',None]:
            hydrophobic_potential = 3
        potential_bias_list.append(str(hydrophobic_potential))

    if pssm not in ['',None]:
        potential_list.append('PSSM')
        potential_bias_list.append('5')
        S.args['PSSM'] = _get_file_path(pssm)

    if len(potential_list) > 0:
        S.args['potentials'] = ','.join(potential_list)
        S.args['potential_scale'] = ','.join(potential_bias_list)

    S.args['secondary_structure'] = secondary_structure
    S.args['helix_bias'] = helix_bias
    S.args['strand_bias'] = strand_bias
    S.args['loop_bias'] = loop_bias

    if num_steps in ['',None]:
        S.args['T'] = 20
    else:
        S.args['T'] = int(num_steps)

    if 'normal' in noise:
        S.args['sample_distribution'] = noise
        S.args['sample_distribution_gmm_means'] = [0]
        S.args['sample_distribution_gmm_variances'] = [1]
    elif 'gmm2' in noise:
        S.args['sample_distribution'] = noise
        S.args['sample_distribution_gmm_means'] = [-1,1]
        S.args['sample_distribution_gmm_variances'] = [1,1]
    elif 'gmm3' in noise:
        S.args['sample_distribution'] = noise
        S.args['sample_distribution_gmm_means'] = [-1,0,1]
        S.args['sample_distribution_gmm_variances'] = [1,1,1]

    if secondary_structure not in ['',None] or helix_bias+strand_bias+loop_bias > 0:
        S.args['checkpoint'] = dssp_checkpoint
        S.args['d_t1d'] = 29
        print('using dssp checkpoint')
    else:
        S.args['checkpoint'] = og_checkpoint
        S.args['d_t1d'] = 24
        print('using og checkpoint')

    for k,v in S.args.items():
        print(f"{k} --> {v}")

    S.model_init()
    S.diffuser_init()
    S.setup()

    plddt_data = []
    for j in range(S.max_t):
        print(f'on step {j}')
        output_seq, output_pdb, plddt = S.take_step_get_outputs(j)
        plddt_data.append(plddt)
        yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)

    output_seq, output_pdb, plddt = S.get_outputs()
    yield output_seq, output_pdb, display_pdb(output_pdb), get_plddt_plot(plddt_data, S.max_t)

def get_plddt_plot(plddt_data, max_t):
    x = [i+1 for i in range(len(plddt_data))]
    fig, ax = plt.subplots(figsize=(15,6))
    ax.plot(x,plddt_data,color='#661dbf', linewidth=3,marker='o')
    ax.set_xticks([i+1 for i in range(max_t)])
    ax.set_yticks([(i+1)/10 for i in range(10)])
    ax.set_ylim([0,1])
    ax.set_ylabel('model confidence (plddt)')
    ax.set_xlabel('diffusion steps (t)')
    return fig

def display_pdb(path_to_pdb):
    pdb = open(path_to_pdb, "r").read()
    view = py3Dmol.view(width=500, height=500)
    view.addModel(pdb, "pdb")
    view.setStyle({'model': -1}, {"cartoon": {'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':1}}})
    view.zoomTo()
    return _iframe(view._make_html().replace("'", '"'))


# MOTIF SCAFFOLDING
def get_motif_preview(pdb_id, contigs):
    input_pdb = fetch_pdb(pdb_id=pdb_id.lower())

    parse = parse_pdb(input_pdb)
    output_name = input_pdb

    pdb = open(output_name, "r").read()
    view = py3Dmol.view(width=500, height=500)
    view.addModel(pdb, "pdb")

    if contigs in ['',0]:
        contigs = ['0']
    else:
        contigs = [contigs]

    pdb_map = get_mappings(ContigMap(parse,contigs))
    roi = [x[1]-1 for x in pdb_map['con_ref_pdb_idx']]

    colormap = {0:'#D3D3D3', 1:'#F74CFF'}
    colors = {i+1: colormap[1] if i in roi else colormap[0] for i in range(parse['xyz'].shape[0])}
    view.setStyle({"cartoon": {"colorscheme": {"prop": "resi", "map": colors}}})
    view.zoomTo()
    return _iframe(view._make_html().replace("'", '"')), output_name

def fetch_pdb(pdb_id=None):
    if pdb_id is None or pdb_id == "":
        return None
    else:
        os.system(f"wget -qnc https://files.rcsb.org/view/{pdb_id}.pdb")
        return f"{pdb_id}.pdb"

# MSA AND PSSM GUIDANCE
def save_pssm(file_upload):
    filename = _get_file_path(file_upload)
    orig_name = file_upload.orig_name if hasattr(file_upload, 'orig_name') else filename
    if filename.split('.')[-1] in ['fasta', 'a3m']:
        return msa_to_pssm(file_upload)
    return filename

def msa_to_pssm(msa_file):
    aa_to_index = {'A': 0, 'R': 1, 'N': 2, 'D': 3, 'C': 4, 'Q': 5, 'E': 6, 'G': 7, 'H': 8, 'I': 9, 'L': 10,
                'K': 11, 'M': 12, 'F': 13, 'P': 14, 'S': 15, 'T': 16, 'W': 17, 'Y': 18, 'V': 19, 'X': 20, '-': 21}
    records = list(SeqIO.parse(_get_file_path(msa_file), "fasta"))

    assert len(records) >= 1, "MSA must contain more than one protein sequence."

    first_seq = str(records[0].seq)
    aligned_seqs = [first_seq]
    aligner = Align.PairwiseAligner()
    aligner.open_gap_score = -0.7
    aligner.extend_gap_score = -0.3
    for record in records[1:]:
        alignment = aligner.align(first_seq, str(record.seq))[0]
        alignment = alignment.format().split("\n")
        al1 = alignment[0]
        al2 = alignment[2]
        al1_fin = ""
        al2_fin = ""
        percent_gap = al2.count('-')/ len(al2)
        if percent_gap > 0.4:
            continue
        for i in range(len(al1)):
            if al1[i] != '-':
                al1_fin += al1[i]
                al2_fin += al2[i]
        aligned_seqs.append(str(al2_fin))
    aligned_seq_length = len(first_seq)
    matrix = np.zeros((22, aligned_seq_length))
    for seq in aligned_seqs:
        for i in range(aligned_seq_length):
            if i == len(seq):
                break
            amino_acid = seq[i]
            if amino_acid.upper() not in aa_to_index.keys():
                continue
            else:
                aa_index = aa_to_index[amino_acid.upper()]
            matrix[aa_index, i] += 1
    matrix /= len(aligned_seqs)
    print(len(aligned_seqs))
    matrix[20:,]=0

    msa_path = _get_file_path(msa_file)
    outdir = ".".join(msa_path.split('.')[:-1]) + ".csv"
    np.savetxt(outdir, matrix[:21,:].T, delimiter=",")
    return outdir

def get_pssm(fasta_msa, input_pssm):
    if input_pssm not in ['',None]:
        outdir = _get_file_path(input_pssm)
    else:
        outdir = save_pssm(fasta_msa)

    pssm = np.loadtxt(outdir, delimiter=",", dtype=float)
    fig, ax = plt.subplots(figsize=(15,6))
    plt.imshow(torch.permute(torch.tensor(pssm),(1,0)))

    return fig, outdir


def toggle_seq_input(choice):
    if choice == "protein length":
        return gr.Slider(visible=True, value=None), gr.Textbox(visible=False, value=None)
    elif choice == "custom sequence":
        return gr.Slider(visible=False, value=None), gr.Textbox(visible=True, value=None)

def toggle_secondary_structure(choice):
    if choice == "sliders":
        return gr.Slider(visible=True, value=None),gr.Slider(visible=True, value=None),gr.Slider(visible=True, value=None),gr.Textbox(visible=False, value=None)
    elif choice == "explicit":
        return gr.Slider(visible=False, value=None),gr.Slider(visible=False, value=None),gr.Slider(visible=False, value=None),gr.Textbox(visible=True, value=None)


# ── protenix structure prediction ─────────────────────────────────────────────

PROTENIX_MODEL = os.environ.get("PROTENIX_MODEL", "protenix-v2")
PROTENIX_TIMEOUT = int(os.environ.get("PROTENIX_TIMEOUT", "1200"))  # first run: JIT compile + weight download

# Checkpoint directory protenix uses: $PROTENIX_ROOT_DIR/checkpoint (default ~/checkpoint)
_PROTENIX_CKPT_DIR = Path(os.environ.get("PROTENIX_ROOT_DIR", Path.home())) / "checkpoint"

# Map of model_name → HF Hub repo/file that mirrors the ByteDance CDN
# (CDN protenix.tos-cn-beijing.volces.com is geo-blocked from HF Spaces)
_HF_WEIGHT_SOURCES = {
    "protenix-v2": ("TMF001/protenix-v2-weights", "protenix-v2.pt"),
}

def _ensure_protenix_weights(model_name: str) -> str | None:
    """Download model weights from HF Hub if not already present. Returns error str or None."""
    ckpt_path = _PROTENIX_CKPT_DIR / f"{model_name}.pt"
    if ckpt_path.exists():
        return None
    if model_name not in _HF_WEIGHT_SOURCES:
        return f"No HF Hub mirror known for model '{model_name}'; ByteDance CDN is geo-blocked from HF Spaces"
    repo_id, filename = _HF_WEIGHT_SOURCES[model_name]
    print(f"[protenix-weights] Downloading {filename} from {repo_id} …")
    try:
        from huggingface_hub import hf_hub_download
        _PROTENIX_CKPT_DIR.mkdir(parents=True, exist_ok=True)
        tmp = hf_hub_download(repo_id=repo_id, filename=filename)
        shutil.copy(tmp, ckpt_path)
        print(f"[protenix-weights] Saved to {ckpt_path}")
        return None
    except Exception as e:
        return f"Weight download failed: {e}"

# pLDDT confidence colour bands (AlphaFold convention, B-factor scale 0-100)
_PLDDT_COLORS = [
    (90, "#0053D6"),   # very high — dark blue
    (70, "#65CBF3"),   # confident — light blue
    (50, "#FFDB13"),   # low — yellow
    (0,  "#FF7D45"),   # very low — orange
]

def _plddt_color(score):
    for threshold, color in _PLDDT_COLORS:
        if score >= threshold:
            return color
    return _PLDDT_COLORS[-1][1]

def _extract_plddt(cif_path):
    """Return per-residue pLDDT scores (0-100) from a Protenix CIF output."""
    from Bio.PDB import MMCIFParser
    parser = MMCIFParser(QUIET=True)
    structure = parser.get_structure("pred", cif_path)
    scores = []
    for model in structure:
        for chain in model:
            for residue in chain:
                ca_atoms = [a for a in residue if a.get_name() == "CA"]
                if ca_atoms:
                    scores.append(ca_atoms[0].get_bfactor())
    return scores

def _plddt_bar_plot(scores):
    residues = list(range(1, len(scores) + 1))
    colors = [_plddt_color(s) for s in scores]
    fig, ax = plt.subplots(figsize=(max(8, len(scores) / 5), 4))
    ax.bar(residues, scores, color=colors, width=1.0, edgecolor="none")
    ax.set_xlim(0.5, len(scores) + 0.5)
    ax.set_ylim(0, 100)
    ax.set_xlabel("Residue")
    ax.set_ylabel("pLDDT")
    ax.set_title("Per-residue confidence (Protenix v2)")
    ax.axhline(90, color="#0053D6", linewidth=0.8, linestyle="--", alpha=0.6)
    ax.axhline(70, color="#65CBF3", linewidth=0.8, linestyle="--", alpha=0.6)
    ax.axhline(50, color="#FFDB13", linewidth=0.8, linestyle="--", alpha=0.6)
    from matplotlib.patches import Patch
    legend = [
        Patch(color="#0053D6", label="Very high (>90)"),
        Patch(color="#65CBF3", label="Confident (70-90)"),
        Patch(color="#FFDB13", label="Low (50-70)"),
        Patch(color="#FF7D45", label="Very low (<50)"),
    ]
    ax.legend(handles=legend, loc="lower right", fontsize=9)
    fig.tight_layout()
    return fig

def display_cif(cif_path):
    content = open(cif_path).read()
    view = py3Dmol.view(width=500, height=500)
    view.addModel(content, "cif")
    view.setStyle(
        {'model': -1},
        {"cartoon": {'colorscheme': {'prop': 'b', 'gradient': 'roygb', 'min': 50, 'max': 100}}},
    )
    view.zoomTo()
    return _iframe(view._make_html().replace("'", '"'))

def fold_with_protenix(sequence, msa_file, model_name):
    try:
        sequence = sequence.strip().upper()
        if not sequence:
            return None, None, None, "Error: sequence is empty"

        valid = set("ACDEFGHIKLMNPQRSTVWXY")
        bad = set(sequence) - valid
        if bad:
            return None, None, None, f"Error: invalid amino acid characters: {sorted(bad)}"

        job_id = secrets.token_hex(8)
        work_dir = Path(tempfile.mkdtemp(prefix=f"protenix_{job_id}_"))

        try:
            protein_chain = {"sequence": sequence, "count": 1}

            if msa_file is not None:
                msa_src = Path(_get_file_path(msa_file))
                msa_dest = work_dir / "input.a3m"
                shutil.copy(msa_src, msa_dest)
                protein_chain["unpairedMsaPath"] = str(msa_dest.resolve())

            input_data = [{
                "name": job_id,
                "sequences": [{"proteinChain": protein_chain}],
                "covalent_bonds": [],
            }]

            input_json = work_dir / "input.json"
            input_json.write_text(json.dumps(input_data, indent=2))

            output_dir = work_dir / "output"
            output_dir.mkdir()

            # Ensure model weights are present (ByteDance CDN is geo-blocked; use HF Hub mirror)
            weight_err = _ensure_protenix_weights(model_name)
            if weight_err:
                return None, None, None, f"Weight download error: {weight_err}"

            # Invoke protenix via Python directly — bypasses the entry-point binary
            # which is sometimes not created during startup installs.
            # Equivalent to: protenix pred -i <input> -o <output> -n <model>
            py_cmd = (
                "import sys; "
                f"sys.argv = ['protenix','pred','-i',{str(input_json)!r},"
                f"'-o',{str(output_dir)!r},'-n',{model_name!r}];\n"
                "from runner.batch_inference import protenix_cli; protenix_cli()"
            )
            result = subprocess.run(
                [sys.executable, "-c", py_cmd],
                capture_output=True,
                text=True,
                timeout=PROTENIX_TIMEOUT,
            )

            if result.returncode != 0:
                return None, None, None, (
                    f"Protenix failed (rc={result.returncode}):\n"
                    f"STDOUT: {result.stdout[-1000:]}\n"
                    f"STDERR: {result.stderr[-1500:]}"
                )

            cif_files = sorted(output_dir.rglob("*.cif")) or sorted(output_dir.rglob("*.pdb"))
            if not cif_files:
                return None, None, None, "Protenix produced no output structure file"

            best = cif_files[0]
            out_path = f"./tmp/{job_id}{best.suffix}"
            shutil.copy(best, out_path)

            try:
                plddt = _extract_plddt(out_path)
                plddt_fig = _plddt_bar_plot(plddt)
                mean_plddt = f"{np.mean(plddt):.1f}" if plddt else "N/A"
            except Exception as pe:
                plddt_fig = None
                mean_plddt = f"pLDDT parse error: {pe}"

            return display_cif(out_path), out_path, plddt_fig, f"Mean pLDDT: {mean_plddt}"

        except subprocess.TimeoutExpired:
            return None, None, None, f"Timed out after {PROTENIX_TIMEOUT}s"
        except Exception as e:
            import traceback
            return None, None, None, f"Error: {e}\n{traceback.format_exc()[-1000:]}"
        finally:
            shutil.rmtree(work_dir, ignore_errors=True)

    except Exception as e:
        import traceback
        return None, None, None, f"Outer error: {e}\n{traceback.format_exc()[-500:]}"


# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(theme=gr.themes.Soft(font="Arial")) as demo:

    gr.Markdown("# Protein Tools")

    with gr.Tabs():

        # ── Tab 1: Sequence Diffusion (original) ─────────────────────────────
        with gr.Tab("Sequence Diffusion"):

            with gr.Row():
                with gr.Column(min_width=500):
                    gr.Markdown("""
                            ## How does it work?\n
                            --- [PREPRINT](https://biorxiv.org/content/10.1101/2023.05.08.539766v1) ---
                            Protein sequence and structure co-generation is a long outstanding problem in the field of protein design. By implementing [ddpm](https://arxiv.org/abs/2006.11239) style diffusion over protein seqeuence space we generate protein sequence and structure pairs. Starting with [RoseTTAFold](https://www.science.org/doi/10.1126/science.abj8754), a protein structure prediction network, we finetuned it to predict sequence and structure given a partially noised sequence. By applying losses to both the predicted sequence and structure the model is forced to generate meaningful pairs. Diffusing in sequence space makes it easy to implement potentials to guide the diffusive process toward particular amino acid composition, net charge, and more! Furthermore, you can sample proteins from a family of sequences or even train a small sequence to function classifier to guide generation toward desired sequences.
                            ![fig1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/diffusion_landscape.png)

                            ## How to use it?\n
                            A user can either design a custom input sequence to diffuse from or specify a length below. To scaffold a sequence use the following format where X represent residues to diffuse: XXXXXXXXSCIENCESCIENCEXXXXXXXXXXXXXXXXXXX. You can even design a protein with your name XXXXXXXXXXXXNAMEHEREXXXXXXXXXXXXX!

                            ### Acknowledgements\n
                            Thank you to Simon Dürr and the Hugging Face team for setting us up with a community GPU grant!
                            """)

                gr.Markdown("""
                ## Model in Action
                ![gif1](http://files.ipd.uw.edu/pub/sequence_diffusion/figs/seqdiff_anim_720p.gif)
                """)

            with gr.Row(equal_height=False):
                with gr.Column():
                    with gr.Tabs():
                        with gr.Tab("Inputs"):
                            gr.Markdown("## INPUTS")
                            gr.Markdown("""#### Start Sequence
                                        Specify the protein length for complete unconditional generation, or scaffold a motif (or your name) using the custom sequence input""")
                            seq_opt = gr.Radio(["protein length","custom sequence"], label="How would you like to specify the starting sequence?", value='protein length')

                            sequence = gr.Textbox(label="custom sequence", lines=1, placeholder='AMINO ACIDS: A,C,D,E,F,G,H,I,K,L,M,N,P,Q,R,S,T,V,W,Y\n  MASK TOKEN: X', visible=False)
                            seq_len = gr.Slider(minimum=5.0, maximum=250.0, label="protein length", value=100, visible=True)

                            seq_opt.change(fn=toggle_seq_input,
                                                    inputs=[seq_opt],
                                                    outputs=[seq_len, sequence])

                            gr.Markdown("### Optional Parameters")
                            with gr.Accordion(label='Secondary Structure',open=True):
                                gr.Markdown("Try changing the sliders or inputing explicit secondary structure conditioning for each residue")
                                sec_str_opt = gr.Radio(["sliders","explicit"], label="How would you like to specify secondary structure?", value='sliders')

                                secondary_structure = gr.Textbox(label="secondary structure", lines=1, placeholder='HELIX = H  STRAND = S  LOOP = L  MASK = X(must be the same length as input sequence)', visible=False)

                                with gr.Column():
                                    helix_bias = gr.Slider(minimum=0.0, maximum=0.05, label="helix bias", visible=True)
                                    strand_bias = gr.Slider(minimum=0.0, maximum=0.05, label="strand bias", visible=True)
                                    loop_bias = gr.Slider(minimum=0.0, maximum=0.20, label="loop bias", visible=True)

                                sec_str_opt.change(fn=toggle_secondary_structure,
                                                        inputs=[sec_str_opt],
                                                        outputs=[helix_bias,strand_bias,loop_bias,secondary_structure])

                            with gr.Accordion(label='Amino Acid Compositional Bias',open=False):
                                gr.Markdown("Bias sequence composition for particular amino acids by specifying the one letter code followed by the fraction to bias. This can be input as a list for example: W0.2,E0.1")
                                with gr.Row():
                                    aa_bias = gr.Textbox(label="aa bias", lines=1, placeholder='specify one letter AA and fraction to bias, for example W0.1 or M0.1,K0.1' )
                                    aa_bias_potential = gr.Textbox(label="aa bias scale", lines=1, placeholder='AA Bias potential scale (recomended range 1.0-5.0)')

                            with gr.Accordion(label='Hydrophobic Bias',open=False):
                                gr.Markdown("Bias for or against hydrophobic composition, to get more soluble proteins, bias away with a negative target score (ex. -5)")
                                with gr.Row():
                                    hydrophobic_target_score = gr.Textbox(label="hydrophobic score", lines=1, placeholder='hydrophobic score to target (negative score is good for solublility)')
                                    hydrophobic_potential = gr.Textbox(label="hydrophobic potential scale", lines=1, placeholder='hydrophobic potential scale (recomended range 1.0-2.0)')

                            with gr.Accordion(label='Diffusion Params',open=False):
                                gr.Markdown("Increasing T to more steps can be helpful for harder design challenges, sampling from different distributions can change the sequence and structural composition")
                                with gr.Row():
                                    num_steps = gr.Textbox(label="T", lines=1, placeholder='number of diffusion steps (25 or less will speed things up)')
                                    noise = gr.Dropdown(['normal','gmm2 [-1,1]','gmm3 [-1,0,1]'], label='noise type', value='normal')

                        with gr.Tab("Motif Selection"):

                            gr.Markdown("### Motif Selection Preview")
                            gr.Markdown('Contigs explained: to grab residues (seq and str) on a pdb chain you will provide the chain letter followed by a range of residues as indexed in the pdb file for example (A3-10) is the syntax to select residues 3-10 on chain A (the chain always needs to be specified). To add diffused residues to either side of this motif you can specify a range or discrete value without a chain letter infront. To add 15 residues before the motif and 20-30 residues (randomly sampled) after use the following syntax: 15,A3-10,20-30 commas are used to separate regions selected from the pdb and designed (diffused) resiudes which will be added. ')
                            pdb_id_code = gr.Textbox(label="PDB ID", lines=1, placeholder='INPUT PDB ID TO FETCH (ex. 1DPX)', visible=True)
                            contigs = gr.Textbox(label="contigs", lines=1, placeholder='specify contigs to grab particular residues from pdb ()', visible=True)
                            gr.Markdown('Using the same contig syntax, seq or str of input motif residues can be masked, allowing the model to hold strucutre fixed and design sequence or vice-versa')
                            with gr.Row():
                                seq_mask = gr.Textbox(label='seq mask',lines=1,placeholder='input residues to mask sequence')
                                str_mask = gr.Textbox(label='str mask',lines=1,placeholder='input residues to mask structure')
                            preview_viewer = gr.HTML()
                            rewrite_pdb = gr.File(label='PDB file')
                            preview_btn = gr.Button("Preview Motif")

                        with gr.Tab("MSA to PSSM"):
                            gr.Markdown("### MSA to PSSM Generation")
                            gr.Markdown('input either an MSA or PSSM to guide the model toward generating samples within your family of interest')
                            with gr.Row():
                                fasta_msa = gr.File(label='MSA')
                                input_pssm = gr.File(label='PSSM (.csv)')
                            pssm = gr.File(label='Generated PSSM')
                            pssm_view = gr.Plot(label='PSSM Viewer')
                            pssm_gen_btn = gr.Button("Generate PSSM")

                    btn = gr.Button("GENERATE")

                with gr.Column():
                    gr.Markdown("## OUTPUTS")
                    gr.Markdown("#### Confidence score for generated structure at each timestep")
                    plddt_plot = gr.Plot(label='plddt at step t')
                    gr.Markdown("#### Output protein sequence")
                    output_seq = gr.Textbox(label="sequence")
                    gr.Markdown("#### Download PDB file")
                    output_pdb = gr.File(label="PDB file")
                    gr.Markdown("#### Structure viewer")
                    output_viewer = gr.HTML()

            preview_btn.click(get_motif_preview, [pdb_id_code, contigs], [preview_viewer, rewrite_pdb])
            pssm_gen_btn.click(get_pssm, [fasta_msa, input_pssm], [pssm_view, pssm])
            btn.click(
                protein_diffusion_model,
                [sequence, seq_len, helix_bias, strand_bias, loop_bias,
                 secondary_structure, aa_bias, aa_bias_potential,
                 num_steps, noise, hydrophobic_target_score, hydrophobic_potential,
                 contigs, pssm, seq_mask, str_mask, rewrite_pdb],
                [output_seq, output_pdb, output_viewer, plddt_plot],
            )

        # ── Tab 2: Structure Prediction (Protenix v2) ─────────────────────────
        with gr.Tab("Structure Prediction (Protenix v2)"):

            gr.Markdown("""
            ## Protein Structure Prediction
            Fold a protein sequence using [Protenix v2](https://github.com/bytedance/Protenix) —
            ByteDance's open-source reimplementation of AlphaFold3 (~464M parameters).
            Optionally supply a pre-computed MSA (`.a3m`) to improve accuracy.
            Structure is coloured by per-residue pLDDT confidence (blue = high, red = low).
            """)

            with gr.Row(equal_height=False):
                with gr.Column():
                    px_sequence = gr.Textbox(
                        label="Protein sequence",
                        lines=4,
                        placeholder="Paste your amino acid sequence here (single-letter codes)",
                    )
                    px_msa = gr.File(
                        label="MSA file (optional, .a3m)",
                        file_types=[".a3m"],
                    )
                    px_model = gr.Dropdown(
                        choices=[
                            "protenix-v2",
                            "protenix_base_default_v1.0.0",
                            "protenix_base_20250630_v1.0.0",
                            "protenix_mini_default_v0.5.0",
                        ],
                        value=PROTENIX_MODEL,
                        label="Model",
                    )
                    gr.Examples(
                        examples=[
                            # Crambin (1CRN) — 46 residues, classic small test protein
                            ["TTCCPSIVARSNFNVCRLPGTPEALCATYTGCIIIPGATCPGDYAN", None, "protenix-v2"],
                            # Villin headpiece (1VII) — 36 residues, ultra-fast folder
                            ["LSDEDFKAVFGMTRSAFANLPLWKQQNLKKEKGLF", None, "protenix-v2"],
                        ],
                        inputs=[px_sequence, px_msa, px_model],
                        label="Example sequences from PDB",
                    )
                    px_btn = gr.Button("Fold", variant="primary")

                with gr.Column():
                    px_mean_plddt = gr.Textbox(label="Mean pLDDT", interactive=False)
                    px_viewer = gr.HTML(label="Structure viewer")
                    px_cif = gr.File(label="Download CIF")
                    px_plddt_plot = gr.Plot(label="Per-residue pLDDT")

            px_btn.click(
                fold_with_protenix,
                inputs=[px_sequence, px_msa, px_model],
                outputs=[px_viewer, px_cif, px_plddt_plot, px_mean_plddt],
            )

demo.launch(debug=True)