import json, math, random
from pathlib import Path
import numpy as np
from scipy.optimize import differential_evolution, minimize_scalar

SEED = 7
np.random.seed(SEED); random.seed(SEED)


def cycle_matrix(weights):
    n = len(weights); A = np.zeros((n,n), dtype=complex)
    # y_i receives from predecessor i-1; product of the directed cycle is alpha
    for i, w in enumerate(weights): A[i, (i-1) % n] = w
    return A


def determinant_identity_check(n=5):
    rng = np.random.default_rng(SEED)
    weights = rng.uniform(.3, 1.4, n)
    alpha = np.prod(weights)
    errs=[]
    for _ in range(30):
        z = rng.normal()+1j*rng.normal()
        h = rng.normal()+1j*rng.normal()
        hs = h*(1 + .3*(rng.normal(size=n)+1j*rng.normal(size=n)))
        D=np.diag(hs); A=cycle_matrix(weights)
        lhs=np.linalg.det(np.eye(n)-D@A)
        rhs=1-alpha*np.prod(hs)
        errs.append(abs(lhs-rhs))
    return float(max(errs)), float(np.mean(errs))

# Stable unit h(z)=b/(z-a), with positive feedback and cycle alpha.
def q_abs(omega, phases, n, alpha=.5, a=.8, b=1., r=0.):
    z=np.exp(1j*omega); h=b/(z-a)
    return abs(1-alpha*(h**n)*np.prod(1+r*np.exp(1j*np.asarray(phases))))

def critical_radius(n, alpha=.5, a=.8, b=1.):
    # Minimize the scalar characteristic residual over unit-circle frequency and
    # independent uncertainty phases. This is a finite-grid approximation to rc.
    def obj(x): return q_abs(x[0], x[1:], n, alpha, a, b, x[-1] if False else 0.)
    # Instead search r jointly and find where min residual crosses 0; phases have
    # an exact useful reduction: product factors are optimized by DE at each r.
    def minres(r):
        def f(x): return q_abs(x[0], x[1:], n, alpha, a, b, r)
        bounds=[(0,2*np.pi)]+[(0,2*np.pi)]*n
        out=differential_evolution(f,bounds,seed=SEED,popsize=8,maxiter=100,tol=1e-8,polish=True)
        return out.fun
    # Direct zeros are numerically ill-conditioned, so locate the first r for
    # which a boundary root exists using scalar minimization of residual.
    rs=np.linspace(0,.99,100)
    vals=np.array([minres(float(r)) for r in rs])
    # report the first local near-zero crossing; residual has a numerical floor
    k=int(np.argmin(vals))
    # A robust scalar prediction for the destabilizing/stabilizing real direction.
    # Root z=a+(alpha*b^n)^(1/n)*(1-r), reaching z=1.
    predicted=1-(1-a)/(alpha**(1/n)*b)
    return predicted, float(rs[k]), float(vals[k]), [(float(rs[i]),float(vals[i])) for i in range(0,len(rs),10)]


def root_radius_for_real_gain(n, r, alpha=.5, a=.8, b=1.):
    # Exact roots under all factors (1-r): (z-a)^n=alpha*b^n*(1-r)^n.
    c=(alpha**(1/n))*b*(1-r)
    roots=np.array([a+c*np.exp(2j*np.pi*k/n) for k in range(n)])
    return float(np.max(np.abs(roots))), roots


def mechanism_sweeps():
    detmax, detmean=determinant_identity_check()
    alpha=.5; a=.8; b=1.
    ns=[2,3,4,6,8]
    scaling=[]
    for n in ns:
        pred=1-(1-a)/(alpha**(1/n)*b)
        # binary search observed real-gain boundary max pole radius <= 1
        lo,hi=0.,.999
        for _ in range(60):
            mid=(lo+hi)/2
            if root_radius_for_real_gain(n,mid,alpha,a,b)[0] > 1: lo=mid
            else: hi=mid
        scaling.append({'n':n,'predicted_rc':pred,'observed_rc':hi,'abs_error':abs(pred-hi)})
    radii=[]
    n=4; pred=scaling[2]['predicted_rc']
    for r in [0., .4, .7, pred-.02, pred+.02, .9]:
        rad,_=root_radius_for_real_gain(n,r,alpha,a,b)
        radii.append({'r':r,'max_pole_radius':rad,'unstable':rad>1+1e-10})
    return {'determinant_max_abs_error':detmax,'determinant_mean_abs_error':detmean,
            'critical_radius_scaling':scaling,'radius_transition':radii}


def tiny_rnn():
    # Small deterministic comparison; cyclic B is the proposed topology, dense B baseline.
    try:
        import torch
        torch.manual_seed(SEED); np.random.seed(SEED)
        device='cuda' if torch.cuda.is_available() else 'cpu'
        class RNN(torch.nn.Module):
            def __init__(self,n,cyclic):
                super().__init__(); self.n=n; self.cyclic=cyclic
                self.B=torch.nn.Parameter(torch.zeros(n,n)); self.C=torch.nn.Parameter(torch.randn(n,1)*.15); self.O=torch.nn.Parameter(torch.randn(1,n)*.15)
                mask=torch.zeros(n,n)
                for i in range(n): mask[i,(i-1)%n]=1
                self.register_buffer('mask',mask)
                torch.nn.init.orthogonal_(self.B)
            def forward(self,u,gain=None):
                h=torch.zeros(u.shape[0],self.n,device=u.device); ys=[]
                B=self.B*self.mask if self.cyclic else self.B
                if gain is not None: B=B*gain
                for t in range(u.shape[1]):
                    h=torch.tanh(h@B.T+u[:,t,:]@self.C.T); ys.append(h@self.O.T)
                return torch.stack(ys,1)
        n=8; T=48; batch=64
        t=torch.arange(T+1,device=device).float()[None,:]
        u=torch.sin(.23*t).repeat(batch,1).unsqueeze(-1); target=u[:,1:,:]; inp=u[:,:-1,:]
        out={}
        for cyclic in [False,True]:
            model=RNN(n,cyclic).to(device); opt=torch.optim.Adam(model.parameters(),lr=.02)
            for _ in range(250):
                opt.zero_grad(); loss=((model(inp)-target)**2).mean(); loss.backward(); opt.step()
            with torch.no_grad():
                clean=float(((model(inp)-target)**2).mean())
                vals=[]
                for r in [.1,.3,.5]:
                    g=torch.exp(torch.randn(n,device=device)*r-r*r/2).view(1,n)
                    vals.append(float(((model(inp,gain=g)-target)**2).mean()))
            out['cyclic' if cyclic else 'dense']={'clean_mse':clean,'gain_noise_mse':vals}
        return out
    except Exception as e:
        return {'error':repr(e)}

if __name__=='__main__':
    result={'seed':SEED,'math_and_mechanism':mechanism_sweeps(),'tiny_rnn':tiny_rnn()}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))