import json, math, random
import numpy as np
import torch
from scipy.linalg import eigh

SEED=2950
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def beam(nel=8, L=1.0, EI=1.0, rhoA=1.0):
    h=L/nel; K=np.zeros((2*(nel+1),2*(nel+1))); M=K.copy()
    ke=EI/h**3*np.array([[12,6*h,-12,6*h],[6*h,4*h*h,-6*h,2*h*h],[-12,-6*h,12,-6*h],[6*h,2*h*h,-6*h,4*h*h]],float)
    me=rhoA*h/420*np.array([[156,22*h,54,-13*h],[22*h,4*h*h,13*h,-3*h*h],[54,13*h,156,-22*h],[-13*h,-3*h*h,-22*h,4*h*h]],float)
    for e in range(nel):
        ix=[2*e,2*e+1,2*e+2,2*e+3]; K[np.ix_(ix,ix)]+=ke; M[np.ix_(ix,ix)]+=me
    # clamped left, free right
    return K[2:,2:], M[2:,2:]

def inertia(K,M,w,tol=1e-12):
    # For positive-definite M, inertia(K-w^2 M) equals count of frequencies below w.
    A=K-(w*w)*M; ev=np.linalg.eigvalsh(A); scale=max(1.,np.max(np.abs(A)))
    return int(np.sum(ev < -tol*scale))

def bracket(K,M,target,hi=None,tol=1e-3):
    vals=eigh(K,M,eigvals_only=True); freqs=np.sqrt(np.maximum(vals,0))
    if hi is None: hi=float(freqs[-1]*1.05)
    lo=0.; upper=hi
    # Find the transition count target-1 -> target; bisection around exact transition.
    grid=np.linspace(0,upper,400)
    for a,b in zip(grid[:-1],grid[1:]):
        if inertia(K,M,a)==target-1 and inertia(K,M,b)>=target:
            lo,upper=float(a),float(b); break
    # narrow an interval whose endpoints have counts target-1 and target
    while upper-lo>tol:
        mid=(lo+upper)/2
        if inertia(K,M,mid)<target: lo=mid
        else: upper=mid
    return lo,upper,freqs

def train(K,M,target,interval=None,steps=2000,seed=0):
    torch.manual_seed(seed); np.random.seed(seed)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        kt=torch.tensor(K,dtype=torch.float64,device=dev); mt=torch.tensor(M,dtype=torch.float64,device=dev)
        n=K.shape[0]; q=torch.randn(n,dtype=torch.float64,device=dev,requires_grad=True)
        # omega starts randomly in a broad range; bracketed version starts inside its certified interval
        exact=np.sqrt(eigh(K,M,eigvals_only=True));
        if interval is None: w=torch.tensor(float(np.random.uniform(.2,exact[-1]*1.15)),dtype=torch.float64,device=dev,requires_grad=True)
        else: w=torch.tensor(float((interval[0]+interval[1])/2),dtype=torch.float64,device=dev,requires_grad=True)
        opt=torch.optim.Adam([q,w],lr=.025)
        for t in range(steps):
            nq=q/(torch.linalg.vector_norm(q)+1e-12)
            r=kt@nq-(w*w)*(mt@nq)
            loss=torch.sum(r*r)
            # This is the stated box penalty; modest normalization avoids trivial q=0.
            if interval is not None:
                lo,up=interval; loss=loss+2000*(torch.relu(torch.tensor(lo,device=dev)-w)**2+torch.relu(w-torch.tensor(up,device=dev))**2)
                if t < steps//5: # freeze frequency during initial eigenvector settling
                    w.grad=None
            loss.backward(); opt.step(); opt.zero_grad()
            if interval is not None and t < steps//5: w.data.fill_((interval[0]+interval[1])/2)
        qn=(q.detach()/torch.linalg.vector_norm(q.detach())).cpu().numpy(); ww=float(w.detach().cpu())
        ev,vec=eigh(K,M); freqs=np.sqrt(ev); # mass-normalized reference vectors
        mac=np.array([(abs(qn@vec[:,i])**2/((qn@qn)*(vec[:,i]@vec[:,i]))) for i in range(len(freqs))])
        return ww, int(np.argmax(mac)), float(mac.max()), float(freqs[np.argmin(abs(freqs-ww))])
    except Exception:
        # CPU retry is required for shared/fragile CUDA environments.
        if dev=='cuda':
            torch.cuda.empty_cache(); torch.cuda.is_available=lambda: False
            return train(K,M,target,interval,steps,seed)
        raise

def main():
    K,M=beam(8); exact=np.sqrt(eigh(K,M,eigvals_only=True)); target=3
    ws=np.linspace(0,exact[-1]*1.1,250)
    counts=np.array([inertia(K,M,float(w)) for w in ws])
    monotone=bool(np.all(np.diff(counts)>=0)); changes=int(np.sum(np.diff(counts)>0))
    lo,up,freqs=bracket(K,M,target)
    contained=int(np.sum((freqs>lo)&(freqs<up)))
    baseline=[train(K,M,target,None,2000,s) for s in range(8)]
    idea=[train(K,M,target,(lo,up),2000,s) for s in range(8)]
    # lower-mode collapse means identified FEM mode index below requested target.
    b_lower=sum(x[1]<target-1 for x in baseline); i_lower=sum(x[1]<target-1 for x in idea)
    out={'exact_frequencies':exact.tolist(),'count_monotone':monotone,'count_grid_changes':changes,'bracket':[lo,up],'bracket_width':up-lo,'frequencies_inside_bracket':contained,'target_mode_zero_based':target-1,
         'baseline':baseline,'idea':idea,'baseline_lower_mode_rate':b_lower/len(baseline),'idea_lower_mode_rate':i_lower/len(idea),
         'baseline_target_rate':sum(x[1]==target-1 for x in baseline)/len(baseline),'idea_target_rate':sum(x[1]==target-1 for x in idea)/len(idea)}
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()