#include "SequenceVisitor.cuh"
#include "MatchUpstreamMuon.h"

#include "SequenceVisitor.cuh"
#include "RunForwardCPU.h"
#include "Tools.h"
#include "MuonFeaturesExtraction.cuh"

template<>
void SequenceVisitor::set_arguments_size<cpu_match_upstream_muon_t>(
  cpu_match_upstream_muon_t::arguments_t arguments,
  const RuntimeOptions& runtime_options,
  const Constants& constants,
  const HostBuffers& host_buffers)
{
  arguments.set_size<dev_muon_hits>(runtime_options.number_of_events);
  arguments.set_size<dev_match_upstream_muon>(host_buffers.host_number_of_selected_events[0]);
}

template<>
void SequenceVisitor::visit<cpu_match_upstream_muon_t>(
  cpu_match_upstream_muon_t& state,
  const cpu_match_upstream_muon_t::arguments_t& arguments,
  const RuntimeOptions& runtime_options,
  const Constants& constants,
  HostBuffers& host_buffers,
  cudaStream_t& cuda_stream,
  cudaEvent_t& cuda_generic_event)
{

  cudaEventRecord(cuda_generic_event, cuda_stream);
  cudaEventSynchronize(cuda_generic_event);

  std::vector<Muon::HitsSoA> muon_hits_events(host_buffers.host_number_of_selected_events[0]);
  //std::vector<Muon::HitsSoA> muon_hits_events(host_buffers.host_number_of_selected_events[0]);

  cudaCheck(cudaMemcpyAsync(
      muon_hits_events.data(),
      arguments.offset<dev_muon_hits>(),
      runtime_options.number_of_events * sizeof(Muon::HitsSoA),
      cudaMemcpyDeviceToHost,
      cuda_stream
  ));

  std::vector<bool> MuonMatch = state.invoke(
    constants.host_inv_clus_res,
    host_buffers.host_atomics_velo,
    host_buffers.host_velo_track_hit_number,
    host_buffers.host_kalmanvelo_states,
    host_buffers.host_atomics_ut,
    host_buffers.host_ut_track_hit_number,
    host_buffers.host_ut_qop,
    host_buffers.host_ut_track_velo_indices,
    muon_hits_events.data(),
    //host_buffers.host_muon_hits_events.data(),
    runtime_options.number_of_events);

  cudaCheck(cudaMemcpyAsync(
      arguments.offset<dev_match_upstream_muon>(),
      MuonMatch.data(),
      host_buffers.host_number_of_selected_events[0] * sizeof(bool),
      cudaMemcpyHostToDevice,
      cuda_stream
  ));
}