"""Synthetic, multi-surface phakic eye with a finite, bored ICL volume.

SI geometry, vector Snell refraction, Fresnel power, real aperture clipping.
ICL local surfaces are searched non-sequentially, including cylindrical bore
and optic boundary. Ordinary reflected branches are accounted as loss, not
rendered ghost images. TIR within the implant is followed up to max_events.
This is an explicit synthetic design, not a manufacturer or patient model.
"""

from dataclasses import asdict, dataclass, replace
import math

import numpy as np

PATHS=("material", "transparent_port", "outside_optic", "port_wall", "optic_wall", "no_implant")


@dataclass(frozen=True)
class Eye:
    pupil_mm: float = 6.5
    optic_mm: float = 5.8
    port_mm: float = .36
    icl_front_mm: float = 3.4
    icl_thickness_mm: float = .1
    icl_front_radius_mm: float = -28.5
    icl_back_radius_mm: float = 28.5
    icl_decenter_x_mm: float = 0.
    icl_decenter_y_mm: float = 0.
    cornea_front_mm: float = 0.
    cornea_thickness_mm: float = .55
    cornea_front_radius_mm: float = 7.8
    cornea_back_radius_mm: float = 6.5
    cornea_front_q: float = -.26
    cornea_back_q: float = -.26
    iris_mm: float = 3.1
    lens_front_mm: float = 4.2
    lens_back_mm: float = 8.0
    lens_front_radius_mm: float = 10.
    lens_back_radius_mm: float = -6.
    lens_front_q: float = -1.
    lens_back_q: float = -4.
    n_cornea: float = 1.376
    n_aqueous: float = 1.336
    n_icl: float = 1.45
    n_lens: float = 1.42
    n_vitreous: float = 1.336
    beam_radius_mm: float = 4.5
    implant: bool = True
    fresnel: bool = True
    retina_mm: float = 28.
    angular_focal_mm: float = 20.
    max_events: int = 12

    def to_dict(self):return asdict(self)

    def validate(self):
        for key,value in self.to_dict().items():
            if isinstance(value,(float,int)) and not math.isfinite(value):raise ValueError(f"nonfinite {key}")
        if not 0<=self.port_mm<self.optic_mm or not 0<self.pupil_mm<8.5:raise ValueError("invalid aperture")
        if not self.cornea_thickness_mm<self.iris_mm<self.icl_front_mm<self.lens_front_mm<self.lens_back_mm<self.retina_mm:raise ValueError("invalid axial anatomy")
        if self.icl_thickness_mm<=0 or self.beam_radius_mm<=0:raise ValueError("positive dimensions required")
        if self.n_icl<1 or min(self.n_cornea,self.n_aqueous,self.n_lens,self.n_vitreous)<1:raise ValueError("invalid refractive index")
        if self.icl_front_radius_mm>=0 or self.icl_back_radius_mm<=0:raise ValueError("this synthetic design is biconcave")
        r=self.optic_mm/2
        front=self.icl_front_mm+sag(r,self.icl_front_radius_mm,0.)
        back=self.icl_front_mm+self.icl_thickness_mm+sag(r,self.icl_back_radius_mm,0.)
        if front<=self.iris_mm or back>=self.lens_front_mm:raise ValueError("implant volume overlaps iris or crystalline lens")


def sag(radius,curvature_radius,q=0.):
    radius=np.asarray(radius)
    root=1-(1+q)*(radius/curvature_radius)**2
    return radius**2/(curvature_radius*(1+np.sqrt(np.maximum(root,0))))


def surface_hit(p,d,vertex,radius,q=0.,offset=(0.,0.),r_min=0.,r_max=np.inf):
    """Nearest positive intersection on the vertex branch of a conic."""
    x=p[:,0]-offset[0];y=p[:,1]-offset[1];z=p[:,2]-vertex
    a=d[:,0]**2+d[:,1]**2+(1+q)*d[:,2]**2
    b=2*(x*d[:,0]+y*d[:,1]-radius*d[:,2]+(1+q)*z*d[:,2])
    c=x*x+y*y-2*radius*z+(1+q)*z*z
    disc=b*b-4*a*c
    root=np.sqrt(np.maximum(disc,0.))
    # Stable quadratic roots plus the paraboloid/axial linear limit.
    qq=-.5*(b+np.where(b>=0,root,-root))
    t1=np.divide(qq,a,out=np.full(len(p),np.inf),where=np.abs(a)>1e-20)
    t2=np.divide(c,qq,out=np.full(len(p),np.inf),where=np.abs(qq)>1e-25)
    linear=np.divide(-c,b,out=np.full(len(p),np.inf),where=np.abs(b)>1e-25)
    t1=np.where(np.abs(a)<1e-20,linear,t1)
    best=np.full(len(p),np.inf)
    for t in (t1,t2):
        safe=np.where(np.isfinite(t),t,0.)
        hx=x+safe*d[:,0];hy=y+safe*d[:,1];hz=z+safe*d[:,2]
        rr=hx*hx+hy*hy
        domain=1-(1+q)*rr/radius**2
        expected=sag(np.sqrt(rr),radius,q)
        valid=(disc>=0)&(t>2e-9)&(rr>=r_min*r_min)&(rr<=r_max*r_max)&(domain>=0)&(np.abs(hz-expected)<1e-7)
        best=np.where(valid& (t<best),t,best)
    hit=p+np.where(np.isfinite(best),best,0.)[:,None]*d
    normal=np.column_stack((hit[:,0]-offset[0],hit[:,1]-offset[1],-radius+(1+q)*(hit[:,2]-vertex)))
    normal/=np.maximum(np.linalg.norm(normal,axis=1)[:,None],1e-30)
    return best,normal


def cylinder_hit(p,d,radius,eye):
    x=p[:,0]-eye.icl_decenter_x_mm*1e-3;y=p[:,1]-eye.icl_decenter_y_mm*1e-3
    a=d[:,0]**2+d[:,1]**2;b=2*(x*d[:,0]+y*d[:,1]);c=x*x+y*y-radius**2
    disc=b*b-4*a*c
    root=np.sqrt(np.maximum(disc,0.))
    best=np.full(len(p),np.inf)
    for sign in (-1,1):
        t=np.divide(-b+sign*root,2*a,out=np.full(len(p),np.inf),where=a>1e-25)
        z=p[:,2]+np.where(np.isfinite(t),t,0.)*d[:,2]
        front=eye.icl_front_mm*1e-3+sag(radius,eye.icl_front_radius_mm*1e-3)
        back=(eye.icl_front_mm+eye.icl_thickness_mm)*1e-3+sag(radius,eye.icl_back_radius_mm*1e-3)
        valid=(disc>=0)&(t>2e-9)&(z>=front-1e-10)&(z<=back+1e-10)
        best=np.where(valid&(t<best),t,best)
    h=p+np.where(np.isfinite(best),best,0.)[:,None]*d
    normal=np.column_stack((h[:,0]-eye.icl_decenter_x_mm*1e-3,h[:,1]-eye.icl_decenter_y_mm*1e-3,np.zeros(len(p))))
    normal/=np.maximum(np.linalg.norm(normal,axis=1)[:,None],1e-30)
    return best,normal


def snell_fresnel(direction,normal,n1,n2):
    """Unit transmitted/TIR-reflected direction, unpolarized power R and TIR."""
    n1=np.broadcast_to(n1,(len(direction),));n2=np.broadcast_to(n2,(len(direction),))
    normal=np.where((np.sum(direction*normal,axis=1)>0)[:,None],-normal,normal)
    ci=np.clip(-np.sum(direction*normal,axis=1),0.,1.)
    eta=n1/n2;k=1-eta**2*(1-ci**2);tir=k<0
    ct=np.sqrt(np.maximum(k,0.))
    refracted=eta[:,None]*direction+(eta*ci-ct)[:,None]*normal
    reflected=direction+2*ci[:,None]*normal
    result=np.where(tir[:,None],reflected,refracted)
    result/=np.maximum(np.linalg.norm(result,axis=1)[:,None],1e-30)
    rs=((n1*ci-n2*ct)/np.maximum(n1*ci+n2*ct,1e-30))**2
    rp=((n2*ci-n1*ct)/np.maximum(n2*ci+n1*ct,1e-30))**2
    reflect=np.where(tir,1.,.5*(rs+rp))
    return result,reflect,tir


def disk_samples(count,radius):
    index=np.arange(count)+.5
    rr=radius*np.sqrt(index/count)
    phi=index*(np.pi*(3-np.sqrt(5.)))
    return np.column_stack((rr*np.cos(phi),rr*np.sin(phi)))


def _implant(p,d,w,eye,ledger,history):
    if not eye.implant:return p,d,w,np.full(len(p),5,dtype=int)
    local=np.array([eye.icl_decenter_x_mm,eye.icl_decenter_y_mm])*1e-3
    front=eye.icl_front_mm*1e-3;back=(eye.icl_front_mm+eye.icl_thickness_mm)*1e-3
    inner=eye.port_mm*.5e-3;outer=eye.optic_mm*.5e-3
    exit_z=back+float(sag(outer,eye.icl_back_radius_mm*1e-3))+1e-7
    t=(.5*(front+back)-p[:,2])/d[:,2]
    cross=p[:,:2]+t[:,None]*d[:,:2]-local
    rr=np.linalg.norm(cross,axis=1)
    labels=np.where(rr<inner,1,2)
    inside=np.zeros(len(p),dtype=bool);active=w>0
    for _ in range(eye.max_events):
        ids=np.where(active)[0]
        if len(ids)==0:break
        pp,dd=p[ids],d[ids]
        af,an=surface_hit(pp,dd,front,eye.icl_front_radius_mm*1e-3,offset=local,r_min=inner,r_max=outer)
        bf,bn=surface_hit(pp,dd,back,eye.icl_back_radius_mm*1e-3,offset=local,r_min=inner,r_max=outer)
        hf,hn=cylinder_hit(pp,dd,inner,eye) if inner>0 else (np.full(len(ids),np.inf),np.zeros_like(pp))
        of,on=cylinder_hit(pp,dd,outer,eye)
        times=np.stack((af,bf,hf,of),axis=1);chosen=np.argmin(times,axis=1)
        distance=times[np.arange(len(ids)),chosen]
        to_exit=np.divide(exit_z-pp[:,2],dd[:,2],out=np.full(len(ids),np.inf),where=dd[:,2]>0)
        done=(to_exit<distance)&(to_exit>=0)&~inside[ids]
        backward=(dd[:,2]<=0)&~inside[ids]
        invalid=(~np.isfinite(distance))&~done&~backward
        p[ids[done]]=pp[done]+to_exit[done,None]*dd[done]
        active[ids[done]]=False
        for mask,key in ((backward,"backward_exit"),(invalid,"unresolved")):
            ledger[key]+=float(w[ids[mask]].sum());w[ids[mask]]=0;active[ids[mask]]=False
        hit=~(done|backward|invalid)
        hs=ids[hit]
        if len(hs):
            pick=chosen[hit]
            normals=np.stack((an,bn,hn,on),axis=1)[np.where(hit)[0],pick]
            p[hs]+=distance[hit,None]*d[hs]
            n1=np.where(inside[hs],eye.n_icl,eye.n_aqueous)
            n2=np.where(inside[hs],eye.n_aqueous,eye.n_icl)
            new,reflect,tir=snell_fresnel(d[hs],normals,n1,n2)
            losses=w[hs]*np.where(tir,0.,reflect if eye.fresnel else 0.)
            ledger["reflected_untraced"]+=float(losses.sum());w[hs]-=losses
            inside[hs]=np.where(tir,inside[hs],~inside[hs])
            # A path can interact with both wall families; keep the last wall
            # tag, and export total wall-event counts separately.
            labels[hs[(pick<2)&(labels[hs]==2)]]=0
            labels[hs[pick==2]]=3;labels[hs[pick==3]]=4
            ledger["port_wall_events"]+=int(np.count_nonzero(pick==2))
            ledger["optic_wall_events"]+=int(np.count_nonzero(pick==3))
            ledger["tir_events"]+=int(np.count_nonzero(tir))
            d[hs]=new;p[hs]+=new*3e-9
        if history is not None:history.append(p.copy())
    ledger["unresolved"]+=float(w[active].sum());w[active]=0.
    return p,d,w,labels


def trace_to_vitreous(eye,field_deg=(0.,0.),count=131072,xy=None,keep_paths=False):
    eye.validate()
    xy=disk_samples(count,eye.beam_radius_mm*1e-3) if xy is None else np.asarray(xy,dtype=float)
    p=np.column_stack((xy,np.full(len(xy),-1e-3)))
    theta=np.radians(field_deg)
    direction=np.array([-np.tan(theta[0]),-np.tan(theta[1]),1.]);direction/=np.linalg.norm(direction)
    d=np.repeat(direction[None,:],len(p),axis=0)
    w=np.full(len(p),1/len(p));history=[p.copy()] if keep_paths else None
    ledger={"incident":1.,"reflected_untraced":0.,"blocked_iris":0.,"missed_surface":0.,"backward_exit":0.,"unresolved":0.,
            "port_wall_events":0,"optic_wall_events":0,"tir_events":0}
    def interface(vertex,radius,q,n1,n2,aperture):
        nonlocal p,d,w
        ids=np.where(w>0)[0]
        if not len(ids):return
        t,normal=surface_hit(p[ids],d[ids],vertex*1e-3,radius*1e-3,q,r_max=aperture*1e-3)
        valid=np.isfinite(t)&(d[ids,2]>0)
        ledger["missed_surface"]+=float(w[ids[~valid]].sum());w[ids[~valid]]=0.
        good=ids[valid]
        p[good]+=t[valid,None]*d[good]
        new,reflect,tir=snell_fresnel(d[good],normal[valid],n1,n2)
        fraction=np.where(tir,1.,reflect if eye.fresnel else 0.)
        loss=w[good]*fraction;ledger["reflected_untraced"]+=float(loss.sum());w[good]-=loss
        d[good]=new;p[good]+=new*3e-9
        if history is not None:history.append(p.copy())
    interface(eye.cornea_front_mm,eye.cornea_front_radius_mm,eye.cornea_front_q,1.,eye.n_cornea,5.)
    interface(eye.cornea_front_mm+eye.cornea_thickness_mm,eye.cornea_back_radius_mm,eye.cornea_back_q,eye.n_cornea,eye.n_aqueous,5.)
    t=(eye.iris_mm*1e-3-p[:,2])/d[:,2];p+=t[:,None]*d
    blocked=np.linalg.norm(p[:,:2],axis=1)>eye.pupil_mm*.5e-3
    ledger["blocked_iris"]+=float(w[blocked].sum());w[blocked]=0
    ledger["admitted_after_iris"]=float(w.sum())
    if history is not None:history.append(p.copy())
    p,d,w,labels=_implant(p,d,w,eye,ledger,history)
    interface(eye.lens_front_mm,eye.lens_front_radius_mm,eye.lens_front_q,eye.n_aqueous,eye.n_lens,5.)
    interface(eye.lens_back_mm,eye.lens_back_radius_mm,eye.lens_back_q,eye.n_lens,eye.n_vitreous,5.)
    return {"position":p,"direction":d,"weights":w,"path_id":labels,"ledger":ledger,"history":history}


def calibrate_eye(eye=Eye()):
    """Choose an axial retina for paraxial focus with a solid centered ICL.

    This constructs a synthetic myopic eye; it does not infer a patient's
    axial length from a prescription. This retina stays fixed in ablations.
    """
    calibration=replace(eye,port_mm=0.,pupil_mm=3.,implant=True,icl_decenter_x_mm=0.,icl_decenter_y_mm=0.)
    xy=disk_samples(512,.25e-3)
    rays=trace_to_vitreous(calibration,xy=xy)
    p,d,w=rays["position"],rays["direction"],rays["weights"]
    slopes=d[:,:2]/d[:,2,None];intercept=p[:,:2]-slopes*p[:,2,None]
    focus=-float(np.sum(w[:,None]*slopes*intercept)/np.sum(w[:,None]*slopes**2))
    xmeans=[]
    for angle in (-.2,.2):
        rays=trace_to_vitreous(calibration,field_deg=(angle,0.),xy=xy)
        p,d,w=rays["position"],rays["direction"],rays["weights"]
        hit=p+(focus-p[:,2])[:,None]*d/d[:,2,None]
        xmeans.append(float(np.sum(w*hit[:,0])/w.sum()))
    focal=-(xmeans[1]-xmeans[0])/(2*np.tan(np.radians(.2)))
    return replace(eye,retina_mm=focus*1e3,angular_focal_mm=focal*1e3)


def trace_eye(eye,field_deg=(0.,0.),count=131072,xy=None,keep_paths=False):
    result=trace_to_vitreous(eye,field_deg,count,xy,keep_paths)
    p,d,w=result["position"],result["direction"],result["weights"]
    distance=(eye.retina_mm*1e-3-p[:,2])/d[:,2]
    hit=p+distance[:,None]*d
    result["retina_xy_m"]=hit[:,:2]
    result["view_deg"]=-np.degrees(np.arctan(hit[:,:2]/(eye.angular_focal_mm*1e-3)))
    if result["history"] is not None:result["history"].append(hit)
    ledger=result["ledger"]
    ledger["retina_arrived"]=float(w.sum())
    terminals=("reflected_untraced","blocked_iris","missed_surface","backward_exit","unresolved","retina_arrived")
    ledger["energy_balance_error"]=1-sum(ledger[key] for key in terminals)
    ledger["path_power"]={name:float(w[result["path_id"]==i].sum()) for i,name in enumerate(PATHS)}
    return result


def retinal_histogram(rays,n=1024,fov_deg=8.):
    """Absolute per-incident-beam energy in a +x right/+y up angular grid."""
    maps=[]
    xy=rays["view_deg"];w=rays["weights"]
    edges=np.arange(n+1)*fov_deg/n-fov_deg/2-fov_deg/(2*n)
    for index in range(len(PATHS)):
        keep=(rays["path_id"]==index)&(w>0)
        hist,_,_=np.histogram2d(xy[keep,1],xy[keep,0],bins=(edges,edges),weights=w[keep])
        maps.append(hist)
    return np.asarray(maps)