//
// reduction.metal
// Metal kernels for NMF reduction operations
//
// This file contains GPU kernels for parallel reduction operations
// used in Non-negative Matrix Factorization on Apple Silicon.
//

#include <metal_stdlib>
using namespace metal;

/// Parameters for reduction operations
struct ReductionParams {
    uint total_elements; // Total number of elements to reduce
    uint block_size;     // Threads per threadgroup
};

/// Threadgroup size for reductions (must match dispatch)
constant uint REDUCTION_BLOCK_SIZE = 256;

/// First-pass reduction: compute squared values and partial sums
/// Used for Frobenius norm calculation: ||A||_F = sqrt(sum(a_ij^2))
///
/// @param data         [in]  Input data array
/// @param partial_sums [out] Partial sum for each threadgroup
/// @param params       [in]  Reduction parameters
/// @param tid          Thread index within threadgroup
/// @param gid          Global thread index
/// @param threadgroup_id Threadgroup index
/// @param shared       Threadgroup shared memory
kernel void nmf_frobenius_norm_squared(
    constant float* data                        [[buffer(0)]],
    device float* partial_sums                  [[buffer(1)]],
    constant ReductionParams& params            [[buffer(2)]],
    uint tid [[thread_index_in_threadgroup]],
    uint gid [[thread_position_in_grid]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Load data and compute square
    float value = 0.0f;
    if (gid < params.total_elements) {
        value = data[gid];
        value = value * value;
    }
    shared[tid] = value;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel reduction in shared memory
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write result for this threadgroup
    if (tid == 0) {
        partial_sums[threadgroup_id] = shared[0];
    }
}

/// Second-pass reduction: sum partial sums from first pass
/// This kernel is used to finalize multi-block reductions
///
/// @param partial_sums [in/out] Partial sums to reduce
/// @param num_blocks   [in]     Number of partial sums
/// @param tid          Thread index within threadgroup
/// @param shared       Threadgroup shared memory
kernel void nmf_reduce_partial_sums(
    device float* partial_sums                  [[buffer(0)]],
    constant uint& num_blocks                   [[buffer(1)]],
    uint tid [[thread_index_in_threadgroup]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Load partial sums
    float value = 0.0f;
    if (tid < num_blocks) {
        value = partial_sums[tid];
    }
    shared[tid] = value;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s && tid + s < num_blocks) {
            shared[tid] += shared[tid + s];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write final result
    if (tid == 0) {
        partial_sums[0] = shared[0];
    }
}

/// Sum reduction (first pass)
///
/// @param data         [in]  Input data array
/// @param partial_sums [out] Partial sum for each threadgroup
/// @param params       [in]  Reduction parameters
/// @param tid          Thread index within threadgroup
/// @param gid          Global thread index
/// @param threadgroup_id Threadgroup index
/// @param shared       Threadgroup shared memory
kernel void nmf_sum_reduction(
    constant float* data                        [[buffer(0)]],
    device float* partial_sums                  [[buffer(1)]],
    constant ReductionParams& params            [[buffer(2)]],
    uint tid [[thread_index_in_threadgroup]],
    uint gid [[thread_position_in_grid]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Load data
    float value = 0.0f;
    if (gid < params.total_elements) {
        value = data[gid];
    }
    shared[tid] = value;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write result for this threadgroup
    if (tid == 0) {
        partial_sums[threadgroup_id] = shared[0];
    }
}

/// Max reduction (first pass)
///
/// @param data         [in]  Input data array
/// @param partial_max  [out] Partial max for each threadgroup
/// @param params       [in]  Reduction parameters
/// @param tid          Thread index within threadgroup
/// @param gid          Global thread index
/// @param threadgroup_id Threadgroup index
/// @param shared       Threadgroup shared memory
kernel void nmf_max_reduction(
    constant float* data                        [[buffer(0)]],
    device float* partial_max                   [[buffer(1)]],
    constant ReductionParams& params            [[buffer(2)]],
    uint tid [[thread_index_in_threadgroup]],
    uint gid [[thread_position_in_grid]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Load data
    float value = -INFINITY;
    if (gid < params.total_elements) {
        value = data[gid];
    }
    shared[tid] = value;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel max reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] = max(shared[tid], shared[tid + s]);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write result for this threadgroup
    if (tid == 0) {
        partial_max[threadgroup_id] = shared[0];
    }
}

/// Max reduction second pass
///
/// @param partial_max [in/out] Partial max values to reduce
/// @param num_blocks  [in]     Number of partial max values
/// @param tid         Thread index within threadgroup
/// @param shared      Threadgroup shared memory
kernel void nmf_reduce_partial_max(
    device float* partial_max                   [[buffer(0)]],
    constant uint& num_blocks                   [[buffer(1)]],
    uint tid [[thread_index_in_threadgroup]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Load partial max values
    float value = -INFINITY;
    if (tid < num_blocks) {
        value = partial_max[tid];
    }
    shared[tid] = value;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel max reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s && tid + s < num_blocks) {
            shared[tid] = max(shared[tid], shared[tid + s]);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write final result
    if (tid == 0) {
        partial_max[0] = shared[0];
    }
}

/// Min reduction (first pass)
///
/// @param data         [in]  Input data array
/// @param partial_min  [out] Partial min for each threadgroup
/// @param params       [in]  Reduction parameters
/// @param tid          Thread index within threadgroup
/// @param gid          Global thread index
/// @param threadgroup_id Threadgroup index
/// @param shared       Threadgroup shared memory
kernel void nmf_min_reduction(
    constant float* data                        [[buffer(0)]],
    device float* partial_min                   [[buffer(1)]],
    constant ReductionParams& params            [[buffer(2)]],
    uint tid [[thread_index_in_threadgroup]],
    uint gid [[thread_position_in_grid]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Load data
    float value = INFINITY;
    if (gid < params.total_elements) {
        value = data[gid];
    }
    shared[tid] = value;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel min reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] = min(shared[tid], shared[tid + s]);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write result for this threadgroup
    if (tid == 0) {
        partial_min[threadgroup_id] = shared[0];
    }
}

/// Convergence check: compute max absolute difference
/// Used to check NMF convergence: max|H_new - H_old| < tolerance
///
/// @param new_data         [in]  New matrix values
/// @param old_data         [in]  Old matrix values
/// @param partial_max_diff [out] Partial max differences
/// @param params           [in]  Reduction parameters
/// @param tid              Thread index within threadgroup
/// @param gid              Global thread index
/// @param threadgroup_id   Threadgroup index
/// @param shared           Threadgroup shared memory
kernel void nmf_max_abs_diff(
    constant float* new_data                    [[buffer(0)]],
    constant float* old_data                    [[buffer(1)]],
    device float* partial_max_diff              [[buffer(2)]],
    constant ReductionParams& params            [[buffer(3)]],
    uint tid [[thread_index_in_threadgroup]],
    uint gid [[thread_position_in_grid]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Compute absolute difference
    float diff = 0.0f;
    if (gid < params.total_elements) {
        diff = abs(new_data[gid] - old_data[gid]);
    }
    shared[tid] = diff;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel max reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] = max(shared[tid], shared[tid + s]);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write result for this threadgroup
    if (tid == 0) {
        partial_max_diff[threadgroup_id] = shared[0];
    }
}

/// Reconstruction error: compute ||V - WH||_F^2
/// Used to measure NMF quality: smaller is better
///
/// @param v_data        [in]  Original matrix V
/// @param wh_data       [in]  Reconstructed matrix W*H
/// @param partial_error [out] Partial squared errors
/// @param params        [in]  Reduction parameters
/// @param tid           Thread index within threadgroup
/// @param gid           Global thread index
/// @param threadgroup_id Threadgroup index
/// @param shared        Threadgroup shared memory
kernel void nmf_reconstruction_error(
    constant float* v_data                      [[buffer(0)]],
    constant float* wh_data                     [[buffer(1)]],
    device float* partial_error                 [[buffer(2)]],
    constant ReductionParams& params            [[buffer(3)]],
    uint tid [[thread_index_in_threadgroup]],
    uint gid [[thread_position_in_grid]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Compute squared difference
    float diff = 0.0f;
    if (gid < params.total_elements) {
        diff = v_data[gid] - wh_data[gid];
        diff = diff * diff;
    }
    shared[tid] = diff;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write result for this threadgroup
    if (tid == 0) {
        partial_error[threadgroup_id] = shared[0];
    }
}

/// KL divergence: compute D_KL(V || WH)
/// Alternative measure for NMF quality
/// D_KL = sum(V * log(V/WH) - V + WH)
///
/// @param v_data        [in]  Original matrix V
/// @param wh_data       [in]  Reconstructed matrix W*H
/// @param partial_div   [out] Partial divergence values
/// @param epsilon       [in]  Small value for numerical stability
/// @param params        [in]  Reduction parameters
/// @param tid           Thread index within threadgroup
/// @param gid           Global thread index
/// @param threadgroup_id Threadgroup index
/// @param shared        Threadgroup shared memory
kernel void nmf_kl_divergence(
    constant float* v_data                      [[buffer(0)]],
    constant float* wh_data                     [[buffer(1)]],
    device float* partial_div                   [[buffer(2)]],
    constant float& epsilon                     [[buffer(3)]],
    constant ReductionParams& params            [[buffer(4)]],
    uint tid [[thread_index_in_threadgroup]],
    uint gid [[thread_position_in_grid]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    threadgroup float* shared [[threadgroup(0)]]
) {
    // Compute KL divergence term
    float div = 0.0f;
    if (gid < params.total_elements) {
        float v = v_data[gid] + epsilon;
        float wh = wh_data[gid] + epsilon;
        div = v * log(v / wh) - v + wh;
    }
    shared[tid] = div;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Parallel reduction
    for (uint s = REDUCTION_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Write result for this threadgroup
    if (tid == 0) {
        partial_div[threadgroup_id] = shared[0];
    }
}