import streamlit as st
import torch
import torch.nn as nn
import torch.nn.functional as F
import pickle
import numpy as np
import re
from sklearn.cluster import KMeans
from transformers import T5Tokenizer, T5EncoderModel
import os
from dotenv import load_dotenv
import google.generativeai as genai
from bio_clip_recommender import BioCLIP, ProteinEmbedder, ProtocolRecommender

load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if not GEMINI_API_KEY:
    st.error("API Key not found! Please check that your .env file exists and is formatted correctly.")
    st.stop()

genai.configure(api_key=GEMINI_API_KEY)

sys_instruct = """
You are a Senior Research Biochemist. Your primary role is to design scientifically rigorous, highly cohesive protein purification pipelines based ONLY on provided laboratory data.

Follow these strict guidelines:
1. Tone & Style: Be concise, professional, and direct. Omit conversational filler and pleasantries. 
2. Formatting: Use standard Markdown. Create clear headers for each step (e.g., '### Step 1: Lysis'). Bold all specific buffer concentrations, proteins, reagents, and pH values so they are easy to read at the bench.
3. Biochemical Rationale: For every step, explicitly state *why* specific reagents are used based on standard biochemical principles.
4. Chemical Guardrails: Never invent or hallucinate protocols or buffers. If you detect chemical incompatibilities in the user's request or retrieved data (e.g., high DTT concentrations applied to standard Ni-NTA columns, or inappropriate detergents for soluble proteins), explicitly flag them with a bold **WARNING**.
"""

model = genai.GenerativeModel('gemini-3.5-flash', 
                              system_instruction=sys_instruct)

st.set_page_config(
    page_title="Bio-CLIP Recommender",
    layout="wide",
    initial_sidebar_state="expanded"
)

@st.cache_resource(show_spinner=False)
def load_ai_engines():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    embedder = ProteinEmbedder(device=device)
    recommender = ProtocolRecommender(device=device)
    return embedder, recommender

st.title("Bio-CLIP: Protein Purification AI")
st.markdown("Zero-shot protocol recommendation directly from 1D amino acid sequences.")

with st.spinner("Booting up ProtT5 and Bio-CLIP Models (this takes a moment)..."):
    try:
        embedder, recommender = load_ai_engines()
        models_loaded = True
    except Exception as e:
        st.error(f"Failed to load model weights. Ensure `.pth` and `.pkl` files are in the same folder. \n\nError: {e}")
        models_loaded = False

if models_loaded:
    with st.sidebar:
        st.header("About")
        st.write("This tool uses **ProtT5** to embed amino acids and a Contrastive Learning (**Bio-CLIP**) model to map them to historical purification protocols.")
        st.markdown("---")
        st.write("**Top K Results**")
        top_k = st.slider("Number of recommendations:", 1, 5, 3)

    st.markdown("### Input Sequence")
    default_gfp = "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"
    sequence_input = st.text_area("Paste Amino Acid Sequence here:", value=default_gfp, height=150)

    user_goal = st.text_input("Optional: What is your specific purification goal? (e.g., 'Prioritize highest yield', 'Are there temperature concerns?')")

    if st.button("Generate Purification Pipeline", type="primary"):
        sequence_input = sequence_input.strip()
        
        if len(sequence_input) < 10:
            st.warning("Please enter a valid amino acid sequence (at least 10 characters).")
        else:
            with st.spinner("1/2: Running ProtT5 Sequence Embedding..."):
                seq_vector = embedder.embed_raw_sequence(sequence_input)
            
            with st.spinner("2/2: Querying Bio-CLIP Latent Space..."):
                results = recommender.search(seq_vector, top_k=top_k)
            
            st.success("Search Complete!")
            st.markdown("---")
            
            col1, col2, col3 = st.columns(3)
            
            columns = {'lysis': col1, 'elution': col2, 'desalting': col3}
            
            for step in ['lysis', 'elution', 'desalting']:
                with columns[step]:
                    st.subheader(f"{step.capitalize()}")
                    for i, res in enumerate(results[step]):
                    
                        conf = max(0.0, min(1.0, (res['confidence'] + 0.1)))
                        st.markdown(f"**Option {i+1}**")
                        st.progress(float(conf), text=f"Confidence: {conf*100:.1f}%")
                        st.info(res['text'])
                        st.write("")

            st.markdown("---")

            st.markdown("### 3. AI Synthesis")
            with st.spinner("3/3: Gemini AI analyzing retrieved protocols..."):
                try:
                    retrieved_context = f"""
                    Lysis Options: {[res['text'] for res in results['lysis']]}
                    Elution Options: {[res['text'] for res in results['elution']]}
                    Desalting Options: {[res['text'] for res in results['desalting']]}
                    """
                    has_his_tag = "HHHHHH" in sequence_input

                    rag_prompt = f"""

                    User Sequence: {sequence_input}
                    Contains His-Tag: {has_his_tag}
                    
                    Synthesize a cohesive, step-by-step protein purification protocol using ONLY the retrieved options below. 
                    If the user provided a specific goal, tailor the synthesis to that goal.

                    User's specific goal/question: {user_goal if user_goal else "Provide a standard, recommended cohesive pipeline from these options."}

                    --- RETRIEVED PROTOCOL OPTIONS ---
                    {retrieved_context}
                    ----------------------------------
                    """

                    response = model.generate_content(rag_prompt)
                    
                    st.write(response.text)
                    
                except Exception as e:
                    st.error(f"An error occurred while generating the AI summary: {e}")