import os
import re
import gradio as gr
import spaces
import torch
from transformers import LlamaForCausalLM, LlamaTokenizer
from huggingface_hub import InferenceClient

MODEL_NAME = "GreatCaptainNemo/ProLLaMA"

print(f"Loading {MODEL_NAME} at startup (only happens once)...")
tokenizer = LlamaTokenizer.from_pretrained(MODEL_NAME)
model = LlamaForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32
)
model.eval()
print("Model loaded.")


KNOWN_SUPERFAMILIES = [
    "Ankyrin repeat-containing domain superfamily",
    "Immunoglobulin-like fold",
    "TIM barrel",
    "Winged helix DNA-binding domain superfamily",
    "SH3-like domain superfamily",
    "Leucine-rich repeat domain superfamily",
    "Zinc finger domain superfamily",
    "Alpha/beta hydrolase fold",
    "P-loop containing nucleoside triphosphate hydrolase",
    "Globin-like superfamily",
]


def extract_sequence(raw_output):
    raw_output = raw_output.upper()

    # Look for Seq=<...>
    match = re.search(r"SEQ=<([ACDEFGHIKLMNPQRSTVWY]+)>", raw_output)
    if match:
        return match.group(1)

    # Look for a line containing only amino acid letters
    for line in raw_output.splitlines():
        line = line.strip()
        if re.fullmatch(r"[ACDEFGHIKLMNPQRSTVWY]{30,}", line):
            return line

    return ""


@spaces.GPU
def generate(description, max_new_tokens):
    model.to("cuda")

    prompt = f"""Design a protein sequence with the following function.

Description:
{description}

Protein sequence:
"""

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=True,
            top_p=0.9,
            temperature=0.8,
            repetition_penalty=1.2,
            pad_token_id=tokenizer.eos_token_id,
        )

    # Decode ONLY the generated text
    generated_ids = output_ids[0][inputs["input_ids"].shape[1]:]
    raw = tokenizer.decode(generated_ids, skip_special_tokens=True)

    print("=" * 80)
    print("RAW MODEL OUTPUT:")
    print(raw)
    print("=" * 80)

    seq = extract_sequence(raw)

    return raw, seq


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Describe the protein you want",
            placeholder="e.g. a protein that binds to cancer cells",
        ),
        gr.Number(
            label="Max new tokens",
            value=200,
        ),
    ],
    outputs=[
        gr.Textbox(label="Raw model output"),
        gr.Textbox(label="Extracted sequence"),
    ],
    description=(
        "IMPORTANT: this generates a sequence belonging to a structurally "
        "similar protein family. It is NOT validated to bind any specific "
        "target. Always verify predicted structures and experimentally "
        "validate any candidates."
    ),
    api_name="generate_protein",
)

demo.launch()