import json
import math
from pathlib import Path
import numpy as np

SEED = 2961
DT = 0.002


def exp_filter(signal, a, dt=DT, z0=0.0):
    """Exact zero-order-hold update for zdot=-a*z+a*q."""
    beta = math.exp(-a * dt)
    z = float(z0)
    out = np.empty(len(signal))
    for i, q in enumerate(signal):
        z = beta * z + (1.0 - beta) * float(q)
        out[i] = z
    return out


def frequency_check():
    rows = []
    for a in [0.3, 1.0, 3.0]:
        for ratio in [2.0, 5.0, 10.0]:
            w = ratio * a
            dt = 0.001
            n = int(max(20 * (2 * math.pi / w) / dt, 12000))
            t = np.arange(n) * dt
            z = exp_filter(np.sin(w * t), a, dt)
            cut = n // 3
            amp = 2 * abs(np.mean(z[cut:] * np.exp(-1j * w * t[cut:])))
            theory = a / math.sqrt(a * a + w * w)
            rows.append({"a": a, "omega": w, "omega_over_a": ratio,
                         "measured": float(amp), "theory": float(theory),
                         "relative_error": float(abs(amp - theory) / theory)})
    return rows


def run_controller(kind, a=3.0, noise_amp=0.08, omega=30.0, episodes=20):
    # Scalar safe system xdot=u, h=x. Observation noise affects the estimated q.
    # The direct h>=delta safeguard is implemented as a one-step viability clamp.
    alpha, eps, umax, kappa, delta = 1.0, 0.015, 1.0, a / 2.0, 0.01
    decay = math.exp(-a * DT)
    all_tv, all_minx, all_viol, all_infeas = [], [], [], []
    for ep in range(episodes):
        phase = 0.37 * ep
        x, z = 0.22, 0.0
        actions, xs = [], []
        for i in range(3000):
            t = i * DT
            obs = x + noise_amp * math.sin(omega * t + phase)
            # Nominal policy is intentionally sensitive to the noisy observation.
            u_pi = -0.42 - 0.75 * obs
            qhat_pi = u_pi + alpha * obs
            if kind == "memoryless":
                lower = -alpha * obs + eps
            else:
                lower = -alpha * obs + (1.0 - kappa / a) * z + eps
            infeasible = lower > umax
            u = float(np.clip(max(u_pi, lower), -umax, umax))
            # Explicit direct margin h>=delta, represented by a one-step clamp.
            if x <= delta:
                u = max(u, 0.0)
            q_exec = u + alpha * x
            if kind == "filtered":
                # Exactly one state update, using the executed residual.
                z = decay * z + (1.0 - decay) * q_exec
            x += DT * u
            actions.append(u)
            xs.append(x)
            all_infeas.append(float(infeasible))
        all_tv.append(float(np.sum(np.abs(np.diff(actions)))))
        all_minx.append(float(np.min(xs)))
        all_viol.append(float(np.mean(np.asarray(xs) < 0.0)))
    return {"tv_mean": float(np.mean(all_tv)), "tv_std": float(np.std(all_tv)),
            "min_x_mean": float(np.mean(all_minx)), "violation_rate": float(np.mean(all_viol)),
            "infeasible_rate": float(np.mean(all_infeas))}


def main():
    freq = frequency_check()
    comparison = {}
    for a in [0.3, 1.0, 3.0, 10.0]:
        comparison[str(a)] = {
            "memoryless": run_controller("memoryless", a=a, omega=30.0),
            "filtered": run_controller("filtered", a=a, omega=30.0),
        }
    result = {"seed": SEED, "dt": DT, "frequency_response": freq,
              "toy_comparison": comparison}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()