#pragma once
#include "clusterer_utils.h"
//#include "cuda_resources/dfu_clustering.cuh"
#include "data_structs/cuda_hit_buffer.cuh"
#include "data_structs/cuda_constants.h"
#include "cuda_utils.cuh"

#include "cuda_runtime.h"
#include "device_launch_parameters.h"

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/generate.h>
#include <thrust/sort.h>
#include <thrust/copy.h>
#include <thrust/random.h>
#include <thrust/gather.h>
#include <thrust/scatter.h>
#include <functional>
namespace clustering{
    namespace SoAClustering
    {

        //high level API of the clustering, expected to be called from CPU processing nodes/threads
        template<typename hit_type>
        struct gpu_interface
        {
            static void copy_constants_to_device(const cuda_constants  & constants);    

            static void copy_data_to_device(host_data<hit_type>* host_data, device_data<hit_type> * dev);
            static void copy_data_from_device(host_data<hit_type> *host_data, device_data<hit_type> *dev, uint32_t buffer_size = 0, bool copy_raw_data = true);
            static void init_dfu(device_data<hit_type> dev, uint32_t total_thread_count, uint64_t matrix_size, uint32_t init_dfu_threads_per_block);
            static device_data<hit_type> *initialize_device_memory(uint64_t max_hit_count, uint64_t total_thread_count, uint64_t matrix_size, uint32_t init_dfu_threads_per_block, const node_args & args);
            
            static void deallocate_device_memory(device_data<hit_type> *dev);
            static void sort_data(device_data<hit_type>& dev);
            static void sort_by_labels(device_data<hit_type>& dev);
            static void sort_by_root_time(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint32_t padded_buffer_size, uint32_t equality_sort_bit_count = 32);
            static void label_dfu_clusters(device_data<hit_type> dev, uint32_t total_thread_count, uint32_t dfu_threads_per_block, uint32_t max_hit_count, uint32_t shared_memory_bytes);
            
            // Tiled clustering:
            static device_data<hit_type> *initialize_device_memory_tiled(uint64_t max_hit_count, uint64_t total_thread_count, uint64_t matrix_size, uint32_t init_dfu_threads_per_block, uint32_t dfu_threads_per_block, uint32_t hpw, const node_args & args);
            static void sort_tiled(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint16_t tile_size, uint16_t tile_offset, bool preserve_pointers, uint32_t hpw);
            static void sort_tiled_cub(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint16_t tile_size, uint16_t tile_offset, bool preserve_pointers, uint32_t hpw);
            static void reorganize_coalesced(device_data<hit_type>& dev, uint32_t block_size, uint32_t hpw);
            
            static void label_dfu_tiled(device_data<hit_type> dev, uint32_t total_thread_count, uint32_t dfu_threads_per_block, uint32_t max_hit_count, uint32_t shared_memory_bytes, uint8_t skip_config, bool atomic_merge,  uint16_t tile_size, uint16_t tile_offset);
            static void add_padding_coalesced(device_data<hit_type> dev, uint32_t max_buffer_size, uint32_t unlimited_threads_per_block);
            
            static void sort_tiled_cub_2pass(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint16_t tile_size, uint16_t tile_offset, bool preserve_pointers, uint32_t hpw, uint32_t padded_buffer_size);
            static void sort_tiled_cub_impl(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint16_t tile_size, uint16_t tile_offset, bool preserve_pointers, uint32_t hpw, bool two_pass, uint32_t padded_buffer_size);

            // general attribute computation - based on configuration, it calls attribute_computation class methods below
            
            static void compute_cluster_attributes(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block);
            
        };

        template<typename num_type>
            struct InRangeSelect
         {
            float min_val, max_val;
            num_type*values;
            __host__ __device__
            InRangeSelect(float minv, float maxv, num_type*values) : min_val(minv), max_val(maxv), values(values) {}
            __device__ __forceinline__ bool operator()(const uint32_t idx) const {
                num_type val = values[idx]; // Load value into a register
                return val > min_val && val < max_val;
            }
        };
        template<typename hit_type>
        struct attribute_computation
        {

            static cudaError_t compute_energy(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block);
            static cudaError_t compute_energy_2d_map(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block);
            static cudaError_t compute_cluster_size(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block);
            static cudaError_t compute_cluster_energy(device_data<hit_type>& dev);
            static cudaError_t compute_cluster_energy_2d_map(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint32_t* filtered_cluster_count_ptr, uint32_t* indices_to_select, uint32_t* nonzero_pixel_matrix_count_clustered_ptr, uint32_t output_offset);
            static cudaError_t compute_cluster_energy_histogram(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block);
            static cudaError_t compute_energy_weighted_centroid(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block);
            static cudaError_t compute_cluster_start_toa(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block);
            
       
            static void compute_cluster_energy_accepted_indices(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, float min_energy, float max_energy, uint32_t* cluster_indices, uint32_t * selected_cluster_count);
            
        private:
            static void overbin_floats(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, float* coord_centroids, uint32_t* coord_overbinned, uint32_t overbinning_factor);
            static void compute_energy_weighted_averages(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, typename hit_type::x_type* coord_in, float* coord_centroids, uint32_t* indices_to_select);
            static void compute_frame_offsets(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint32_t* indices, uint32_t* filtered_cluster_indices, uint32_t* filtered_cluster_count_ptr);
            static void compute_frame_indices(device_data<hit_type>& dev, uint32_t unlimited_threads_per_block, uint32_t* offsets, uint32_t* frame_indices);
            

        };

        template<typename hit_type>
        struct kernel_implementation
        {
            //implementation of the kernels - hidden in templated class because
            // - global functions should not be hidden within class
            // - if we only used templated global functions (outside class) we would need to explicitly
            // instantiate each method for all data types
            // This way we only instantiate the class once for each type and can also hide sensitive implementation in the .cpp/.cu file 
            using compact_toa_type = typename hit_type::compact_toa_type;
            using x_type = typename hit_type::x_type;
            using y_type = typename hit_type::y_type;
            using toa_type = typename hit_type::toa_type;
            using tot_type = typename hit_type::tot_type;
            
            //Base clustering
            __device__ static void compute_dfu_no_clean(device_data<hit_type> dev);
            __device__ static void connect_dfu_border_path_no_clean(device_data<hit_type> dev);
            __device__ static void clean_dfu_hit_matrices(device_data<hit_type> dev);
            __device__ static void init_dfu_hit_matrices(device_data<hit_type> dev);
            __device__ static void shorten_dfu_paths(device_data<hit_type> dev);
            
            //Tiled clustering
            __device__ static void compute_dfu_tiled(device_data<hit_type> dev, bool atomic_merge, uint16_t tile_size, uint16_t tile_offset);
            __device__ static void connect_dfu_border_path_tiled(device_data<hit_type> dev, bool atomic_merge, uint16_t tile_size, uint16_t tile_offset);
            __device__ static void shorten_dfu_paths_tiled(device_data<hit_type> dev);// possibly we do not need this
            __device__ static void remove_tile_from_prefix(device_data<hit_type> dev, uint32_t hpw);
            __device__ static void add_padding(device_data<hit_type> dev, uint32_t max_hit_count);
            
            __device__ static void add_tile_to_prefix(device_data<hit_type> dev, uint16_t tile_size, uint16_t tile_offset, uint32_t hpw);

            __device__ static void compute_root_toas(device_data<hit_type> dev, typename hit_type::compact_toa_type* root_toas, uint32_t padded_buffer_size);

            __device__ static void compute_tiles(device_data<hit_type> dev, uint16_t tile_size, uint16_t tile_offset, typename  hit_type::tile_id_type * tile_ids, uint32_t padded_buffer_size);

            // attribute computation
            
            __device__ static void compute_cluster_start_indicators(device_data<hit_type> dev, bool* start_bits);
            
            __device__ static void compute_energy(device_data<hit_type> dev, extra_attribute_data_wrapper<hit_type> attr);

        };

        //global functions do not like to be a member function
        //global functions also should not be implemented elsewhere
        //see notes in "kernel imeplmentation" class on why kernels are designed this way 

        // Tiled clustering kernels
        template<typename hit_type>
        __global__ void kernel_add_padding_coalesced(device_data<hit_type> dev, uint32_t max_hit_count)
        {
            kernel_implementation<hit_type>::add_padding(dev, max_hit_count);
        }
        template<typename hit_type>
        __global__ void kernel_remove_tile_from_prefix(device_data<hit_type> dev, uint32_t hpw)
        {
            kernel_implementation<hit_type>::remove_tile_from_prefix(dev, hpw);
        }
        template<typename hit_type>
        __global__ void kernel_add_tile_to_prefix(device_data<hit_type> dev, uint16_t tile_size, uint16_t tile_offset, uint32_t hpw)
        {
            kernel_implementation<hit_type>::add_tile_to_prefix(dev, tile_size, tile_offset, hpw);
        }

        template<typename hit_type>
        __global__ void kernel_compute_dfu_tiled(device_data<hit_type> dev, bool atomic_merge, uint16_t tile_size, uint16_t tile_offset)
        {
            kernel_implementation<hit_type>::compute_dfu_tiled(dev, atomic_merge, tile_size, tile_offset);
        }

        template<typename hit_type>
        __global__ void kernel_connect_dfu_border_path_tiled(device_data<hit_type> dev, bool atomic_merge, uint16_t tile_size, uint16_t tile_offset)
        {
            kernel_implementation<hit_type>::connect_dfu_border_path_tiled(dev, atomic_merge, tile_size, tile_offset);
        }
        
        template<typename hit_type>
        __global__ void kernel_shorten_dfu_paths_tiled(device_data<hit_type> dev)
        {
            kernel_implementation<hit_type>::shorten_dfu_paths_tiled(dev);
        }

        template<typename hit_type>
        __global__ void kernel_compute_tiles(device_data<hit_type> dev, uint16_t tile_size, uint16_t tile_offset, typename  hit_type::tile_id_type * tile_ids, uint32_t padded_buffer_size)
        {
            kernel_implementation<hit_type>::compute_tiles(dev, tile_size, tile_offset, tile_ids, padded_buffer_size);
        }

        template<typename hit_type>
        __global__ void kernel_compute_root_toas(device_data<hit_type> dev, typename hit_type::compact_toa_type* root_toas, uint32_t padded_buffer_size)
        {
            kernel_implementation<hit_type>::compute_root_toas(dev, root_toas, padded_buffer_size);
        }

        // attribute computation kernels

        template<typename hit_type>
        __global__ static void kernel_compute_cluster_start_indicators(device_data<hit_type> dev, bool* start_bits)
        {
            kernel_implementation<hit_type>::compute_cluster_start_indicators(dev, start_bits);
        }
        
        template<typename hit_type>
        __global__ static void kernel_compute_energy(device_data<hit_type> dev, extra_attribute_data_wrapper<hit_type> attr)
        {
            kernel_implementation<hit_type>::compute_energy(dev, attr);
        }

        template<typename hit_type, typename x_type, typename y_type>
        __global__ static void kernel_compute_linear_coord(device_data<hit_type> dev, x_type * xs, y_type* ys, uint32_t* target, uint32_t* valid_size)
        {
            //we implement this kernel in the header due to its similicity and because due to 3
            // template parameters it would be pain to explicitly instantiate
            // Note: xs and ys can have also different type than just hit_type::x_type because of overbinning
            
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
            if (dev.hit_count <= tid)
                return;
            if(tid >= *valid_size)
            {
                target[tid] = max_value_unsigned<uint32_t>;
                return;
            }
            auto x = xs[tid];
            auto y = ys[tid];
            target[tid] = x * static_cast<uint32_t>(hit_type::detector_type::y_size) + y;
        }
        
        // Non tiled kernels

        template<typename hit_type>
        __global__ void kernel_compute_dfu_no_clean(device_data<hit_type> dev)
        {
            kernel_implementation<hit_type>::compute_dfu_no_clean(dev);
        }

        template<typename hit_type>
        __global__ void kernel_connect_dfu_border_path_no_clean(device_data<hit_type> dev)
        {
            kernel_implementation<hit_type>::connect_dfu_border_path_no_clean(dev);
        }

        template<typename hit_type>
        __global__ void kernel_clean_dfu_hit_matrices(device_data<hit_type> dev)
        {
            kernel_implementation<hit_type>::clean_dfu_hit_matrices(dev);
        }

        template<typename hit_type>
        __global__ void kernel_init_dfu_hit_matrices(device_data<hit_type> dev)
        {
            kernel_implementation<hit_type>::init_dfu_hit_matrices(dev);
        }

        template<typename hit_type>
        __global__ void kernel_shorten_dfu_paths(device_data<hit_type> dev)
        {
            kernel_implementation<hit_type>::shorten_dfu_paths(dev);
        }

        // simple arithmetics kernels, independent of hit_type, can be implemented directly, no instantiation needed
        template<typename num_type, typename denom_type, typename result_type, typename size_type>
        __global__ void kernel_divide_elementwise(num_type*num, denom_type* denom, result_type* result, size_type* size_ptr)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
            
            if (*size_ptr <= tid)
                return;    
            
            result[tid] = num[tid] / denom[tid];
        }
        
        template<typename num1_type, typename num2_type, typename result_type>
        __global__ void kernel_multiply_elementwise(num1_type*num1, num2_type* num2, result_type* result, uint32_t* size_ptr)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
        
            if (*size_ptr <= tid)
                return;       
            result[tid] = num1[tid] * num2[tid];
        }

        template<typename num_type, typename const_type, typename result_type, typename size_type>
        __global__ void kernel_multiply_by_constant(num_type* val, const_type constant, result_type* output, size_type* size_ptr)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
        
            if (*size_ptr <= tid)
                return;       
            output[tid] = val[tid] * constant;
        }

        template<typename num_type, typename const_type, typename result_type, typename size_type>
        __global__ void kernel_add_constant(num_type* val, const_type constant, result_type* output, size_type* size_ptr)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
        
            if (*size_ptr <= tid)
                return;       
            output[tid] = val[tid] + constant;
        }

        template<typename value1_type, typename value2_type, typename result_type, typename size_type>
        __global__ void kernel_sum(value1_type* left, value2_type *right, result_type* output, size_type* size_ptr)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
        
            if (*size_ptr <= tid)
                return; 
            output[tid] = left[tid] + right[tid];
        }

        template<typename array_type, typename offset_type, typename diff_type, typename upper_bound_type>
        __global__ void kernel_generate_offseted_sequence(array_type* arr, offset_type offset, diff_type diff, upper_bound_type upper_bound)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
            if (upper_bound <= tid)
                return; 
            uint64_t new_value = offset + tid * diff;
            if(max_value_unsigned<array_type> <= new_value)
                arr[tid] = max_value_unsigned<array_type>;  
            else       
                arr[tid] = new_value; 
        }
        
        template<typename float_type, typename int_type, typename size_type>
        __global__ void kernel_round_to_int(float_type* val_float, int_type* val_uint, size_type* size_ptr)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
            
            if (*size_ptr <= tid)
                return;       
            val_uint[tid] = static_cast<int_type>(roundf(val_float[tid]));
        }

        template<typename val_type, typename pad_type, typename start_pad_idx_type, typename end_pad_idx_type>
        __global__ void kernel_pad_array(val_type* val, pad_type val_pad, start_pad_idx_type* pad_start_ptr, end_pad_idx_type* pad_end_ptr)
        {
            size_t tid = threadIdx.x + blockDim.x * blockIdx.x;
            
            if (*pad_end_ptr <= tid || tid < *pad_start_ptr)
                return;
            val[tid] = val_pad;
        }

        template<typename data_type>
        struct pointers_to_swap {
            data_type** first;
            data_type** second;
        };
        
        void CUDART_CB host_function_trampoline(void* userData);
        template<typename data_type>
        std::function<void()>* create_swap_lambda(data_type** first, data_type** second)
        {
            return new std::function<void()>(
            [first, second]() {
                // This is the logic that will run on the CPU when the callback is triggered.
                std::swap(*first, *second);
            }
            );
        }
    
    }
}
