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);
    }

    /**
     * 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)
                int coarse_vx = (int) Math.round((double) vx / velocity_decimate);
                int coarse_vy = (int) Math.round((double) vy / velocity_decimate);
               
                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;
    }
}
