import gradio as gr
import torch
import yaml
import json
from tokenizers import Tokenizer

# --- 1. Load Custom Model Code ---
# This dynamically loads your corrected HRM source code.
# Since the files are in the root, we import directly.
from hrm_act_v1 import HierarchicalReasoningModel_ACTV1

# --- 2. Load Artifacts ---
print("Loading artifacts...")
# Load the tokenizer
tokenizer = Tokenizer.from_file("tokenizer.json")
# Load the model configuration
with open('config.yaml', 'r') as f:
    config_data = yaml.safe_load(f)
model_config = config_data['arch']
# Load the grant type mapping
with open('activity_code_map.json', 'r') as f:
    activity_code_map = json.load(f)

# --- 3. Initialize the Model ---
print("Initializing model...")
# The model expects a dict, so we pass the Pydantic model's dict representation
# We also need to add other required keys from the root of the config
model_config.update({
    'batch_size': config_data['global_batch_size'],
    'seq_len': 512, # You may need to get this from your dataset metadata
    'num_puzzle_identifiers': len(activity_code_map) + 1,
    'vocab_size': tokenizer.get_vocab_size()
})
model = HierarchicalReasoningModel_ACTV1(config_dict=model_config)
# Load the fine-tuned weights
model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
model.eval() # Set the model to evaluation mode
print("Model loaded successfully!")

# --- 4. Define the Inference Function ---
def optimize_abstract(draft_abstract, grant_type):
    """
    Takes a draft abstract and grant type, runs the model, and returns the optimized text.
    """
    if not draft_abstract or not grant_type:
        return "Please provide both a draft abstract and a grant type."

    try:
        # Prepare inputs
        tokenizer.enable_padding(length=512)
        tokenizer.enable_truncation(max_length=512)
        
        input_ids = tokenizer.encode(draft_abstract).ids
        grant_type_id = activity_code_map.get(grant_type, 0) # Default to 0 if unknown

        # Convert to PyTorch tensors
        input_tensor = torch.tensor([input_ids], dtype=torch.long)
        grant_tensor = torch.tensor([grant_type_id], dtype=torch.long)

        # Create the batch dictionary that the model expects
        batch = {
            "inputs": input_tensor,
            "puzzle_identifiers": grant_tensor,
            # The model requires a 'labels' field, even for inference, so we provide a dummy one
            "labels": torch.zeros_like(input_tensor)
        }

        # Run inference
        with torch.no_grad():
            carry = model.initial_carry(batch)
            # The model runs in a loop; for inference, we run it for the max steps
            for _ in range(model_config['halt_max_steps']):
                carry, _ = model(carry=carry, batch=batch)
            
            # Get the final logits from the carry state
            final_logits = model.inner.lm_head(carry.inner_carry.z_H)[:, model.inner.puzzle_emb_len:]
            predicted_ids = torch.argmax(final_logits, dim=-1).squeeze().tolist()
        
        # Decode the output
        optimized_text = tokenizer.decode(predicted_ids, skip_special_tokens=True)
        
        return optimized_text

    except Exception as e:
        print(f"An error occurred during inference: {e}")
        return f"Error: Could not process the abstract. Details: {e}"

# --- 5. Create the Gradio Interface ---
grant_type_choices = list(activity_code_map.keys())

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 HRM Grant Abstract Optimizer")
    gr.Markdown("Enter a draft abstract and select the grant type to get a version optimized by the fine-tuned Hierarchical Reasoning Model.")
    
    with gr.Row():
        with gr.Column():
            draft_input = gr.Textbox(label="Draft Abstract", lines=15, placeholder="Paste your draft abstract here...")
            grant_type = gr.Dropdown(label="Grant Type", choices=grant_type_choices, value=grant_type_choices[0] if grant_type_choices else None)
            optimize_btn = gr.Button("Optimize Abstract", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="Optimized Abstract", lines=17, interactive=False)
            
    optimize_btn.click(
        fn=optimize_abstract,
        inputs=[draft_input, grant_type],
        outputs=output_text
    )

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