#!/usr/bin/env python
# -*- coding: utf-8 -*-

# app.py - RLAnOxPeptide Gradio Web Application


import os
import torch
import torch.nn as nn
import pandas as pd
import joblib
import numpy as np
import gradio as gr
from sklearn.cluster import KMeans
from tqdm import tqdm
import transformers
import time
import copy  # ✅ ADDED: For deep copying the base model

# NEW DEPENDENCY: peft library for LoRA
from peft import PeftModel

# Suppress verbose logging from transformers
transformers.logging.set_verbosity_error()

# --------------------------------------------------------------------------
# SECTION 1: CORE CLASS AND FUNCTION DEFINITIONS
# --------------------------------------------------------------------------

# --- Vocabulary Definition ---
AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
token2id = {aa: i + 2 for i, aa in enumerate(AMINO_ACIDS)}
token2id["<PAD>"] = 0
token2id["<EOS>"] = 1
id2token = {i: t for t, i in token2id.items()}
VOCAB_SIZE = len(token2id)


# --- Validator's Feature Extractor Class ---
# ✅ MODIFIED: Accepts a pre-loaded model instead of loading its own.
class LoRAProtT5Extractor:
    def __init__(self, preloaded_base_model, preloaded_tokenizer, lora_adapter_path):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Initializing Validator Feature Extractor on device: {self.device}")

        base_model = preloaded_base_model
        self.tokenizer = preloaded_tokenizer
        
        if not os.path.exists(lora_adapter_path):
            raise FileNotFoundError(f"Error: Validator LoRA adapter directory not found at: {lora_adapter_path}")
            
        print(f"  - [Validator] Applying LoRA adapter from: {lora_adapter_path}")
        lora_model = PeftModel.from_pretrained(base_model, lora_adapter_path)
        
        print("  - [Validator] Merging LoRA weights for faster inference...")
        self.model = lora_model.merge_and_unload().to(self.device)
        self.model.eval()
        print("  - Validator feature extractor is ready.")

    def encode(self, sequence):
        if not sequence or not isinstance(sequence, str):
            return np.zeros((1, 1024), dtype=np.float32)
        seq_spaced = " ".join(list(sequence))
        encoded_input = self.tokenizer(seq_spaced, return_tensors='pt', padding=True, truncation=True)
        encoded_input = {k: v.to(self.device) for k, v in encoded_input.items()}
        with torch.no_grad():
            embedding = self.model(**encoded_input).last_hidden_state
        emb_np = embedding.squeeze(0).cpu().numpy()
        return emb_np if emb_np.shape[0] > 0 else np.zeros((1, 1024), dtype=np.float32)


# --- Predictor Model Head Architecture (Unchanged) ---
class AntioxidantPredictor(nn.Module):
    def __init__(self, input_dim=1914, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1):
        super(AntioxidantPredictor, self).__init__()
        self.prott5_dim = 1024
        self.handcrafted_dim = input_dim - self.prott5_dim
        self.seq_len = 16
        self.prott5_feature_dim = 64
        encoder_layer = nn.TransformerEncoderLayer(d_model=self.prott5_feature_dim, nhead=transformer_heads, dropout=transformer_dropout, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=transformer_layers)
        fused_dim = self.prott5_feature_dim + self.handcrafted_dim
        self.fusion_fc = nn.Sequential(nn.Linear(fused_dim, 1024), nn.ReLU(), nn.Dropout(0.3), nn.Linear(1024, 512), nn.ReLU(), nn.Dropout(0.3))
        self.classifier = nn.Sequential(nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, 1))
        self.temperature = nn.Parameter(torch.ones(1), requires_grad=False)
    def forward(self, x):
        batch_size = x.size(0)
        prot_t5_features = x[:, :self.prott5_dim]
        handcrafted_features = x[:, self.prott5_dim:]
        prot_t5_seq = prot_t5_features.view(batch_size, self.seq_len, self.prott5_feature_dim)
        encoded_seq = self.transformer_encoder(prot_t5_seq)
        refined_prott5 = encoded_seq.mean(dim=1)
        fused_features = torch.cat([refined_prott5, handcrafted_features], dim=1)
        fused_output = self.fusion_fc(fused_features)
        logits = self.classifier(fused_output)
        return logits / self.temperature
    def get_temperature(self):
        return self.temperature.item()


# --- Generator Model Architecture ---
# ✅ MODIFIED: Accepts a pre-loaded model instead of loading its own.
class AdvancedProtT5Generator(nn.Module):
    def __init__(self, preloaded_base_model, lora_adapter_path, vocab_size):
        super(AdvancedProtT5Generator, self).__init__()
        
        base_model = preloaded_base_model
        print(f"  - [Generator] Applying LoRA adapter from: {lora_adapter_path}")
        self.backbone = PeftModel.from_pretrained(base_model, lora_adapter_path)
        
        self.embed_tokens = self.backbone.get_input_embeddings()
        
        embed_dim = self.backbone.config.d_model
        self.lm_head = nn.Linear(embed_dim, vocab_size)
        
        self.vocab_size = vocab_size
        self.eos_token_id = token2id["<EOS>"]
        self.pad_token_id = token2id["<PAD>"]
        print("  - Advanced Generator framework initialized.")

    def forward(self, input_ids):
        attention_mask = (input_ids != self.pad_token_id).int()
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        sequence_output = outputs.last_hidden_state
        logits = self.lm_head(sequence_output)
        return logits

    def sample(self, batch_size, max_length=20, device="cpu", temperature=2.5, min_decoded_length=3):
        start_token = torch.randint(2, self.vocab_size, (batch_size, 1), device=device)
        generated = start_token
        for _ in range(max_length - 1):
            logits = self.forward(generated)
            next_logits = logits[:, -1, :] / temperature
            if generated.size(1) < min_decoded_length:
                next_logits[:, self.eos_token_id] = -float("inf")
            probs = torch.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            generated = torch.cat((generated, next_token), dim=1)
            if (generated == self.eos_token_id).any(dim=1).all():
                break
        return generated

    def decode(self, token_ids_batch):
        sequences = []
        for ids_tensor in token_ids_batch:
            seq = ""
            for token_id in ids_tensor.tolist()[1:]:
                if token_id == self.eos_token_id: break
                if token_id == token2id["<PAD>"]: continue
                seq += id2token.get(token_id, "?")
            sequences.append(seq)
        return sequences

# --- CRITICAL DEPENDENCY: feature_extract.py (Unchanged) ---
try:
    from feature_extract import extract_features
except ImportError:
    raise gr.Error("Fatal Error: `feature_extract.py` not found. This file is required. Please upload it to your repository.")

# --- Clustering Logic (Unchanged) ---
def cluster_sequences(generator, sequences, num_clusters, device):
    if not sequences or len(sequences) < num_clusters:
        return sequences[:num_clusters]
    with torch.no_grad():
        token_ids_list = []
        max_len = max(len(seq) for seq in sequences) + 2
        for seq in sequences:
            ids = [np.random.randint(2, VOCAB_SIZE)] + [token2id.get(aa, 0) for aa in seq] + [generator.eos_token_id]
            ids += [token2id["<PAD>"]] * (max_len - len(ids))
            token_ids_list.append(ids)
        input_ids = torch.tensor(token_ids_list, dtype=torch.long, device=device)
        embeddings = generator.embed_tokens(input_ids)
        mask = (input_ids != token2id["<PAD>"]).unsqueeze(-1).float()
        seq_embeds = (embeddings * mask).sum(dim=1) / (mask.sum(dim=1) + 1e-9)
        seq_embeds_np = seq_embeds.cpu().numpy()
    kmeans = KMeans(n_clusters=int(num_clusters), random_state=42, n_init='auto').fit(seq_embeds_np)
    reps = []
    for i in range(int(num_clusters)):
        idxs = np.where(kmeans.labels_ == i)[0]
        if len(idxs) == 0: continue
        center = kmeans.cluster_centers_[i]
        distances = np.linalg.norm(seq_embeds_np[idxs] - center, axis=1)
        rep_idx = idxs[np.argmin(distances)]
        reps.append(sequences[rep_idx])
    return reps

# --------------------------------------------------------------------------
# SECTION 2: GLOBAL MODEL AND DEPENDENCY LOADING
# --------------------------------------------------------------------------

print("--- Starting Application: Loading all models and dependencies ---")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

try:
    # --- Define file paths ---
    PROTT5_BASE_MODEL_ID = "Rostlab/prot_t5_xl_uniref50"
    VALIDATOR_LORA_PATH = "./lora_finetuned_prott5"
    PREDICTOR_HEAD_CHECKPOINT_PATH = "./predictor_with_lora_checkpoints/final_predictor_with_lora.pth"
    SCALER_PATH = "./predictor_with_lora_checkpoints/scaler_lora.pkl"
    GENERATOR_LORA_DIR = "./generator_with_lora_output/final_lora_generator"
    GENERATOR_LM_HEAD_PATH = os.path.join(GENERATOR_LORA_DIR, "lm_head.pth")

    # ✅ OPTIMIZED: Load the base model and tokenizer only ONCE
    print(f"--- Loading Base ProtT5 Model ({PROTT5_BASE_MODEL_ID}) just once... ---")
    base_prot_t5_model = transformers.T5EncoderModel.from_pretrained(PROTT5_BASE_MODEL_ID)
    base_tokenizer = transformers.T5Tokenizer.from_pretrained(PROTT5_BASE_MODEL_ID)
    print("✅ Base ProtT5 Model loaded.")

    # --- Load Validator System ---
    print("\n--- Initializing Validator System ---")
    VALIDATOR_SCALER = joblib.load(SCALER_PATH)
    # Pass a deep copy of the base model to prevent modification conflicts
    VALIDATOR_EXTRACTOR = LoRAProtT5Extractor(
        preloaded_base_model=copy.deepcopy(base_prot_t5_model),
        preloaded_tokenizer=base_tokenizer,
        lora_adapter_path=VALIDATOR_LORA_PATH
    )
    PREDICTOR_MODEL = AntioxidantPredictor(input_dim=1914)
    PREDICTOR_MODEL.load_state_dict(torch.load(PREDICTOR_HEAD_CHECKPOINT_PATH, map_location=DEVICE))
    PREDICTOR_MODEL.to(DEVICE)
    PREDICTOR_MODEL.eval()
    print("✅ Validator System loaded successfully.")

    # --- Load Generator System ---
    print("\n--- Initializing Generator System ---")
    # Pass a deep copy of the base model here as well
    GENERATOR_MODEL = AdvancedProtT5Generator(
        preloaded_base_model=copy.deepcopy(base_prot_t5_model),
        lora_adapter_path=GENERATOR_LORA_DIR,
        vocab_size=VOCAB_SIZE
    )
    if not os.path.exists(GENERATOR_LM_HEAD_PATH):
        raise FileNotFoundError(f"Generator's lm_head weights not found at: {GENERATOR_LM_HEAD_PATH}")
    GENERATOR_MODEL.lm_head.load_state_dict(torch.load(GENERATOR_LM_HEAD_PATH, map_location=DEVICE))
    GENERATOR_MODEL.to(DEVICE)
    GENERATOR_MODEL.eval()
    print("✅ Generator System loaded successfully.")
    
    print("\n--- All models loaded! Gradio app is ready. ---\n")

except Exception as e:
    print(f"💥 FATAL ERROR during model loading: {e}")
    raise gr.Error(f"A required model or file could not be loaded. Please check your repository file structure and paths. Error details: {e}")

# --------------------------------------------------------------------------
# SECTION 3: WRAPPER FUNCTIONS FOR GRADIO UI
# --------------------------------------------------------------------------

def predict_peptide_wrapper(sequence_str):
    if not sequence_str or not isinstance(sequence_str, str) or any(c not in AMINO_ACIDS for c in sequence_str.upper()):
        return "0.0000", "Error: Please enter a valid peptide sequence using standard amino acids."
    
    try:
        features = extract_features(sequence_str.upper(), VALIDATOR_EXTRACTOR, L_fixed=29, d_model_pe=16)
        scaled_features = VALIDATOR_SCALER.transform(features.reshape(1, -1))
        
        with torch.no_grad():
            features_tensor = torch.tensor(scaled_features, dtype=torch.float32).to(DEVICE)
            logits = PREDICTOR_MODEL(features_tensor)
            probability = torch.sigmoid(logits).squeeze().item()
        
        classification = "Antioxidant" if probability >= 0.5 else "Non-Antioxidant"
        return f"{probability:.4f}", classification

    except Exception as e:
        print(f"Prediction Error for sequence '{sequence_str}': {e}")
        return "N/A", f"An error occurred during prediction: {e}"

def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, diversity_factor, progress=gr.Progress()):
    num_to_generate = int(num_to_generate)
    min_len = int(min_len)
    max_len = int(max_len)

    if min_len > max_len:
        gr.Warning("Minimum Length cannot be greater than Maximum Length. Adjusting min_len = max_len.")
        min_len = max_len
    
    try:
        validated_pool = {}
        attempts = 0
        max_attempts = 20
        generation_batch_size = 10

        while len(validated_pool) < num_to_generate and attempts < max_attempts:
            progress(len(validated_pool) / num_to_generate, desc=f"Found {len(validated_pool)} / {num_to_generate} peptides. (Attempt {attempts+1}/{max_attempts})")

            with torch.no_grad():
                generated_tokens = GENERATOR_MODEL.sample(
                    batch_size=generation_batch_size, max_length=max_len, device=DEVICE,
                    temperature=temperature, min_decoded_length=min_len
                )
            decoded_sequences = GENERATOR_MODEL.decode(generated_tokens)
            
            new_candidates = []
            for seq in decoded_sequences:
                if min_len <= len(seq) <= max_len:
                    if seq not in validated_pool:
                        new_candidates.append(seq)
            
            for seq in new_candidates:
                prob_str, _ = predict_peptide_wrapper(seq)
                try:
                    prob = float(prob_str)
                    if prob > 0.90:
                        validated_pool[seq] = prob
                        if len(validated_pool) >= num_to_generate:
                            break 
                except (ValueError, TypeError):
                    continue
            
            attempts += 1
            if len(validated_pool) >= num_to_generate:
                break 
        
        progress(1.0, desc=f"Collected {len(validated_pool)} high-quality peptides. Clustering for diversity...")
        time.sleep(1)

        if not validated_pool:
            return pd.DataFrame([{"Sequence": "Could not generate any high-activity peptides (>0.9 prob) with current settings.", "Predicted Probability": "N/A"}])
        
        high_quality_sequences = list(validated_pool.keys())
        final_diverse_seqs = cluster_sequences(GENERATOR_MODEL, high_quality_sequences, num_to_generate, DEVICE)
        
        final_results = [(seq, f"{validated_pool[seq]:.4f}") for seq in final_diverse_seqs]
        final_results.sort(key=lambda x: float(x[1]), reverse=True)
        
        return pd.DataFrame(final_results, columns=["Sequence", "Predicted Probability"])

    except Exception as e:
        print(f"Generation Pipeline Error: {e}")
        return pd.DataFrame([{"Sequence": f"An error occurred during generation: {e}", "Predicted Probability": "N/A"}])

# --------------------------------------------------------------------------
# SECTION 4: GRADIO UI CONSTRUCTION (Unchanged)
# --------------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="RLAnOxPeptide") as demo:
    gr.Markdown("# RLAnOxPeptide: Intelligent Peptide Design and Prediction")
    gr.Markdown("An integrated framework combining reinforcement learning and a Transformer model for the efficient prediction and innovative design of antioxidant peptides.")

    with gr.Tabs():
        # --- PREDICTION TAB ---
        with gr.TabItem("Peptide Activity Predictor"):
            gr.Markdown("### Enter an amino acid sequence to predict its antioxidant activity.")
            with gr.Row():
                peptide_input = gr.Textbox(label="Peptide Sequence", placeholder="e.g., WHYHDYKY", scale=3)
                predict_button = gr.Button("Predict", variant="primary", scale=1)
            with gr.Row():
                probability_output = gr.Textbox(label="Predicted Probability", interactive=False)
                class_output = gr.Textbox(label="Predicted Class", interactive=False)
            
            predict_button.click(
                fn=predict_peptide_wrapper,
                inputs=peptide_input,
                outputs=[probability_output, class_output]
            )
            gr.Examples(
                examples=[["WHYHDYKY"], ["YPGG"], ["LVLHEHGGN"], ["WKYG"]],
                inputs=peptide_input,
                fn=predict_peptide_wrapper,
                outputs=[probability_output, class_output],
                cache_examples=True
            )

        # --- GENERATION TAB ---
        with gr.TabItem("Novel Sequence Generator"):
            gr.Markdown("### Set parameters to generate novel, high-activity antioxidant peptides.")
            with gr.Column():
                with gr.Row():
                    num_input = gr.Slider(minimum=1, maximum=10, value=10, step=1, label="Number of Final Peptides to Generate")
                    min_len_input = gr.Slider(minimum=2, maximum=20, value=3, step=1, label="Minimum Length")
                    max_len_input = gr.Slider(minimum=2, maximum=20, value=20, step=1, label="Maximum Length")
                with gr.Row():
                    temp_input = gr.Slider(minimum=0.5, maximum=3.0, value=2.5, step=0.1, label="Temperature (Higher = More random)")
                    diversity_input = gr.Slider(minimum=1.1, maximum=5.0, value=1.5, step=0.1, label="Diversity Factor (Larger initial pool for clustering)")
            
            generate_button = gr.Button("Generate Peptides", variant="primary")
            results_output = gr.DataFrame(headers=["Sequence", "Predicted Probability"], label="Generated & Validated Peptides (>90% Probability)", wrap=True)

            def update_min_len_range(max_len):
                return gr.Slider(maximum=max_len)
            max_len_input.change(fn=update_min_len_range, inputs=max_len_input, outputs=max_len_input)

            def update_max_len_range(min_len):
                return gr.Slider(minimum=min_len)
            min_len_input.change(fn=update_max_len_range, inputs=min_len_input, outputs=max_len_input)

            generate_button.click(
                fn=generate_peptide_wrapper,
                inputs=[num_input, min_len_input, max_len_input, temp_input, diversity_input],
                outputs=results_output
            )

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