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

# app.py - RLAnOxPeptide Gradio Web Application
# Final version incorporating user feedback on generator logic and UI controls.

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

# 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)


# --- Feature Extractor Model Class (For ProtT5) ---
class FeatureProtT5Model:
    def __init__(self, base_model_id, finetuned_weights_path=None):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Initializing ProtT5 for feature extraction on device: {self.device}")

        print(f"Loading base model and tokenizer from '{base_model_id}'...")
        self.tokenizer = transformers.T5Tokenizer.from_pretrained(base_model_id, do_lower_case=False)
        self.model = transformers.T5EncoderModel.from_pretrained(base_model_id)

        if finetuned_weights_path and os.path.exists(finetuned_weights_path):
            print(f"Applying local fine-tuned weights from: {finetuned_weights_path}")
            state_dict = torch.load(finetuned_weights_path, map_location=self.device)
            self.model.load_state_dict(state_dict, strict=False)
            print("Successfully applied fine-tuned weights.")
        else:
            print("Warning: Fine-tuned weights not found or not provided. Using base ProtT5 weights.")

        self.model.to(self.device)
        self.model.eval()

    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 Architecture ---
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 ---
class ProtT5Generator(nn.Module):
    def __init__(self, vocab_size, embed_dim=512, num_layers=6, num_heads=8, dropout=0.1):
        super(ProtT5Generator, self).__init__()
        self.embed_tokens = nn.Embedding(vocab_size, embed_dim, padding_idx=token2id["<PAD>"])
        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        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>"]

    def forward(self, input_ids):
        embeddings = self.embed_tokens(input_ids)
        encoder_output = self.encoder(embeddings)
        logits = self.lm_head(encoder_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)
        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 == self.pad_token_id: continue
                seq += id2token.get(token_id, "")
            sequences.append(seq)
        return sequences

# --- CRITICAL DEPENDENCY: feature_extract.py ---
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 ---
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), default=0) + 2
        for seq in sequences:
            ids = [token2id.get(aa, 0) for aa in seq] + [generator.eos_token_id]
            ids = [np.random.randint(2, VOCAB_SIZE)] + ids
            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)
    representatives = []
    for i in range(int(num_clusters)):
        indices = np.where(kmeans.labels_ == i)[0]
        if len(indices) == 0: continue
        cluster_center = kmeans.cluster_centers_[i]
        cluster_embeddings = seq_embeds_np[indices]
        distances = np.linalg.norm(cluster_embeddings - cluster_center, axis=1)
        representative_index = indices[np.argmin(distances)]
        representatives.append(sequences[representative_index])
    return representatives

# --------------------------------------------------------------------------
# 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 ---
    PREDICTOR_CHECKPOINT_PATH = "checkpoints/final_rl_model_logitp0.1_calibrated_FINETUNED_PROTT5.pth"
    SCALER_PATH = "checkpoints/scaler_FINETUNED_PROTT5.pkl"
    GENERATOR_CHECKPOINT_PATH = "generator_checkpoints_v3.6/final_generator_model.pth"
    PROTT5_BASE_MODEL_ID = "Rostlab/prot_t5_xl_uniref50"
    FINETUNED_PROTT5_FOR_FEATURES_PATH = "prott5/model/finetuned_prott5.bin"

    # --- Load Predictor Model ---
    print(f"Loading Predictor from: {PREDICTOR_CHECKPOINT_PATH}")
    PREDICTOR_MODEL = AntioxidantPredictor(input_dim=1914)
    PREDICTOR_MODEL.load_state_dict(torch.load(PREDICTOR_CHECKPOINT_PATH, map_location=DEVICE))
    PREDICTOR_MODEL.to(DEVICE)
    PREDICTOR_MODEL.eval()
    print(f"✅ Predictor model loaded (Temp: {PREDICTOR_MODEL.get_temperature():.4f}).")

    # --- Load Scaler & Feature Extractor ---
    print(f"Loading Scaler from: {SCALER_PATH}")
    SCALER = joblib.load(SCALER_PATH)
    print("Loading ProtT5 Feature Extractor...")
    PROTT5_EXTRACTOR = FeatureProtT5Model(
        base_model_id=PROTT5_BASE_MODEL_ID,
        finetuned_weights_path=FINETUNED_PROTT5_FOR_FEATURES_PATH
    )
    print("✅ Scaler and Feature Extractor loaded.")

    # --- Load Generator Model ---
    print(f"Loading Generator from: {GENERATOR_CHECKPOINT_PATH}")
    GENERATOR_MODEL = ProtT5Generator(vocab_size=VOCAB_SIZE)
    GENERATOR_MODEL.load_state_dict(torch.load(GENERATOR_CHECKPOINT_PATH, map_location=DEVICE))
    GENERATOR_MODEL.to(DEVICE)
    GENERATOR_MODEL.eval()
    print("✅ Generator model loaded.")
    
    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 (ACDEFGHIKLMNPQRSTVWY)."
    
    try:
        features = extract_features(sequence_str.upper(), PROTT5_EXTRACTOR, L_fixed=29, d_model_pe=16)
        scaled_features = 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()):
    """
    Handles the full generation-validation-clustering pipeline with a loop to ensure
    the target number of peptides is generated.
    """
    num_to_generate = int(num_to_generate)
    min_len = int(min_len)
    max_len = int(max_len)

    # Safety check for length
    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 = {}  # Use a dictionary to store unique sequences and their probabilities
        attempts = 0
        max_attempts = 20  # Safety break to prevent infinite loops
        generation_batch_size = 200 # Number of sequences to generate in each attempt

        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})")

            # Generate a batch of candidate sequences
            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)
            
            # Filter for length and uniqueness
            new_candidates = []
            for seq in decoded_sequences:
                if min_len <= len(seq) <= max_len:
                    if seq not in validated_pool:
                        new_candidates.append(seq)
            
            # Validate the new, unique candidates
            for seq in new_candidates:
                prob_str, _ = predict_peptide_wrapper(seq)
                try:
                    prob = float(prob_str)
                    if prob > 0.90:
                        validated_pool[seq] = prob
                        # Check if we have reached the target
                        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 the current settings. Try different parameters.", "Predicted Probability": "N/A"}])
        
        # --- Final Processing ---
        high_quality_sequences = list(validated_pool.keys())
        
        # Cluster to ensure diversity, selecting up to the target number
        final_diverse_seqs = cluster_sequences(GENERATOR_MODEL, high_quality_sequences, num_to_generate, DEVICE)
        
        # Format final results into a DataFrame
        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
# --------------------------------------------------------------------------
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=5, maximum=50, value=10, step=1, label="Number of Final Peptides to Generate")
                    # ✅ MODIFIED: Length sliders both have a range of 2-20
                    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)

            # ✅ ADDED: Dynamic linking of min and max length sliders for better UX
            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=min_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()