"""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=group.calibrated_points(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
