# app.py
# Coherent_Compute_Engine — RFTSystems
# Real, on-machine coherent throughput benchmark + tamper-evident receipt download (SHA-256).
#
# What an “item” is:
#   1 item = one per-oscillator coherent state update of [Psi, E, L] per step.
#
# Notes:
# - Measurements reflect the Hugging Face Space runtime hardware (not the visitor’s local machine).
# - Includes optional baselines (NumPy vectorised + tiny Python loop) computed live.
# - Produces a downloadable receipt JSON with canonical hashing for verification.

import os
import json
import time
import math
import csv
import hashlib
import platform
import datetime as dt
from pathlib import Path

import numpy as np
import gradio as gr

# Optional: Numba acceleration (used if available)
try:
    import numba as nb
    NUMBA_OK = True
except Exception:
    nb = None
    NUMBA_OK = False

APP_NAME = "Coherent_Compute_Engine"
APP_VERSION = "v1.0.0"
RESULTS_DIR = Path("results")
RESULTS_DIR.mkdir(exist_ok=True)

# ----------------------------
# Canonical JSON + integrity
# ----------------------------
def canon_json_bytes(obj) -> bytes:
    return json.dumps(
        obj,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")

def sha256_hex(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()

def utc_now_iso() -> str:
    # Avoid deprecated utcnow warning in newer Python versions
    return dt.datetime.now(dt.timezone.utc).isoformat().replace("+00:00", "Z")

def write_receipt(payload: dict) -> str:
    """
    Writes a JSON receipt to disk and returns the filepath for Gradio download.
    Receipt includes its own SHA-256 of the canonical JSON (tamper-evident).
    """
    # hash without integrity field
    payload = dict(payload)  # copy
    payload.pop("integrity", None)

    b0 = canon_json_bytes(payload)
    h = sha256_hex(b0)

    payload["integrity"] = {
        "sha256": h,
        "receipt_id": h[:12],
        "canonical_json": "sorted_keys + compact_separators",
    }

    b1 = canon_json_bytes(payload)

    ts = payload.get("timestamp_utc", utc_now_iso())
    safe_ts = ts.replace(":", "").replace(".", "").replace("Z", "")
    fname = f"receipt_{safe_ts}_{h[:12]}.json"
    path = RESULTS_DIR / fname
    path.write_bytes(b1)
    return str(path)

def append_csv_row(row: dict, csv_path: Path):
    headers = [
        "timestamp_utc","engine","device_note","oscillators","steps",
        "elapsed_s","throughput_Bps","coherence_abs","mean_energy",
        "python","platform","cpu_count_logical","numba_available","seed","scale"
    ]
    need_header = not csv_path.exists()
    with csv_path.open("a", newline="") as f:
        w = csv.DictWriter(f, fieldnames=headers)
        if need_header:
            w.writeheader()
        w.writerow({k: row.get(k, "") for k in headers})

# ----------------------------
# RFT-lite coherent kernel
# ----------------------------
def _np_step(Psi, E, L, scale=1.0):
    phase = 0.997 * Psi + 0.003 * E
    drive = np.tanh(phase * scale)
    Psi_n = 0.999 * Psi + 0.001 * drive
    E_n   = 0.995 * E   + 0.004 * Psi_n
    L_n   = 0.998 * L   + 0.001 * (Psi_n * E_n)
    return Psi_n, E_n, L_n

def coherence_abs(Psi0: np.ndarray, Psi1: np.ndarray) -> float:
    v0 = Psi0.astype(np.float64, copy=False)
    v1 = Psi1.astype(np.float64, copy=False)
    num = float(np.dot(v0, v1))
    den = float(np.linalg.norm(v0) * np.linalg.norm(v1)) + 1e-12
    return abs(num / den)

def mean_energy(E: np.ndarray) -> float:
    return float(np.mean(np.clip(E, 0.0, 1.5)))

def run_engine_numpy(n: int, steps: int, seed: int, scale: float):
    rng = np.random.default_rng(seed)
    Psi = rng.random(n, dtype=np.float32)
    E   = rng.random(n, dtype=np.float32)
    L   = rng.random(n, dtype=np.float32)

    sample = min(n, 200_000)
    Psi0 = Psi[:sample].copy()

    t0 = time.perf_counter()
    for _ in range(steps):
        Psi, E, L = _np_step(Psi, E, L, scale=scale)
    t1 = time.perf_counter()

    elapsed = t1 - t0
    Psi1 = Psi[:sample].copy()

    items = int(n) * int(steps)
    throughput_Bps = (items / elapsed) / 1e9

    return {
        "engine": "numpy",
        "oscillators": int(n),
        "steps": int(steps),
        "elapsed_s": float(elapsed),
        "throughput_Bps": float(throughput_Bps),
        "coherence_abs": float(coherence_abs(Psi0, Psi1)),
        "mean_energy": float(mean_energy(E[:sample])),
    }

# ----------------------------
# Baselines (optional, live)
# ----------------------------
def run_baseline_python(n: int, steps: int, seed: int):
    # Safety caps for hosted runtimes
    n = min(n, 200_000)
    steps = min(steps, 10)

    rng = np.random.default_rng(seed)
    Psi = rng.random(n).tolist()
    E   = rng.random(n).tolist()
    L   = rng.random(n).tolist()

    def step(Psi, E, L):
        outPsi = [0.0]*n
        outE   = [0.0]*n
        outL   = [0.0]*n
        for i in range(n):
            phase = 0.997*Psi[i] + 0.003*E[i]
            drive = math.tanh(phase)
            p = 0.999*Psi[i] + 0.001*drive
            e = 0.995*E[i] + 0.004*p
            l = 0.998*L[i] + 0.001*(p*e)
            outPsi[i], outE[i], outL[i] = p, e, l
        return outPsi, outE, outL

    t0 = time.perf_counter()
    for _ in range(steps):
        Psi, E, L = step(Psi, E, L)
    t1 = time.perf_counter()

    elapsed = t1 - t0
    items = int(n) * int(steps)
    throughput_Bps = (items / elapsed) / 1e9

    return {
        "engine": "python_loop",
        "oscillators": int(n),
        "steps": int(steps),
        "elapsed_s": float(elapsed),
        "throughput_Bps": float(throughput_Bps),
        "coherence_abs": 1.0,
        "mean_energy": float(np.mean(np.clip(np.array(E[:min(n, 50_000)], dtype=np.float32), 0.0, 1.5))),
        "note": "Python baseline is safety-capped (n<=200k, steps<=10).",
    }

# ----------------------------
# Numba engine (if available)
# ----------------------------
if NUMBA_OK:
    @nb.njit(fastmath=True)
    def _numba_kernel(Psi, E, L, scale):
        n = Psi.shape[0]
        for i in range(n):
            phase = 0.997 * Psi[i] + 0.003 * E[i]
            drive = math.tanh(phase * scale)
            p = 0.999 * Psi[i] + 0.001 * drive
            e = 0.995 * E[i] + 0.004 * p
            l = 0.998 * L[i] + 0.001 * (p * e)
            Psi[i] = p
            E[i] = e
            L[i] = l

    def run_engine_numba(n: int, steps: int, seed: int, scale: float):
        rng = np.random.default_rng(seed)
        Psi = rng.random(n, dtype=np.float32)
        E   = rng.random(n, dtype=np.float32)
        L   = rng.random(n, dtype=np.float32)

        sample = min(n, 200_000)
        Psi0 = Psi[:sample].copy()

        # compile warmup on tiny arrays
        tiny = min(n, 1024)
        _numba_kernel(Psi[:tiny], E[:tiny], L[:tiny], scale)

        t0 = time.perf_counter()
        for _ in range(steps):
            _numba_kernel(Psi, E, L, scale)
        t1 = time.perf_counter()

        elapsed = t1 - t0
        Psi1 = Psi[:sample].copy()

        items = int(n) * int(steps)
        throughput_Bps = (items / elapsed) / 1e9

        return {
            "engine": "numba",
            "oscillators": int(n),
            "steps": int(steps),
            "elapsed_s": float(elapsed),
            "throughput_Bps": float(throughput_Bps),
            "coherence_abs": float(coherence_abs(Psi0, Psi1)),
            "mean_energy": float(mean_energy(E[:sample])),
        }

# ----------------------------
# Runner + UI helpers
# ----------------------------
def format_ui(primary: dict, baselines: dict | None):
    ui = {
        "Throughput (B/s)": f'{primary["throughput_Bps"]:.3f} B/s',
        "Coherence (|C|)": f'{primary["coherence_abs"]:.5f}',
        "Mean Energy": f'{primary["mean_energy"]:.5f}',
        "Elapsed Time (s)": f'{primary["elapsed_s"]:.2f}',
        "Oscillators": f'{primary["oscillators"]:,}',
        "Steps": f'{primary["steps"]:,}',
        "Engine": primary["engine"],
        "CPU Cores Available": os.cpu_count(),
    }
    if baselines:
        for name, b in baselines.items():
            ui[f"Baseline: {name} (B/s)"] = f'{b["throughput_Bps"]:.3f}'
            ui[f"Baseline: {name} Engine"] = b["engine"]
    return ui

def run_and_receipt(n_oscillators, steps, seed, scale, include_baselines):
    n = int(n_oscillators)
    s = int(steps)
    seed = int(seed)
    scale = float(scale)

    # Safety rails (prevents accidental crash / OOM on Space)
    n = max(100_000, min(n, 40_000_000))
    s = max(10, min(s, 5000))

    # primary engine
    if NUMBA_OK:
        primary = run_engine_numba(n, s, seed, scale)
    else:
        primary = run_engine_numpy(n, s, seed, scale)

    # optional baselines
    baselines = {}
    if include_baselines:
        baselines["numpy"] = run_engine_numpy(min(n, 8_000_000), min(s, 2000), seed, scale)
        baselines["python_loop"] = run_baseline_python(min(n, 500_000), min(s, 200), seed)

    # build receipt payload
    payload = {
        "timestamp_utc": utc_now_iso(),
        "app": {"name": APP_NAME, "version": APP_VERSION},
        "definition": {
            "item": "1 item = one per-oscillator coherent state update of [Psi, E, L] per step (as implemented here)."
        },
        "runtime": {
            "device_note": "Hugging Face Space runtime machine (results are not from the visitor's local CPU).",
            "platform": platform.platform(),
            "python": platform.python_version(),
            "cpu_count_logical": os.cpu_count(),
            "numba_available": bool(NUMBA_OK),
        },
        "inputs": {
            "oscillators": int(n),
            "steps": int(s),
            "seed": int(seed),
            "scale": float(scale),
            "include_baselines": bool(include_baselines),
        },
        "results": {
            "primary": primary,
            "baselines": baselines,
        },
    }

    receipt_path = write_receipt(payload)

    # append csv row (primary only)
    csv_row = {
        "timestamp_utc": payload["timestamp_utc"],
        "engine": primary["engine"],
        "device_note": payload["runtime"]["device_note"],
        "oscillators": primary["oscillators"],
        "steps": primary["steps"],
        "elapsed_s": primary["elapsed_s"],
        "throughput_Bps": primary["throughput_Bps"],
        "coherence_abs": primary["coherence_abs"],
        "mean_energy": primary["mean_energy"],
        "python": payload["runtime"]["python"],
        "platform": payload["runtime"]["platform"],
        "cpu_count_logical": payload["runtime"]["cpu_count_logical"],
        "numba_available": payload["runtime"]["numba_available"],
        "seed": seed,
        "scale": scale,
    }
    append_csv_row(csv_row, RESULTS_DIR / "runs.csv")

    # UI output
    ui = format_ui(primary, baselines if include_baselines else None)
    ui["Receipt SHA-256 (in file)"] = payload["integrity"]["sha256"] if "integrity" in payload else "written in receipt"

    return ui, receipt_path

# ----------------------------
# UI
# ----------------------------
CSS = """
:root { --rft-accent: #ff7a18; }
.gradio-container { max-width: 980px !important; }
#titlebar h1 { font-size: 2.05rem; letter-spacing: -0.02em; }
.rft-card { border-radius: 16px !important; border: 1px solid rgba(255,255,255,0.08) !important; }
"""

with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
<div id="titlebar"><h1>Coherent Compute Engine</h1></div>

**What this Space does**  
Runs a real coherent state-update benchmark and reports measured throughput, stability proxy, and energy behavior. No precomputed results.

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

**Verification**  
Every run generates a tamper-evident receipt (JSON) with a SHA-256 hash you can download.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group(elem_classes=["rft-card"]):
                n_slider = gr.Slider(
                    minimum=100_000, maximum=40_000_000, step=100_000,
                    value=6_400_000, label="Number of Oscillators"
                )
                steps_slider = gr.Slider(
                    minimum=10, maximum=5000, step=10,
                    value=650, label="Simulation Steps"
                )
                seed_box = gr.Number(value=7, precision=0, label="Seed")
                scale_box = gr.Number(value=1.0, precision=3, label="Scale (stability knob)")
                include_baselines = gr.Checkbox(
                    value=False,
                    label="Include baselines (NumPy + tiny Python loop)",
                    info="Baselines are measured live too. Python loop is safety-capped."
                )
                run_btn = gr.Button("Run Engine", variant="primary")

        with gr.Column(scale=1):
            results_json = gr.JSON(label="Results")
            receipt_file = gr.File(label="Receipt (download)")

    run_btn.click(
        fn=run_and_receipt,
        inputs=[n_slider, steps_slider, seed_box, scale_box, include_baselines],
        outputs=[results_json, receipt_file],
    )

    gr.Markdown(
        """
**Notes**
- This runs on the Hugging Face Space runtime machine. Your browser is just the UI.
- If the Space is under load, throughput will vary; the receipt captures the environment at run time.
        """
    )

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