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

# --- 1. Load Custom Model Code ---
# This import now works because we have the correct models/hrm/ structure
from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1

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

# --- 3. Initialize the Model ---
print("Initializing model...")
model_config.update({
    'batch_size': config_data['global_batch_size'],
    'seq_len': 512,
    'num_puzzle_identifiers': len(activity_code_map) + 1,
    'vocab_size': tokenizer.get_vocab_size()
})
model = HierarchicalReasoningModel_ACTV1(config_dict=model_config)

# --- MODIFICATION: Clean the state dict keys before loading ---
# Load the original state dict
original_state_dict = torch.load('pytorch_model.bin', map_location='cpu')
# Create a new state dict with the corrected keys
new_state_dict = OrderedDict()
for k, v in original_state_dict.items():
    if k.startswith('_orig_mod.model.'):
        name = k[len('_orig_mod.model.'):] # remove the prefix
        new_state_dict[name] = v
    else:
        new_state_dict[k] = v

# Load the cleaned state dict
model.load_state_dict(new_state_dict)
# --- END MODIFICATION ---

model.eval() # Set the model to evaluation mode
print("Model loaded successfully!")

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

    try:
        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)

        input_tensor = torch.tensor([input_ids], dtype=torch.long)
        grant_tensor = torch.tensor([grant_type_id], dtype=torch.long)

        batch = {
            "inputs": input_tensor,
            "puzzle_identifiers": grant_tensor,
            "labels": torch.zeros_like(input_tensor)
        }

        with torch.no_grad():
            carry = model.initial_carry(batch)
            for _ in range(model_config['halt_max_steps']):
                carry, _ = model(carry=carry, batch=batch)
            
            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()
        
        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()