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

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

#include <vector>
#include <iostream>

namespace clustering {
namespace SoAClustering
{

    //register friendly 3x3 array struct
    template <typename data_type>
    class array3x3
    {
    public:
        data_type a00;
        data_type a01;
        data_type a02;
        data_type a10;
        data_type a11;
        data_type a12;
        data_type a20;
        data_type a21;
        data_type a22;
        
        __device__ array3x3():
        a00(), a01(), a02(), a10(), a11(), a12(), a20(), a21(), a22()
        {}
        __device__ const data_type & get(uint8_t idx) const
        {
            data_type return_val = 0;
            return_val += (idx==0) * a00;
            return_val += (idx==1) * a01;
            return_val += (idx==2) * a02;
            return_val += (idx==3) * a10;
            return_val += (idx==4) * a11;
            return_val += (idx==5) * a12;
            return_val += (idx==6) * a20;
            return_val += (idx==7) * a21;
            return_val += (idx==8) * a22;
            return return_val;
        }
        __device__ void set(uint8_t idx, data_type value)
        {
            a00 = (idx == 0) ? value : a00;
            a01 = (idx == 1) ? value : a01;
            a02 = (idx == 2) ? value : a02;
            a10 = (idx == 3) ? value : a10;
            a11 = (idx == 4) ? value : a11;
            a12 = (idx == 5) ? value : a12;
            a20 = (idx == 6) ? value : a20;
            a21 = (idx == 7) ? value : a21;
            a22 = (idx == 8) ? value : a22;
            /*
            switch(idx)
            {
                case 0:
                    a00 = value;
                    break;
                case 1:
                    a01 = value;
                    break;
                case 2:
                    a02 = value;
                    break;
                case 3:
                    a10 = value;
                    break;
                case 4:
                    a11 = value;
                    break;
                case 5:
                    a12 = value;
                    break;
                case 6:
                    a20 = value;
                    break;
                case 7:
                    a21 = value;
                    break;
                default:
                    a22 = value;
                    break;
            }*/
        }
    };

    template<typename hit_type>
    class dfu_tiled_clusterer
    {
        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;

    public:
        static constexpr uint16_t initial_magic_tile_id = 65535;
    private:
        const uint32_t id;
        const bool processing_safeguard;
        device_data<hit_type> *device;
        uint32_t* neighbor_matrices;//TODO: try to store roots directly in the hit index matrices
        uint16_t last_tile_id = initial_magic_tile_id;
        uint16_t tile_size;
        uint16_t tile_offset;
        array3x3<compact_toa_type> toa_matrix; 
    public:
        
        __device__ dfu_tiled_clusterer(uint32_t tid, device_data<hit_type> *device, uint16_t tile_size = 3, uint16_t tile_offset = 0, bool processing_safeguard = false, uint32_t * shared_mem = nullptr, uint16_t init_tile_id = initial_magic_tile_id);
        __device__ bool process_hit(uint32_t hit_idx, bool atomic_merge);
        __device__ bool process_hit_toa_registers3_idxs_shared(uint32_t hit_idx, bool atomic_merge);
        __device__ uint32_t find_root(uint32_t representative_index, bool force_contract = false, bool atomic_contraction = false);
        
        __device__ uint32_t to_linear_coordinate(int16_t x, int16_t y) const;
        __device__ void load_neighbor_matrix(device_data<hit_type> device);
        __device__ void store_neighbor_matrix(device_data<hit_type> device) const;
        __device__ void clean_hit_index_matrix();
        static __device__ uint16_t identify_tile(int16_t x, int16_t y, uint16_t tile_size, uint16_t tile_offset)
        {
            //Note: x and y both should have at most value of 1024 
            return ( (uint16_t)((x + tile_offset) / tile_size) << 8) | ((y+tile_offset) / tile_size);
        }
        __device__ uint32_t x_tiled_coordinate(int16_t x) const;
        __device__ uint32_t y_tiled_coordinate(int16_t x) const;
        
        
    private:
        __device__ uint16_t identify_tile(int16_t x, int16_t y) const;
        
        __device__ void assign_label(uint32_t hit_idx, uint32_t label, bool atomic, uint32_t* new_representative = nullptr, compact_toa_type* new_representative_toa = nullptr);
        __device__ uint32_t get_label(uint32_t hit_idx);
        __device__ uint32_t to_neighbor_matrix_idx(uint32_t idx) const;
        __device__ void compute_neighbor_ids(int16_t x, int16_t y, uint32_t* neighbor_ids);
        __device__ int32_t compute_neighbor_id(int16_t x, int16_t y, uint8_t idx) const;
        __device__ bool has_nonzero_neighbors( typename hit_type::x_type x, typename hit_type::y_type y, typename hit_type::compact_toa_type toa);
        __device__ array3x3<uint32_t> get_nonzero_neighbors( typename hit_type::x_type x, typename hit_type::y_type y, typename hit_type::compact_toa_type toa, uint8_t * neighbor_count, array3x3<uint8_t> * valid_neighbor_idxs);
        
    };
    }
}