import json, random
from pathlib import Path
import numpy as np
import torch
from fisher_observable import fisher_information

SEED = 2387
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)

def rollout(z0, k, dt=0.5):
    z = z0
    for _ in range(k):
        z = torch.stack((z[0] + dt*z[2], z[1] + dt*z[3], z[2], z[3]))
    return z

def obs(z, use_brightness=True, c=1.0):
    x, y = z[0], z[1]
    bearing = torch.atan2(y, x)
    if use_brightness:
        return torch.stack((bearing, c/(x*x+y*y)))
    return bearing.reshape(1)

def jacobians(z0, K, use_brightness, dt=0.5):
    return [torch.autograd.functional.jacobian(lambda q: obs(rollout(q,k,dt),use_brightness), z0) for k in range(K+1)]

def fisher(z0, K, use_brightness, sig_b=0.01, sig_l=0.01):
    js = jacobians(z0,K,use_brightness)
    vars = [sig_b**2] if not use_brightness else [sig_b**2,sig_l**2]
    return fisher_information(js, vars).detach().numpy()

def recover(z_init, observations, K, use_brightness, sig_b, sig_l, steps=500):
    q=torch.tensor(z_init, requires_grad=True); opt=torch.optim.Adam([q],lr=0.04)
    target=torch.tensor(observations)
    for _ in range(steps):
        opt.zero_grad(); pred=torch.cat([obs(rollout(q,k),use_brightness) for k in range(K+1)])
        if use_brightness:
            rb=torch.atan2(torch.sin(pred[0::2]-target[0::2]), torch.cos(pred[0::2]-target[0::2]))
            rl=pred[1::2]-target[1::2]; loss=(rb/sig_b).square().mean()+(rl/sig_l).square().mean()
        else:
            rr=torch.atan2(torch.sin(pred-target),torch.cos(pred-target)); loss=(rr/sig_b).square().mean()
        loss.backward(); opt.step()
    return q.detach().numpy()

def main():
    z=torch.tensor([3.0,1.0,0.4,-0.2]); out={"seed":SEED}
    rank_rows=[]
    for K in [0,1,2,4,8]:
        row={"K":K}
        for name,use in [("bearing",False),("bearing_brightness",True)]:
            e=np.linalg.eigvalsh(fisher(z,K,use)); row[name+"_rank"]=int(np.sum(e>1e-8)); row[name+"_eigenvalues"]=e.tolist()
        rank_rows.append(row)
    out["rank_sweep"]=rank_rows
    # At K=4 both position and velocity directions can be observed.
    scaling=[]
    for s in [0.005,0.01,0.02,0.04,0.08]:
        e=np.linalg.eigvalsh(fisher(z,4,True,sig_b=0.01,sig_l=s))
        scaling.append({"sigma_l":s,"sigma_inv_sq":1/s**2,"lambda_min":float(e[0]),"lambda_max":float(e[-1])})
    out["noise_scaling"]=scaling
    # Explicit transition sweep around the predicted equal-whitened-sensitivity noise.
    transition=[]
    for s in [0.05, 0.10, 0.15, 0.158113883, 0.20, 0.30, 0.60]:
        e=np.linalg.eigvalsh(fisher(z,4,True,sig_b=0.01,sig_l=s))
        transition.append({"sigma_l":s, "lambda_min":float(e[0]), "condition":float(e[-1]/max(e[0],1e-15))})
    out["noise_transition"]=transition
    # Independent central-difference check of the autodiff chain-rule Jacobian.
    kcheck=2; h=1e-6; q=z.detach().numpy(); J=torch.cat(jacobians(z,kcheck,True),dim=0).numpy()
    def flat_obs(qn):
        qq=torch.tensor(qn); return np.concatenate([obs(rollout(qq,k),True).detach().numpy() for k in range(kcheck+1)])
    Jfd=np.column_stack([(flat_obs(q+np.eye(4)[j]*h)-flat_obs(q-np.eye(4)[j]*h))/(2*h) for j in range(4)])
    out["jacobian_check_max_abs_error"]=float(np.max(np.abs(J-Jfd)))
    out["first_full_rank_K"]={name:next((r["K"] for r in rank_rows if r[name+"_rank"]==4),None) for name in ["bearing","bearing_brightness"]}
    jb=jacobians(z,0,True)[0].numpy(); bearing_norm=np.linalg.norm(jb[0]/0.01)
    threshold=1.0/np.linalg.norm(jb[1]/0.01)
    out["predicted_equal_whitened_sensitivity_sigma_brightness"]=float(threshold)
    K=4; sb=0.01; sl=0.01
    clean=np.concatenate([obs(rollout(z,k),True).detach().numpy() for k in range(K+1)])
    rng=np.random.default_rng(SEED); noisy=clean.copy(); noisy[0::2]+=rng.normal(0,sb,K+1); noisy[1::2]+=rng.normal(0,sl,K+1)
    init=np.array([2.7,1.3,0.25,-0.05]); base=recover(init,noisy[0::2],K,False,sb,sl); idea=recover(init,noisy,K,True,sb,sl)
    true=z.numpy()
    out["recovery"]={"true":true.tolist(),"bearing_only_estimate":base.tolist(),"bearing_plus_brightness_estimate":idea.tolist(),"bearing_only_rmse":float(np.sqrt(np.mean((base-true)**2))),"idea_rmse":float(np.sqrt(np.mean((idea-true)**2)))}
    out["checks"]={"bearing_rank_at_K8":rank_rows[-1]["bearing_rank"],"combined_rank_at_K1":rank_rows[1]["bearing_brightness_rank"],"lambda_min_ratio_sigma_.01_to_.02":scaling[1]["lambda_min"]/scaling[2]["lambda_min"],"expected_inverse_variance_ratio":4.0,
      "bearing_condition_at_K4":float(rank_rows[3]["bearing_eigenvalues"][-1]/max(rank_rows[3]["bearing_eigenvalues"][1],1e-15)),
      "combined_condition_at_K4":float(rank_rows[3]["bearing_brightness_eigenvalues"][-1]/rank_rows[3]["bearing_brightness_eigenvalues"][0]),"bearing_whitened_norm":bearing_norm,"brightness_equal_norm_sigma":float(threshold)}
    Path("results.json").write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__ == '__main__': main()