import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from rich.console import Console
import time

# Initialize rich console for better logging
console = Console()

# Load the model and tokenizer with the same configuration as training
console.print("[bold green]Loading model and tokenizer...[/bold green]")

# Load model with memory optimizations
model = AutoModelForCausalLM.from_pretrained(
    "./fine-tuned-model",
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.float16,  # Use float16 for memory efficiency
    low_cpu_mem_usage=True,  # Add this for better memory handling
)
tokenizer = AutoTokenizer.from_pretrained("./fine-tuned-model")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'left'

# Load base model for before/after comparison
console.print("[bold green]Loading base model for comparison...[/bold green]")
base_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/phi-2",
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,  # Add this for better memory handling
)

def generate_response(
    prompt,
    max_length=128,  # Match training max_length
    temperature=0.7,
    top_p=0.9,
    num_generations=2,  # Match training num_generations
    repetition_penalty=1.1,
    do_sample=True,
    show_comparison=True,  # New parameter for comparison toggle
):
    try:
        # Get the device of the model
        device = next(model.parameters()).device
        
        # Tokenize the input
        inputs = tokenizer(prompt, return_tensors="pt", padding=True)
        
        # Move inputs to the same device as the model
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        # Generate response from fine-tuned model
        with torch.no_grad():  # Disable gradient computation
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_length,
                do_sample=do_sample,
                temperature=temperature,
                top_p=top_p,
                num_return_sequences=num_generations,
                repetition_penalty=repetition_penalty,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        
        # Decode and return the responses
        responses = []
        for output in outputs:
            response = tokenizer.decode(output, skip_special_tokens=True)
            responses.append(response)
        
        fine_tuned_response = "\n\n---\n\n".join(responses)
        
        if show_comparison:
            # Generate response from base model
            with torch.no_grad():
                base_outputs = base_model.generate(
                    **inputs,
                    max_new_tokens=max_length,
                    do_sample=do_sample,
                    temperature=temperature,
                    top_p=top_p,
                    num_return_sequences=1,  # Only one for comparison
                    repetition_penalty=repetition_penalty,
                    pad_token_id=tokenizer.eos_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                )
            
            base_response = tokenizer.decode(base_outputs[0], skip_special_tokens=True)
            
            return f"""
### Before Fine-tuning (Base Model)
{base_response}

### After Fine-tuning
{fine_tuned_response}
"""
        else:
            return fine_tuned_response
            
    except Exception as e:
        console.print(f"[bold red]Error during generation: {str(e)}[/bold red]")
        return f"Error: {str(e)}"

# Create custom CSS for better UI
custom_css = """
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
.container {
    max-width: 800px;
    margin: auto;
    padding: 20px;
}
.title {
    text-align: center;
    color: #2c3e50;
    margin-bottom: 20px;
}
.description {
    color: #34495e;
    line-height: 1.6;
    margin-bottom: 20px;
}
.comparison {
    background-color: #f8f9fa;
    padding: 15px;
    border-radius: 8px;
    margin: 10px 0;
}
.prompt-box {
    background-color: #ffffff;
    border: 2px solid #3498db;
    border-radius: 8px;
    padding: 15px;
    margin-bottom: 20px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.prompt-box label {
    font-size: 1.1em;
    font-weight: bold;
    color: #2c3e50;
    margin-bottom: 10px;
    display: block;
}
.prompt-box textarea {
    width: 100%;
    min-height: 100px;
    padding: 10px;
    border: 1px solid #bdc3c7;
    border-radius: 4px;
    font-size: 1em;
    line-height: 1.5;
}
.output-box {
    background-color: #ffffff;
    border: 2px solid #2ecc71;
    border-radius: 8px;
    padding: 20px;
    margin-top: 20px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.output-box label {
    font-size: 1.1em;
    font-weight: bold;
    color: #2c3e50;
    margin-bottom: 15px;
    display: block;
}
.output-box .markdown {
    background-color: #f8f9fa;
    padding: 15px;
    border-radius: 6px;
    border: 1px solid #e9ecef;
}
.output-box h3 {
    color: #2c3e50;
    border-bottom: 2px solid #3498db;
    padding-bottom: 8px;
    margin-top: 20px;
}
.output-box p {
    line-height: 1.6;
    color: #34495e;
    margin: 10px 0;
}
.loading {
    display: flex;
    align-items: center;
    justify-content: center;
    padding: 20px;
    background-color: #f8f9fa;
    border-radius: 8px;
    margin: 10px 0;
}
.loading-spinner {
    width: 40px;
    height: 40px;
    border: 4px solid #f3f3f3;
    border-top: 4px solid #3498db;
    border-radius: 50%;
    animation: spin 1s linear infinite;
    margin-right: 15px;
}
@keyframes spin {
    0% { transform: rotate(0deg); }
    100% { transform: rotate(360deg); }
}
.loading-text {
    color: #2c3e50;
    font-size: 1.1em;
    font-weight: 500;
}
"""

# Create the Gradio interface with enhanced UI
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Phi-2 Fine-tuned with GRPO and qLoRA
        This model has been fine-tuned using GRPO (Generative Reward-Penalized Optimization) and compressed using qLoRA.
        Try it out with different prompts and generation parameters!
        """,
        elem_classes="title"
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            with gr.Column(elem_classes="prompt-box"):
                prompt = gr.Textbox(
                    label="Enter Your Prompt Here",
                    placeholder="Type your prompt here... (e.g., 'What is machine learning?' or 'Write a story about a robot learning to paint')",
                    lines=5,
                    show_label=True,
                )
            
            with gr.Row():
                with gr.Column():
                    max_length = gr.Slider(
                        minimum=32,
                        maximum=256,
                        value=128,
                        step=32,
                        label="Max Length",
                        info="Maximum number of tokens to generate"
                    )
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.7,
                        step=0.1,
                        label="Temperature",
                        info="Higher values make output more random, lower values more deterministic"
                    )
                with gr.Column():
                    top_p = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.9,
                        step=0.1,
                        label="Top-p",
                        info="Nucleus sampling parameter"
                    )
                    num_generations = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=2,
                        step=1,
                        label="Number of Generations",
                        info="Number of different responses to generate"
                    )
            
            with gr.Row():
                with gr.Column():
                    repetition_penalty = gr.Slider(
                        minimum=1.0,
                        maximum=2.0,
                        value=1.1,
                        step=0.1,
                        label="Repetition Penalty",
                        info="Higher values prevent repetition"
                    )
                with gr.Column():
                    do_sample = gr.Checkbox(
                        value=True,
                        label="Enable Sampling",
                        info="Enable/disable sampling for deterministic output"
                    )
            
            show_comparison = gr.Checkbox(
                value=True,
                label="Show Before/After Comparison",
                info="Toggle to show responses from both base and fine-tuned models"
            )
            
            generate_btn = gr.Button("Generate", variant="primary", size="large")
        
        with gr.Column(scale=3):
            with gr.Column(elem_classes="output-box"):
                output = gr.Markdown(
                    label="Generated Response(s)",
                    show_label=True,
                    value="Your generated responses will appear here...",  # Add default value
                )
                loading_status = gr.Markdown(
                    value="",
                    show_label=False,
                    elem_classes="loading"
                )
    
    gr.Markdown(
        """
        ### Example Prompts
        Try these example prompts to test the model:
        
        1. **Technical Questions**:
           - "What is machine learning?"
           - "What is deep learning?"
           - "What is the difference between supervised and unsupervised learning?"
        
        2. **Creative Writing**:
           - "Write a short story about a robot learning to paint."
           - "Write a story about a time-traveling smartphone."
           - "Write a fairy tale about a computer learning to dream."
           - "Create a story about an AI becoming an artist."
        
        3. **Technical Explanations**:
           - "How does neural network training work?"
           - "Explain quantum computing in simple terms."
           - "What is transfer learning?"
        
        4. **Creative Tasks**:
           - "Write a poem about artificial intelligence."
           - "Write a poem about the future of technology."
           - "Create a story about a robot learning to dream."
        """,
        elem_classes="description"
    )
    
    def generate_with_status(*args):
        # Show loading status
        loading_status.value = """
        <div class="loading">
            <div class="loading-spinner"></div>
            <div class="loading-text">Generating responses... Please wait...</div>
        </div>
        """
        # Generate response
        result = generate_response(*args)
        # Clear loading status
        loading_status.value = ""
        return result
    
    # Connect the interface
    generate_btn.click(
        fn=generate_with_status,
        inputs=[
            prompt,
            max_length,
            temperature,
            top_p,
            num_generations,
            repetition_penalty,
            do_sample,
            show_comparison
        ],
        outputs=output
    )

if __name__ == "__main__":
    console.print("[bold green]Starting Gradio interface...[/bold green]")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True  # Enable sharing for HuggingFace Spaces
    )