import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Points to your newly renamed repo
MODEL_ID = "assix-research/genomic-transformer-aphasia-recovery"

# Load model and tokenizer
print("Loading model from Hub...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_ID, 
    trust_remote_code=True
)

def predict_recovery(dna_sequence):
    """
    Function to process DNA sequence and return a clinical prediction.
    """
    # Basic validation for DNA characters
    dna_sequence = dna_sequence.strip().upper()
    if not dna_sequence or not all(base in "ATCG" for base in dna_sequence):
        return "Error", "Invalid DNA sequence. Please use only A, T, C, and G."

    # Tokenize and predict
    inputs = tokenizer(dna_sequence, return_tensors="pt", truncation=True, max_length=1024)
    with torch.no_grad():
        logits = model(**inputs).logits
    
    # Scale result for WAB-AQ (0-100)
    score = max(0, min(100, logits.item()))
    
    # Clinical Categorization
    if score > 75:
        category = "High Plasticity Potential"
        insight = "Genomic markers suggest a high capacity for spontaneous neural reorganization."
    elif score > 40:
        category = "Moderate Plasticity"
        insight = "Standard recovery path; likely to benefit significantly from consistent SLT."
    else:
        category = "Structural Reliance"
        insight = "Recovery may depend more on physical white matter integrity than genomic plasticity."

    return f"{score:.2f}/100", f"{category}: {insight}"

# Build the Interface
demo = gr.Interface(
    fn=predict_recovery, # Function name now matches definition above
    inputs=gr.Textbox(
        label="Input Genomic Sequence (Nucleotides)", 
        placeholder="Enter 1024bp DNA sequence...",
        lines=8
    ),
    outputs=[
        gr.Label(label="Predicted WAB-AQ Potential"),
        gr.Textbox(label="Clinical Interpretation")
    ],
    title="🧬 Stroke-Recovery-Analyser (v1.0)",
    description="Analyze the biological plasticity potential for post-stroke aphasia recovery using a fine-tuned Nucleotide Transformer v2.",
    examples=[
        ["ATGACCATCCTTTTCCTTACTATGGTTATTTCATACTTTGGTTGCATGAAGGCTGCCCCCATGAAAGAAGCAAACATCCGAGGACAAGGTGGCTTGGCCTACCCAGGTGTGCGGACCCATGGGACTCTGGAGAGCGTGAATGGGCCCAAGGCAGGTTCAAGAGGCTTGACATCATTGGCTGACACTTTCGAACACGTGATAGAAGAGCTGTTGGATGAGGACCAGAAAGTTCGGCCCAATGAAGAAAACAATAAGGACGCAGACTTGTACACGTCCAGGGTGATGCTCAGTAGTCAAGTGCCTTTGGAGCCTCCTCTTCTCTTTCTGCTGGAGGAATACAAAAATTACCTAGATGCTGCAAACATGTCCATGAGGGTCCGGCGCCACTCTGACCCTGCCCGCCGAGGGGAGCTGAGCGTGTGTGACAGTATTAGTGAGTGGGTAACGGCGGCAGACAAAAAGACTGCAGTGGACATGTCGGGCGGGACGGTCACAGTCCTTGAAAAGGTCCCTGTATCAAAAGGCCAACTGAAGCAATACTTCTACGAGACCAAGTGCAATCCCATGGGTTACACAAAAGAAGGCTGCAGGGGCATAGACAAAAGGCATTGGAACTCCCAGTGCCGAACTACCCAGTCGTACGTGCGGGCCCTTACCATGGATAGCAAAAAGAGAATTGGCTGGCGATTCATAAGGATAGACACTTCTTGTGTATGTACATTGACCATTAAAAGGGGAAGATAG"], # Val/Val
        ["ATGACCATCCTTTTCCTTACTATGGTTATTTCATACTTTGGTTGCATGAAGGCTGCCCCCATGAAAGAAGCAAACATCCGAGGACAAGGTGGCTTGGCCTACCCAGGTGTGCGGACCCATGGGACTCTGGAGAGCGTGAATGGGCCCAAGGCAGGTTCAAGAGGCTTGACATCATTGGCTGACACTTTCGAACACATGATAGAAGAGCTGTTGGATGAGGACCAGAAAGTTCGGCCCAATGAAGAAAACAATAAGGACGCAGACTTGTACACGTCCAGGGTGATGCTCAGTAGTCAAGTGCCTTTGGAGCCTCCTCTTCTCTTTCTGCTGGAGGAATACAAAAATTACCTAGATGCTGCAAACATGTCCATGAGGGTCCGGCGCCACTCTGACCCTGCCCGCCGAGGGGAGCTGAGCGTGTGTGACAGTATTAGTGAGTGGGTAACGGCGGCAGACAAAAAGACTGCAGTGGACATGTCGGGCGGGACGGTCACAGTCCTTGAAAAGGTCCCTGTATCAAAAGGCCAACTGAAGCAATACTTCTACGAGACCAAGTGCAATCCCATGGGTTACACAAAAGAAGGCTGCAGGGGCATAGACAAAAGGCATTGGAACTCCCAGTGCCGAACTACCCAGTCGTACGTGCGGGCCCTTACCATGGATAGCAAAAAGAGAATTGGCTGGCGATTCATAAGGATAGACACTTCTTGTGTATGTACATTGACCATTAAAAGGGGAAGATAG"]  # Val/Met
    ]
)

if __name__ == "__main__":
    demo.launch()