
#pragma once
#include "data_flow/dataflow_package.h"

#include "data_structs/calibration.h"
#include "data_structs/cuda_hit_buffer.cuh"
#include "clusterer_utils.h"
#include "cuda_resources/cuda_kernels.cuh"
// #include "../../cuda_resources/cuda_kernels_transposed.cuh"
#include "cuda_resources/dfu_clustering.cuh"
#include "data_structs/clustered_data.h"
// #include "cuda_runtime.h"
// #include "device_launch_parameters.h"

#include <algorithm>
#include <cassert>
#include <queue>
#include <type_traits>
#include <chrono>
#include <limits>
#include <cstdint>
// converts the burda hit to mm_hit
// computes energy, and time in nanoseconds

namespace clustering
{


  using namespace std::chrono_literals;
  namespace SoAClusteringType = SoAClustering;
  /*template <typename hit_type>
  class external_dataflow_controller;
  
  
  */


  struct temporary_merge_result
  {
    int32_t* first_start_idx = nullptr;
    float* first_start_energy = nullptr;
    uint64_t first_size = 0;
     
    int32_t* second_start_idx = nullptr;
    float* second_start_energy = nullptr;
    uint64_t second_size = 0;

    bool is_empty()
    {
      return (first_start_idx == nullptr && first_start_energy == nullptr &&
        first_size == 0 && second_start_idx == nullptr && second_start_energy == nullptr
        && second_size == 0);
    }
    
  };

  template <typename hit_type>
  class external_cuda_clusterer
      : virtual public i_data_node
  {
  public:
    using output_func_type = typename std::function<void(clustered_data<hit_type>)>;
    external_cuda_clusterer(const node_args &args, output_func_type output_fn, calibration *calib = nullptr);
    
    std::string name() override { return "cuda_clusterer"; }

    void receive_data_callback(SoAClusteringType::host_data<hit_type> *host_data);

    void done_callback();

  protected:
    struct parameters
    {
      static constexpr uint32_t matrix_padding_size = 2;
      uint64_t buffer_size;
      uint64_t max_unsortedness_mus;
      uint32_t chip_sensor_size_x;
      uint32_t chip_sensor_size_y;
      uint32_t cuda_streams_per_worker;
      uint32_t cuda_thread_count;
      uint16_t cuda_clustering_threads_per_block;
      uint32_t cuda_min_launch_stream_count;
      uint32_t cuda_sleep_duration_empty_ms;
      uint32_t cuda_aux_matrix_size;
      uint64_t toa_ns_multiplier;
      bool use_label_cache_normal;
      bool use_label_cache_border;
      uint32_t label_cache_size;
      uint32_t hit_data_cache_size;
      uint16_t init_threads_per_block;
      uint64_t max_cluster_join_time_ns;
      uint32_t shared_memory_bytes;
      uint16_t tile_size;
      bool discard_raw_data;
    };
    bool calibrate;
    std::unique_ptr<calibration> calibrator;
    output_func_type output_fn;
    std::vector<SoAClusteringType::device_data<hit_type> *> devs;
    std::deque<SoAClusteringType::host_data<hit_type> *> host_buffers_to_process;
    parameters params;
    bool done = false;
    std::chrono::milliseconds empty_sleep_time{10ms};
    const bool flush_on_dt = false;
    double last_t = 0;
    SoAClustering::host_data<hit_type> * waiting_buffer = nullptr;
    std::vector<int32_t> waiting_frame_idx;
    std::vector<float> waiting_frame_energy;
    std::vector<temporary_merge_result> last_merge_results;
    std::vector<temporary_merge_result>::iterator last_merge_result;
    std::vector<cudaStream_t> streams_used = {0}; //always using the default stream
    
    // std::vector<int32_t> cluster_id_map = std::vector(SoAClusteringType::MAX_HIT_COUNT, -1);
    friend class external_dataflow_controller<hit_type>;

    static constexpr uint64_t US_IN_MS = 1000;
    uint32_t hits_per_thread();
    void initialize_cuda_constants();
    mm_hit convert_hit(uint8_t x, uint8_t y, typename hit_type::toa_type int_toa, uint16_t tot);
    void output_clusters(SoAClusteringType::host_data<hit_type> *host);
    virtual void start() override;
    virtual ~external_cuda_clusterer();
    void run_base_clustering();
    void run_tile_clustering();
    void merge_buffers(SoAClustering::host_data<hit_type> * first_buffer, SoAClustering::host_data<hit_type> * second_buffer);
    void merge_cluster_pixel_matrices(SoAClustering::host_data<hit_type> * first_buffer, SoAClustering::host_data<hit_type> * second_buffer, uint32_t frame_offset, uint32_t sensor_offset);
    std::vector<cudaStream_t>& streams();
    private:
    void postprocess_host_buffer(SoAClustering::host_data<hit_type>& buffer);  
    std::pair<std::vector<int32_t>, std::vector<float>> union_sparse_pixel_matrices(SoAClustering::host_data<hit_type> * first_buffer, SoAClustering::host_data<hit_type> * second_buffer,
      int64_t first_current_idx, int64_t first_frame_end_idx, 
      int64_t second_current_idx, int64_t second_frame_end_idx,
    int64_t first_index_modifier, int64_t second_index_modifier, uint64_t frame_filter_offset);
    void find_frame_stop_indices(SoAClustering::host_data<hit_type> * first_buffer, SoAClustering::host_data<hit_type> * second_buffer,
       int64_t first_stop_val, int64_t second_stop_val, int64_t& first_start_idx, int64_t& second_end_idx, uint32_t frame_filter_offset);
  clustered_data<hit_type> create_output_data(SoAClustering::host_data<hit_type> * buffer, bool done = false);
       
  };
}