"""Private bounded geometry reuse experiment, separate from optimizer logic."""
from __future__ import annotations

from collections import OrderedDict
from dataclasses import replace
import importlib.util
import math
from pathlib import Path

import numpy as np


def install(backend):
    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, hits=0, misses=0, evictions=0, peak_bytes=0,
                 download_bytes=0, cuda_calls=0, fallback_calls=0)
    budget = 128*1024**2
    cache = OrderedDict()
    state = {}

    def capture(artifact):
        if state:
            assert state['artifact'] is artifact
            return
        settings = replace(m.ProbabilisticCameraLidarRefinementOptions(), minimum_confidence=0.0)
        prepared = base(artifact.frames, settings)
        assert not prepared.fallback and len(prepared.frames) == len(artifact.frames)
        state.update(artifact=artifact, frames=prepared.frames, offsets={}, live_bytes=0)
        total = 0
        for original, frame in zip(artifact.frames, prepared.frames, strict=True):
            state['offsets'][original.frame_id] = total
            total += len(frame.points)
        state['total'] = total
        if backend == 'cuda':
            import cupy as cp
            spec = importlib.util.spec_from_file_location('cuda_screen_kernel', Path(__file__).with_name('.tmp-probabilistic-cuda-screen-v1.py'))
            extension = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(extension)
            rows, ids, cameras = [], [], []
            for index, frame in enumerate(prepared.frames):
                intr = frame.intrinsics
                assert not (intr.distortion_model == 'radial-tangential' and len(intr.distortion) == 5 and intr.distortion[4] != 0.0)
                rows.append(np.column_stack((frame.points, frame.means, frame.covariance, frame.weights)))
                ids.append(np.full(len(frame.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')])
            state['device'] = (cp.asarray(np.concatenate(rows)),cp.asarray(np.concatenate(ids)),
                               cp.asarray(np.asarray(cameras,dtype=np.float64)),cp.empty((3,total),dtype=cp.float64))
            state['kernel'] = cp.RawKernel(extension.KERNEL,'evaluate',options=('--std=c++11','--fmad=false'))

    def cpu_points(transform):
        lx,ly,lz,lw = normalize_quaternion_xyzw(transform.rotation_quat_xyzw)
        rx,ry,rz,rw = quaternion_conjugate_xyzw((lx,ly,lz,lw))
        out = np.zeros((3,state['total']),dtype=float)
        offset = 0
        with np.errstate(over='ignore', invalid='ignore', divide='ignore'):
            for frame in state['frames']:
                x,y,z = frame.points.T
                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+transform.translation_m[0]
                yy=aw*ry-ax*rz+ay*rw+az*rx+transform.translation_m[1]
                zz=aw*rz+ax*ry-ay*rx+az*rw+transform.translation_m[2]
                rows=np.flatnonzero(zz>1e-9)
                x,y,z=xx[rows],yy[rows],zz[rows]
                intr=frame.intrinsics
                if intr.projection in {'mei','double_sphere'}:
                    assert intr.xi is not None
                    norm=np.sqrt(x*x+y*y+z*z)
                    if intr.projection=='mei':
                        denominator=z+intr.xi*norm
                    else:
                        assert intr.alpha is not None
                        zeta=intr.xi*norm+z
                        second=np.sqrt(x*x+y*y+zeta*zeta)
                        denominator=intr.alpha*second+(1.0-intr.alpha)*zeta
                    valid=denominator>1e-12
                    rows,x,y,denominator=rows[valid],x[valid],y[valid],denominator[valid]
                    x,y=x/denominator,y/denominator
                else:
                    x,y=x/z,y/z
                if intr.distortion_model=='radial-tangential' and len(intr.distortion) in {4,5}:
                    k1,k2,p1,p2=intr.distortion[:4]
                    k3=intr.distortion[4] if len(intr.distortion)==5 else 0.0
                    r=x*x+y*y
                    square,cube=m._distortion_powers(r)
                    radial=1.0+k1*r+k2*square+k3*cube
                    dx=x*radial+2.0*p1*x*y+p2*(r+2.0*x*x)
                    dy=y*radial+p1*(r+2.0*y*y)+2.0*p2*x*y
                    x,y=dx,dy
                u,v=intr.fx*x+intr.cx,intr.fy*y+intr.cy
                valid=(u>=0.0)&(u<intr.width)&(v>=0.0)&(v<intr.height)
                rows,u,v=rows[valid],u[valid],v[valid]
                du,dv=u-frame.means[rows,0],v-frame.means[rows,1]
                pixel=du*du+dv*dv
                a,b,c,d=frame.covariance[rows].T
                square=(d*du*du-(b+c)*du*dv+a*dv*dv)/(a*d-b*c)
                out[0,offset+rows]=pixel
                out[1,offset+rows]=np.sqrt(np.where(square>0.0,square,0.0))
                out[2,offset+rows]=1.0
                offset+=len(frame.points)
        return out

    def values(transform):
        stats['calls']+=1
        key=np.asarray((*transform.translation_m,*transform.rotation_quat_xyzw),dtype=np.float64).tobytes()
        if key in cache:
            stats['hits']+=1
            cache.move_to_end(key)
            return cache[key]
        stats['misses']+=1
        if backend=='cuda':
            left=normalize_quaternion_xyzw(transform.rotation_quat_xyzw)
            right=quaternion_conjugate_xyzw(left)
            data,ids,cameras,out=state['device']
            n=state['total']
            state['kernel'](((n+255)//256,),(256,),
                            (data,ids,cameras,np.int32(n),np.int32(1),
                             *(np.float64(v) for v in (*left,*right,*transform.translation_m)),out))
            result=out.get()
            stats['cuda_calls']+=1
            stats['download_bytes']+=result.nbytes
            assert not np.any(result[2]==2.0)
        else:
            result=cpu_points(transform)
        if result.nbytes<=budget:
            while state['live_bytes']+result.nbytes>budget:
                _,old=cache.popitem(last=False)
                state['live_bytes']-=old.nbytes
                stats['evictions']+=1
            cache[key]=result
            state['live_bytes']+=result.nbytes
            stats['peak_bytes']=max(stats['peak_bytes'],state['live_bytes'])
        return result

    class Prepared(base):
        def __init__(self,frames,settings):
            super().__init__(frames,settings)
            assert state and not self.fallback
            rows,weights=[],[]
            for frame in frames:
                selected=[(i,m._confidence(item,settings)) for i,item in enumerate(frame.correspondences)
                          if m._confidence(item,settings)>=settings.minimum_confidence]
                rows.extend(state['offsets'][frame.frame_id]+i for i,_ in selected)
                weights.extend(w for _,w in selected)
            self.rows=np.asarray(rows,dtype=np.int64)
            self.weights=np.asarray(weights,dtype=float)

        def evaluate(self,transform):
            result=values(transform)[:,self.rows]
            valid=result[2]==1.0
            count=int(np.count_nonzero(valid))
            if not count:
                return m.ProbabilisticPoseEvaluation(math.inf,None,None,0,self.total)
            pixels=result[0,valid]
            errors=result[1,valid] if self.settings.use_covariance else np.sqrt(np.where(pixels>0.0,pixels,0.0))
            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,capture
