import json, math, random
import numpy as np
import torch

SEED = 381
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)


def determinant_lemma_check():
    n = 6
    R = np.random.randn(n, n)
    A = 0.65 * R / max(abs(np.linalg.eigvals(R)))
    b, c = np.random.randn(n), np.random.randn(n)
    J = A + np.outer(b, c)
    pts = [.7+.4j, -1.1+.2j, .3-.9j]
    errors = []
    for z in pts:
        lhs = np.linalg.det(z*np.eye(n)-J)
        q = np.linalg.solve(z*np.eye(n)-A, b)
        rhs = np.linalg.det(z*np.eye(n)-A) * (1-c@q)
        errors.append(abs(lhs-rhs)/max(1., abs(lhs), abs(rhs)))
    eigJ = np.linalg.eigvals(J); eigA = np.linalg.eigvals(A)
    mode_errors = []
    for z in eigJ:
        if min(abs(z-eigA)) > 1e-6:
            mode_errors.append(abs(1-c@np.linalg.solve(z*np.eye(n)-A,b)))
    return {"determinant_relative_error": float(max(errors)),
            "mode_characteristic_error": float(max(mode_errors)),
            "open_loop_radius": float(max(abs(eigA))),
            "closed_loop_radius": float(max(abs(eigJ)))}


def contour_penalty(g, a=.92, margin=.10, radii=(.98, 1.0, 1.02), ntheta=16):
    # A=aI, b=c=sqrt(g), hence H(z)=g/(z-a); this is still evaluated by a solve.
    A = torch.eye(1, dtype=torch.complex64) * a
    b = torch.sqrt(torch.clamp(g, min=1e-8)).to(torch.complex64).reshape(1)
    c = b
    vals = []
    for r in radii:
        for k in range(ntheta):
            th = 2*math.pi*k/ntheta
            z = torch.tensor(r*math.cos(th)+1j*r*math.sin(th), dtype=torch.complex64)
            q = torch.linalg.solve(z*torch.eye(1, dtype=torch.complex64)-A, b)
            vals.append(torch.relu(torch.abs(torch.sum(c*q))-(1-margin))**2)
    return torch.stack(vals).mean()


def train(reg, steps=240, lr=.035):
    # The task optimum deliberately asks for a strong positive feedback gain.
    # Regularization should prevent the associated pole a+g from crossing one.
    g = torch.tensor(0.02, requires_grad=True)
    opt = torch.optim.SGD([g], lr=lr)
    trace=[]
    for step in range(steps):
        opt.zero_grad()
        task = (g-0.18)**2
        pen = contour_penalty(g) if reg else 0*g
        loss = task + 1.8*pen
        loss.backward(); opt.step()
        with torch.no_grad(): g.clamp_(0, .35)
        if step % 20 == 0 or step == steps-1:
            trace.append((step, float(g), float(task), float(pen)))
    gv=float(g)
    pole=.92+gv
    # impulse response of the scalar closed-loop recurrence
    impulse=[pole**k for k in range(20)]
    return {"g":gv, "task_mse":(gv-.18)**2, "max_contour_gain":max(
        float(torch.abs(torch.tensor(gv/(r*math.cos(2*math.pi*k/16)+1j*r*math.sin(2*math.pi*k/16)-.92))))
        for r in (.98,1.,1.02) for k in range(16)),
        "closed_loop_pole":pole, "impulse_abs_k19":abs(impulse[-1]), "trace":trace}


def main():
    check=determinant_lemma_check()
    base=train(False); idea=train(True)
    result={"seed":SEED, "math_check":check, "baseline":base, "idea":idea,
            "claim_signal": {"lower_pole": idea["closed_loop_pole"] < base["closed_loop_pole"],
                             "lower_contour_gain": idea["max_contour_gain"] < base["max_contour_gain"],
                             "faster_impulse_decay": idea["impulse_abs_k19"] < base["impulse_abs_k19"]}}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__ == "__main__": main()