"""Private ordered finite-difference batch prototype; not benchmark evidence."""
import numpy as np


def weighted(weights, offsets):
    products = weights[None, :, :] * np.asarray(offsets, dtype=np.float64)[:, None, :]
    import sys
    if products.shape[0] * products.shape[1] < 512 or sys.implementation.name != 'cpython' or sys.version_info[:2] not in ((3, 10), (3, 11), (3, 12)):
        rows = products.reshape((-1, weights.shape[1])).tolist()
        return np.asarray([sum(row) for row in rows], dtype=np.float64).reshape(products.shape[:2])
    total = np.zeros(products.shape[:2], dtype=np.float64) + products[..., 0]
    if sys.version_info[:2] < (3, 12):
        for index in range(1, products.shape[-1]):
            total = total + products[..., index]
        return total
    correction = np.zeros_like(total)
    for index in range(1, products.shape[-1]):
        term = products[..., index]
        updated = total + term
        correction += np.where(np.abs(total) >= np.abs(term), (total - updated) + term, (term - updated) + total)
        total = updated
    np.copyto(total, total + correction, where=(correction != 0.0) & np.isfinite(correction))
    return total


def unique_weighted(weights, offsets):
    unique = []
    positions = []
    known = {}
    for row in offsets:
        key = row.tobytes()
        position = known.get(key)
        if position is None:
            position = len(unique)
            known[key] = position
            unique.append(row)
        positions.append(position)
    return weighted(weights, np.asarray(unique))[positions]


def points(group, probes, selected):
    parameters = [tuple(values[name] for name in group.names[2:]) for values in probes]
    source = group.sensor_points[selected]
    weights = group.weights[selected]
    target = None
    if group.kind in ('xyz', 'pair'):
        offsets = np.asarray([parameters[0] for parameters in parameters])
        result = np.stack([
            source[None, :, axis] + unique_weighted(weights, offsets[:, axis::3])
            for axis in range(3)
        ], axis=-1)
        if group.kind == 'pair':
            target = np.stack([
                group.target_points[None, selected, axis]
                + unique_weighted(group.target_weights[selected], offsets[:, axis::3])
                for axis in range(3)
            ], axis=-1)
    elif group.kind == 'depth':
        import math
        depth = np.asarray([
            math.exp(parameter[0][0]) * group.depths[selected] + parameter[0][1]
            for parameter in parameters
        ])
        result = source[None, :, :] * depth[:, :, None]
    else:
        offsets = [parameter[-1] for parameter in parameters]
        nominal = group.depths[None, selected]
        if group.kind == 'centered':
            nominal = nominal + np.asarray([parameter[0][0] for parameter in parameters])[:, None]
        depth = nominal + weighted(weights, offsets)
        result = source[None, :, :] * depth[:, :, None]
    return result, target


def field_columns(plan, values, name, block, options):
    from calibrex.graph.joint_factors import se3_from_tangent
    from calibrex.graph.joint_optimization import _perturb
    from calibrex.graph._point_plane_batches import _transform_points

    steps = [block.finite_difference_steps[index] if block.finite_difference_steps
             else options.default_finite_difference_step for index in range(block.dimension)]
    probes = [_perturb(values, name, index, sign * step)
              for index, step in enumerate(steps) for sign in (1.0, -1.0)]
    residuals = np.zeros((len(probes), plan.size), dtype=np.float64)
    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)
        width = max(1, 4096 // len(probes))
        for start in range(0, len(group.rows), width):
            selected = slice(start, start + width)
            source, target = points(group, probes, selected)
            shape = source.shape
            world = _transform_points(source.reshape((-1, 3)), first).reshape(shape)
            difference = world - (
                _transform_points(target.reshape((-1, 3)), second).reshape(shape)
                if target is not None else group.planes[None, selected, :]
            )
            normal = group.normals[None, selected, :]
            raw = (normal[..., 0] * difference[..., 0]
                   + normal[..., 1] * difference[..., 1]
                   + normal[..., 2] * difference[..., 2])
            residuals[:, group.rows[selected]] = group.sqrt_weights[None, selected] * raw
    for probe_index, probe in enumerate(probes):
        for row, dimension, factor in plan.scalars:
            if name in factor.variable_names:
                result = factor.residuals(probe)
                if len(result) != dimension:
                    raise ValueError('joint factor residual dimension changed during linearization')
                residuals[probe_index, row:row + dimension] = result
    if not np.isfinite(residuals).all():
        raise ValueError('joint factor residuals must be non-empty and finite')
    return ((residuals[0::2] - residuals[1::2]) / (2.0 * np.asarray(steps)[:, None])).T


def linearize(self, values, blocks, layout, options):
    from calibrex.graph.joint_optimization import _perturb

    baseline = self.evaluate(values)
    jacobian = np.zeros((self.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 group.names[2:] for group in self.groups)
            and not any(name in group.names[:2] for group in self.groups)):
            jacobian[:, column] = field_columns(self, values, name, block, options)
            continue
        affected = np.zeros(self.size, dtype=np.bool_)
        for group in self.groups:
            if name in group.names:
                affected[group.rows] = True
        for row, dimension, factor in self.scalars:
            if name in factor.variable_names:
                affected[row:row + dimension] = True
        for index in range(block.dimension):
            step = (block.finite_difference_steps[index] if block.finite_difference_steps
                    else options.default_finite_difference_step)
            plus = self.evaluate(_perturb(values, name, index, step), name)
            minus = self.evaluate(_perturb(values, name, index, -step), name)
            jacobian[affected, column.start + index] = (
                plus[affected] - minus[affected]
            ) / (2.0 * step)
    return baseline, jacobian


def install():
    from calibrex.graph import _rgbd_batches
    _rgbd_batches._RGBDPlan.linearize = linearize
