import json
import math
import numpy as np
from pathlib import Path

# Adjoint-weak fractional residual MVP.
# The Grünwald matrix is used only to make the discrete integration-by-parts
# identity explicit; weak measurements use its exact transpose.

def grunwald_matrix(n, alpha, T=1.0):
    h = T / (n - 1)
    # left-sided GL derivative: (D f)_i = h^-alpha sum_{k<=i} (-1)^k C(alpha,k) f_{i-k}
    w = np.empty(n)
    w[0] = 1.0
    for k in range(1, n):
        w[k] = w[k-1] * (-(alpha-k+1)/k)
    A = np.zeros((n, n))
    for i in range(n):
        A[i, :i+1] = h**(-alpha) * w[:i+1][::-1]
    return A

def trapezoid_weights(n, T=1.0):
    h = T/(n-1)
    q = np.full(n, h)
    q[[0,-1]] = h/2
    return q

def gaussian_phi(t, center=.62, width=.16):
    return np.exp(-0.5*((t-center)/width)**2)

def one_check(n, alpha, sigma, rng):
    hq = trapezoid_weights(n)
    A = grunwald_matrix(n, alpha)
    phi = gaussian_phi(np.linspace(0,1,n))
    # Integral phi^T W D f = f^T D^T W phi.
    lhs = phi @ (hq * (A @ rng.normal(size=n)))
    qweak = A.T @ (hq * phi)
    rhs = qweak @ rng.normal(size=n) # independent draw, used only for norm check below
    # exact coefficient vector acting on sampled f
    q = qweak
    return float(np.dot(q,q)), float(np.dot(hq*phi, hq*phi)), A, q, phi, hq

def noisy_variance_sweep():
    rng = np.random.default_rng(11)
    n, alpha = 257, .8
    A = grunwald_matrix(n, alpha)
    t = np.linspace(0,1,n); hq = trapezoid_weights(n)
    phi = gaussian_phi(t)
    qweak = A.T @ (hq*phi)
    # A pointwise derivative at an interior location, versus one weak scalar.
    i = n//2
    qstrong = A[i].copy()
    rows=[]
    for sigma in [.01,.03,.1,.3]:
        vals_w=[]; vals_s=[]
        for _ in range(5000):
            e=rng.normal(0,sigma,n)
            vals_w.append(qweak@e); vals_s.append(qstrong@e)
        vw=np.var(vals_w, ddof=1); vs=np.var(vals_s, ddof=1)
        rows.append({'sigma':sigma,'weak_var':vw,'strong_var':vs,
                     'weak_over_sigma2':vw/sigma**2,'strong_over_sigma2':vs/sigma**2})
    return rows

def resolution_sweep():
    alpha=.8; rng=np.random.default_rng(12); out=[]
    for n in [65,129,257,513]:
        t=np.linspace(0,1,n); hq=trapezoid_weights(n)
        A=grunwald_matrix(n,alpha); phi=gaussian_phi(t)
        qw=A.T@(hq*phi); qs=A[n//2]
        out.append({'n':n,'weak_norm2':float(qw@qw),'strong_norm2':float(qs@qs)})
    # consecutive ratios are the directly testable scaling predictions
    for i in range(1,len(out)):
        out[i]['weak_ratio_prev']=out[i]['weak_norm2']/out[i-1]['weak_norm2']
        out[i]['strong_ratio_prev']=out[i]['strong_norm2']/out[i-1]['strong_norm2']
    return out

def coefficient_fit(seed, sigma, n=129, alpha=.8, trials=200):
    rng=np.random.default_rng(seed); t=np.linspace(0,1,n); hq=trapezoid_weights(n)
    A=grunwald_matrix(n,alpha); phi=gaussian_phi(t)
    # Synthetic relation D^alpha u = c u on a fixed smooth trajectory. We estimate
    # c from noisy sampled u; weak form transfers D to phi.
    u=np.exp(-.8*t)*(1+.2*np.sin(2*np.pi*t)); ctrue=.7
    # use a generated target equal to the discrete operator, avoiding continuum mismatch
    target=A@u
    qweak=A.T@(hq*phi); qid=hq*phi
    # Compare noise-induced error to each estimator's noiseless value; this
    # isolates robustness without claiming the arbitrary trajectory obeys a
    # pointwise constant-coefficient fractional ODE.
    weak0=(qweak@u)/(qid@u)
    inds=np.arange(n//2-3,n//2+4)
    strong0=np.sum(u[inds]*(A[inds]@u))/np.sum(u[inds]**2)
    csw=[]; css=[]
    for _ in range(trials):
        un=u+rng.normal(0,sigma,n)
        # weak scalar equation: int phi D u = c int phi u
        den=qid@un; csw.append((qweak@un)/den)
        # strong local equation, average a central neighborhood for fair scalar noise
        inds=np.arange(n//2-3,n//2+4)
        d=A[inds]@un
        css.append(np.sum(un[inds]*d)/np.sum(un[inds]**2))
    return {'sigma':sigma,'weak_rmse':float(np.sqrt(np.mean((np.array(csw)-weak0)**2))),
            'strong_rmse':float(np.sqrt(np.mean((np.array(css)-strong0)**2))),
            'weak_noiseless':float(weak0),'strong_noiseless':float(strong0)}

def identity_check():
    rng=np.random.default_rng(4); n=97; alpha=.65
    A=grunwald_matrix(n,alpha); q=trapezoid_weights(n); f=rng.normal(size=n); phi=rng.normal(size=n)
    lhs=phi@(q*(A@f)); rhs=f@(A.T@(q*phi))
    return {'absolute_error':float(abs(lhs-rhs)), 'relative_error':float(abs(lhs-rhs)/(1+abs(lhs)))}

def main():
    result={'identity':identity_check(), 'noise_sweep':noisy_variance_sweep(),
            'resolution_sweep':resolution_sweep(),
            'coefficient_fit':[coefficient_fit(100+i,s) for i,s in enumerate([.01,.03,.1,.3])]}
    # Predictions from the mechanism: variance is sigma^2 ||q||^2; weak q norm
    # scales approximately h^(1/2) for a fixed smooth localized kernel. The
    # GL point-row norm is dominated by summable near-diagonal weights, so its
    # squared norm scales as h^(-2 alpha), giving ratio 2^(2 alpha).
    r=result['resolution_sweep']
    result['predictions']={
      'noise_linearity':'weak variance / sigma^2 should be constant',
      'weak_resolution_ratio_predicted':0.5,
      'strong_resolution_ratio_predicted':2**(2*.8),
      'observed_weak_ratios':[x.get('weak_ratio_prev') for x in r[1:]],
      'observed_strong_ratios':[x.get('strong_ratio_prev') for x in r[1:]],
      'identity_tolerance':'machine precision'}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()