"""Exploratory strict-FP64 CUDA kernels for native RGB-D differences."""
import math
import sys
import numpy as np
import cupy as cp

CUDA = r'''
__device__ double weighted(const double* weights,const double* offsets,int row,int probe,
                          int controls,int width,int start,int stride,int compensated) {
    double total=0.0 + weights[row*controls]*offsets[probe*width+start];
    double correction=0.0;
    for(int j=1;j<controls;++j) {
        double term=weights[row*controls+j]*offsets[probe*width+start+stride*j];
        double updated=total+term;
        if(compensated) {
            if(fabs(total)>=fabs(term)) correction+=(total-updated)+term;
            else correction+=(term-updated)+total;
        }
        total=updated;
    }
    if(compensated && correction!=0.0 && isfinite(correction)) total+=correction;
    return total;
}
__device__ void transformed(double x,double y,double z,const double* f,double* out) {
    double qx=f[3],qy=f[4],qz=f[5],qw=f[6],cx=f[7],cy=f[8],cz=f[9],cw=f[10];
    double ax=qw*x+qx*0.0+qy*z-qz*y;
    double ay=qw*y-qx*z+qy*0.0+qz*x;
    double az=qw*z+qx*y-qy*x+qz*0.0;
    double aw=qw*0.0-qx*x-qy*y-qz*z;
    out[0]=aw*cx+ax*cw+ay*cz-az*cy+f[0];
    out[1]=aw*cy-ax*cz+ay*cw+az*cx+f[1];
    out[2]=aw*cz+ax*cy-ay*cx+az*cw+f[2];
}
extern "C" __global__ void fields(int kind,int n,int d,int controls,int width,
 const double* source,const double* depths,const double* weights,const double* target,
 const double* target_weights,const double* planes,const double* normals,const double* sqrt_weights,
 const double* parameters,const double* first,const double* second,const double* steps,
 int compensated,double* output) {
    int index=blockDim.x*blockIdx.x+threadIdx.x;
    if(index>=n*d) return;
    int row=index/d,column=index%d;
    double positive=0.0;
    for(int sign=0;sign<2;++sign) {
        int probe=2*column+sign;
        double x,y,z;
        if(kind==0 || kind==1) {
            x=source[3*row]+weighted(weights,parameters,row,probe,controls,width,0,3,compensated);
            y=source[3*row+1]+weighted(weights,parameters,row,probe,controls,width,1,3,compensated);
            z=source[3*row+2]+weighted(weights,parameters,row,probe,controls,width,2,3,compensated);
        } else {
            double depth;
            if(kind==2) depth=parameters[probe*width]*depths[row]+parameters[probe*width+1];
            else {
                double nominal=depths[row];
                if(kind==4) nominal=nominal+parameters[probe*width];
                depth=nominal+weighted(weights,parameters,row,probe,controls,width,kind==4?1:0,1,compensated);
            }
            x=source[3*row]*depth;y=source[3*row+1]*depth;z=source[3*row+2]*depth;
        }
        double a[3],b[3];transformed(x,y,z,first,a);
        if(kind==1) {
            x=target[3*row]+weighted(target_weights,parameters,row,probe,controls,width,0,3,compensated);
            y=target[3*row+1]+weighted(target_weights,parameters,row,probe,controls,width,1,3,compensated);
            z=target[3*row+2]+weighted(target_weights,parameters,row,probe,controls,width,2,3,compensated);
            transformed(x,y,z,second,b);
        } else {b[0]=planes[3*row];b[1]=planes[3*row+1];b[2]=planes[3*row+2];}
        double dx=a[0]-b[0],dy=a[1]-b[1],dz=a[2]-b[2];
        double residual=sqrt_weights[row]*(normals[3*row]*dx+normals[3*row+1]*dy+normals[3*row+2]*dz);
        if(sign==0) positive=residual;
        else output[index]=(positive-residual)/(2.0*steps[column]);
    }
}
extern "C" __global__ void poses(int pair,int n,int d,const double* points,
 const double* target,const double* planes,const double* normals,const double* sqrt_weights,
 const double* first,const double* second,const double* steps,double* output) {
    int index=blockDim.x*blockIdx.x+threadIdx.x;
    if(index>=n*d) return;
    int row=index/d,column=index%d;
    double positive=0.0;
    for(int sign=0;sign<2;++sign) {
        int probe=2*column+sign;
        double a[3],b[3];transformed(points[3*row],points[3*row+1],points[3*row+2],first+11*probe,a);
        if(pair) transformed(target[3*row],target[3*row+1],target[3*row+2],second+11*probe,b);
        else {b[0]=planes[3*row];b[1]=planes[3*row+1];b[2]=planes[3*row+2];}
        double dx=a[0]-b[0],dy=a[1]-b[1],dz=a[2]-b[2];
        double residual=sqrt_weights[row]*(normals[3*row]*dx+normals[3*row+1]*dy+normals[3*row+2]*dz);
        if(sign==0) positive=residual;
        else output[index]=(positive-residual)/(2.0*steps[column]);
    }
}
'''
_field=cp.RawKernel(CUDA,'fields',options=('--fmad=false',))
_pose=cp.RawKernel(CUDA,'poses',options=('--fmad=false',))
_cache={}

def uploaded(array):
    found=_cache.get(id(array))
    if found is not None:
        assert found[0] is array
        return found[1]
    result=cp.asarray(np.ascontiguousarray(array))
    _cache[id(array)]=(array,result)
    return result

def kernel(kind,source,depths,weights,target,target_weights,planes,normals,sqrt_weights,
           parameters,baseline,first,second,steps,compensated):
    n,d=len(source),len(steps)
    output=cp.empty((n,d),dtype=cp.float64)
    arrays=[uploaded(a) for a in [source,depths,weights,target,target_weights,planes,normals,sqrt_weights]]
    changing=[cp.asarray(a) for a in [parameters,first,second,steps]]
    _field(((n*d+255)//256,),(256,), (np.int32(kind),np.int32(n),np.int32(d),np.int32(weights.shape[1]),np.int32(parameters.shape[1]),*arrays,*changing,np.int32(compensated),output))
    result=output.get()
    assert np.isfinite(result).all()
    return result

def pose_kernel(pair,points,target,planes,normals,sqrt_weights,first,second,steps):
    n,d=len(points),len(steps)
    output=cp.empty((n,d),dtype=cp.float64)
    arrays=[uploaded(a) for a in [points,target,planes,normals,sqrt_weights]]
    changing=[cp.asarray(a) for a in [first,second,steps]]
    _pose(((n*d+255)//256,),(256,), (np.int32(pair),np.int32(n),np.int32(d),*arrays,*changing,output))
    result=output.get()
    assert np.isfinite(result).all()
    return result

def install(variant):
    if variant != 'fused':
        return {}
    from calibrex.graph import _rgbd_field_columns as fields
    from calibrex.graph.joint_factors import se3_from_tangent
    from calibrex.graph.joint_optimization import _perturb
    from calibrex.core.geometry import normalize_quaternion_xyzw, quaternion_conjugate_xyzw
    def frame(transform):
        q=normalize_quaternion_xyzw(transform.rotation_quat_xyzw)
        return np.asarray((*transform.translation_m,*q,*quaternion_conjugate_xyzw(q)),dtype=np.float64)
    def compiled(plan, values, name, block, options, output):
        steps=np.asarray(block.finite_difference_steps or (options.default_finite_difference_step,)*block.dimension,dtype=np.float64)
        probes=[_perturb(values,name,index,sign*float(step)) for index,step in enumerate(steps) for sign in (1.0,-1.0)]
        for group in plan.groups:
            if name not in group.names:
                continue
            first=se3_from_tangent(values[group.names[0]]).compose(group.world)
            second=se3_from_tangent(values[group.names[1]]).compose(group.sensor)
            if group.kind!='pair':
                first=first.compose(second)
            kind={'xyz':0,'pair':1,'depth':2,'trilinear':3,'centered':4}[group.kind]
            if kind==2:
                parameters=np.asarray([(math.exp(probe[group.names[2]][0]),probe[group.names[2]][1]) for probe in probes],dtype=np.float64)
            elif kind==4:
                parameters=np.asarray([(probe[group.names[2]][0],*probe[group.names[3]]) for probe in probes],dtype=np.float64)
            else:
                parameters=np.asarray([probe[group.names[2]] for probe in probes],dtype=np.float64)
            output[group.rows]=kernel(kind,group.sensor_points,group.depths,group.weights,group.target_points,group.target_weights,group.planes,group.normals,group.sqrt_weights,parameters,np.asarray([values[group.names[2]]],dtype=np.float64),frame(first),frame(second),steps,sys.version_info[:2]==(3,12))
        for row,dimension,factor in plan.scalars:
            if name not in factor.variable_names:
                continue
            for index,step in enumerate(steps):
                plus=np.asarray(factor.residuals(probes[2*index]),dtype=np.float64)
                minus=np.asarray(factor.residuals(probes[2*index+1]),dtype=np.float64)
                if len(plus)!=dimension or len(minus)!=dimension:
                    raise ValueError('joint factor residual dimension changed during linearization')
                output[row:row+dimension,index]=(plus-minus)/(2.0*step)
    fields._field_columns=compiled
    return {'fused_cuda_fields':True,'fmad':False,'fp64':True}

_install_fields = install
def install(variant):
    changes=_install_fields(variant)
    if variant!='fused': return changes
    from calibrex.graph import _rgbd_batches as batches
    from calibrex.graph import _rgbd_field_columns as fields
    from calibrex.graph.joint_factors import se3_from_tangent
    from calibrex.graph.joint_optimization import _perturb
    from calibrex.core.geometry import normalize_quaternion_xyzw,quaternion_conjugate_xyzw
    def frame(t):
        q=normalize_quaternion_xyzw(t.rotation_quat_xyzw)
        return np.asarray((*t.translation_m,*q,*quaternion_conjugate_xyzw(q)),dtype=np.float64)
    def linearize(plan,values,blocks,layout,options):
        baseline=plan.evaluate(values)
        jacobian=np.zeros((plan.size,sum(s.stop-s.start for s in layout.values())))
        by_name={block.name:block for block in blocks}
        for name,column in layout.items():
            block=by_name[name]
            if (block.dimension>=8 and any(name in g.names[2:] for g in plan.groups)
                    and not any(name in g.names[:2] for g in plan.groups)):
                fields._field_columns(plan,values,name,block,options,jacobian[:,column])
                continue
            steps=np.asarray(block.finite_difference_steps or (options.default_finite_difference_step,)*block.dimension,dtype=np.float64)
            probes=[_perturb(values,name,index,sign*float(step)) for index,step in enumerate(steps) for sign in (1.0,-1.0)]
            if (any(name in g.names[:2] for g in plan.groups)
                    and not any(name in g.names[2:] for g in plan.groups)):
                for group in plan.groups:
                    if name not in group.names: continue
                    points,target=gpu_calibrated(group,values)
                    transforms=[]
                    for probe in probes:
                        a=se3_from_tangent(probe[group.names[0]]).compose(group.world)
                        b=se3_from_tangent(probe[group.names[1]]).compose(group.sensor)
                        if group.kind!='pair': a=a.compose(b)
                        transforms.append((frame(a),frame(b)))
                    first=np.asarray([a for a,b in transforms],dtype=np.float64)
                    second=np.asarray([b for a,b in transforms],dtype=np.float64)
                    jacobian[group.rows,column]=pose_kernel(group.kind=='pair',points,target if target is not None else group.target_points,group.planes,group.normals,group.sqrt_weights,first,second,steps)
                for row,dimension,factor in plan.scalars:
                    if name not in factor.variable_names: continue
                    for index,step in enumerate(steps):
                        plus=np.asarray(factor.residuals(probes[2*index]),dtype=np.float64)
                        minus=np.asarray(factor.residuals(probes[2*index+1]),dtype=np.float64)
                        if len(plus)!=dimension or len(minus)!=dimension:
                            raise ValueError('joint factor residual dimension changed during linearization')
                        jacobian[row:row+dimension,column.start+index]=(plus-minus)/(2.0*step)
                continue
            affected=np.zeros(plan.size,dtype=np.bool_)
            for group in plan.groups:
                if name in group.names: affected[group.rows]=True
            for row,dimension,factor in plan.scalars:
                if name in factor.variable_names: affected[row:row+dimension]=True
            for index,step in enumerate(steps):
                plus=plan.evaluate(probes[2*index],name)
                minus=plan.evaluate(probes[2*index+1],name)
                jacobian[affected,column.start+index]=(plus[affected]-minus[affected])/(2.0*step)
        return baseline,jacobian
    batches._RGBDPlan.linearize=linearize
    changes['fused_pose_columns']=True
    return changes


CUDA += r'''
extern "C" __global__ void calibrate(int kind,int n,int controls,int width,const double* source,
 const double* depths,const double* weights,const double* parameters,int compensated,double* output) {
    int row=blockDim.x*blockIdx.x+threadIdx.x;
    if(row>=n) return;
    if(kind==0 || kind==1) {
        for(int axis=0;axis<3;++axis)
            output[3*row+axis]=source[3*row+axis]+weighted(weights,parameters,row,0,controls,width,axis,3,compensated);
    } else {
        double depth;
        if(kind==2) depth=parameters[0]*depths[row]+parameters[1];
        else {
            double nominal=depths[row];
            if(kind==4) nominal=nominal+parameters[0];
            depth=nominal+weighted(weights,parameters,row,0,controls,width,kind==4?1:0,1,compensated);
        }
        for(int axis=0;axis<3;++axis) output[3*row+axis]=source[3*row+axis]*depth;
    }
}
extern "C" __global__ void evaluate(int pair,int n,const double* points,const double* target,
 const double* planes,const double* normals,const double* sqrt_weights,const double* first,
 const double* second,double* output) {
    int row=blockDim.x*blockIdx.x+threadIdx.x;
    if(row>=n) return;
    double a[3],b[3];transformed(points[3*row],points[3*row+1],points[3*row+2],first,a);
    if(pair) transformed(target[3*row],target[3*row+1],target[3*row+2],second,b);
    else {b[0]=planes[3*row];b[1]=planes[3*row+1];b[2]=planes[3*row+2];}
    double dx=a[0]-b[0],dy=a[1]-b[1],dz=a[2]-b[2];
    output[row]=sqrt_weights[row]*(normals[3*row]*dx+normals[3*row+1]*dy+normals[3*row+2]*dz);
}
'''
_field=cp.RawKernel(CUDA,'fields',options=('--fmad=false',))
_pose=cp.RawKernel(CUDA,'poses',options=('--fmad=false',))
_calibrate=cp.RawKernel(CUDA,'calibrate',options=('--fmad=false',))
_evaluate=cp.RawKernel(CUDA,'evaluate',options=('--fmad=false',))
_upload_host=uploaded
def uploaded(array):
    return array if isinstance(array,cp.ndarray) else _upload_host(array)

_calibrated={}
def gpu_calibrated(group,values):
    from collections import OrderedDict
    kind={'xyz':0,'pair':1,'depth':2,'trilinear':3,'centered':4}[group.kind]
    parameters=tuple(values[name] for name in group.names[2:])
    key=tuple(id(p) for p in parameters)
    entry=_calibrated.setdefault(id(group),(group,OrderedDict()))
    assert entry[0] is group
    cache=entry[1]
    if key in cache:
        value=cache[key]
        assert all(a is b for a,b in zip(parameters,value[0],strict=True))
        cache.move_to_end(key)
        return value[1],value[2]
    if kind==2:
        array=np.asarray([(math.exp(parameters[0][0]),parameters[0][1])],dtype=np.float64)
    elif kind==4:
        array=np.asarray([(parameters[0][0],*parameters[1])],dtype=np.float64)
    else:
        array=np.asarray([parameters[0]],dtype=np.float64)
    packed=cp.asarray(array)
    n=len(group.rows)
    def one(points,weights):
        result=cp.empty((n,3),dtype=cp.float64)
        _calibrate(((n+255)//256,),(256,), (np.int32(kind),np.int32(n),np.int32(weights.shape[1]),np.int32(array.shape[1]),uploaded(points),uploaded(group.depths),uploaded(weights),packed,np.int32(sys.version_info[:2]==(3,12)),result))
        return result
    points=one(group.sensor_points,group.weights)
    target=one(group.target_points,group.target_weights) if kind==1 else None
    cache[key]=(parameters,points,target)
    if len(cache)>2: cache.popitem(last=False)
    return points,target

_install_columns=install
def install(variant):
    changes=_install_columns(variant)
    if variant!='fused': return changes
    from calibrex.graph import _rgbd_batches as batches
    from calibrex.graph import joint_optimization as graph
    from calibrex.graph.joint_factors import se3_from_tangent
    from calibrex.core.geometry import normalize_quaternion_xyzw,quaternion_conjugate_xyzw
    from calibrex.graph._point_plane_batches import _OwnedPointFactors
    def frame(t):
        q=normalize_quaternion_xyzw(t.rotation_quat_xyzw)
        return np.asarray((*t.translation_m,*q,*quaternion_conjugate_xyzw(q)),dtype=np.float64)
    def evaluate(plan,values,active=None):
        result=np.zeros(plan.size,dtype=np.float64)
        for group in plan.groups:
            if active is not None and active not in group.names: continue
            points,target=gpu_calibrated(group,values)
            first=se3_from_tangent(values[group.names[0]]).compose(group.world)
            second=se3_from_tangent(values[group.names[1]]).compose(group.sensor)
            if group.kind!='pair': first=first.compose(second)
            n=len(points)
            output=cp.empty(n,dtype=cp.float64)
            _evaluate(((n+255)//256,),(256,), (np.int32(group.kind=='pair'),np.int32(n),points,target if target is not None else uploaded(group.target_points),uploaded(group.planes),uploaded(group.normals),uploaded(group.sqrt_weights),cp.asarray(frame(first)),cp.asarray(frame(second)),output))
            result[group.rows]=output.get()
        for row,dimension,factor in plan.scalars:
            if active is None or active in factor.variable_names:
                residual=factor.residuals(values)
                if len(residual)!=dimension: raise ValueError('joint factor residual dimension changed during linearization')
                result[row:row+dimension]=residual
        if not np.isfinite(result).all(): raise ValueError('joint factor residuals must be non-empty and finite')
        return result
    batches._RGBDPlan.evaluate=evaluate
    original=graph._factor_family_diagnostics
    class Factor:
        def __init__(self,factor,dimension): self.factor=factor;self.dimension=dimension
        def __getattr__(self,name): return getattr(self.factor,name)
        def residuals(self,values): return (0.0,)*self.dimension
    def diagnostics(factors,values,residual,jacobian,options):
        plan=batches._prepare_rgbd_plan(factors,values) if type(factors) is _OwnedPointFactors else None
        if plan is None: return original(factors,values,residual,jacobian,options)
        scalar={row:dimension for row,dimension,_ in plan.scalars}
        row,proxies=0,[]
        for factor in factors:
            dimension=scalar.get(row,1);proxies.append(Factor(factor,dimension));row+=dimension
        assert row==plan.size
        return original(proxies,values,residual,jacobian,options)
    graph._factor_family_diagnostics=diagnostics
    changes.update(resident_calibrated_points=True,native_diagnostic_dimensions=True)
    return changes
