"""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)
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)
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, first, second, steps, compensated):
    output = np.empty((len(source), len(steps)), dtype=np.float64)
    for row in range(len(source)):
        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 = source[row, 0] + weighted(weights, parameters, row, probe, 0, 3, compensated)
                    y = source[row, 1] + weighted(weights, parameters, row, probe, 1, 3, compensated)
                    z = source[row, 2] + weighted(weights, parameters, row, probe, 2, 3, compensated)
                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 = target[row, 0] + weighted(target_weights, parameters, row, probe, 0, 3, compensated)
                    ty = target[row, 1] + weighted(target_weights, parameters, row, probe, 1, 3, compensated)
                    tz = target[row, 2] + weighted(target_weights, parameters, row, probe, 2, 3, compensated)
                    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,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}
