"""Reproducible teaching calculations, not a calibrated OPC/PDK model.

Requires NumPy and Matplotlib. All optical lengths use lambda/NA; the
gradient example uses pixels. No wafer measurements are used.
"""
import json
from pathlib import Path
import numpy as np

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

def convolve_same_batch(a,k):
    n=a.shape[-1]+len(k)-1
    size=1 << (n-1).bit_length()
    full=np.fft.ifft(np.fft.fft(a,size,axis=-1)*np.fft.fft(k,size),axis=-1)
    start=(len(k)-1)//2
    return full[...,start:start+a.shape[-1]]

def illumination(kind,count):
    q,w=np.polynomial.legendre.leggauss(count)
    if kind=='coherent': return np.array([0.]),np.array([1.])
    if kind=='broad': return .7*q,w/2
    # Two rectangular poles, each of half-width .08, centered at +/-.65.
    return np.r_[.65+.08*q,-.65+.08*q],np.r_[w/4,w/4]

RADII=np.array([.5,.75,1,1.5,2,2.5,3,4,5,6,8,10,12,16,20,24,28,32.])
REMOTE=[(1.5,2.2),(3.4,4.5),(7.,11.),(14.,20.),(-2.6,-2.),(-6.,-4.3),(-18.,-13.)]

def mask(x):
    dx=x[1]-x[0]
    m=np.zeros_like(x)
    for a,b in [(-.4,.4)]+REMOTE:
        # Exact area fraction for every mask pixel, so refinement does not
        # move the intended rectangle boundary by a whole grid point.
        m+=np.maximum(0,np.minimum(x+dx/2,b)-np.maximum(x-dx/2,a))/dx
    return m

def central_edges(x,image):
    ids=np.flatnonzero((image[:-1]-.25)*(image[1:]-.25)<=0)
    roots=[]
    for j in ids:
        if image[j+1]!=image[j]:
            roots.append(float(x[j]+(.25-image[j])*(x[j+1]-x[j])/(image[j+1]-image[j])))
    left=[r for r in roots if -.95<r<0]
    right=[r for r in roots if 0<r<.95]
    if not left or not right: raise ValueError('central threshold contour lost')
    return np.array([max(left),min(right)])

def optical_case(kind,dx=1/128,count=81,radii=RADII):
    x=np.arange(-32,32+dx/2,dx)
    s,w=illumination(kind,count)
    mod=mask(x)[None,:]*np.exp(2j*np.pi*s[:,None]*x[None,:])
    # Ideal slit pupil P(f)=1 for |f|<=1: h(x)=2*sinc(2x).
    half=max(32.,float(np.max(radii)))
    hx=np.arange(-half,half+dx/2,dx)
    h=2*np.sinc(2*hx)*dx
    keep=np.abs(x)<.95
    images=[]; edges=[]
    for radius in radii:
        window=np.clip((radius+dx/2-np.abs(hx))/dx,0,1)
        field=convolve_same_batch(mod,h*window)
        image=np.sum(w[:,None]*np.abs(field)**2,axis=0)
        images.append(image[keep]); edges.append(central_edges(x,image))
    edges=np.array(edges)
    reference=edges[-1]
    err=np.max(np.abs(edges-reference),axis=1)
    return dict(kind=kind,dx=dx,source_points=len(s),radii=list(map(float,radii)),
        edge_positions=edges.tolist(),reference_edges=reference.tolist(),
        max_edge_error=err.tolist(),sampled_tail_envelope=np.maximum.accumulate(err[::-1])[::-1].tolist(),
        core_max_intensity_error=np.max(np.abs(np.array(images)-images[-1]),axis=1).tolist()),(x[keep],np.array(images))

def gradient_experiment():
    n,r,tile,steps=512,8,128,80
    xx=np.arange(-r,r+1); h=np.exp(-.5*(xx/3.)**2); h/=h.sum()
    desired=np.zeros(n)
    for a,b in [(40,63),(116,140),(240,261),(376,397),(440,464)]: desired[a:b]=1
    initial=.15+.65*desired
    def gradient(m,t):
        u=np.convolve(m,h,'same')
        return 2*np.convolve((u*u-t)*u,h,'same')
    def tiled(m,halo,frozen=False):
        g=np.empty_like(m)
        for lo in range(0,n,tile):
            hi=lo+tile; a=max(0,lo-halo); b=min(n,hi+halo)
            local=(initial if frozen else m)[a:b].copy()
            local[lo-a:hi-a]=m[lo:hi]
            g[lo:hi]=gradient(local,desired[a:b])[lo-a:hi-a]
        return g
    true=gradient(initial,desired)
    g1=tiled(initial,r); g2=tiled(initial,2*r)
    state={key:initial.copy() for key in ['global','synchronized','frozen']}
    for _ in range(steps):
        for key,m in state.items():
            g=gradient(m,desired) if key=='global' else tiled(m,2*r,key=='frozen')
            state[key]=np.clip(m-.1*g,0,1)
    result=dict(kernel_radius_pixels=r,tile_pixels=tile,iterations=steps,
        gradient_error_R=float(np.max(abs(g1-true))),gradient_error_2R=float(np.max(abs(g2-true))),
        final_mask_error_synchronized=float(np.max(abs(state['synchronized']-state['global']))),
        final_mask_error_frozen=float(np.max(abs(state['frozen']-state['global']))))
    assert result['gradient_error_2R']<1e-12
    assert result['final_mask_error_synchronized']<1e-12
    assert result['gradient_error_R']>1e-5
    assert result['final_mask_error_frozen']>1e-3
    # A separate finite-difference check tests the analytic derivative.
    j=127; eps=1e-6
    def loss(m): return .5*np.sum((np.convolve(m,h,'same')**2-desired)**2)
    a=initial.copy(); b=initial.copy(); a[j]+=eps; b[j]-=eps
    result['finite_difference_gradient_error']=float(abs((loss(a)-loss(b))/(2*eps)-true[j]))
    assert result['finite_difference_gradient_error']<1e-7
    return result,dict(g1=g1,g2=g2,true=true,state=state)

def run():
    cases=[]; raw={}
    for kind in ['coherent','broad','two_poles']:
        result,curves=optical_case(kind); cases.append(result); raw[kind]=curves
    checks={}
    for kind in ['coherent','broad','two_poles']:
        base=np.array(next(i for i in cases if i['kind']==kind)['reference_edges'])
        fine,_=optical_case(kind,dx=1/256,count=161,radii=np.array([32.,40.]))
        checks[kind]=dict(fine_reference_edges=fine['reference_edges'],
            base_to_refined_max_edge_change=float(np.max(abs(base-np.array(fine['reference_edges'])))),
            radius_32_to_40_change=float(np.max(abs(np.array(fine['edge_positions'][0])-np.array(fine['edge_positions'][1])))))
        assert checks[kind]['radius_32_to_40_change']<1e-12
        assert checks[kind]['base_to_refined_max_edge_change']<5e-5
    grad,gdata=gradient_experiment()
    data=dict(scope='1D scalar slit-pupil teaching calculation; not a process recommendation',
        numpy_version=np.__version__,
        coordinate_unit='lambda/NA',pupil='P(f)=1 for |f|<=1; zero otherwise',
        mask=dict(central_open_interval=[-.4,.4],remote_open_intervals=REMOTE,elsewhere=0),
        threshold=.25,domain=[-32,32],core=[-.95,.95],
        method='linear FFT convolution with zero padding; no kernel renormalization; source integration by Gauss-Legendre',
        cases=cases,refinement_checks=checks,gradient_demo=grad)
    (ROOT/'simulation-results.json').write_text(json.dumps(data,indent=2)+'\n')
    figures(data,raw,gdata)
    print(json.dumps(dict(refinement=checks,gradient=grad),indent=2))

def figures(data,raw,gdata):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    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'})
    colors=['#166b74','#c25d2d','#7667a8']
    def save(fig,name):
        fig.savefig(OUT/(name+'.svg')); fig.savefig(OUT/(name+'.png'),dpi=180); plt.close(fig)
    # These are analytic source Fourier transforms, not reproduced source figures.
    d=np.linspace(0,12,1600)
    fig,ax=plt.subplots(figsize=(9,4.5))
    ax.plot(d,np.ones_like(d),label='Single source point',c=colors[0])
    ax.plot(d,np.abs(np.sinc(2*.7*d)),label='Uniform source [-0.7, 0.7]',c=colors[1])
    ax.plot(d,np.abs(np.cos(2*np.pi*.65*d)*np.sinc(2*.08*d)),label='Two poles: centers +/-0.65, half-width 0.08',c=colors[2])
    ax.set(xlabel=r'Separation $d/(\lambda/NA)$',ylabel=r'Normalized mutual coherence $|\mu(d)|$',ylim=(-.03,1.12))
    ax.legend(fontsize=9,loc='upper right'); fig.tight_layout(); save(fig,'source-coherence')
    fig,axes=plt.subplots(1,2,figsize=(10.5,4.2))
    for i,case in enumerate(data['cases']):
        r=np.array(case['radii']); e=np.array(case['max_edge_error'])
        axes[0].semilogy(r,np.where(e>1e-12,e,np.nan),'-o',markersize=3,c=colors[i],label=case['kind'])
        envelope=np.array(case['sampled_tail_envelope'])
        axes[1].semilogy(r,np.where(envelope>1e-12,envelope,np.nan),'-o',markersize=3,c=colors[i])
    for ax in axes:
        ax.set(xlabel=r'Kernel radius $R/(\lambda/NA)$',ylim=(1e-5,.2),xlim=(0,25))
        ax.axhline(3e-5,c='#777777',ls=':',lw=1)
        ax.axvspan(21,25,color='#eeeeee',zorder=-1)
        ax.text(23,3e-5,'All mask\nincluded',ha='center',fontsize=8)
    axes[0].set(ylabel=r'Max central-edge error / $(\lambda/NA)$',title='Error at each sampled radius')
    axes[1].set(title='Envelope over sampled larger radii')
    axes[0].legend(fontsize=9); fig.tight_layout(); save(fig,'radius-convergence')
    fig,ax=plt.subplots(figsize=(6.5,5.2))
    for i,case in enumerate(data['cases']):
        e=np.array(case['max_edge_error'])
        ax.semilogy(case['radii'],np.where(e>1e-12,e,np.nan),'-o',markersize=4,c=colors[i],label=case['kind'])
    ax.axhline(3e-5,c='#777777',ls=':',lw=1)
    ax.axvspan(21,25,color='#eeeeee',zorder=-1)
    ax.text(23,1e-3,'All mask\nincluded',ha='center',fontsize=8,rotation=90)
    ax.set(xlabel=r'Radius $R/(\lambda/NA)$',ylabel=r'Edge error / $(\lambda/NA)$',xlim=(0,25),ylim=(1e-5,.2))
    ax.legend(fontsize=10); fig.tight_layout(); save(fig,'radius-convergence-mobile')
    fig,ax=plt.subplots(figsize=(9,4.6))
    ratio=np.linspace(0,1,200)
    ax.plot(ratio,(1+2*ratio)**2,label='Forward area: $(1+2R/L)^2$',c=colors[0],lw=2)
    ax.plot(ratio,(1+4*ratio)**2,label='Independent gradient area: $(1+4R/L)^2$',c=colors[1],lw=2)
    ax.set(xlabel='Halo radius / core side length (R/L)',ylabel='Read area / core area',ylim=(0,26))
    ax.legend(); fig.tight_layout(); save(fig,'halo-cost')
    fig,axes=plt.subplots(2,1,figsize=(9,6),sharex=True)
    t=np.arange(512)
    axes[0].plot(t,abs(gdata['g1']-gdata['true']),c=colors[1],label='Gradient error, halo R')
    axes[0].plot(t,abs(gdata['g2']-gdata['true']),c=colors[0],label='Gradient error, halo 2R')
    axes[0].set(ylabel='Absolute gradient error'); axes[0].legend(fontsize=9)
    state=gdata['state']
    axes[1].plot(t,abs(state['frozen']-state['global']),c=colors[2],label='Final-mask error: frozen halo')
    axes[1].plot(t,abs(state['synchronized']-state['global']),c=colors[0],label='Final-mask error: synchronized')
    for ax in axes:
        for b in [128,256,384]: ax.axvline(b,c='#777',ls=':',lw=1)
        ax.set_xlim(220,300)
    axes[1].set(xlabel='Pixel position (dotted line = tile boundary)',ylabel='Absolute mask-value error')
    axes[1].legend(fontsize=9); fig.tight_layout(); save(fig,'gradient-and-sync')

if __name__=='__main__': run()