"""Analytic 2D Airy-tail and linear inverse-problem teaching calculations.

All optical lengths are in lambda/NA. The Airy bound applies only to a
normalized, unaberrated, scalar, spatially incoherent circular-pupil model.
The matrix example is dimensionless and is not a calibrated OPC model.
"""
import json
from pathlib import Path
import numpy as np
from scipy.special import j0, j1, jn_zeros
from scipy.integrate import quad
from scipy.optimize import brentq
import scipy
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

ROOT=Path(__file__).resolve().parent
OUT=ROOT/'figures'
OUT.mkdir(exist_ok=True)

def tail(r):
    v=2*np.pi*np.asarray(r)
    return j0(v)**2+j1(v)**2

def radial_energy(v):
    return 0.0 if v==0 else 2*j1(v)**2/v

def run():
    first_zero=float(jn_zeros(1,1)[0]/(2*np.pi))
    budget=.25*2*.2/40
    radius=float(brentq(lambda r: tail(r)-budget,1,100,xtol=1e-12))
    checks=[]
    for r in [first_zero,1,5,10,20,radius,100]:
        v=2*np.pi*r
        # Integrate only a finite interval; add the analytic remaining tail.
        value,err=quad(radial_energy,0,v,epsabs=2e-12,epsrel=2e-12,limit=1000)
        checks.append(dict(radius=r,analytic_tail=float(tail(r)),
            finite_integral=value,quadrature_error_estimate=err,
            normalization_residual=float(value+tail(r)-1)))
    assert max(abs(c['normalization_residual']) for c in checks)<1e-10
    assert abs(tail(radius)-budget)<1e-12
    n=81; mid=n//2
    matrix_cases=[]
    for alpha in [.45,.49]:
        J=np.eye(n)-alpha*(np.eye(n,k=1)+np.eye(n,k=-1))
        b=np.eye(n)[:,mid]
        delta=np.linalg.solve(J,b)
        residual=float(np.max(abs(J@delta-b)))
        matrix_cases.append(dict(alpha=alpha,size=n,condition_number=float(np.linalg.cond(J)),
            solution=delta.tolist(),normalized_solution=(delta/delta[mid]).tolist(),
            equation_residual=residual))
        assert residual<1e-12
    data=dict(description=__doc__,versions=dict(numpy=np.__version__,scipy=scipy.__version__),
        airy=dict(first_zero=first_zero,first_zero_tail=float(tail(first_zero)),
            edge_budget=dict(width_nm=40,nils=2,threshold=.25,edge_nm=.2,intensity=budget),
            radius_for_budget=radius,checks=checks),matrix_cases=matrix_cases)
    (ROOT/'first-principles-results.json').write_text(json.dumps(data,indent=2)+'\n')
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':11,'axes.spines.top':False,
        'axes.spines.right':False,'figure.facecolor':'white','axes.grid':True,
        'grid.alpha':.18,'svg.fonttype':'none','savefig.bbox':'tight'})
    def save(fig,name):
        fig.savefig(OUT/(name+'.svg')); fig.savefig(OUT/(name+'.png'),dpi=180); plt.close(fig)
    r=np.linspace(0,10,8001); v=2*np.pi*r
    q=np.ones_like(r); q[1:]=(2*j1(v[1:])/v[1:])**2
    fig,axes=plt.subplots(1,2,figsize=(11,4.6))
    axes[0].semilogy(r,q,color='#176d77',lw=1.3)
    axes[0].axvline(first_zero,ls=':',color='#777')
    axes[0].set(xlabel=r'Radius $R/(\lambda/NA)$',ylabel=r'Local PSF / peak',ylim=(1e-8,2),
        title='A weak local sidelobe is not the total tail')
    r=np.geomspace(.2,100,4000)
    axes[1].loglog(r,tail(r),color='#176d77',label=r'$T(R)=J_0^2+J_1^2$')
    axes[1].loglog(r,1/(np.pi**2*r),'--',color='#c25d2d',label=r'Asymptote $1/(\pi^2 R)$')
    axes[1].axhline(budget,ls=':',color='#777',label='Illustrative intensity budget 0.0025')
    axes[1].plot(radius,budget,'o',color='#7667a8')
    axes[1].annotate(f'{radius:.2f}',(radius,budget),xytext=(-25,15),textcoords='offset points')
    axes[1].set(xlabel=r'Radius $R/(\lambda/NA)$',ylabel='Integrated energy outside radius',
        title='Ideal circular pupil, incoherent intensity bound')
    axes[1].legend(fontsize=8.4); fig.tight_layout(); save(fig,'airy-tail-budget')
    # Portrait version for the WeChat reading column.
    fig,ax=plt.subplots(figsize=(6.5,4.9))
    ax.loglog(r,tail(r),lw=2,color='#176d77',label='Integrated energy outside R')
    ax.axhline(budget,ls=':',color='#777',label='Illustrative budget 0.0025')
    ax.plot(radius,budget,'o',color='#c25d2d')
    ax.annotate(f'R = {radius:.2f}',(radius,budget),xytext=(-90,18),textcoords='offset points')
    ax.set(xlabel=r'Radius $R/(\lambda/NA)$',ylabel='Normalized energy',ylim=(.0008,1))
    ax.legend(fontsize=10);fig.tight_layout();save(fig,'airy-tail-budget-mobile')
    fig,ax=plt.subplots(figsize=(9,4.8))
    x=np.arange(n)-mid
    ax.stem([-1,0,1],[.45,1,.45],linefmt='#aeb8bd',markerfmt='o',basefmt=' ',label=r'Nonzero $|J_{i0}|$, $\alpha=0.45$')
    for case,c in zip(matrix_cases,['#176d77','#c25d2d']):
        ax.semilogy(x,np.abs(case['normalized_solution']),color=c,
            label=rf'Normalized $J^{{-1}}e_0$, $\alpha={case["alpha"]}$')
    ax.set(xlabel='Segment index relative to perturbed target',ylabel='Normalized magnitude',
        xlim=(-25,25),ylim=(1e-6,2),title='A local Jacobian can have a nonlocal inverse')
    ax.legend(fontsize=9);fig.tight_layout();save(fig,'inverse-coupling')
    print(json.dumps(dict(airy=data['airy'],matrix_residuals=[c['equation_residual'] for c in matrix_cases]),indent=2))

if __name__=='__main__': run()