package com.elphel.imagej.cuas.rt;

public class TemporalKernelGenerator {

    public static class ConsolidationKernel {
        public final int spatialRadius;
        public final int spatialDim; // 2 * spatialRadius + 1
       
        // Maps high-res velocity index -> coarse 9-velocity index (0-8).
        // Contains -1 if the high-res velocity falls outside the target tracking limit.
        public final int[] coarseVIdx;
       
        // The Convolution Tensor
        // Dimensions: [historic_t][v_out_idx][sub_idx][dy][dx]
        // dy and dx are indices 0 to spatialDim-1, representing offsets from -spatialRadius to +spatialRadius
        public final double[][][][][] weights;

        public ConsolidationKernel(int numHistory, int numVOut, int numSub, int sRad) {
            this.spatialRadius = sRad;
            this.spatialDim = 2 * sRad + 1;
            this.coarseVIdx = new int[numVOut];
            this.weights = new double[numHistory][numVOut][numSub][spatialDim][spatialDim];
        }
        public int getSpatialRadius() {
        	return spatialRadius;
        }
        public int getSpatialDim() {
        	return spatialDim;
        }
        public double [][][][][] getKernel(){
        	return weights;
        }
        public int [] getCoarseVIdx() {
        	return coarseVIdx;
        }
    }

    /**
     * 1D Cumulative Distribution Function for a Trapezoid centered at 0.
     */
    private static double trapezoidCDF(double x, double w_top, double w_bot) {
        double x0 = -w_bot / 2.0;
        double x1 = -w_top / 2.0;
        double x2 =  w_top / 2.0;
        double x3 =  w_bot / 2.0;

        if (x <= x0) return 0.0;

        if (Math.abs(w_bot - w_top) < 1e-9) { // t=0 Boxcar
            if (x >= x3) return w_bot;
            return (x - x0);
        }

        double m = 1.0 / (x1 - x0);
        double area_rising = 0.5 * m * (x1 - x0) * (x1 - x0);
        double area_flat = 1.0 * (x2 - x1);

        if (x <= x1) return 0.5 * m * (x - x0) * (x - x0);
        if (x <= x2) return area_rising + 1.0 * (x - x1);
        if (x <= x3) {
            double dx = x - x2;
            return area_rising + area_flat + (1.0 * dx - 0.5 * m * dx * dx);
        }
        return area_rising + area_flat + 0.5 * m * (x3 - x2) * (x3 - x2);
    }

    // Canonical target bump (separable half-cosine, 3-px half-period): the same // By Claude on 06/12/2026
    // shape used by G, CuasRecurrentBuffer.condition() and the synthetic generator.
    private static double halfcos(double d) { // By Claude on 06/12/2026
        double a = Math.abs(d);
        return (a >= 1.5) ? 0.0 : Math.cos(Math.PI / 3.0 * a);
    }

    /**
     * Matched-filter variant: the velocity-channel kernel IS the canonical       // By Claude on 06/12/2026
     * half-cosine target bump swept along the velocity trajectory - the optimal
     * linear detector for a clean bump input, with no optical-blur model (the
     * trapezoid in generateKernel/Direct was an analogy, found inaccurate). LoG/
     * LReLU are deliberately out of the template (LoG on a zero-background bump ~
     * passes it; tune LoG radius separately). Same projection convention as
     * generateKernelDirect (proj = dx + 0.5 + (v/vd)*t) so it plugs into
     * convolve3dDirect unchanged; trapezoid CDF replaced by a bump SAMPLE at the
     * source-pixel offset. Each frame's spatial weights are normalized to the
     * frame's temporal weight, so a constant-amplitude target reads out its
     * amplitude at its velocity channel (energy-preserving across sub-pixel phase).
     */
    public static ConsolidationKernel generateKernelMatched( // By Claude on 06/12/2026
            double[] temporal_weights,
            int pixel_decimate,
            int velocity_decimate,
            int velocity_radius
    ) {
        int num_history = temporal_weights.length;
        int num_subpixels = pixel_decimate * pixel_decimate;
        int v_dim = 2 * velocity_radius + 1;
        int num_velocities = v_dim * v_dim;

        double max_time = num_history > 0 ? num_history - 1 : 0;
        double max_shift = ((double) velocity_radius / velocity_decimate) * max_time;
        int spatial_radius = (int) Math.ceil(max_shift + 1.5); // + bump half-extent

        ConsolidationKernel kernel = new ConsolidationKernel(num_history, num_velocities, num_subpixels, spatial_radius);
        int sdim = 2 * spatial_radius + 1;
        double[][] tmp = new double[sdim][sdim];

        for (int vx = -velocity_radius; vx <= velocity_radius; vx++) {
            for (int vy = -velocity_radius; vy <= velocity_radius; vy++) {
                int v_out_idx = (vy + velocity_radius) * v_dim + (vx + velocity_radius);
                kernel.coarseVIdx[v_out_idx] = 0; // unused in the direct (scalar frames) path

                for (int t = 0; t < num_history; t++) {
                    double time_weight = temporal_weights[t];
                    if (time_weight == 0) continue;

                    for (int sy = 0; sy < pixel_decimate; sy++) {
                        for (int sx = 0; sx < pixel_decimate; sx++) {
                            int sub_idx = sy * pixel_decimate + sx;
                            double sub_cx = (sx + 0.5) / pixel_decimate;
                            double sub_cy = (sy + 0.5) / pixel_decimate;
                            double sum_sp = 0;
                            for (int dy = -spatial_radius; dy <= spatial_radius; dy++) {
                                double dist_y = (dy + 0.5 + ((double) vy / velocity_decimate) * t) - sub_cy;
                                double by = halfcos(dist_y);
                                for (int dx = -spatial_radius; dx <= spatial_radius; dx++) {
                                    double dist_x = (dx + 0.5 + ((double) vx / velocity_decimate) * t) - sub_cx;
                                    double w = by * halfcos(dist_x);
                                    tmp[dy + spatial_radius][dx + spatial_radius] = w;
                                    sum_sp += w;
                                }
                            }
                            if (sum_sp <= 0) continue;
                            double scale = time_weight / sum_sp; // per-frame normalization
                            for (int dy = 0; dy < sdim; dy++) {
                                for (int dx = 0; dx < sdim; dx++) {
                                    double w = tmp[dy][dx] * scale;
                                    if (w > 0) kernel.weights[t][v_out_idx][sub_idx][dy][dx] = w;
                                }
                            }
                        }
                    }
                }
            }
        }
        return kernel;
    }

    /**
     * Direct pixel-domain variant: the kernel convolves scalar (LoG+LReLU) frames // By Claude on 06/11/2026
     * directly, bypassing the 3x3 coarse-velocity stage - a single-stage velocity-
     * matched filter (optimal in the low-SNR linear regime where convolve3D3LReLU
     * only adds noise). Identical trapezoid-CDF projection math; differences from
     * generateKernel(): no coarse 3x3 plane mapping (coarseVIdx unused, set to 0,
     * no tracking-limit exclusion) and t counts FRAMES of the level sequence
     * (pass temporal_weights one entry longer to span the same temporal extent
     * as num_hist 3d3 pair-steps).
     */
    public static ConsolidationKernel generateKernelDirect( // By Claude on 06/11/2026
            double[] temporal_weights,
            int pixel_decimate,
            int velocity_decimate,
            int velocity_radius
    ) {
        int num_history = temporal_weights.length;
        int num_subpixels = pixel_decimate * pixel_decimate;
        int v_dim = 2 * velocity_radius + 1;
        int num_velocities = v_dim * v_dim;

        double max_time = num_history > 0 ? num_history - 1 : 0;
        double max_vel = (double) velocity_radius / velocity_decimate;
        double max_shift = max_vel * max_time;
        double max_width = 1.0 + max_time / velocity_decimate;
        int spatial_radius = (int) Math.ceil(max_shift + max_width / 2.0);

        ConsolidationKernel kernel = new ConsolidationKernel(num_history, num_velocities, num_subpixels, spatial_radius);

        for (int vx = -velocity_radius; vx <= velocity_radius; vx++) {
            for (int vy = -velocity_radius; vy <= velocity_radius; vy++) {

                int v_out_idx = (vy + velocity_radius) * v_dim + (vx + velocity_radius);
                kernel.coarseVIdx[v_out_idx] = 0; // unused in the direct (scalar frames) path

                for (int t = 0; t < num_history; t++) {
                    double time_weight = temporal_weights[t];
                    if (time_weight == 0) continue;

                    double blur_width = (double) t / velocity_decimate;
                    double w_bot = 1.0 + blur_width;
                    double w_top = Math.abs(1.0 - blur_width);
                    double total_area = 0.5 * (w_bot + w_top);

                    for (int dy = -spatial_radius; dy <= spatial_radius; dy++) {
                        for (int dx = -spatial_radius; dx <= spatial_radius; dx++) {

                            double proj_x = dx + 0.5 + ((double) vx / velocity_decimate) * t;
                            double proj_y = dy + 0.5 + ((double) vy / velocity_decimate) * t;

                            for (int sy = 0; sy < pixel_decimate; sy++) {
                                for (int sx = 0; sx < pixel_decimate; sx++) {

                                    int sub_idx = sy * pixel_decimate + sx;

                                    double left_x = (double) sx / pixel_decimate;
                                    double right_x = (double) (sx + 1) / pixel_decimate;
                                    double top_y = (double) sy / pixel_decimate;
                                    double bot_y = (double) (sy + 1) / pixel_decimate;

                                    double weight_x = (trapezoidCDF(right_x - proj_x, w_top, w_bot) -
                                                       trapezoidCDF(left_x - proj_x, w_top, w_bot)) / total_area;

                                    double weight_y = (trapezoidCDF(bot_y - proj_y, w_top, w_bot) -
                                                       trapezoidCDF(top_y - proj_y, w_top, w_bot)) / total_area;

                                    double final_weight = weight_x * weight_y * time_weight;

                                    if (final_weight > 0) {
                                        kernel.weights[t][v_out_idx][sub_idx][dy + spatial_radius][dx + spatial_radius] = final_weight;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        return kernel;
    }

    /**
     * Generates the multi-dimensional consolidation tensor.
     */
    public static ConsolidationKernel generateKernel(
            double[] temporal_weights,
            int pixel_decimate,
            int velocity_decimate,
            int velocity_radius
    ) {
        int num_history = temporal_weights.length;
        int num_subpixels = pixel_decimate * pixel_decimate;
        int v_dim = 2 * velocity_radius + 1;
        int num_velocities = v_dim * v_dim;

        // Calculate the maximum required spatial neighborhood to catch all trapezoid spillover
        double max_time = num_history > 0 ? num_history - 1 : 0;
        double max_vel = (double) velocity_radius / velocity_decimate;
        double max_shift = max_vel * max_time;
        double max_width = 1.0 + max_time / velocity_decimate;
        int spatial_radius = (int) Math.ceil(max_shift + max_width / 2.0);

        ConsolidationKernel kernel = new ConsolidationKernel(num_history, num_velocities, num_subpixels, spatial_radius);

        for (int vx = -velocity_radius; vx <= velocity_radius; vx++) {
            for (int vy = -velocity_radius; vy <= velocity_radius; vy++) {
               
                int v_out_idx = (vy + velocity_radius) * v_dim + (vx + velocity_radius);
               
                // Map to coarse 3x3 velocity grid (-1, 0, 1)
                // round half AWAY FROM ZERO: Java Math.round(-0.5) = 0 while round(0.5) = 1, // By Claude on 06/12/2026
                // which mapped fine velocity +-vel_decimate/2 (e.g. +-2 at VD4) to DIFFERENT
                // coarse planes (+1 vs 0) - an asymmetric response to symmetric targets
                int coarse_vx = (vx >= 0) ? (int) Math.round((double) vx / velocity_decimate) :
                                           -(int) Math.round((double) -vx / velocity_decimate); // By Claude on 06/12/2026
                int coarse_vy = (vy >= 0) ? (int) Math.round((double) vy / velocity_decimate) :
                                           -(int) Math.round((double) -vy / velocity_decimate); // By Claude on 06/12/2026
               
                if (Math.abs(coarse_vx) > 1 || Math.abs(coarse_vy) > 1) {
                    kernel.coarseVIdx[v_out_idx] = -1; // Outside tracking limit
                    continue;
                }
               
                kernel.coarseVIdx[v_out_idx] = (coarse_vy + 1) * 3 + (coarse_vx + 1);

                for (int t = 0; t < num_history; t++) {
                    double time_weight = temporal_weights[t];
                    if (time_weight == 0) continue;

                    double blur_width = (double) t / velocity_decimate;
                    double w_bot = 1.0 + blur_width;
                    double w_top = Math.abs(1.0 - blur_width);
                    double total_area = 0.5 * (w_bot + w_top);

                    // Loop over the spatial neighborhood (source macro-pixels)
                    for (int dy = -spatial_radius; dy <= spatial_radius; dy++) {
                        for (int dx = -spatial_radius; dx <= spatial_radius; dx++) {
                           
                            // Center of projected trapezoid from source pixel (dx, dy)
                            double proj_x = dx + 0.5 + ((double) vx / velocity_decimate) * t;
                            double proj_y = dy + 0.5 + ((double) vy / velocity_decimate) * t;

                            for (int sy = 0; sy < pixel_decimate; sy++) {
                                for (int sx = 0; sx < pixel_decimate; sx++) {
                                   
                                    int sub_idx = sy * pixel_decimate + sx;
                                   
                                    // Target subpixel bounds (relative to target macro-pixel 0,0)
                                    double left_x = (double) sx / pixel_decimate;
                                    double right_x = (double) (sx + 1) / pixel_decimate;
                                    double top_y = (double) sy / pixel_decimate;
                                    double bot_y = (double) (sy + 1) / pixel_decimate;

                                    // Shift so trapezoid is centered at 0 for CDF
                                    double weight_x = (trapezoidCDF(right_x - proj_x, w_top, w_bot) -
                                                       trapezoidCDF(left_x - proj_x, w_top, w_bot)) / total_area;
                                   
                                    double weight_y = (trapezoidCDF(bot_y - proj_y, w_top, w_bot) -
                                                       trapezoidCDF(top_y - proj_y, w_top, w_bot)) / total_area;

                                    double final_weight = weight_x * weight_y * time_weight;

                                    if (final_weight > 0) {
                                        kernel.weights[t][v_out_idx][sub_idx][dy + spatial_radius][dx + spatial_radius] = final_weight;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        return kernel;
    }
}
