import gradio as gr
import torch
import yaml
import os

def load_model():
    """Load the HRM model and config"""
    try:
        # Load config
        with open('all_config.yaml', 'r') as f:
            config = yaml.safe_load(f)
        
        # Load checkpoint
        checkpoint = torch.load('pytorch_model.bin', map_location='cpu')
        
        return config, checkpoint, "✅ Model loaded successfully!"
    except Exception as e:
        return None, None, f"❌ Error loading model: {str(e)}"

def test_model_info(config, checkpoint):
    """Display model information"""
    if config is None or checkpoint is None:
        return "Model not loaded"
    
    info = f"""
    **Model Architecture**: {config['arch']['name']}
    **Hidden Size**: {config['arch']['hidden_size']}
    **H Layers**: {config['arch']['H_layers']}
    **L Layers**: {config['arch']['L_layers']}
    **Parameters in Checkpoint**: {len(checkpoint)}
    **Model Purpose**: Grant Abstract Optimization
    
    **Training Details**:
    - Steps: 492,500 (final checkpoint)
    - Batch Size: {config['global_batch_size']}
    - Learning Rate: {config['lr']}
    """
    return info

def placeholder_inference(draft_abstract, grant_type):
    """Placeholder for actual inference (requires full training pipeline)"""
    return f"""
    **Input Abstract**: {draft_abstract[:100]}...
    
    **Grant Type**: {grant_type}
    
    **Status**: Model checkpoint loaded successfully! 
    
    ⚠️ **Note**: Full inference requires the original training pipeline with tokenizer and preprocessing code. 
    This demo shows that the model weights are accessible and the architecture is properly configured.
    
    **Next Steps**: 
    1. Integrate with original training codebase
    2. Load tokenizer and preprocessing pipeline  
    3. Implement full inference function
    """

# Load model on startup
config, checkpoint, load_status = load_model()

# Create Gradio interface
with gr.Blocks(title="HRM Grant Abstract Optimizer") as demo:
    gr.Markdown("# 🎯 Hierarchical Reasoning Model for Grant Abstract Optimization")
    gr.Markdown("A specialized 27M-parameter model for transforming draft grant abstracts into funding-worthy versions.")
    
    with gr.Tab("Model Info"):
        gr.Markdown("## Model Status")
        gr.Markdown(load_status)
        
        if config is not None:
            model_info = test_model_info(config, checkpoint)
            gr.Markdown(model_info)
    
    with gr.Tab("Test Interface"):
        gr.Markdown("## Abstract Optimization Demo")
        gr.Markdown("*Note: This is a demonstration interface. Full inference requires integration with the training pipeline.*")
        
        with gr.Row():
            with gr.Column():
                draft_input = gr.Textbox(
                    label="Draft Abstract",
                    placeholder="Enter your sub-optimal grant abstract here...",
                    lines=8,
                    value="Our study will investigate protein interactions in cancer cells. We believe this research could be important for understanding disease mechanisms."
                )
                grant_type = gr.Dropdown(
                    choices=["R01", "F32", "K99", "R21", "R15"],
                    label="Grant Type",
                    value="R01"
                )
                optimize_btn = gr.Button("Optimize Abstract", variant="primary")
            
            with gr.Column():
                output = gr.Textbox(
                    label="Optimized Abstract",
                    lines=10,
                    interactive=False
                )
        
        optimize_btn.click(
            fn=placeholder_inference,
            inputs=[draft_input, grant_type],
            outputs=output
        )

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