/*//////////////////////////////////////////////////////////////////////////////
// CMT Cosmic Muon Tomography project //////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

  Copyright (c) 2014, Universita' degli Studi di Padova, INFN sez. di Padova
  All rights reserved

  Authors: Andrea Rigoni Garola < andrea.rigoni@pd.infn.it >

//////////////////////////////////////////////////////////////////////////////*/

#include "testing-prototype.h"
#include <Core/Vector.h>

#ifdef USE_CUDA
#include <thrust/device_ptr.h>
#include <thrust/transform.h>

struct DoubleFunctor {
  __host__ __device__ int operator()(int x) const { return x * 2; }
};
#endif

int main() {
  BEGIN_TESTING(VectorMetaAllocator);

  uLib::Vector<int> v;
  std::cout << "Pushing elements...\n";
  v << 1, 2, 3, 4, 5;

  std::cout << "Initial RAM contents: ";
  for (size_t i = 0; i < v.size(); ++i) {
    std::cout << v[i] << " ";
    if (v[i] != (int)(i + 1)) {
      std::cout << "\nError: Value mismatch at index " << i << "\n";
      exit(1);
    }
  }
  std::cout << "\n";

#ifdef USE_CUDA
  std::cout << "Moving to VRAM...\n";
  v.MoveToVRAM();

  int *vram_ptr = v.GetVRAMData();
  if (vram_ptr) {
    std::cout << "Successfully obtained VRAM pointer: " << vram_ptr << "\n";
  } else {
    std::cout << "Error: Failed to obtain VRAM pointer!\n";
    exit(1);
  }

  // Verify DeviceData() matches GetVRAMData()
  {
    thrust::device_ptr<int> dev_ptr = v.DeviceData();
    if (dev_ptr.get() != vram_ptr) {
      std::cout << "Error: DeviceData() does not match GetVRAMData()!\n";
      exit(1);
    }
    std::cout << "DeviceData() matches GetVRAMData(). OK\n";
  }

  // Use thrust::transform via DeviceBegin()/DeviceEnd() to double all elements
  // on device
  std::cout << "Doubling elements on device via thrust::transform...\n";
  thrust::transform(v.DeviceBegin(), v.DeviceEnd(), v.DeviceBegin(),
                    DoubleFunctor{});

  std::cout << "Moving back to RAM...\n";
  v.MoveToRAM();

  std::cout << "RAM contents after VRAM trip + thrust transform: ";
  for (size_t i = 0; i < v.size(); ++i) {
    std::cout << v[i] << " ";
    if (v[i] != (int)((i + 1) * 2)) {
      std::cout << "\nError: Data corrupted after RAM->VRAM->thrust->RAM trip "
                   "at index "
                << i << "\n";
      exit(1);
    }
  }
  std::cout << "\n";
#else
  std::cout << "USE_CUDA not defined, skipping VRAM tests.\n";
#endif

  std::cout << "Scaling vector...\n";
  for (size_t i = 0; i < v.size(); ++i)
    v[i] *= 10;

  std::cout << "Final contents: ";
  for (size_t i = 0; i < v.size(); ++i)
    std::cout << v[i] << " ";
  std::cout << "\n";

  END_TESTING;
}