"""Edge Runtime Example for Sofia Engine.

Demonstrates:
1. Memory-bounded SignalFrame telemetry packaging
2. Hardened Sofia Assembly VM bytecode execution
3. In-situ neural network forward pass & analytical backpropagation
4. Safe, zero-pickle model serialization (JSON manifest + NPZ + SHA-256)

Run:
    python examples/edge_runtime.py
"""

from __future__ import annotations

import tempfile
import time
from pathlib import Path
import numpy as np

from sofia_ai.core.contracts import DataQuality, Sample, SignalFrame, SignalMetadata
from sofia_ai.learning.asm import AssemblyNeuralNetwork
from sofia_ai.learning.asm.vm import (
    AsmInstruction,
    OpCode,
    Register,
    SofiaAsmVM,
    VMStatus,
)
from sofia_ai.learning.manifest import load_assembly_model, save_assembly_model


def run_edge_runtime() -> None:
    print("=== Sofia Engine: Edge Runtime & Embedded VM ===")

    # -------------------------------------------------------------------------
    # 1. Bounded SignalFrame Telemetry Packaging
    # -------------------------------------------------------------------------
    fs = 1000.0
    meta = SignalMetadata(
        sample_rate=fs,
        channel_name="vibration_radial",
        physical_unit="m/s^2",
        sensor_id="piezo_accel_42",
    )

    # Telemetry buffer strictly respects runtime capacity ceilings (<= 65536)
    base_time = time.time()
    samples = tuple(
        Sample(
            timestamp=base_time + i / fs,
            device_id="edge-gateway-42",
            channel="vibration_radial",
            value=float(np.sin(2 * np.pi * 10.0 * (i / fs))),
            unit="m/s^2",
            quality=DataQuality.GOOD,
        )
        for i in range(256)
    )
    frame = SignalFrame(samples=samples, metadata=meta)

    print(f"\n[1] SignalFrame Ingested: {frame.size} samples, Quality={frame.samples[0].quality.name}")
    print(f"    Channel: {frame.metadata.channel_name}, Sensor: {frame.metadata.sensor_id}")

    # -------------------------------------------------------------------------
    # 2. Hardened Sofia Assembly Virtual Machine
    # -------------------------------------------------------------------------
    print("\n[2] Sofia Assembly Virtual Machine Execution")
    vm = SofiaAsmVM(memory_size=1024)

    # Assemble a bounded test program:
    # R1 = 12.5, R2 = 3.5, R1 = R1 + R2 (16.0), Mem[100] = R1, HALT
    program = [
        AsmInstruction(OpCode.LOAD_CONST, Register.R1, imm=12.5),
        AsmInstruction(OpCode.LOAD_CONST, Register.R2, imm=3.5),
        AsmInstruction(OpCode.ADD, Register.R1, Register.R2),
        AsmInstruction(OpCode.LOAD_CONST, Register.R3, imm=100.0),
        AsmInstruction(OpCode.STORE_MEM, Register.R3, Register.R1),
        AsmInstruction(OpCode.HALT),
    ]

    result = vm.run(program, max_cycles=1000)
    print(f"    VM Status:       {result.status.value}")
    print(f"    Executed Cycles: {result.cycles}")
    print(f"    Register R1:     {vm.registers[Register.R1]:.2f}")
    print(f"    Memory[100]:     {vm.memory[100]:.2f}")
    assert result.status == VMStatus.HALTED
    assert vm.memory[100] == 16.0

    # -------------------------------------------------------------------------
    # 3. In-Situ Neural Training with Analytical Backpropagation
    # -------------------------------------------------------------------------
    print("\n[3] In-Situ Neural Network (2-Layer Bytecode Backpropagation)")
    input_dim, hidden_dim, output_dim = 4, 8, 1
    net = AssemblyNeuralNetwork(
        input_dim=input_dim,
        hidden_dim=hidden_dim,
        output_dim=output_dim,
        learning_rate=0.05,
    )

    # Synthetic training sample (e.g. 4 normalized vibration features -> target health)
    x = np.array([0.25, 0.40, 0.10, 0.60], dtype=np.float64)
    target = np.array([0.85], dtype=np.float64)

    initial_pred = net.forward(x)
    initial_loss = float(np.mean((initial_pred - target) ** 2))
    print(f"    Initial Prediction: {initial_pred[0]:.4f} (Target: {target[0]:.4f}, Loss: {initial_loss:.4f})")

    # Perform 15 in-situ gradient update steps
    for _ in range(15):
        net.train_step(x, target)

    final_pred = net.forward(x)
    final_loss = float(np.mean((final_pred - target) ** 2))
    print(f"    Updated Prediction: {final_pred[0]:.4f} (Target: {target[0]:.4f}, Loss: {final_loss:.4f})")
    assert final_loss < initial_loss, "Loss should decrease after training steps"

    # -------------------------------------------------------------------------
    # 4. Safe Model Manifest Export & SHA-256 Checksum Verification
    # -------------------------------------------------------------------------
    print("\n[4] Safe Model Serialization (Zero-Pickle sofia.model.v1)")
    with tempfile.TemporaryDirectory() as tmpdir:
        manifest_path = save_assembly_model(
            model=net,
            destination_dir=tmpdir,
            model_name="edge_demo_model",
            feature_schema="vibration-v3",
            extra_metadata={"trained_on": "synthetic_edge_sample"},
        )
        print(f"    Saved Manifest: {manifest_path.name}")

        # Reload model and verify SHA-256 parameter integrity
        loaded_net, loaded_manifest = load_assembly_model(manifest_path)
        print(f"    Verified SHA-256: {loaded_manifest.parameters_sha256[:16]}... (Valid)")

        reloaded_pred = loaded_net.forward(x)
        print(f"    Reloaded Model Prediction: {reloaded_pred[0]:.4f}")
        assert np.isclose(reloaded_pred[0], final_pred[0]), "Reloaded model must match original weights"

    print("\nAll edge runtime checks passed successfully.")


if __name__ == "__main__":
    run_edge_runtime()