import time
import numpy as np
import gradio as gr
import platform
import psutil

# -------------------------------
# Coherent Compute Core
# -------------------------------

def rft_step(Psi, E, L):
    # Stable, branchless update
    phase = 0.997 * Psi + 0.003 * E
    drive = np.tanh(phase)
    Psi = 0.999 * Psi + 0.001 * drive
    E   = 0.995 * E   + 0.004 * Psi
    L   = 0.998 * L   + 0.001 * (Psi * E)
    return Psi, E, L


def coherence_metric(prev, curr):
    num = np.dot(prev, curr)
    den = (np.linalg.norm(prev) * np.linalg.norm(curr)) + 1e-9
    return float(num / den)


# -------------------------------
# Benchmark Runner
# -------------------------------

def run_engine(oscillators: int, steps: int):
    oscillators = int(oscillators)
    steps = int(steps)

    rng = np.random.default_rng(42)
    Psi = rng.random(oscillators, dtype=np.float32)
    E   = rng.random(oscillators, dtype=np.float32)
    L   = rng.random(oscillators, dtype=np.float32)

    sample = min(200_000, oscillators)
    prev_snapshot = Psi[:sample].copy()

    t0 = time.time()

    for _ in range(steps):
        Psi, E, L = rft_step(Psi, E, L)

    elapsed = time.time() - t0
    curr_snapshot = Psi[:sample].copy()

    coherence = abs(coherence_metric(prev_snapshot, curr_snapshot))
    energy = float(np.mean(E))

    items = oscillators * steps
    throughput = items / elapsed

    return {
        "Throughput (updates/sec)": f"{throughput/1e9:.3f} B/s",
        "Coherence (|C|)": f"{coherence:.5f}",
        "Mean Energy": f"{energy:.5f}",
        "Elapsed Time (s)": f"{elapsed:.2f}",
        "Oscillators": f"{oscillators:,}",
        "Steps": f"{steps:,}",
        "CPU": platform.processor(),
        "Cores Available": psutil.cpu_count(logical=True)
    }


# -------------------------------
# Gradio UI
# -------------------------------

with gr.Blocks(title="Coherent Compute Engine") as demo:
    gr.Markdown(
        """
        # Coherent Compute Engine

        **What this is**
        - A live, CPU-first coherent compute benchmark  
        - No precomputed results  
        - No GPUs required  
        - Measures real throughput, stability, and energy behavior  

        **What an “item” is**
        - One coherent state update of `[Ψ, E, L]` per oscillator per step  

        Everything you see below is computed **right now**, on this machine.
        """
    )

    with gr.Row():
        oscillators = gr.Slider(
            minimum=100_000,
            maximum=10_000_000,
            value=2_000_000,
            step=100_000,
            label="Number of Oscillators"
        )
        steps = gr.Slider(
            minimum=50,
            maximum=1000,
            value=250,
            step=50,
            label="Simulation Steps"
        )

    run_btn = gr.Button("Run Engine")

    output = gr.JSON(label="Results")

    run_btn.click(
        fn=run_engine,
        inputs=[oscillators, steps],
        outputs=output
    )

demo.launch()