import json
import numpy as np
from scipy.linalg import expm

SEED = 2080
A = np.diag([0.4, 0.4])
C1 = np.array([[1.0, 0.0]])
C2 = np.array([[0.0, 1.0]])
I = np.eye(2)


def obs(A, C):
    n = A.shape[0]
    return np.vstack([C @ np.linalg.matrix_power(A, k) for k in range(n)])


def matrices(gain, tau1=0.5, tau2=0.5):
    L1 = np.array([[gain], [0.0]])
    L2 = np.array([[0.0], [gain]])
    F1 = A - L1 @ C1
    F2 = A - L2 @ C2
    P1, P2 = expm(F1 * tau1), expm(F2 * tau2)
    return F1, F2, P1, P2, P2 @ P1


def spectral_radius(M):
    return float(np.max(np.abs(np.linalg.eigvals(M))))


def exact_boundary():
    # For this diagonal complementary system, rho(Phi)=exp((.4-g) tau)
    # and the predicted boundary is gain=.8 for equal dwell time.
    gains = np.linspace(0.0, 1.2, 241)
    values = np.array([spectral_radius(matrices(g)[-1]) for g in gains])
    crossing = gains[np.argmin(np.abs(values - 1.0))]
    return crossing, gains, values


def empirical_decay(gain, tau=0.5, cycles=25):
    Phi = matrices(gain, tau, tau)[-1]
    e = np.array([1.0, -0.7])
    norms = []
    for _ in range(cycles):
        norms.append(np.linalg.norm(e))
        e = Phi @ e
    norms = np.maximum(np.asarray(norms), 1e-300)
    slope = np.polyfit(np.arange(cycles), np.log(norms), 1)[0]
    return slope, norms[-1]


def noisy_switched(gain, tau=0.5, noise=0.15, cycles=80, trials=200):
    # Discrete exact observer with measurement noise. For each dwell, use
    # midpoint-equivalent correction z-C h at the start of the dwell.
    rng = np.random.default_rng(SEED + 9)
    F1, F2, P1, P2, _ = matrices(gain, tau, tau)
    L1 = np.array([[gain], [0.0]])
    L2 = np.array([[0.0], [gain]])
    # Integral exp(F(t-s)) L ds, computed robustly by augmented exponential.
    def noise_map(F, L):
        aug = np.zeros((3, 3)); aug[:2, :2] = F; aug[:2, 2:] = L
        return expm(aug * tau)[:2, 2:]
    Q1, Q2 = noise_map(F1, L1), noise_map(F2, L2)
    vals = []
    for _ in range(trials):
        e = np.array([1.0, -0.7])
        for _ in range(cycles):
            e = P1 @ e + Q1 @ rng.normal(0, noise, size=(1,))
            e = P2 @ e + Q2 @ rng.normal(0, noise, size=(1,))
        vals.append(np.dot(e, e))
    return float(np.mean(vals))


def residual_fusion(gain, tau=0.5, noise=0.15, cycles=80, trials=200):
    # Baseline: each observed coordinate is independently corrected, but its
    # residual is held at zero during the missing-channel dwell (ordinary
    # modality-specific residual fusion, no cycle-aware gain design).
    rng = np.random.default_rng(SEED + 9)
    F1, F2, P1, P2, _ = matrices(gain, tau, tau)
    L1 = np.array([[gain], [0.0]]); L2 = np.array([[0.0], [gain]])
    def noise_map(F, L):
        aug = np.zeros((3, 3)); aug[:2, :2] = F; aug[:2, 2:] = L
        return expm(aug * tau)[:2, 2:]
    Q1, Q2 = noise_map(F1, L1), noise_map(F2, L2)
    vals = []
    for _ in range(trials):
        e = np.array([1.0, -0.7])
        for _ in range(cycles):
            # baseline uses only currently available channel but resets the
            # unobserved residual branch; this is equivalent here to no
            # complementary correction during that half-cycle.
            e = P1 @ e + Q1 @ rng.normal(0, noise, size=(1,))
            e = expm(A * tau) @ e + Q2 @ rng.normal(0, noise, size=(1,))
        vals.append(np.dot(e, e))
    return float(np.mean(vals))


def main():
    O1, O2 = obs(A, C1), obs(A, C2)
    Ostack = np.vstack([O1, O2])
    crossing, gains, rhos = exact_boundary()
    # Prediction 1: equal-dwell stability boundary at gain=1.2.
    boundary = {"predicted": 0.8, "observed_grid": float(crossing),
                "rho_at_0_7": float(spectral_radius(matrices(.7)[-1])),
                "rho_at_0_9": float(spectral_radius(matrices(.9)[-1]))}
    # Prediction 2: log decay slope equals log spectral radius per cycle.
    decay = []
    for g in (0.9, 1.0, 1.2):
        slope, final = empirical_decay(g)
        predicted = np.log(spectral_radius(matrices(g)[-1]))
        decay.append({"gain": g, "predicted_log_slope": float(predicted),
                      "observed_log_slope": float(slope), "final_norm": float(final)})
    # Prediction 3: stable switched cycle can tolerate an individually expanding
    # channel; scan dwell ratios while keeping gain=1.2.
    dwell = []
    for t1, t2 in ((0.2,0.8),(0.3,0.7),(0.4,0.6),(0.5,0.5),(0.6,0.4),(0.7,0.3),(0.8,0.2)):
        F1,F2,P1,P2,Phi = matrices(1.2,t1,t2)
        dwell.append({"tau1":t1,"tau2":t2,"rho_cycle":spectral_radius(Phi),
                      "rho_channel1":spectral_radius(P1),"rho_channel2":spectral_radius(P2)})
    gain = 1.2
    comparison = {"switched_mse": noisy_switched(gain),
                  "residual_fusion_mse": residual_fusion(gain),
                  "gain": gain}
    out = {"observability_rank_channel1": int(np.linalg.matrix_rank(O1)),
           "observability_rank_channel2": int(np.linalg.matrix_rank(O2)),
           "observability_rank_stacked": int(np.linalg.matrix_rank(Ostack)),
           "predictions": {"boundary": boundary, "decay": decay, "dwell_sweep": dwell},
           "comparison": comparison}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()