from __future__ import annotations

import numpy as np
import gradio as gr
import matplotlib.pyplot as plt

from model import HolographicMasterCodeTransformer, synthetic_loss_curve, alpha_to_label


def build_training_figure(alpha_value: float, epochs: int):
    nominal_losses, nominal_regs = synthetic_loss_curve(1 / 137.035, epochs=epochs)
    current_losses, current_regs = synthetic_loss_curve(alpha_value, epochs=epochs)

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(nominal_losses, label="Loss — nominal α", linewidth=2.2)
    ax.plot(current_losses, label="Loss — selected α", linewidth=2.2)
    ax.plot(nominal_regs, label="Reg — nominal α", linestyle="--")
    ax.plot(current_regs, label="Reg — selected α", linestyle="--")
    ax.set_title("OmegaCode training dynamics")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Value")
    ax.grid(True, alpha=0.25)
    ax.legend(loc="upper right")
    fig.tight_layout()
    return fig


def build_interference_figure(alpha_value: float):
    x = np.linspace(-10, 10, 1200)
    alpha_0 = 1 / 137.035
    delta = alpha_value - alpha_0

    # Coherence and phase shift are synthetic and only for visualization.
    visibility = float(np.exp(-25000.0 * abs(delta)))
    visibility = max(0.08, min(0.98, visibility))
    phase = float(0.9 * np.sign(delta) * min(1.0, abs(delta) * 5e4))

    coherent = 1.0 + 0.95 * np.cos(1.15 * x)
    perturbed = 1.0 + visibility * np.cos(1.15 * x + phase)

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(x, coherent, label="Nominal α: coherent pattern", linewidth=2.2)
    ax.plot(x, perturbed, label="Selected α: phase-modulated pattern", linewidth=2.2)
    ax.set_title("Double-slit analogue: interference visibility vs α")
    ax.set_xlabel("Screen coordinate")
    ax.set_ylabel("Normalized intensity")
    ax.grid(True, alpha=0.25)
    ax.legend(loc="upper right")
    fig.tight_layout()
    return fig


def run_demo(alpha_value: float, epochs: int, noise: float):
    model = HolographicMasterCodeTransformer()

    x = np.random.randn(1, 8).astype(np.float32) * max(noise, 1e-6)
    pred, psi, security, delta_phi = model.forward(x, alpha_value)

    training_fig = build_training_figure(alpha_value, epochs)
    interference_fig = build_interference_figure(alpha_value)

    alpha_label = alpha_to_label(alpha_value)
    alpha_0 = model.alpha_0
    delta_alpha = alpha_value - alpha_0

    summary = {
        "alpha": alpha_value,
        "alpha_label": alpha_label,
        "delta_alpha": delta_alpha,
        "prediction": float(pred.squeeze()),
        "psi_mean": float(np.mean(psi)),
        "security_factor": float(security),
        "phase_shift": float(delta_phi),
        "stability_score": float(max(0.0, 1.0 - abs(delta_alpha) * 5e4)),
    }

    return training_fig, interference_fig, summary


with gr.Blocks(title="OmegaCode Holographic Demo") as demo:
    gr.Markdown(
        """
# OmegaCode — Holographic Master Code Demo

A research-oriented, physics-inspired neural prototype with an α-sensitive stability control and interference-based visualization.
It is designed for exploration and does **not** claim to be a validated physical theory.

## What to try

- Move **α** around the nominal value
- Watch the **loss curve** shift
- Compare the **double-slit analogue** as phase coherence changes
        """
    )

    with gr.Row():
        alpha_value = gr.Slider(
            minimum=1 / 137.08,
            maximum=1 / 137.00,
            value=1 / 137.035,
            step=1e-7,
            label="α / Fine-structure constant (effective control parameter)",
        )
        epochs = gr.Slider(50, 300, value=150, step=10, label="Simulated epochs")
        noise = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Input noise")

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

    with gr.Row():
        training_plot = gr.Plot(label="Training dynamics")
        interference_plot = gr.Plot(label="Interference pattern")

    out = gr.JSON(label="Model diagnostics")

    run_btn.click(
        run_demo,
        inputs=[alpha_value, epochs, noise],
        outputs=[training_plot, interference_plot, out],
    )
    alpha_value.change(
        run_demo,
        inputs=[alpha_value, epochs, noise],
        outputs=[training_plot, interference_plot, out],
    )
    epochs.change(
        run_demo,
        inputs=[alpha_value, epochs, noise],
        outputs=[training_plot, interference_plot, out],
    )
    noise.change(
        run_demo,
        inputs=[alpha_value, epochs, noise],
        outputs=[training_plot, interference_plot, out],
    )

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