"""Private fused CPU evaluation screen. No fast-math or changed reductions."""
from pathlib import Path
import importlib.util
import math
import numpy as np
from numba import njit

@njit(cache=True, nogil=True, fastmath=False, error_model='numpy')
def point_kernel(data, frame, cameras, covariance, pose, out):
    lx,ly,lz,lw,rx,ry,rz,rw,tx,ty,tz=pose
    for i in range(len(data)):
        out[2,i]=0.0
        p=data[i]
        c=cameras[frame[i]]
        x,y,z=p[0],p[1],p[2]
        ax=lw*x+lx*0.0+ly*z-lz*y
        ay=lw*y-lx*z+ly*0.0+lz*x
        az=lw*z+lx*y-ly*x+lz*0.0
        aw=lw*0.0-lx*x-ly*y-lz*z
        xx=aw*rx+ax*rw+ay*rz-az*ry+tx
        yy=aw*ry-ax*rz+ay*rw+az*rx+ty
        zz=aw*rz+ax*ry-ay*rx+az*rw+tz
        if not zz>1e-9:
            continue
        denominator=zz
        if c[0]!=0.0:
            norm=math.sqrt(xx*xx+yy*yy+zz*zz)
            if c[0]==1.0:
                denominator=zz+c[7]*norm
            else:
                zeta=c[7]*norm+zz
                second=math.sqrt(xx*xx+yy*yy+zeta*zeta)
                denominator=c[8]*second+(1.0-c[8])*zeta
            if not denominator>1e-12:
                continue
        xx,yy=xx/denominator,yy/denominator
        if c[13]!=0.0:
            radius=xx*xx+yy*yy
            if math.isfinite(radius) and abs(radius)>1e100:
                out[2,i]=2.0
                continue
            squared=radius*radius
            cubed=squared*radius
            radial=1.0+c[9]*radius+c[10]*squared+0.0*cubed
            dx=xx*radial+2.0*c[11]*xx*yy+c[12]*(radius+2.0*xx*xx)
            dy=yy*radial+c[11]*(radius+2.0*yy*yy)+2.0*c[12]*xx*yy
            xx,yy=dx,dy
        u,v=c[1]*xx+c[3],c[2]*yy+c[4]
        if not (u>=0.0 and u<c[5] and v>=0.0 and v<c[6]):
            continue
        du,dv=u-p[3],v-p[4]
        pixel=du*du+dv*dv
        squared_error=pixel
        if covariance:
            a,b,cc,d=p[5],p[6],p[7],p[8]
            squared_error=(d*du*du-(b+cc)*du*dv+a*dv*dv)/(a*d-b*cc)
        out[0,i]=pixel
        out[1,i]=math.sqrt(squared_error if squared_error>0.0 else 0.0)
        out[2,i]=1.0


def install():
    from calibrex.core.geometry import normalize_quaternion_xyzw,quaternion_conjugate_xyzw
    from calibrex.solvers import probabilistic_camera_lidar_refiner as m
    base=m._PreparedPoseEvaluation
    stats=dict(calls=0,compiled_calls=0,fallback_calls=0)
    class Prepared(base):
        def __init__(self,frames,settings):
            super().__init__(frames,settings)
            self.packed=None
            if self.fallback or not self.frames:
                return
            for f in self.frames:
                intr=f.intrinsics
                if (intr.projection!='pinhole' and intr.xi is None) or (intr.projection=='double_sphere' and intr.alpha is None):
                    return
                if intr.distortion_model=='radial-tangential' and (len(intr.distortion) not in {4,5} or (len(intr.distortion)==5 and intr.distortion[4]!=0.0)):
                    return
            rows,ids,cameras=[],[],[]
            for index,f in enumerate(self.frames):
                intr=f.intrinsics
                rows.append(np.column_stack((f.points,f.means,f.covariance,f.weights)))
                ids.append(np.full(len(f.points),index,dtype=np.int32))
                distortion=list(intr.distortion[:4])+[0.0]*4
                cameras.append([{'pinhole':0.,'mei':1.,'double_sphere':2.}[intr.projection],
                                intr.fx,intr.fy,intr.cx,intr.cy,intr.width,intr.height,
                                intr.xi or 0.,intr.alpha or 0.,*distortion[:4],float(intr.distortion_model=='radial-tangential')])
            self.packed=np.concatenate(rows)
            self.ids=np.concatenate(ids)
            self.cameras=np.asarray(cameras,dtype=np.float64)
            self.out=np.empty((3,len(self.packed)),dtype=float)
            self.weights=self.packed[:,9].copy()
        def evaluate(self,transform):
            stats['calls']+=1
            if self.packed is None:
                stats['fallback_calls']+=1
                return super().evaluate(transform)
            left=normalize_quaternion_xyzw(transform.rotation_quat_xyzw)
            right=quaternion_conjugate_xyzw(left)
            pose=np.asarray((*left,*right,*transform.translation_m),dtype=float)
            point_kernel(self.packed,self.ids,self.cameras,self.settings.use_covariance,pose,self.out)
            stats['compiled_calls']+=1
            if np.any(self.out[2]==2.0):
                stats['fallback_calls']+=1
                return super().evaluate(transform)
            valid=self.out[2]==1.0
            count=int(np.count_nonzero(valid))
            if not count:
                return m.ProbabilisticPoseEvaluation(math.inf,None,None,0,self.total)
            pixels,errors=self.out[0,valid],self.out[1,valid]
            weights=self.weights[valid]
            robust=np.log1p((errors/self.settings.cauchy_scale)**2)
            return m.ProbabilisticPoseEvaluation(float(np.sum(weights*robust)/np.sum(weights)),
                     math.sqrt(float(np.sum(weights*pixels))/float(np.sum(weights))),
                     float(np.sum(weights*errors)/np.sum(weights)),count,self.total-count)
    m._PreparedPoseEvaluation=Prepared
    return stats
