#pragma once
#include "node_args.h"
#include "cuda_resources/cuda_utils.cuh"
#include "data_structs/calibration.h"

#include <iostream>

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "devices/current_device.h"
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/sequence.h>
#include <cub/cub.cuh>
#include <vector>
#include <set>
#include <algorithm>
#include <functional>

namespace clustering
{
  /*extern __constant__ uint64_t TOA_NS_MULTIPLIER;
  extern __constant__ uint64_t MAX_CLUSTER_JOIN_TIME;
  extern __constant__ uint64_t MAX_HIT_COUNT;
  extern __constant__ uint32_t APPROX_HITS_PER_WORKER;
  extern __constant__ uint32_t MATRIX_SIZE;
  extern __constant__ uint32_t MATRIX_HEIGHT;
  extern __constant__ uint32_t MATRIX_WIDTH;
  extern __constant__ uint32_t TOTAL_THREAD_COUNT;
  extern __constant__ uint32_t PADDED_MATRIX_WIDTH;
  extern __constant__ uint16_t DFU_THREADS_PER_BLOCK;
  extern __constant__ uint16_t INIT_DFU_THREADS_PER_BLOCK;
  extern __constant__ uint32_t DFU_BLOCK_COUNT;
  extern __constant__ uint32_t LABEL_CACHE_SIZE;
  extern __constant__ uint32_t LABEL_CACHE_MIDPOINT;
  extern __constant__ uint32_t HIT_DATA_CACHE_SIZE;
  
  __constant__ uint16_t UNRESTRICTED_THREADS_PER_BLOCK;
  __constant__ uint32_t UNRESTRICTED_THREAD_COUNT;
  __constant__ uint32_t UNRESTRICTED_HITS_PER_WORKER;
  __constant__ uint32_t UNRESTRICTED_BLOCK_COUNT;
  
  extern __constant__ bool ONLINE_CONTRACT_DFU_PATH;
  extern __constant__ bool USE_LABEL_CACHE_NORMAL;
  extern __constant__ bool USE_LABEL_CACHE_BORDER;
  */
   
  
  
  
namespace SoAClustering{
  template <typename hit_type>
  struct hit_data_buffer
  {
    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 compact_toa_type = typename hit_type::compact_toa_type;
    using tot_type = typename hit_type::tot_type;
    x_type * x;
    y_type* y;
    tot_type* tot;
    compact_toa_type * toa;
    uint64_t current_index_ = 0;
    uint64_t max_size;
    const uint64_t end_buffer_free_space;// = unsortedness_window_mus * max_supported_hitrate_mhz;
    const uint64_t border_toa_offset_to_max_toa_mult_ns;// = 5000 *  TOA_NS_MULTIPLIER; //TODO use
    const uint64_t toa_ns_multiplier;
    const uint64_t toa_precision_digits_ns;
    const uint64_t unsortedness_window_mus;
    const uint32_t sensor_size_y;
    uint64_t offset_timestamp;
    uint64_t max_buffer_toa = 0;
    uint64_t frame_offset = 0;
    bool can_push = true;
    static constexpr uint64_t ns_in_mus = 1000;
    double frame_length_s;
    uint64_t frame_length_int_ns;
    uint64_t buffer_dt;
    uint32_t frame_count = 1;
    //hit_data_buffer();
    hit_data_buffer(uint64_t max_size, uint64_t end_buffer_free_space, uint64_t border_toa_offset_to_max_toa_ns, uint64_t toa_ns_multiplier, uint64_t unsortedness_window_mus, uint32_t sensor_size_y, double frame_length_s);

    uint32_t size();
    bool can_switch_buffer(const hit_type & hit);
    int64_t border_time() const;
    bool can_push_back(const hit_type & hit);

    bool push_back(hit_type && hit);
    
    void clear();
    bool can_reuse();
    
    uint64_t compute_skipped_frame_offset(uint64_t last_frame_offset, uint64_t last_buffer_end_time);

    uint64_t seal(uint64_t last_frame_offset, uint64_t last_border_time);
    const x_type * xs() const {return x;}
    const y_type * ys() const{return y;}
    const compact_toa_type * toas() const {return toa;}
    const tot_type * tots() const {return tot;}
    bool overflow_risk(toa_type toa) const;
    bool can_push_back_phase1(toa_type toa_int) const;
    bool can_push_back_phase2(toa_type toa_int) const;
    void set_start_offset(uint64_t time_offset, const hit_type & hit);
    bool first_frame_open() const;
    bool last_frame_open() const;
    bool entire_frame_open() const;
    uint64_t full_frame_count() const;
    
  };

  static constexpr uint32_t SAFE_THREAD_BLOCK = 1;
  template<typename hit_type>
  struct device_metadata
  {
    //TODO: START WITH REPLACING ARGS WITH __constant__ -> see if we need them on host, if not
    //then just replace and copy data to __constant__
    //IMPORTANT: DO COPYING DURING BUFFER DATA ALLOCATION iN CONSTRUCTOR OF EXTERNAL_CUDA_CLUSTERER 
    //static constexpr uint32_t max_cluster_join_time = 200 * hit_data_buffer::TOA_NS_MULTIPLIER;
/*
    static constexpr uint32_t max_hit_count = 38400000; //or 38400000 * 3
    static constexpr uint32_t total_thread_count = 3840;//3840;// or 8848
    static constexpr uint32_t approx_hits_per_worker = max_hit_count / total_thread_count;
    
    //FIXME: temporarily useless
    //static constexpr uint32_t max_cluster_timespan = 1000 * hit_data_buffer::TOA_NS_MULTIPLIER;

    static constexpr uint32_t matrix_padding_size = 2;
    static constexpr uint32_t matrix_size = (256 + matrix_padding_size) * (256 + matrix_padding_size);
    static constexpr uint16_t matrix_height = 2 << 7;
    static constexpr uint16_t matrix_width = 2 << 7;
    static constexpr uint32_t padded_matrix_width = (2 << 7) + matrix_padding_size;
    static constexpr uint16_t dfu_threads_per_block = 64;
    static constexpr uint16_t init_dfu_threads_per_block = 256;
    static constexpr uint16_t dfu_block_count = (device_metadata::total_thread_count / DFU_THREADS_PER_BLOCK) + SAFE_THREAD_BLOCK;

    static constexpr uint16_t unrestricted_threads_per_block = 256;
    static constexpr uint32_t unrestricted_thread_count = 3840 * 16;
    static constexpr uint16_t unrestricted_hits_per_worker = max_hit_count / unrestricted_thread_count;
    static constexpr uint16_t unrestricted_block_count = (unrestricted_thread_count / unrestricted_threads_per_block) + 1;
    
    
    static constexpr bool online_contract_dfu_path = true;
    static constexpr bool use_global_neighborhood = true;
    
    static constexpr bool use_label_cache_normal = true;
    static constexpr bool use_label_cache_border = true;
    */
    //uint64_t* tile_border_data;
    //initialize to -1s
    //need to cleanup to -1 after clusters are identified
    uint16_t* hit_index_matrices = nullptr;
    uint32_t* labels = nullptr;
    uint32_t* staging_matrices_idx = nullptr;
    typename hit_type::compact_toa_type * staging_matrices_toa = nullptr;
    uint32_t* cluster_offsets;
    //uint16_t* offset_to_last_toa; 
  };

  struct host_metadata
  {
    uint32_t* labels;
    const uint32_t* label() const {return labels;}
  };


  struct extra_data_config
  {
    
    enum class attribute
    {
      cluster_offset,
      cluster_size,
      energy,
      energy_2d_map,
      cluster_energy,
      cluster_energy_spectrum,
      cluster_energy_2d_map,
      cluster_weighted_centroid,
      cluster_start_toa
    };

    enum class memory_location {
    device, // GPU memory
    host    // Pinned host memory
};

    extra_data_config(memory_location mem_loc, const node_args & args):
    memory_space(mem_loc){
        if(args.exists("attribute_frame_length_s"))
          frame_length_s = args.get_arg<double>("attribute_frame_length_s"); 
        if(args.get_arg<bool>("attribute_energy"))
          set(attribute::energy);
        if(args.get_arg<bool>("attribute_energy_2d_map"))
          set(attribute::energy_2d_map);
        if(args.get_arg<bool>("attribute_cluster_energy"))
          set(attribute::cluster_energy);
        if(args.get_arg<bool>("attribute_cluster_size"))
          set(attribute::cluster_size);
        if(args.get_arg<bool>("attribute_cluster_energy_2d_map"))
        {
          if(!args.exists("attribute_frame_length_s"))
            throw std::invalid_argument("Cannot compute cluster energy 2d map when frame length is not defined. Make sure to set attribute_frame_length_s");
          set(attribute::cluster_energy_2d_map);
        }
        if(args.get_arg<bool>("attribute_cluster_energy_spectrum"))
          set(attribute::cluster_energy_spectrum);
        if(args.get_arg<bool>("attribute_cluster_weighted_centroid"))
          set(attribute::cluster_weighted_centroid);
        if(args.get_arg<bool>("attribute_cluster_start_time"))
          set(attribute::cluster_start_toa);

        if(args.exists("attribute_cluster_energy_spectrum_min"))
          spectrum_min_val = args.get_arg<float>("attribute_cluster_energy_spectrum_min");
        if(args.exists("attribute_cluster_energy_spectrum_max"))
          spectrum_max_val = args.get_arg<float>("attribute_cluster_energy_spectrum_max");
        if(args.exists("attribute_cluster_energy_spectrum_bin_count"))
          spectrum_bin_count = args.get_arg<uint32_t>("attribute_cluster_energy_spectrum_bin_count");     

        uint16_t filter_index = 1;
        while(args.exists("attribute_cluster_energy_filter_"+ std::to_string(filter_index)))
        {
          if(!args.exists("attribute_cluster_energy_2d_map"))
            throw std::invalid_argument("Cannot compute cluster energy filter when cluster_energy_2d_map attribute is turned off, make sure to use attribute_cluster_energy_2d_map");
          
          //TODO parse attribute filter interval
          energy_filter_intervals.push_back(args.get_arg<std::pair<double,double>>("attribute_cluster_energy_filter_"+ std::to_string(filter_index)));
          filter_index ++;
        }
          
        if(requires_calibration() && !args.exists("calibration_folder"))
        {
          throw std::invalid_argument("Error: missing calibration folder");
        }    
      
    }
    memory_location memory_space;
    uint64_t config_flags = 0;
    bool contains(attribute value)
    {
      return (config_flags >> (static_cast<int>(value))) & 1 > 0;
    }
    void set(attribute value)
    {
      config_flags |=(1 << static_cast<int>(value ));
    }
    bool requires_calibration()
    {

      return (contains(extra_data_config::attribute::cluster_energy) || 
        contains(extra_data_config::attribute::energy_2d_map) ||
        contains(extra_data_config::attribute::cluster_energy_2d_map) ||
        contains(extra_data_config::attribute::energy) ||
        contains(extra_data_config::attribute::cluster_energy_spectrum) ||
        contains(extra_data_config::attribute::cluster_weighted_centroid)
      );
    }
    double frame_length_s = 0.;
    uint32_t spectrum_bin_count = 100;
    float spectrum_min_val = 0.;
    float spectrum_max_val = 500.;
    std::vector<std::pair<double, double>> energy_filter_intervals;
    uint32_t output_count()
    {
      return std::max<uint32_t>(energy_filter_intervals.size(), 1);
    }
    uint32_t filter_count()
    {
      return energy_filter_intervals.size();
    }
  };

  template<typename hit_type>
  struct extra_attribute_data_wrapper
  {
    //extra_data_config config;
    uint32_t * cluster_offsets;
    uint32_t * cluster_sizes;
    float * calib_a;
    float * calib_b;
    float * calib_c;
    float * calib_t;
    float * energy;
    float * cluster_energy;
    
    float * cluster_energy_2d_map;
    uint32_t* cluster_energy_2d_map_indices;
    
    float * energy_2d_map;
    uint32_t* energy_2d_map_indices;

    uint32_t* cluster_energy_spectrum;
    float* x_weigthed_centroids;
    float* y_weigthed_centroids;
    typename hit_type::compact_toa_type* cluster_start_toa;
    

  };
  
  template<typename hit_type>
  struct extra_attribute_data
  {
    extra_data_config config;
    uint32_t * cluster_offsets;
    uint32_t * cluster_sizes;
    float * calib_a;
    float * calib_b;
    float * calib_c;
    float * calib_t;
    float * energy;
    float * cluster_energy;
    float * cluster_energy_2d_map;
    uint32_t* cluster_energy_2d_map_indices;
    
    float * energy_2d_map;
    uint32_t* energy_2d_map_indices;

    uint32_t* cluster_count_ptr = nullptr;
    uint32_t* nonzero_pixel_matrix_count_ptr = nullptr;
    uint32_t* nonzero_pixel_matrix_count_clustered_ptr = nullptr;
    
    uint32_t*  filtered_nonzero_pixel_matrix_count_clustered_ptr = nullptr;
    uint32_t* filtered_cluster_count_ptr = nullptr;
    
    
    uint32_t* cluster_energy_spectrum;
    float* x_centroid;
    float* y_centroid;
    typename hit_type::compact_toa_type* cluster_start_toa;
    uint64_t * auxiliary_dev_buffer64;
    bool has_ownership;



    extra_attribute_data_wrapper<hit_type> wrapped()
    {
        return { cluster_offsets, cluster_sizes, calib_a, calib_b, 
        calib_c, calib_t, energy, cluster_energy, cluster_energy_2d_map, cluster_energy_2d_map_indices,
         energy_2d_map, energy_2d_map_indices, cluster_energy_spectrum, x_centroid, y_centroid, cluster_start_toa};
    }
    void load_calibration(calibration& calib)
    {
      auto calib_a_flat = calib.a_flat();
      auto calib_b_flat = calib.b_flat();
      auto calib_c_flat = calib.c_flat();
      auto calib_t_flat = calib.t_flat();
       
      std::memcpy(calib_a, calib_a_flat.data(), calib_a_flat.size() * sizeof(*calib_a));
      std::memcpy(calib_b, calib_b_flat.data(), calib_b_flat.size() * sizeof(*calib_b));
      std::memcpy(calib_c, calib_c_flat.data(), calib_c_flat.size() * sizeof(*calib_c));
      std::memcpy(calib_t, calib_t_flat.data(), calib_t_flat.size() * sizeof(*calib_t));
       
    }
    static __host__ __device__ uint32_t linear_coordinate(const typename hit_type::x_type & x, const typename hit_type::y_type & y)
    {
      return x * hit_type::detector_type::size_y() + y; 
    }
    void allocate_host(uint32_t max_hit_count)
    {
      //allocate only the calibration maps and the desired attributes

      //TODO allocate calib a b c t
      
      auto err = cudaMallocHost((void**)&cluster_count_ptr, sizeof(uint32_t));
      check_cuda_error(err, "cudaMallocHost for cluster count");
      err = cudaMallocHost((void**)&nonzero_pixel_matrix_count_ptr, sizeof(uint32_t));
      check_cuda_error(err, "cudaMallocHost for nonzero pixel count");
      err = cudaMallocHost((void**)&nonzero_pixel_matrix_count_clustered_ptr, sizeof(uint32_t));
      check_cuda_error(err, "cudaMallocHost for nonzero pixel count clustered");
      err = cudaMemset(nonzero_pixel_matrix_count_clustered_ptr, 0, sizeof(uint32_t));
      check_cuda_error(err, "cudaMemset for nonzero pixel count clustered");
      err = cudaMemset(nonzero_pixel_matrix_count_ptr, 0, sizeof(uint32_t));
      
      
        
      uint64_t detector_size = hit_type::detector_type::size_x() * hit_type::detector_type::size_y();
      if(config.requires_calibration())
      {
        auto error = cudaMallocHost((void **)&(calib_a), detector_size * sizeof(*(calib_a)));
        check_cuda_error(error, "cudaMallocHost for calib a");
        error = cudaMallocHost((void **)&(calib_b), detector_size * sizeof(*(calib_b)));
        check_cuda_error(error, "cudaMallocHost for calib b");
        error = cudaMallocHost((void **)&(calib_c), detector_size * sizeof(*(calib_c)));
        check_cuda_error(error, "cudaMallocHost for calib c");
        error = cudaMallocHost((void **)&(calib_t), detector_size * sizeof(*(calib_t)));
        check_cuda_error(error, "cudaMallocHost for calib t");
        
      }

      if(config.contains(extra_data_config::attribute::cluster_energy))
      {
        auto error = cudaMallocHost((void **)&(cluster_energy), max_hit_count * sizeof(*(cluster_energy)));
        check_cuda_error(error, "cudaMallocHost for cluster energy");
      }
      if(config.contains(extra_data_config::attribute::energy_2d_map))
      {
        auto error = cudaMallocHost((void **)&(energy_2d_map), detector_size * sizeof(*(energy_2d_map)));
        check_cuda_error(error, "cudaMallocHost for energy 2d map");
        error = cudaMallocHost((void **)&(energy_2d_map_indices), detector_size * sizeof(*(energy_2d_map_indices)));
        check_cuda_error(error, "cudaMallocHost for energy 2d map indices");
      }
      
      if(config.contains(extra_data_config::attribute::cluster_energy_2d_map))
      {
        auto alloc_size = (max_hit_count + detector_size) * config.output_count();
        //here we assume sparse hit encoding was used
        auto error = cudaMallocHost((void **)&(cluster_energy_2d_map), alloc_size * sizeof(*(cluster_energy_2d_map)));
        check_cuda_error(error, "cudaMallocHost for cluster energy 2d map");
        error = cudaMallocHost((void **)&(cluster_energy_2d_map_indices), alloc_size * sizeof(*(cluster_energy_2d_map_indices)));
        check_cuda_error(error, "cudaMallocHost for cluster energy 2d map indices");

        if(config.filter_count() > 0)
        {
          error = cudaMallocHost((void **)&(filtered_nonzero_pixel_matrix_count_clustered_ptr), config.filter_count() * sizeof(*(filtered_nonzero_pixel_matrix_count_clustered_ptr)));
          check_cuda_error(error, "cudaMallocHost for filtered pixel matrix sizes");
          error = cudaMallocHost((void **)&(filtered_cluster_count_ptr), sizeof(*(filtered_cluster_count_ptr)));
          check_cuda_error(error, "cudaMallocHost for filtered cluster count");
          
        }
      
      }

      if(config.contains(extra_data_config::attribute::energy))
      {
        auto error = cudaMallocHost((void **)&(energy), max_hit_count * sizeof(*(energy)));
        check_cuda_error(error, "cudaMallocHost for energy per pixel");
      }
      if(config.contains(extra_data_config::attribute::cluster_size))
      {
        
        auto error = cudaMallocHost((void **)&(cluster_sizes), max_hit_count * sizeof(*(cluster_sizes)));
        check_cuda_error(error, "cudaMallocHost for cluster sizes");
      }
      if(config.contains(extra_data_config::attribute::cluster_weighted_centroid))
      {
        auto error = cudaMallocHost((void **)&(x_centroid), max_hit_count * sizeof(*(x_centroid)));
        check_cuda_error(error, "cudaMallocHost for cluster x centroids");
        error = cudaMallocHost((void **)&(y_centroid), max_hit_count * sizeof(*(y_centroid)));
        check_cuda_error(error, "cudaMallocHost for cluster y centroids");
      }
      if(config.contains(extra_data_config::attribute::cluster_energy_spectrum))
      {
        
        auto error = cudaMallocHost((void **)&(cluster_energy_spectrum), config.spectrum_bin_count * sizeof(*(cluster_energy_spectrum)));
        check_cuda_error(error, "cudaMallocHost for cluster energy spectrum");
        //TODO
      }
      if(config.contains(extra_data_config::attribute::cluster_start_toa))
      {
        auto error = cudaMallocHost((void **)&(cluster_start_toa), max_hit_count * sizeof(*(cluster_start_toa)));
        check_cuda_error(error, "cudaMallocHost for cluster start toa");        
      }
    }

    void allocate_device(uint32_t max_hit_count)
    {
      auto err = cudaMalloc((void**)&cluster_count_ptr, sizeof(uint32_t));
      check_cuda_error(err, "cudaMalloc for cluster count");
      err = cudaMalloc((void**)&nonzero_pixel_matrix_count_ptr, sizeof(uint32_t));
      check_cuda_error(err, "cudaMalloc for non zero pixel count");
      err = cudaMalloc((void**)&nonzero_pixel_matrix_count_clustered_ptr, sizeof(uint32_t));
      check_cuda_error(err, "cudaMalloc for non zero pixel count clustered");
      err = cudaMemset(nonzero_pixel_matrix_count_clustered_ptr, 0, sizeof(uint32_t));
      check_cuda_error(err, "cudaMemset for nonzero pixel count clustered");
      err = cudaMemset(nonzero_pixel_matrix_count_ptr, 0, sizeof(uint32_t));
      

      uint64_t detector_size = hit_type::detector_type::size_x() * hit_type::detector_type::size_y();
      if(config.requires_calibration())
      {
        auto error = cudaMalloc((void **)&(calib_a), detector_size * sizeof(*(calib_a)));
        check_cuda_error(error, "cudaMalloc for calib a");
        error = cudaMalloc((void **)&(calib_b), detector_size * sizeof(*(calib_b)));
        check_cuda_error(error, "cudaMalloc for calib b");
        error = cudaMalloc((void **)&(calib_c), detector_size * sizeof(*(calib_c)));
        check_cuda_error(error, "cudaMalloc for calib c");
        error = cudaMalloc((void **)&(calib_t), detector_size * sizeof(*(calib_t)));
        check_cuda_error(error, "cudaMalloc for calib t");
        
        error = cudaMalloc((void **)&(energy), max_hit_count * sizeof(*(energy)));
        check_cuda_error(error, "cudaMalloc for pixel energies");

      }
      if(config.contains(extra_data_config::attribute::energy_2d_map))
      {
        auto error = cudaMalloc((void **)&(energy_2d_map), detector_size * sizeof(*(energy_2d_map)));
        check_cuda_error(error, "cudaMalloc for 2d pixel energy map");
        error = cudaMalloc((void **)&(energy_2d_map_indices), detector_size * sizeof(*(energy_2d_map_indices)));
        check_cuda_error(error, "cudaMalloc for energy 2d map indices");
      } 
      
      if(config.contains(extra_data_config::attribute::cluster_energy) || 
        config.contains(extra_data_config::attribute::cluster_energy_2d_map) ||
        config.contains(extra_data_config::attribute::cluster_energy_spectrum)
      )
      {
        auto error = cudaMalloc((void **)&(cluster_energy), max_hit_count * sizeof(*(cluster_energy)));
        check_cuda_error(error, "cudaMalloc for cluster energies"); 
      }
      if(config.contains(extra_data_config::attribute::cluster_energy_2d_map))
      {
        uint64_t max_frame_count = 1000;
        auto alloc_size =  (max_hit_count + detector_size) * config.output_count();

        auto error = cudaMalloc((void **)&(cluster_energy_2d_map), alloc_size * sizeof(*(cluster_energy)));
        check_cuda_error(error, "cudaMalloc for cluster energy 2d map");  
        error = cudaMalloc((void **)&(cluster_energy_2d_map_indices), alloc_size * sizeof(*(cluster_energy_2d_map_indices)));
        check_cuda_error(error, "cudaMalloc for energy 2d map indices");
        

        //error = cudaMalloc((void **)&(auxiliary_dev_buffer64), max_hit_count * sizeof(*(auxiliary_dev_buffer64)));
        //check_cuda_error(error, "cudaMalloc for auxiliary device 64bit buffer");  
        
      }
      if(config.contains(extra_data_config::attribute::cluster_size))
      {
        auto error = cudaMalloc((void **)&(cluster_sizes), max_hit_count * sizeof(*(cluster_sizes)));
        check_cuda_error(error, "cudaMalloc for cluster sizes");
      }
      if(config.contains(extra_data_config::attribute::cluster_weighted_centroid) || 
        config.contains(extra_data_config::attribute::cluster_energy_2d_map))
      {
        auto error = cudaMalloc((void **)&(x_centroid), max_hit_count * sizeof(*(x_centroid)));
        check_cuda_error(error, "cudaMalloc for cluster x centroids");
        error = cudaMalloc((void **)&(y_centroid), max_hit_count * sizeof(*(y_centroid)));
        check_cuda_error(error, "cudaMalloc for cluster y centroids");
        if(config.filter_count() > 0)
        {
          error = cudaMalloc((void **)&(filtered_nonzero_pixel_matrix_count_clustered_ptr), config.filter_count() * sizeof(*(filtered_nonzero_pixel_matrix_count_clustered_ptr)));
          check_cuda_error(error, "cudaMalloc for filtered pixel matrix sizes");
          error = cudaMalloc((void **)&(filtered_cluster_count_ptr), sizeof(*(filtered_cluster_count_ptr)));
          check_cuda_error(error, "cudaMalloc for filtered cluster count");
          
          
        }
      }
      if(config.contains(extra_data_config::attribute::cluster_energy_spectrum))
      {
        auto error = cudaMalloc((void **)&(cluster_energy_spectrum), config.spectrum_bin_count * sizeof(*(cluster_energy_spectrum)));
        check_cuda_error(error, "cudaMalloc for cluster energy spectrum"); 
      }
      if(config.contains(extra_data_config::attribute::cluster_start_toa) || config.frame_length_s > 0)
      {
        auto error = cudaMalloc((void **)&(cluster_start_toa), max_hit_count * sizeof(*(cluster_start_toa)));
        check_cuda_error(error, "cudaMalloc for cluster start toa"); 
      }
    }

    void async_copy_to_device(cudaStream_t& stream, extra_attribute_data & target, uint32_t hit_count)
    {
      if(config.requires_calibration())
      {
        uint64_t detector_size = hit_type::detector_type::size_x() * hit_type::detector_type::size_y();
        cudaMemcpyAsync(target.calib_a, calib_a, detector_size * sizeof(*calib_a), cudaMemcpyHostToDevice, stream);
        cudaMemcpyAsync(target.calib_b, calib_b, detector_size * sizeof(*calib_b), cudaMemcpyHostToDevice, stream);
        cudaMemcpyAsync(target.calib_c, calib_c, detector_size * sizeof(*calib_c), cudaMemcpyHostToDevice, stream);
        cudaMemcpyAsync(target.calib_t, calib_t, detector_size * sizeof(*calib_t), cudaMemcpyHostToDevice, stream);
      }
    }
    void async_copy_from_device(cudaStream_t& stream, extra_attribute_data & target, uint32_t hit_count, uint32_t max_hit_count)
    {
      //FIXME
      hit_count = max_hit_count;
      uint64_t detector_size = hit_type::detector_type::size_x() * hit_type::detector_type::size_y();
      uint32_t cluster_count = hit_count; 
      cudaMemcpyAsync(target.cluster_count_ptr, cluster_count_ptr, sizeof(*cluster_count_ptr), cudaMemcpyDeviceToHost, stream);
      cudaMemcpyAsync(target.nonzero_pixel_matrix_count_ptr, nonzero_pixel_matrix_count_ptr, sizeof(*nonzero_pixel_matrix_count_ptr), cudaMemcpyDeviceToHost, stream);
      cudaMemcpyAsync(target.nonzero_pixel_matrix_count_clustered_ptr, nonzero_pixel_matrix_count_clustered_ptr, sizeof(*nonzero_pixel_matrix_count_clustered_ptr), cudaMemcpyDeviceToHost, stream);
        
      if(config.contains(extra_data_config::attribute::energy))
      {
        cudaMemcpyAsync(target.energy, energy, hit_count * sizeof(*energy), cudaMemcpyDeviceToHost, stream);
        //DEBUG_EXEC_CHECK(true, cudaDeviceSynchronize);
        
      } 
      if(config.contains(extra_data_config::attribute::cluster_energy))
      {
        cudaMemcpyAsync(target.cluster_energy, cluster_energy, cluster_count * sizeof(*cluster_energy), cudaMemcpyDeviceToHost, stream);
      }
      if(config.contains(extra_data_config::attribute::energy_2d_map))
      {
        cudaMemcpyAsync(target.energy_2d_map, energy_2d_map, detector_size * sizeof(*energy_2d_map), cudaMemcpyDeviceToHost, stream);
        cudaMemcpyAsync(target.energy_2d_map_indices, energy_2d_map_indices, detector_size * sizeof(*energy_2d_map_indices), cudaMemcpyDeviceToHost, stream);
        
      }
      if(config.contains(extra_data_config::attribute::cluster_energy_2d_map))
      {
        auto copy_size = (max_hit_count + detector_size) * config.output_count();
        
        cudaMemcpyAsync(target.cluster_energy_2d_map, cluster_energy_2d_map, copy_size * sizeof(*cluster_energy_2d_map), cudaMemcpyDeviceToHost, stream);
        cudaMemcpyAsync(target.cluster_energy_2d_map_indices, cluster_energy_2d_map_indices, copy_size * sizeof(*cluster_energy_2d_map_indices), cudaMemcpyDeviceToHost, stream);
        if(config.filter_count() > 0)
        {
          cudaMemcpyAsync(target.filtered_nonzero_pixel_matrix_count_clustered_ptr, filtered_nonzero_pixel_matrix_count_clustered_ptr, config.filter_count() * sizeof(*filtered_nonzero_pixel_matrix_count_clustered_ptr), cudaMemcpyDeviceToHost, stream);
          cudaMemcpyAsync(target.filtered_cluster_count_ptr, filtered_cluster_count_ptr, sizeof(*filtered_nonzero_pixel_matrix_count_clustered_ptr), cudaMemcpyDeviceToHost, stream);
          
        }
      }
      if(config.contains(extra_data_config::attribute::cluster_size))
      {
        cudaMemcpyAsync(target.cluster_sizes, cluster_sizes, hit_count * sizeof(*cluster_sizes), cudaMemcpyDeviceToHost, stream);
      }
      if(config.contains(extra_data_config::attribute::cluster_weighted_centroid))
      {
        cudaMemcpyAsync(target.x_centroid, x_centroid, hit_count * sizeof(*x_centroid), cudaMemcpyDeviceToHost, stream);
        cudaMemcpyAsync(target.y_centroid, y_centroid, hit_count * sizeof(*y_centroid), cudaMemcpyDeviceToHost, stream);
      }
      if(config.contains(extra_data_config::attribute::cluster_energy_spectrum))
      {
        cudaMemcpyAsync(target.cluster_energy_spectrum, cluster_energy_spectrum, config.spectrum_bin_count * sizeof(*cluster_energy_spectrum), cudaMemcpyDeviceToHost, stream);
      }
      if(config.contains(extra_data_config::attribute::cluster_start_toa))
      {
        cudaMemcpyAsync(target.cluster_start_toa, cluster_start_toa, hit_count * sizeof(*cluster_start_toa), cudaMemcpyDeviceToHost, stream);
      }
      //DEBUG_EXEC_CHECK(true, cudaDeviceSynchronize);
    }

    extra_attribute_data(const extra_attribute_data &other):
    config(other.config),
      cluster_offsets(other.cluster_offsets),
      cluster_sizes(other.cluster_sizes),
      calib_a(other.calib_a),
      calib_b(other.calib_b),
      calib_c(other.calib_c),
      calib_t(other.calib_t),
      energy(other.energy),
      cluster_energy(other.cluster_energy),
      cluster_energy_2d_map(other.cluster_energy_2d_map),
      energy_2d_map(other.energy_2d_map),
      cluster_energy_2d_map_indices(other.cluster_energy_2d_map_indices),
      energy_2d_map_indices(other.energy_2d_map_indices),
      x_centroid(other.x_centroid),
      y_centroid(other.y_centroid),
      cluster_energy_spectrum(other.cluster_energy_spectrum),
      cluster_start_toa(other.cluster_start_toa),
      has_ownership(false),
      auxiliary_dev_buffer64(other.auxiliary_dev_buffer64)
    {
      
    }
     extra_attribute_data(extra_attribute_data &&other):
    config(other.config),
      cluster_offsets(other.cluster_offsets),
      cluster_sizes(other.cluster_sizes),
      calib_a(other.calib_a),
      calib_b(other.calib_b),
      calib_c(other.calib_c),
      calib_t(other.calib_t),
      energy(other.energy),
      cluster_energy(other.cluster_energy),
      cluster_energy_2d_map(other.cluster_energy_2d_map),
      energy_2d_map(other.energy_2d_map),
      cluster_energy_2d_map_indices(other.cluster_energy_2d_map_indices),
      energy_2d_map_indices(other.energy_2d_map_indices),
      x_centroid(other.x_centroid),
      y_centroid(other.y_centroid),
      cluster_energy_spectrum(other.cluster_energy_spectrum),
      cluster_start_toa(other.cluster_start_toa),
      has_ownership(false),
      auxiliary_dev_buffer64(other.auxiliary_dev_buffer64)
    {
      
    }

    extra_attribute_data(uint32_t max_hit_count, extra_data_config config):
      config(config),
      cluster_offsets(),
      cluster_sizes(),
      calib_a(),
      calib_b(),
      calib_c(),
      calib_t(),
      energy(),
      cluster_energy(),
      cluster_energy_2d_map(),
      energy_2d_map(),
      cluster_energy_2d_map_indices(),
      energy_2d_map_indices(),
      cluster_energy_spectrum(),
      cluster_start_toa(),
      has_ownership(true),
      auxiliary_dev_buffer64()
    {
      switch (config.memory_space) {
        case extra_data_config::memory_location::device:
            allocate_device(max_hit_count);
            break;
        case extra_data_config::memory_location::host:
            allocate_host(max_hit_count);
            break;
    }
      //allocate space either on host or device based on config.device
    }

    __host__ ~extra_attribute_data()
    {
      if(!has_ownership)
        return;
      std::function<cudaError_t(void*)> free_func = [this](void* data)
      {
        if (data != nullptr)
        {
          if (config.memory_space == extra_data_config::memory_location::host)
          {
            return cudaFreeHost(data);
          }
          else
          {
            return cudaFree(data);
          }
        }
        return cudaSuccess;
      };
      
    free_func(cluster_offsets);
    free_func(cluster_sizes);
    free_func(calib_a);
    free_func(calib_b);
    free_func(calib_c);
    free_func(calib_t);
    free_func(energy);
    free_func(cluster_energy);
    free_func(cluster_energy_2d_map);
    free_func(cluster_energy_2d_map_indices);
    
    free_func(energy_2d_map);
    free_func(energy_2d_map_indices);
    
    free_func(cluster_count_ptr);
    free_func(nonzero_pixel_matrix_count_ptr);
    free_func(nonzero_pixel_matrix_count_clustered_ptr);
    free_func(cluster_energy_spectrum);
    free_func(auxiliary_dev_buffer64);
    }
  };

  template <typename hit_type>
  struct host_data
  {
    using data_type = hit_type;
    extra_attribute_data<hit_type>* extra_data;
    host_metadata metadata;
    hit_data_buffer<hit_type> data;
    uint64_t hit_count;
    bool is_final = false;
    host_data(uint64_t max_size, uint64_t end_buffer_free_space, uint64_t border_toa_offset_to_max_toa_ns, uint64_t toa_ns_multiplier, uint64_t unsortedness_window_mus, uint32_t sensor_size_y, double frame_size_s);
  };

  struct coalescence_permutation_functor
{
    const uint32_t block_size; // Total number of threads
    const uint32_t hpw; // Elements per thread
    // Constructor to pass in N and K
    coalescence_permutation_functor(uint32_t block_size, uint32_t hpw) : hpw(hpw), block_size(block_size) {}

    // The core logic, executed in parallel for each element on the GPU.
    // The 'new_index' is provided by the counting_iterator.
    __host__ __device__
    uint32_t operator()(uint32_t idx) const
    {

      const uint32_t block_idx = idx / (block_size * hpw);
      const uint32_t within_block_idx = idx % (block_size * hpw);
      const uint32_t thread_id = within_block_idx  % hpw;
      const uint32_t k_element = within_block_idx / hpw;
      return block_idx *hpw*block_size +  thread_id * block_size + k_element;
    }
};

  template <typename hit_type>
  struct auxiliary_sorting_buffers
  {
    uint32_t* indices;
    typename hit_type::x_type* sorted_xs;
    typename hit_type::y_type* sorted_ys;
    typename hit_type::tot_type* sorted_tots;
    typename hit_type::compact_toa_type* sorted_toas;
    uint32_t* inverse_sorted_indices;
    uint32_t* new_to_old_parent_pointers;
    uint32_t* sorted_indices;
    uint8_t* cub_temp_storage;
    uint64_t cub_temp_storage_bytes;

    thrust::host_vector<uint32_t> coalescence_permute_indices_host;
    uint32_t* coalescence_permute_indices;
    uint32_t max_size;

    uint16_t* tile_ids;
    typename hit_type::compact_toa_type* root_toas;
    
    auxiliary_sorting_buffers(uint32_t N, uint32_t block_size, uint32_t hpw) :
    max_size(N),  
    indices(),
    sorted_xs(),
    sorted_ys(),
    sorted_tots(),
    sorted_toas(),
    inverse_sorted_indices(),
    new_to_old_parent_pointers(),
    sorted_indices(),
    cub_temp_storage(),
    coalescence_permute_indices(),
    tile_ids(),
    root_toas(),
    coalescence_permute_indices_host(N)
    {
      uint32_t padded_N = N + block_size * hpw;
      uint64_t sort_temp_buffer_bytes;
      cub::DeviceRadixSort::SortPairs(
      nullptr, sort_temp_buffer_bytes,
      (typename hit_type::compact_toa_type*)nullptr, (typename hit_type::compact_toa_type*)nullptr,
      (uint32_t*)nullptr, (uint32_t*)nullptr,
      padded_N);

      uint64_t sort_large_temp_buffer_bytes;
      cub::DeviceRadixSort::SortPairs(
      nullptr, sort_large_temp_buffer_bytes,
      (uint64_t*)nullptr, (uint64_t*)nullptr,
      (uint32_t*)nullptr, (uint32_t*)nullptr,
      padded_N);
        
      
      uint64_t select_temp_buffer_bytes;
      cub::DeviceSelect::Flagged(nullptr, select_temp_buffer_bytes, (uint32_t*)nullptr, (bool *)nullptr,
      (uint32_t*)nullptr, (uint32_t*)nullptr, N  /*...select args...*/);

      uint64_t segmented_temp_sum_buffer_bytes;
      cub::DeviceSegmentedReduce::Sum(nullptr, segmented_temp_sum_buffer_bytes, (float*)nullptr, (float*)nullptr,
        1, (uint32_t*)nullptr, (uint32_t*)nullptr); 
      cub::Sum sum_op;
        cub::DeviceReduce::ReduceByKey(
            nullptr, 
                segmented_temp_sum_buffer_bytes, 
                //keys (linear coords)
                (uint32_t*)nullptr, (uint32_t*)nullptr,
                // values (pixel energies)
                (float*)nullptr, (float*)nullptr,
                //remaining params
                (uint32_t*)nullptr, 
                sum_op, padded_N); 
        

      cub_temp_storage_bytes = std::max(std::max(sort_temp_buffer_bytes, select_temp_buffer_bytes), std::max(sort_large_temp_buffer_bytes, segmented_temp_sum_buffer_bytes));

      //cub_temp_storage = thrust::device_vector<uint8_t>(cub_temp_storage_bytes);
      
      auto error = cudaMalloc((void **)&(sorted_xs), padded_N * sizeof(*(sorted_xs)));
      check_cuda_error(error, "cudaMalloc for sorted xs");
      error = cudaMalloc((void **)&(sorted_ys), padded_N * sizeof(*(sorted_ys)));
      check_cuda_error(error, "cudaMalloc for sorted ys");
      error = cudaMalloc((void **)&(sorted_tots), padded_N * sizeof(*(sorted_tots)));
      check_cuda_error(error, "cudaMalloc for sorted tots");
      error = cudaMalloc((void **)&(sorted_toas), padded_N * sizeof(*(sorted_toas)));
      check_cuda_error(error, "cudaMalloc for sorted toas");

      
      error = cudaMalloc((void **)&(indices), padded_N * sizeof(*(indices)));
      check_cuda_error(error, "cudaMalloc for indices");
      
      error = cudaMalloc((void **)&(inverse_sorted_indices), padded_N * sizeof(*(inverse_sorted_indices)));
      check_cuda_error(error, "cudaMalloc for inverse sorted indices");
      
      error = cudaMalloc((void **)&(new_to_old_parent_pointers), padded_N * sizeof(*(new_to_old_parent_pointers)));
      check_cuda_error(error, "cudaMalloc for new to old parent pointers");
      
      error = cudaMalloc((void **)&(sorted_indices),padded_N * sizeof(*(sorted_indices)));
      check_cuda_error(error, "cudaMalloc for sorted indices");
      
      error = cudaMalloc((void **)&(coalescence_permute_indices), padded_N * sizeof(*(coalescence_permute_indices)));
      check_cuda_error(error, "cudaMalloc for coalescence permute indices");

      error = cudaMalloc((void **)&(cub_temp_storage), cub_temp_storage_bytes);
      check_cuda_error(error, "cudaMalloc for cub temporary storage");

      thrust::transform(
            thrust::counting_iterator<int>(0),      
            thrust::counting_iterator<int>(padded_N),      
            thrust::device_pointer_cast(coalescence_permute_indices),                     
            coalescence_permutation_functor(block_size, hpw)               
        );

      error = cudaMalloc((void **)&(root_toas), padded_N * sizeof(*(root_toas)));
      check_cuda_error(error, "cudaMalloc for auxiliary root toas");

      error = cudaMalloc((void **)&(tile_ids), padded_N * sizeof(*(tile_ids)));
      check_cuda_error(error, "cudaMalloc for auxiliary tile_ids");
      thrust::sequence(thrust::device_pointer_cast(indices), thrust::device_pointer_cast(indices + N));
/*
      thrust::transform(
            thrust::counting_iterator<int>(0),      // Start of virtual input (0)
            thrust::counting_iterator<int>(N),      // End of virtual input (T)
            coalescence_permute_indices_host.begin(),                      // Beginning of the output vector
            coalescence_permutation_functor(block_size, hpw)               // The functor with N and K
        );
      uint32_t* data = coalescence_permute_indices_host.data();
      std::vector<uint32_t> vec(data, data + N);
      //datavec.max()
      
      std::cout <<"Max val:" <<  *std::max_element(vec.begin(), vec.end()) << std::endl;
      std::set<uint32_t> unique_elements(vec.begin(), vec.end());*/


    }
    ~auxiliary_sorting_buffers()
    {

      cudaFree(indices);
      cudaFree(sorted_xs);
      cudaFree(sorted_ys);
      cudaFree(sorted_tots);
      cudaFree(sorted_toas);
      cudaFree(inverse_sorted_indices);
      cudaFree(new_to_old_parent_pointers);
      cudaFree(sorted_indices);
      cudaFree(cub_temp_storage);
      cudaFree(coalescence_permute_indices);
      cudaFree(tile_ids);
      cudaFree(root_toas);

    }
  };

  template <typename hit_data>
  struct device_data
  {
    extra_attribute_data<hit_data>* extra_data;
    device_metadata<hit_data> metadata;
    hit_data_buffer<hit_data> data;
    uint32_t hit_count;
    uint32_t* hit_count_ptr = nullptr;
    cudaStream_t stream;
    bool is_final = false;
    auxiliary_sorting_buffers<hit_data>* aux_sort_buffers = nullptr;
    //device_data(bool create_custom_stream  = false);
    device_data(/*cudaStream_t cudaStream*/ uint64_t max_hit_count);
  };

  class  cuda_coord
  {

public:
    __device__ cuda_coord() {}
    __device__ cuda_coord(int16_t x, int16_t y) : m_x(x),
                                             m_y(y) {}

    __device__ int16_t x() const
    {
        return m_x;
    }
    __device__ int16_t y() const
    {
        return m_y;
    }

    __device__ cuda_coord operator+(const cuda_coord &right)
    {
        return cuda_coord(this->x() + right.x(), this->y() + right.y());
    }
    int16_t m_x;
    int16_t m_y;
};

}
template<typename T>
__device__ T max(T left, T right)
{
  return left < right ? right : left;
}

template<typename T>
__device__ T min(T left, T right)
{
  return left < right ? left : right;
}

/*
namespace SoATClustering{
  struct hit_data_buffer
  {
    uint8_t* x;
    uint8_t* y;
    uint16_t* tot;
    uint64_t* toa;
    uint32_t current_index_ = 0;
    uint64_t max_size;
    const uint64_t end_buffer_free_space; // = unsortedness_window_mus * max_supported_hitrate_mhz;
    const uint64_t border_toa_offset_to_max_toa_mult_ns;// = 5000 *  TOA_NS_MULTIPLIER; //TODO use
    const uint64_t toa_ns_multiplier;
    const uint64_t unsortedness_window_mus;
    static constexpr uint64_t ns_in_mus = 1000;
    uint64_t max_buffer_toa = 0;
    bool can_push = true;
    hit_data_buffer();
    hit_data_buffer(uint64_t max_size, uint64_t end_buffer_free_space, uint64_t border_toa_offset_to_max_toa_ns, uint64_t toa_ns_multiplier, uint64_t unsortedness_window_mus);
    uint32_t size();
    bool can_switch_buffer(const tpx3_hit & hit);
    uint64_t border_time();
    bool can_push_back(const tpx3_hit & hit);

    bool push_back(tpx3_hit && hit);

    void clear();
    bool can_reuse();
    void seal();
  };

  static constexpr uint32_t SAFE_THREAD_BLOCK = 1;
  struct device_metadata
  {
    //static constexpr uint32_t max_cluster_join_time = 200 * hit_data_buffer::TOA_NS_MULTIPLIER;
/*
    static constexpr uint32_t max_hit_count = 38400000; //or 38400000 * 3
    static constexpr uint32_t total_thread_count = 3840;//3840;// or 8848
    static constexpr uint32_t approx_hits_per_worker = max_hit_count / total_thread_count;
    static constexpr uint32_t max_cluster_timespan = 1000 * hit_data_buffer::TOA_NS_MULTIPLIER;

    static constexpr uint32_t matrix_padding_size = 2;
    static constexpr uint32_t matrix_size = (256 + matrix_padding_size) * (256 + matrix_padding_size);
    static constexpr uint16_t matrix_height = 2 << 7;
    static constexpr uint16_t matrix_width = 2 << 7;
    static constexpr uint32_t padded_matrix_width = (2 << 7) + matrix_padding_size;
    static constexpr uint16_t dfu_threads_per_block = 64;
    static constexpr uint16_t init_dfu_threads_per_block = 256;
    static constexpr uint16_t dfu_block_count = (device_metadata::total_thread_count / device_metadata::dfu_threads_per_block) + SAFE_THREAD_BLOCK;

    static constexpr uint16_t unrestricted_threads_per_block = 256;
    static constexpr uint32_t unrestricted_thread_count = 3840 * 16;
    static constexpr uint16_t unrestricted_hits_per_worker = max_hit_count / unrestricted_thread_count;
    static constexpr uint16_t unrestricted_block_count = (unrestricted_thread_count / unrestricted_threads_per_block) + 1;
    
    
    static constexpr bool online_contract_dfu_path = true;
    static constexpr bool use_global_neighborhood = true;
    
    static constexpr bool use_label_cache_normal = true;
    static constexpr bool use_label_cache_border = true;
    *//*
    //uint64_t* tile_border_data;
    //initialize to -1s
    //need to cleanup to -1 after clusters are identified
    uint16_t* hit_index_matrices;
    uint32_t* labels;
    //uint16_t* offset_to_last_toa; 
  };

 
  struct device_data
  {
    device_metadata metadata;
    hit_data_buffer data;
    uint32_t hit_count;
    cudaStream_t stream;
    bool is_final = false;
    device_data(uint64_t max_hit_count);
    //device_data(cudaStream_t cudaStream);
  };

  class  cuda_coord
  {

public:
    __device__ cuda_coord() {}
    __device__ cuda_coord(int16_t x, int16_t y) : m_x(x),
                                             m_y(y) {}

    __device__ int16_t x() const
    {
        return m_x;
    }
    __device__ int16_t y() const
    {
        return m_y;
    }

    __device__ cuda_coord operator+(const cuda_coord &right)
    {
        return cuda_coord(this->x() + right.x(), this->y() + right.y());
    }
    int16_t m_x;
    int16_t m_y;
};*/
}

