"""Exploratory compiled native field probes; not a production backend."""
import math
import sys
import numpy as np
from numba import njit

@njit(cache=False, nogil=True, fastmath=False, inline="always")
def weighted(weights, offsets, row, probe, start, stride, compensated):
    total = 0.0 + weights[row, 0] * offsets[probe, start]
    correction = 0.0
    for index in range(1, weights.shape[1]):
        term = weights[row, index] * offsets[probe, start + stride * index]
        updated = total + term
        if compensated:
            if abs(total) >= abs(term):
                correction += (total - updated) + term
            else:
                correction += (term - updated) + total
        total = updated
    if compensated and correction != 0.0 and math.isfinite(correction):
        total += correction
    return total

@njit(cache=False, nogil=True, fastmath=False, inline="always")
def transformed(x, y, z, frame):
    qx, qy, qz, qw = frame[3], frame[4], frame[5], frame[6]
    cx, cy, cz, cw = frame[7], frame[8], frame[9], frame[10]
    ax = qw*x + qx*0.0 + qy*z - qz*y
    ay = qw*y - qx*z + qy*0.0 + qz*x
    az = qw*z + qx*y - qy*x + qz*0.0
    aw = qw*0.0 - qx*x - qy*y - qz*z
    return (aw*cx + ax*cw + ay*cz - az*cy + frame[0],
            aw*cy - ax*cz + ay*cw + az*cx + frame[1],
            aw*cz + ax*cy - ay*cx + az*cw + frame[2])

@njit(cache=False, nogil=True, fastmath=False)
def kernel(kind, source, depths, weights, target, target_weights, planes, normals,
           sqrt_weights, parameters, baseline, first, second, steps, compensated):
    output = np.empty((len(source), len(steps)), dtype=np.float64)
    for row in range(len(source)):
        bx = by = bz = btx = bty = btz = 0.0
        if kind == 0 or kind == 1:
            bx = source[row, 0] + weighted(weights, baseline, row, 0, 0, 3, compensated)
            by = source[row, 1] + weighted(weights, baseline, row, 0, 1, 3, compensated)
            bz = source[row, 2] + weighted(weights, baseline, row, 0, 2, 3, compensated)
            if kind == 1:
                btx = target[row, 0] + weighted(target_weights, baseline, row, 0, 0, 3, compensated)
                bty = target[row, 1] + weighted(target_weights, baseline, row, 0, 1, 3, compensated)
                btz = target[row, 2] + weighted(target_weights, baseline, row, 0, 2, 3, compensated)
        for column in range(len(steps)):
            positive = 0.0
            for sign in range(2):
                probe = 2*column + sign
                if kind == 0 or kind == 1:
                    x, y, z = bx, by, bz
                    axis = column % 3
                    value = source[row, axis] + weighted(weights, parameters, row, probe, axis, 3, compensated)
                    if axis == 0: x = value
                    elif axis == 1: y = value
                    else: z = value
                else:
                    if kind == 2:
                        depth = parameters[probe, 0]*depths[row] + parameters[probe, 1]
                    else:
                        nominal = depths[row]
                        if kind == 4:
                            nominal = nominal + parameters[probe, 0]
                        depth = nominal + weighted(weights, parameters, row, probe, 1 if kind == 4 else 0, 1, compensated)
                    x, y, z = source[row, 0]*depth, source[row, 1]*depth, source[row, 2]*depth
                a, b, c = transformed(x,y,z,first)
                if kind == 1:
                    tx, ty, tz = btx, bty, btz
                    axis = column % 3
                    value = target[row, axis] + weighted(target_weights, parameters, row, probe, axis, 3, compensated)
                    if axis == 0: tx = value
                    elif axis == 1: ty = value
                    else: tz = value
                    tx, ty, tz = transformed(tx,ty,tz,second)
                else:
                    tx, ty, tz = planes[row, 0], planes[row, 1], planes[row, 2]
                dx,dy,dz = a-tx,b-ty,c-tz
                residual = sqrt_weights[row]*(normals[row,0]*dx + normals[row,1]*dy + normals[row,2]*dz)
                if not math.isfinite(residual):
                    raise ValueError("joint factor residuals must be non-empty and finite")
                if sign == 0:
                    positive = residual
                else:
                    output[row,column] = (positive-residual)/(2.0*steps[column])
    return output

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_numba_fields':True,'cache':False,'fastmath':False}


_install_fields = install

@njit(cache=False, nogil=True, fastmath=False)
def pose_kernel(pair, points, target, planes, normals, sqrt_weights, first, second, steps):
    output=np.empty((len(points),len(steps)),dtype=np.float64)
    for row in range(len(points)):
        for column in range(len(steps)):
            positive=0.0
            for sign in range(2):
                probe=2*column+sign
                a,b,c=transformed(points[row,0],points[row,1],points[row,2],first[probe])
                if pair:
                    tx,ty,tz=transformed(target[row,0],target[row,1],target[row,2],second[probe])
                else:
                    tx,ty,tz=planes[row,0],planes[row,1],planes[row,2]
                residual=sqrt_weights[row]*(normals[row,0]*(a-tx)+normals[row,1]*(b-ty)+normals[row,2]*(c-tz))
                if not math.isfinite(residual):
                    raise ValueError('joint factor residuals must be non-empty and finite')
                if sign==0: positive=residual
                else: output[row,column]=(positive-residual)/(2.0*steps[column])
    return output

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
