import os

# 🛡️ Ensure model is saved before importing anything that loads it
if not os.path.exists("trained_model/config.json"):
    print("⚙️ No model found. Initializing and saving EvoTransformer...")
    from init_model import initialize_and_save_model
    initialize_and_save_model()
else:
    print("✅ EvoTransformer already initialized.")

import gradio as gr
import random
from inference import generate_response
from logger import log_user_feedback
from dashboard import update_dashboard_plot, evolution_accuracy_plot
from watchdog import retrain_model
from init_model import load_model

# === Load model to extract architecture
model = load_model()

def get_architecture_summary(model):
    summary = {
        "Layers": getattr(model, "num_layers", "N/A"),
        "Attention Heads": getattr(model, "num_heads", "N/A"),
        "FFN Dim": getattr(model, "ffn_dim", "N/A"),
        "Memory Enabled": getattr(model, "use_memory", "N/A"),
    }
    return "\n".join(f"{k}: {v}" for k, v in summary.items())

# 🎲 Random examples
examples = [
    {
        "goal": "Escape from a burning house",
        "option1": "Run out through the front door",
        "option2": "Hide in the bathroom"
    },
    {
        "goal": "Improve sleep quality",
        "option1": "Use phone in bed",
        "option2": "Turn off screens 1 hour before bed"
    },
    {
        "goal": "Increase productivity at work",
        "option1": "Multitask all day",
        "option2": "Use Pomodoro technique"
    },
    {
        "goal": "Lose weight safely",
        "option1": "Skip meals",
        "option2": "Exercise regularly and eat balanced meals"
    },
    {
        "goal": "Ace an exam",
        "option1": "Cram the night before",
        "option2": "Study consistently for 2 weeks"
    },
]

def load_random_example():
    example = random.choice(examples)
    return example["goal"], example["option1"], example["option2"]

# 🧠 Model suggestion UI
evo_output = gr.Textbox(label="🧠 EvoTransformer Suggestion")
gpt_output = gr.Textbox(label="💬 GPT-3.5 Suggestion")
feedback_output = gr.Textbox(visible=False)

def evo_chat(goal, sol1, sol2):
    response = generate_response(goal, sol1, sol2)
    evo = response.get("evo_suggestion", "Error")
    gpt = response.get("gpt_suggestion", "Error")
    return evo, gpt

def handle_feedback(goal, sol1, sol2, winner):
    try:
        log_user_feedback(goal, sol1, sol2, winner)
        return "✅ Feedback logged. Thank you!"
    except Exception as e:
        return f"❌ Failed to log: {e}"

with gr.Blocks(title="EvoTransformer v2.1 – Compare Options and Learn") as demo:
    gr.Markdown("""
# 🤖 EvoTransformer v2.1  
*Built Different. Learns Live.*

### 🧬 Why EvoTransformer is Special

Unlike traditional AI models that are fixed once deployed, **EvoTransformer evolves**.

- ✅ Learns from **real-time feedback**
- ✅ **Retrains** on your input (every time you click “♻️ Retrain”)
- ✅ Tracks its **own architecture and accuracy** as it grows
- ✅ Competes with **GPT-3.5** on real reasoning tasks
- ✅ Runs light — only ~13M parameters

This isn't just a chatbot. It's a living AI experiment.  
Every interaction makes it smarter. You're shaping it.

---
""")

    
    gr.Markdown("## 🧠 EvoTransformer v2.1 – Compare Options and Learn")

    with gr.Row():
        goal_input = gr.Textbox(label="Goal", placeholder="e.g. Escape from house on fire")
    with gr.Row():
        option1_input = gr.Textbox(label="Option 1", placeholder="e.g. Exit house through main door")
        option2_input = gr.Textbox(label="Option 2", placeholder="e.g. Hide under bed")

    with gr.Row():
        compare_btn = gr.Button("🔍 Compare")
        random_btn = gr.Button("🎲 Load Random Example")

    with gr.Row():
        evo_output.render()
        gpt_output.render()

    with gr.Row():
        winner_dropdown = gr.Radio(["Solution 1", "Solution 2"], label="Which was better?")
        feedback_btn = gr.Button("✅ Log Feedback")

    feedback_output.render()

    compare_btn.click(
        fn=evo_chat,
        inputs=[goal_input, option1_input, option2_input],
        outputs=[evo_output, gpt_output]
    )

    random_btn.click(
        fn=load_random_example,
        inputs=[],
        outputs=[goal_input, option1_input, option2_input]
    )

    feedback_btn.click(
        fn=handle_feedback,
        inputs=[goal_input, option1_input, option2_input, winner_dropdown],
        outputs=[feedback_output]
    )

    # === Define UI components to allow later refresh ===
    retrain_status = gr.Textbox(label="Retrain Status")
    arch_box = gr.Textbox(label="Model Configuration", value=get_architecture_summary(model), lines=5, interactive=False)
    evo_plot = gr.Plot(label="Accuracy per Generation", value=evolution_accuracy_plot())
    dashboard_img = gr.Image(value=update_dashboard_plot(), label="Solution Votes")

    with gr.Row():
        retrain_button = gr.Button("♻️ Retrain Evo", variant="primary")

    retrain_button.click(
        fn=retrain_model,
        inputs=[],
        outputs=[arch_box, evo_plot, retrain_status],
        show_progress=True
    )

    # === Render inside Accordions ===
    with gr.Accordion("🧬 EvoTransformer Architecture", open=False):
        arch_box

    with gr.Accordion("📈 Evolution Accuracy Plot", open=False):
        evo_plot

    with gr.Accordion("📊 User Feedback Summary", open=False):
        dashboard_img

    with gr.Accordion("🔁 Retrain Status", open=False):
        retrain_status


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