import gradio as gr
import pandas as pd
import traceback
import sys

from cognitive_mapping_probe.orchestrator_seismograph import run_seismic_analysis
from cognitive_mapping_probe.prompts import RESONANCE_PROMPTS
from cognitive_mapping_probe.utils import dbg

# --- UI Theme and Layout ---
theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue").set(
    body_background_fill="#f0f4f9",
    block_background_fill="white",
)

def run_and_display(
    model_id: str,
    prompt_type: str,
    seed: int,
    num_steps: int,
    concept_to_inject: str,
    injection_strength: float,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Führt die seismische Analyse durch, inklusive der optionalen Konzeptinjektion.
    """
    try:
        results = run_seismic_analysis(
            model_id=model_id,
            prompt_type=prompt_type,
            seed=int(seed),
            num_steps=int(num_steps),
            concept_to_inject=concept_to_inject,
            injection_strength=float(injection_strength),
            progress_callback=progress
        )

        verdict = results.get("verdict", "Analysis complete.")
        stats = results.get("stats", {})
        deltas = results.get("state_deltas", [])

        df = pd.DataFrame({
            "Internal Step": range(len(deltas)),
            "State Change (Delta)": deltas
        })

        stats_md = f"### Statistical Signature\n"
        stats_md += f"- **Mean Delta:** {stats.get('mean_delta', 0):.4f}\n"
        stats_md += f"- **Std Dev Delta:** {stats.get('std_delta', 0):.4f}\n"
        stats_md += f"- **Max Delta:** {stats.get('max_delta', 0):.4f}\n"

        return f"{verdict}\n\n{stats_md}", df, results

    except Exception:
        error_str = traceback.format_exc()
        return f"### ❌ Analysis Failed\nAn unexpected error occurred:\n\n```\n{error_str}\n```", pd.DataFrame(), {}

# --- Gradio App Definition ---
with gr.Blocks(theme=theme, title="Cognitive Seismograph 2.0") as demo:
    gr.Markdown("# 🧠 Cognitive Seismograph 2.0: Modulating Internal Dynamics")
    gr.Markdown(
        "**Neues Paradigma:** Wir messen nicht nur die intrinsische Dynamik, sondern versuchen sie aktiv durch **Konzeptinjektionen** zu modulieren. Vergleiche die 'seismische Signatur' mit und ohne Injektion."
    )
    with gr.Row(variant='panel'):
        with gr.Column(scale=1):
            gr.Markdown("### 1. General Parameters")
            model_id_input = gr.Textbox(value="google/gemma-3-1b-it", label="Model ID")
            prompt_type_input = gr.Radio(choices=list(RESONANCE_PROMPTS.keys()), value="resonance_prompt", label="Prompt Type")
            seed_input = gr.Slider(1, 1000, 42, step=1, label="Seed")
            num_steps_input = gr.Slider(50, 1000, 300, step=10, label="Number of Internal Steps")

            gr.Markdown("### 2. Modulation Parameters")
            concept_input = gr.Textbox(label="Concept to Inject", placeholder="e.g., 'solitude' or 'ocean' (leave blank for baseline)")
            strength_input = gr.Slider(0.0, 5.0, 1.0, step=0.1, label="Injection Strength")

            run_btn = gr.Button("Run Seismic Analysis", variant="primary")

        with gr.Column(scale=2):
            gr.Markdown("### Results")
            verdict_output = gr.Markdown("Die Analyse der Dynamik erscheint hier.")
            plot_output = gr.LinePlot(
                x="Internal Step",
                y="State Change (Delta)",
                title="Internal State Dynamics (Cognitive EKG)",
                show_label=True,
                height=400,
            )
            with gr.Accordion("Raw JSON Output", open=False):
                raw_json_output = gr.JSON()

    run_btn.click(
        fn=run_and_display,
        inputs=[model_id_input, prompt_type_input, seed_input, num_steps_input, concept_input, strength_input],
        outputs=[verdict_output, plot_output, raw_json_output]
    )

if __name__ == "__main__":
    print("="*80)
    print("🔬 COGNITIVE SEISMOGRAPH 2.0 (MODULATION-ENABLED) INITIALIZED")
    print("="*80)
    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)