# app.py
import gradio as gr
import torch
import torch.nn.functional as F # <--- Added import
import pytorch_lightning as pl   # <--- Added import (needed for type hints, model access)
import os
import json
import logging
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
import gc # For garbage collection on potential OOM
import math # Needed for PositionalEncoding if moved here (or keep in enhanced_trainer)

# --- Configuration ---
MODEL_REPO_ID = "AdrianM0/smiles-to-iupac-translator"
CHECKPOINT_FILENAME = "last.ckpt"
SMILES_TOKENIZER_FILENAME = "smiles_bytelevel_bpe_tokenizer_scaled.json"
IUPAC_TOKENIZER_FILENAME = "iupac_unigram_tokenizer_scaled.json"
CONFIG_FILENAME = "config.json"
# --- End Configuration ---

# --- Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- Load Helper Code (Only Model Definition Needed) ---
try:
    # We only need the LightningModule definition and the mask function now
    from enhanced_trainer import SmilesIupacLitModule, generate_square_subsequent_mask
    logging.info("Successfully imported from enhanced_trainer.py.")

    # We will define beam_search_decode and translate locally in this file
    # REMOVED: from test_ckpt import beam_search_decode, translate

except ImportError as e:
    logging.error(f"Failed to import helper code from enhanced_trainer.py: {e}. Make sure enhanced_trainer.py is in the root of the Hugging Face repo '{MODEL_REPO_ID}'.")
    gr.Error(f"Initialization Error: Could not load necessary Python modules (enhanced_trainer.py). Check Space logs. Error: {e}")
    exit()
except Exception as e:
    logging.error(f"An unexpected error occurred during helper code import: {e}", exc_info=True)
    gr.Error(f"Initialization Error: An unexpected error occurred loading helper modules. Check Space logs. Error: {e}")
    exit()

# --- Global Variables (Load Model Once) ---
model: pl.LightningModule | None = None # Added type hint
smiles_tokenizer: Tokenizer | None = None
iupac_tokenizer: Tokenizer | None = None
device: torch.device | None = None
config: dict | None = None

# --- Beam Search Decoding Logic (Moved from test_ckpt.py) ---

def beam_search_decode(
    model: pl.LightningModule,
    src: torch.Tensor,
    src_padding_mask: torch.Tensor,
    max_len: int,
    sos_idx: int,
    eos_idx: int,
    pad_idx: int, # Needed for padding mask check if src has padding
    device: torch.device,
    beam_width: int = 5,
    n_best: int = 5, # Number of top sequences to return
    length_penalty: float = 0.6 # Alpha for length normalization (0=no penalty, 1=full penalty)
) -> list[torch.Tensor]:
    """
    Performs beam search decoding using the LightningModule's model.
    (Code copied and pasted from test_ckpt.py)
    """
    # Ensure model is in eval mode (redundant if called after model.eval(), but safe)
    model.eval()
    transformer_model = model.model # Access the underlying Seq2SeqTransformer
    n_best = min(n_best, beam_width) # Cannot return more than beam_width sequences

    try:
        with torch.no_grad():
            # --- Encode Source ---
            memory = transformer_model.encode(src, src_padding_mask) # [1, src_len, emb_size]
            memory = memory.to(device)
            # Ensure memory_key_padding_mask is also on the correct device for decode
            memory_key_padding_mask = src_padding_mask.to(memory.device) # [1, src_len]

            # --- Initialize Beams ---
            initial_beam_seq = torch.ones(1, 1, dtype=torch.long, device=device).fill_(sos_idx) # [1, 1]
            initial_beam_score = torch.zeros(1, dtype=torch.float, device=device) # [1]
            active_beams = [(initial_beam_seq, initial_beam_score)]
            finished_beams = []

            # --- Decoding Loop ---
            for step in range(max_len - 1):
                if not active_beams:
                    break

                potential_next_beams = []
                for current_seq, current_score in active_beams:
                    if current_seq[0, -1].item() == eos_idx:
                        finished_beams.append((current_seq, current_score))
                        continue

                    tgt_input = current_seq # [1, current_len]
                    tgt_seq_len = tgt_input.shape[1]
                    tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device).to(device) # [curr_len, curr_len]
                    tgt_padding_mask = torch.zeros(tgt_input.shape, dtype=torch.bool, device=device) # [1, curr_len]

                    decoder_output = transformer_model.decode(
                        tgt=tgt_input,
                        memory=memory,
                        tgt_mask=tgt_mask,
                        tgt_padding_mask=tgt_padding_mask,
                        memory_key_padding_mask=memory_key_padding_mask
                    ) # [1, curr_len, emb_size]

                    next_token_logits = transformer_model.generator(decoder_output[:, -1, :]) # [1, tgt_vocab_size]
                    log_probs = F.log_softmax(next_token_logits, dim=-1) # [1, tgt_vocab_size]

                    topk_log_probs, topk_indices = torch.topk(log_probs + current_score, beam_width, dim=-1)

                    for i in range(beam_width):
                        next_token_id = topk_indices[0, i].item()
                        next_score = topk_log_probs[0, i].reshape(1) # Keep as tensor [1]
                        next_token_tensor = torch.tensor([[next_token_id]], dtype=torch.long, device=device) # [1, 1]
                        new_seq = torch.cat([current_seq, next_token_tensor], dim=1) # [1, current_len + 1]
                        potential_next_beams.append((new_seq, next_score))

                potential_next_beams.sort(key=lambda x: x[1].item(), reverse=True)

                active_beams = []
                added_count = 0
                for seq, score in potential_next_beams:
                     is_finished = seq[0, -1].item() == eos_idx
                     if is_finished:
                         finished_beams.append((seq, score))
                     elif added_count < beam_width:
                         active_beams.append((seq, score))
                         added_count += 1
                     elif added_count >= beam_width:
                         break

            finished_beams.extend(active_beams)

            # Apply length penalty and sort
            # Handle potential division by zero if sequence length is 1 (or 0?)
            def get_score(beam_tuple):
                seq, score = beam_tuple
                seq_len = seq.shape[1]
                if length_penalty == 0.0 or seq_len <= 1:
                    return score.item()
                else:
                    # Ensure seq_len is float for pow
                    return score.item() / (float(seq_len) ** length_penalty)

            finished_beams.sort(key=get_score, reverse=True) # Higher score is better

            top_sequences = [seq[:, 1:] for seq, score in finished_beams[:n_best]] # seq shape [1, len] -> [1, len-1]
            return top_sequences

    except RuntimeError as e:
        logging.error(f"Runtime error during beam search decode: {e}")
        if "CUDA out of memory" in str(e):
            gc.collect(); torch.cuda.empty_cache()
        return [] # Return empty list on error
    except Exception as e:
        logging.error(f"Unexpected error during beam search decode: {e}", exc_info=True)
        return []

# --- Translation Function (Moved from test_ckpt.py) ---

def translate(
    model: pl.LightningModule,
    src_sentence: str,
    smiles_tokenizer: Tokenizer,
    iupac_tokenizer: Tokenizer,
    device: torch.device,
    max_len: int,
    sos_idx: int,
    eos_idx: int,
    pad_idx: int,
    beam_width: int = 5,
    n_best: int = 5,
    length_penalty: float = 0.6
) -> list[str]:
    """
    Translates a single SMILES string using beam search.
    (Code copied and pasted from test_ckpt.py)
    """
    model.eval() # Ensure model is in eval mode
    translations = []

    # --- Tokenize Source ---
    try:
        src_encoded = smiles_tokenizer.encode(src_sentence)
        if not src_encoded or not src_encoded.ids:
            logging.warning(f"Encoding failed or empty for SMILES: {src_sentence}")
            return ["[Encoding Error]"] * n_best
        src_ids = src_encoded.ids[:max_len] # Truncate source
        if not src_ids:
             logging.warning(f"Source empty after truncation: {src_sentence}")
             return ["[Encoding Error - Empty Src]"] * n_best
    except Exception as e:
        logging.error(f"Error tokenizing SMILES '{src_sentence}': {e}")
        return ["[Encoding Error]"] * n_best

    # --- Prepare Input Tensor and Mask ---
    src = torch.tensor(src_ids, dtype=torch.long).unsqueeze(0).to(device) # [1, src_len]
    src_padding_mask = (src == pad_idx).to(device) # [1, src_len]

    # --- Perform Beam Search Decoding ---
    # Calls the beam_search_decode function defined above in this file
    tgt_tokens_list = beam_search_decode(
        model=model,
        src=src,
        src_padding_mask=src_padding_mask,
        max_len=max_len,
        sos_idx=sos_idx,
        eos_idx=eos_idx,
        pad_idx=pad_idx,
        device=device,
        beam_width=beam_width,
        n_best=n_best,
        length_penalty=length_penalty
    ) # Returns list of tensors

    # --- Decode Generated Tokens ---
    if not tgt_tokens_list:
         logging.warning(f"Beam search returned empty list for SMILES: {src_sentence}")
         return ["[Decoding Error - Empty Output]"] * n_best

    for tgt_tokens_tensor in tgt_tokens_list:
        if tgt_tokens_tensor.numel() > 0:
            tgt_tokens = tgt_tokens_tensor.flatten().cpu().numpy().tolist()
            try:
                translation = iupac_tokenizer.decode(tgt_tokens, skip_special_tokens=True)
                translations.append(translation)
            except Exception as e:
                logging.error(f"Error decoding target tokens {tgt_tokens}: {e}")
                translations.append("[Decoding Error]")
        else:
            translations.append("[Decoding Error - Empty Tensor]")

    # Pad with error messages if fewer than n_best results were generated
    while len(translations) < n_best:
        translations.append("[Decoding Error - Fewer Results]")

    return translations


# --- Model/Tokenizer Loading Function (Unchanged) ---
def load_model_and_tokenizers():
    """Loads tokenizers, config, and model from Hugging Face Hub."""
    global model, smiles_tokenizer, iupac_tokenizer, device, config
    if model is not None: # Already loaded
        logging.info("Model and tokenizers already loaded.")
        return

    logging.info(f"Starting model and tokenizer loading from {MODEL_REPO_ID}...")
    try:
        device = torch.device("cpu")
        logging.info(f"Using device: {device}")

        # Download files from HF Hub
        logging.info("Downloading files from Hugging Face Hub...")
        try:
            checkpoint_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=CHECKPOINT_FILENAME)
            smiles_tokenizer_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=SMILES_TOKENIZER_FILENAME)
            iupac_tokenizer_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=IUPAC_TOKENIZER_FILENAME)
            config_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=CONFIG_FILENAME)
            logging.info("Files downloaded successfully.")
        except Exception as e:
            logging.error(f"Failed to download files from {MODEL_REPO_ID}. Check filenames and repo status. Error: {e}", exc_info=True)
            raise gr.Error(f"Download Error: Could not download required files from {MODEL_REPO_ID}. Check Space logs. Error: {e}")

        # Load config
        logging.info("Loading configuration...")
        try:
            with open(config_path, 'r') as f:
                config = json.load(f)
            logging.info("Configuration loaded.")
            # --- Validate essential config keys ---
            required_keys = [
                'src_vocab_size', 'tgt_vocab_size', 'emb_size', 'nhead',
                'ffn_hid_dim', 'num_encoder_layers', 'num_decoder_layers',
                'dropout', 'max_len', 'bos_token_id', 'eos_token_id', 'pad_token_id'
            ]
            missing_keys = [key for key in required_keys if key not in config]
            if missing_keys:
                raise ValueError(f"Config file '{CONFIG_FILENAME}' is missing required keys: {missing_keys}")
            # --- End Validation ---
        except FileNotFoundError:
             logging.error(f"Config file not found locally after download attempt: {config_path}")
             raise gr.Error(f"Config Error: Config file '{CONFIG_FILENAME}' not found. Check file exists in repo.")
        except json.JSONDecodeError as e:
            logging.error(f"Error decoding JSON from config file {config_path}: {e}")
            raise gr.Error(f"Config Error: Could not parse '{CONFIG_FILENAME}'. Check its format. Error: {e}")
        except ValueError as e:
            logging.error(f"Config validation error: {e}")
            raise gr.Error(f"Config Error: {e}")


        # Load tokenizers
        logging.info("Loading tokenizers...")
        try:
            smiles_tokenizer = Tokenizer.from_file(smiles_tokenizer_path)
            iupac_tokenizer = Tokenizer.from_file(iupac_tokenizer_path)
            logging.info("Tokenizers loaded.")
            # --- Validate Tokenizer Special Tokens ---
            # Add more robust checks if necessary
            if smiles_tokenizer.token_to_id("<pad>") != config['pad_token_id'] or \
               smiles_tokenizer.token_to_id("<unk>") is None:
                 logging.warning("SMILES tokenizer special tokens might not match config or are missing.")
            if iupac_tokenizer.token_to_id("<pad>") != config['pad_token_id'] or \
               iupac_tokenizer.token_to_id("<sos>") != config['bos_token_id'] or \
               iupac_tokenizer.token_to_id("<eos>") != config['eos_token_id'] or \
               iupac_tokenizer.token_to_id("<unk>") is None:
                 logging.warning("IUPAC tokenizer special tokens might not match config or are missing.")
            # --- End Validation ---
        except Exception as e:
            logging.error(f"Failed to load tokenizers from {smiles_tokenizer_path} or {iupac_tokenizer_path}: {e}", exc_info=True)
            raise gr.Error(f"Tokenizer Error: Could not load tokenizer files. Check Space logs. Error: {e}")

        # Load model
        logging.info("Loading model from checkpoint...")
        try:
            model = SmilesIupacLitModule.load_from_checkpoint(
                checkpoint_path,
                src_vocab_size=config['src_vocab_size'],
                tgt_vocab_size=config['tgt_vocab_size'],
                map_location=device,
                hparams_dict=config,
                strict=False,
                device="cpu"
            )
            model.to(device)
            model.eval()
            model.freeze()
            logging.info("Model loaded successfully, set to eval mode, frozen, and moved to device.")

        except FileNotFoundError:
            logging.error(f"Checkpoint file not found locally after download attempt: {checkpoint_path}")
            raise gr.Error(f"Model Error: Checkpoint file '{CHECKPOINT_FILENAME}' not found.")
        except Exception as e:
            logging.error(f"Error loading model from checkpoint {checkpoint_path}: {e}", exc_info=True)
            if "memory" in str(e).lower():
                gc.collect()
                if device == torch.device("cuda"):
                    torch.cuda.empty_cache()
            raise gr.Error(f"Model Error: Failed to load model checkpoint. Check Space logs. Error: {e}")

    except gr.Error:
         raise
    except Exception as e:
        logging.error(f"Unexpected error during model/tokenizer loading: {e}", exc_info=True)
        raise gr.Error(f"Initialization Error: An unexpected error occurred. Check Space logs. Error: {e}")


# --- Inference Function for Gradio (Unchanged, calls local translate) ---
def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
    """
    Performs SMILES to IUPAC translation using the loaded model and beam search.
    """
    global model, smiles_tokenizer, iupac_tokenizer, device, config

    if not all([model, smiles_tokenizer, iupac_tokenizer, device, config]):
        error_msg = "Error: Model or tokenizers not loaded properly. Check Space logs."
        # Ensure n_best is int for range, default to 1 if conversion fails early
        try: n_best_int = int(n_best)
        except: n_best_int = 1
        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best_int)])

    if not smiles_string or not smiles_string.strip():
        error_msg = "Error: Please enter a valid SMILES string."
        try: n_best_int = int(n_best)
        except: n_best_int = 1
        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best_int)])

    smiles_input = smiles_string.strip()
    try:
        beam_width = int(beam_width)
        n_best = int(n_best)
        length_penalty = float(length_penalty)
    except ValueError as e:
         error_msg = f"Error: Invalid input parameter type ({e})."
         return f"1. {error_msg}" # Cannot determine n_best here

    logging.info(f"Translating SMILES: '{smiles_input}' (Beam={beam_width}, N={n_best}, Penalty={length_penalty})")

    try:
        # Calls the translate function defined *above in this file*
        predicted_names = translate(
            model=model,
            src_sentence=smiles_input,
            smiles_tokenizer=smiles_tokenizer,
            iupac_tokenizer=iupac_tokenizer,
            device=device,
            max_len=config['max_len'],
            sos_idx=config['bos_token_id'],
            eos_idx=config['eos_token_id'],
            pad_idx=config['pad_token_id'],
            beam_width=beam_width,
            n_best=n_best,
            length_penalty=length_penalty
        )
        logging.info(f"Predictions returned: {predicted_names}")

        if not predicted_names:
             output_text = f"Input SMILES: {smiles_input}\n\nNo predictions generated."
        else:
            output_text = f"Input SMILES: {smiles_input}\n\nTop {len(predicted_names)} Predictions (Beam Width={beam_width}, Length Penalty={length_penalty:.2f}):\n"
            output_text += "\n".join([f"{i+1}. {name}" for i, name in enumerate(predicted_names)])

        return output_text

    except RuntimeError as e:
        logging.error(f"Runtime error during translation: {e}", exc_info=True)
        error_msg = f"Runtime Error during translation: {e}"
        if "memory" in str(e).lower():
            gc.collect()
            if device == torch.device("cuda"):
                torch.cuda.empty_cache()
            error_msg += " (Potential OOM)"
        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best)])

    except Exception as e:
        logging.error(f"Unexpected error during translation: {e}", exc_info=True)
        error_msg = f"Unexpected Error during translation: {e}"
        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best)])


# --- Load Model on App Start (Unchanged) ---
try:
    load_model_and_tokenizers()
except gr.Error:
    pass # Error already raised for Gradio UI
except Exception as e:
    logging.error(f"Critical error during initial model loading sequence: {e}", exc_info=True)
    gr.Error(f"Fatal Initialization Error: {e}. Check Space logs.")


# --- Create Gradio Interface (Unchanged) ---
title = "SMILES to IUPAC Name Translator"
description = f"""
Enter a SMILES string to translate it into its IUPAC chemical name using a Transformer model and beam search decoding.
Model repository: <a href='https://huggingface.co/{MODEL_REPO_ID}' target='_blank'>{MODEL_REPO_ID}</a>.
Adjust beam search parameters below. Higher beam width explores more possibilities but is slower. Length penalty influences the preference for shorter/longer names.
"""

examples = [
    ["CCO", 5, 3, 0.6],
    ["C1=CC=CC=C1", 5, 3, 0.6],
    ["CC(=O)Oc1ccccc1C(=O)O", 5, 3, 0.6], # Aspirin
    ["CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", 5, 3, 0.6], # Ibuprofen
    ["CC(=O)O[C@@H]1C[C@@H]2[C@]3(CCCC([C@@H]3CC[C@]2([C@H]4[C@]1([C@H]5[C@@H](OC(=O)C5=CC4)OC)C)C)(C)C)C", 5, 1, 0.6], # Complex example
    ["INVALID_SMILES", 5, 1, 0.6],
]

smiles_input = gr.Textbox(
    label="SMILES String",
    placeholder="Enter SMILES string here (e.g., CCO for Ethanol)",
    lines=1
)
beam_width_input = gr.Slider(
    minimum=1,
    maximum=10,
    value=5,
    step=1,
    label="Beam Width (k)",
    info="Number of sequences to keep at each decoding step (higher = more exploration, slower)."
)
n_best_input = gr.Slider(
    minimum=1,
    maximum=10,
    value=3,
    step=1,
    label="Number of Results (n_best)",
    info="How many top-scoring sequences to return (must be <= Beam Width)."
)
length_penalty_input = gr.Slider(
    minimum=0.0,
    maximum=2.0,
    value=0.6,
    step=0.1,
    label="Length Penalty (alpha)",
    info="Controls preference for sequence length. >1 prefers longer, <1 prefers shorter, 0 no penalty."
)
output_text = gr.Textbox(
    label="Predicted IUPAC Name(s)",
    lines=5,
    show_copy_button=True
)

iface = gr.Interface(
    fn=predict_iupac,
    inputs=[smiles_input, beam_width_input, n_best_input, length_penalty_input],
    outputs=output_text,
    title=title,
    description=description,
    examples=examples,
    allow_flagging="never",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"),
    article="Note: Translation quality depends on the training data and model size. Complex molecules might yield less accurate results."
)

# --- Launch the App (Unchanged) ---
if __name__ == "__main__":
    iface.launch(share=True)