#include <functional>
#include <iomanip>
#include <iostream>
#include <string>
#include <tuple>
#include <utility/include_all.h>
#include <utility/iterator.h>
#include <utility/math.h>
#include <utility/unit_math.h>
#include <vector>

#ifndef TESTING
struct virtual_particles {
  int32_t requiredSlices_x, requiredSlices_y, requiredSlices_z;
  float r;
  value_unit<float4, SI::m> p;
  value_unit<float4, SI::m> c;
  bool valid = false;

  template <typename V, typename U> virtual_particles(U &&position, V &&volume, float spacing) {
    p = position;
    c = position; // float4_u<SI::m>(0., 0., 0., p.val.w);
    float radius = pow(volume.val / (4.f / 3.f * CUDART_PI_F), 1.f / 3.f);
    float h = p.val.w * Kernel<kernel_kind::spline4>::kernel_size();
    valid = true;

    r = spacing * radius;

    auto k = static_cast<int32_t>(floor(c.val.z / r * 3. / (2. * sqrt(6.))));
    auto j = static_cast<int32_t>(floor(c.val.y / r / sqrt(3.) - 1. / 3. * (k % 2)));
    auto i = static_cast<int32_t>(floor((c.val.x / r - (j + k) % 2) / 2.));
    c = value_unit<float4, SI::m>{2.f * i + ((j + k) % 2), sqrtf(3.f) * (j + 1.f / 3.f * (k % 2)),
                                  2.f * sqrtf(6.f) / 3.f * k, 0.f};
    c = value_unit<float4, SI::m>{c.val.x * r, c.val.y * r, c.val.z * r, 0.f};
    // std::cout << c.val.x << ", " << c.val.y << ", " << c.val.z << std::endl;
    c = value_unit<float4, SI::m>{0.f, 0.f, 0.f, 0.f};
    requiredSlices_x = 2 * static_cast<int32_t>(ceil(h / r));
    requiredSlices_y = 2 * static_cast<int32_t>(ceil(h / (sqrtf(3.f) * r)));
    requiredSlices_z = 2 * static_cast<int32_t>(ceil(h / r * 3.f / (sqrtf(6.f) * 2.f)));
  }
  struct neighbor_it {
    const int32_t requiredSlices_x, requiredSlices_y, requiredSlices_z;
    const float r;
    value_unit<float4, SI::m> p;
    value_unit<float4, SI::m> s;
    value_unit<float4, SI::m> c;
    int32_t i, j, k;

    neighbor_it(int32_t x_len, int32_t y_len, int32_t z_len, float r, value_unit<float4, SI::m> pos,
                value_unit<float4, SI::m> center, int32_t _x, int32_t _y, int32_t _z)
        : requiredSlices_x(x_len), requiredSlices_y(y_len), requiredSlices_z(z_len), r(r), p(pos), c(center), i(_x),
          j(_y), k(_z) {
      increment();
    }

    value_unit<float4, SI::m> operator*() { return s; }
    bool operator==(const neighbor_it &rawIterator) const { return (k == rawIterator.k); }
    bool operator!=(const neighbor_it &rawIterator) const { return (k != rawIterator.k); }

    void increment() {
      //using namespace math::ops;
      bool in_range = false;
      do {
        i++;
        if (i == requiredSlices_x + 1) {
          i = -requiredSlices_x;
          j++;
          if (j == requiredSlices_y + 1) {
            j = -requiredSlices_y;
            k++;
          }
        }
        float4 initial{2.f * i + ((j + k) % 2), sqrtf(3.f) * (j + 1.f / 3.f * (k % 2)), 2.f * sqrtf(6.f) / 3.f * k,
                       0.f};

        s = value_unit<float4, SI::m>{c.val.x + initial.x * r, c.val.y + initial.y * r, c.val.z + initial.z * r,
                                      p.val.w};
        //float dist = math::sqdistance3(p.val, s.val);
        //dist = sqrt(dist);
        in_range = true; // dist < (p.val.w *
                         // Kernel<kernel_kind::spline4>::kernel_size());
                         // s.val.w = 1. - d / (p.val.w *
        // Kernel<kernel_kind::spline4>::kernel_size());
      } while (k != requiredSlices_z + 1 && !(in_range));
    }

    neighbor_it &operator++() {
      increment();
      return (*this);
    }
    const neighbor_it operator++(int) {
      auto temp(*this);
      increment();
      return temp;
    }
  };
  neighbor_it begin() const {
    return neighbor_it{requiredSlices_x,
                       requiredSlices_y,
                       requiredSlices_z,
                       r,
                       p,
                       c,
                       -requiredSlices_x,
                       -requiredSlices_y,
                       !valid ? requiredSlices_z + 1 : -requiredSlices_z};
  }
  neighbor_it end() const {
    return neighbor_it{requiredSlices_x, requiredSlices_y, requiredSlices_z,    r, p, c,
                       requiredSlices_x, requiredSlices_y, requiredSlices_z + 1};
  }
  neighbor_it cbegin() const {
    return neighbor_it{requiredSlices_x,
                       requiredSlices_y,
                       requiredSlices_z,
                       r,
                       p,
                       c,
                       -requiredSlices_x,
                       -requiredSlices_y,
                       !valid ? requiredSlices_z + 1 : -requiredSlices_z};
  }
  neighbor_it cend() const {
    return neighbor_it{requiredSlices_x, requiredSlices_y, requiredSlices_z,    r, p, c,
                       requiredSlices_x, requiredSlices_y, requiredSlices_z + 1};
  }
};

template <typename T> auto print_val(const std::string& str, T val) {
  std::cout << std::setw(32) << str << " = " << std::setw(24)
            << std::setprecision(std::numeric_limits<long double>::digits10 + 1) << val.val << " -> " << std::hexfloat
            << val.val << std::defaultfloat << std::endl;
};
template <> auto print_val<float>(const std::string& str, float val) {
  std::cout << std::setw(32) << str << " = " << std::setw(24)
            << std::setprecision(std::numeric_limits<long double>::digits10 + 1) << val << " -> " << std::hexfloat
            << val << std::defaultfloat << std::endl;
};
template <> auto print_val<int32_t>(const std::string& str, int32_t val) {
  std::cout << std::setw(32) << str << " = " << val << std::endl;
};

#define PRINT(x) print_val(#x, x);

template <kernel_kind K = kernel_kind::spline4, typename T> hostDeviceInline auto supportFromVol(T volume) {
  auto target_neighbors = Kernel<K>::neighbor_number;
  auto kernel_epsilon = (1.f) * powf((target_neighbors)*PI4O3_1, 1.f / 3.f) / Kernel<K>::kernel_size();
  auto h = kernel_epsilon * math::power<ratio<1, 3>>(volume);
  return h;
}

template <typename C> auto createLUT(C func, int32_t steps, float spacing) {
  float_u<SI::volume> volume(1.f);
  auto h = supportFromVol(volume);
  auto H = h * Kernel<kernel_kind::spline4>::kernel_size<float>();
  auto step = H / (static_cast<float>(steps)) * 2.f;

  auto current_position = float4_u<SI::m>{0.f, 0.f, 0.f, h.val};
  auto center_position = float4_u<SI::m>{0.f, 0.f, 0.f, h.val};

  using res_t = decltype(math::unit_get<1>(func(current_position, current_position)) * volume);
  std::vector<res_t> LUT;

  for (int32_t i = 0; i < steps; ++i) {
    math::unit_assign<1>(current_position, H - step * static_cast<float>(i));
    res_t density{0.f};
    for (const auto &var : virtual_particles(center_position, volume, spacing)) {
      if (math::unit_get<1>(var) > 0._m){
        continue;
	  }
      density += volume * math::unit_get<1>(func(current_position, var));
    }
    LUT.push_back(density);
  }
  return LUT;
}

auto createCounterLUT(int32_t steps, float spacing) {
  float_u<SI::volume> volume(1.f);
  auto h = supportFromVol(volume);
  auto H = h * Kernel<kernel_kind::spline4>::kernel_size<float>();
  auto step = H / (static_cast<float>(steps)) * 2.f;

  auto current_position = float4_u<SI::m>{0.f, 0.f, 0.f, h.val};
  auto center_position = float4_u<SI::m>{0.f, 0.f, 0.f, h.val};

  std::vector<int32_t> LUT;

  for (int32_t i = 0; i < steps; ++i) {
    math::unit_assign<1>(current_position, H - step * static_cast<float>(i));
    int32_t counter = 0;
    for (const auto &var : virtual_particles(center_position, volume, spacing)) {
      if (math::unit_get<1>(var) > 0._m){
        continue;
	  }
      if (spline4_kernel(current_position, var) > 0.f){
        counter++;
	  }
    }
    LUT.push_back(counter);
  }
  return LUT;
}

template <typename C> auto create4DLUT(C func, int32_t steps, float spacing) {
  float_u<SI::volume> volume(1.f);
  auto h = supportFromVol(volume);
  auto H = h * Kernel<kernel_kind::spline4>::kernel_size<float>();
  auto step = H / (static_cast<float>(steps)) * 2.f;

  auto current_position = float4_u<SI::m>{0.f, 0.f, 0.f, h.val};
  auto center_position = float4_u<SI::m>{0.f, 0.f, 0.f, h.val};

  using res_t = decltype((func(current_position, current_position)) * volume);
  std::vector<res_t> LUT;

  for (int32_t i = 0; i < steps; ++i) {
    math::unit_assign<1>(current_position, H - step * static_cast<float>(i));
    res_t density{0.f};
    for (const auto &var : virtual_particles(center_position, volume, spacing)) {
      if (math::unit_get<1>(var) > 0._m){
        continue;
	  }
      density += volume * (func(current_position, var));
    }
    LUT.push_back(density);
  }
  return LUT;
}

template <typename T> auto printLUT(const std::string& name, const std::string& type, T LUT) {
  std::cout << "std::vector<" << type << "> " << name << "{ ";
  for (auto v : LUT) {
    std::cout << std::scientific << std::setprecision(std::numeric_limits<float>::digits10 + 1)
              << static_cast<float>(math::unit_get<1>(v).val) << "f, ";
  }
  std::cout << "};" << std::endl;
}

auto printLUT(const std::string& name, const std::vector<int32_t>& LUT) {
  std::cout << "std::vector<int32_t> " << name << "{ ";
  for (auto v : LUT) {
    std::cout << v << ", ";
  }
  std::cout << "};" << std::endl;
}
#include <config/config.h>
#include <fstream>
#ifdef _WIN32
#include <experimental/filesystem>
namespace fs = std::experimental::filesystem;
#else
#include <boost/filesystem.hpp>
namespace fs = boost::filesystem;
#endif

template <typename T> auto writeLUT(const std::string& name, const std::string& type, const std::vector<T>& LUT) {
  fs::path bin_dir(binaryDirectory);
  auto file = bin_dir / "config" / name;
  file.replace_extension("h");

  if (fs::exists(file)) {
    if (fs::exists(__FILE__)) {
      auto input_ts = fs::last_write_time(__FILE__);
      auto output_ts = fs::last_write_time(file);
      if (input_ts <= output_ts){
        return;
	  }
    }
  }
  std::cout << "Writing " << file.string() << std::endl;

  std::ofstream output(file.string());
  output << "std::vector<" << type << "> " << name << "{ ";
  int32_t ctr = 0;
  for (auto v : LUT) {
    output << std::scientific << std::setprecision(std::numeric_limits<float>::digits10 + 1)
           << static_cast<float>(math::unit_get<1>(v).val) << "f, " << (ctr++ % 5 == 0 ? "\n" : "");
  }
  output << "};" << std::endl;
  output.close();
}
auto writeLUT(const std::string& name, const std::vector<int32_t>& LUT) {
  fs::path bin_dir(binaryDirectory);
  auto file = bin_dir / "config" / name;
  file.replace_extension("h");
  if (fs::exists(file)) {
    if (fs::exists(__FILE__)) {
      auto input_ts = fs::last_write_time(__FILE__);
      auto output_ts = fs::last_write_time(file);
      if (input_ts <= output_ts){
        return;
	  }
    }
  }
  std::cout << "Writing " << file.string() << std::endl;
  std::ofstream output(file.string());
  int32_t ctr = 0;
  output << "std::vector<int32_t> " << name << "{ ";
  for (auto v : LUT) {
    output << v << ", " << (ctr++ % 10 == 0 ? "\n" : "");
  }
  output << "};" << std::endl;
  output.close();
}

int main() {
  std::cout << "Running LUT generation program" << std::endl;
  using pos_t = float4_u<SI::m>;
  bool test_spline4Kernel_LUT = false;
  bool test_spline4Gradient_LUT = false;
  bool test_estimateSpline4Gradient_LUT = false;
  bool test_testSpikyGradient_LUT = false;

  float_u<SI::volume> volume(1.f);
  auto h = supportFromVol(volume);
  auto H = h * Kernel<kernel_kind::spline4>::kernel_size<float>();

  auto position = float4_u<SI::m>{H.val / 2.f, 0.f, 0.f, h.val};

  float r = math::brentsMethod(
      [=](float_u<> radius) {
        float_u<> density = -1.f;
        for (const auto &var : virtual_particles(position, volume, math::getValue(radius))) {
          auto kernel = spline4_kernel(position, var);
          density += volume * kernel;
        }
        return density;
      },
      0.7f, 1.1f, 1e-7f, 10000);

  auto LUT = createLUT([](pos_t a, pos_t b) { return spline4_kernel(a, b); }, 2048, r);
  auto spline4gradientLUT = createLUT([](pos_t a, pos_t b) { return spline4_gradient(a, b); }, 2048, r);
  auto spikygradientLUT =
      createLUT([](pos_t a, pos_t b) { return PressureKernel<kernel_kind::spline4>::gradient(a, b); }, 2048, r);
  auto xbarLUT = createLUT([](pos_t a, pos_t b) { return b * spline4_kernel(a, b); }, 2048, r * 0.9f);
  auto ctrLUT = createCounterLUT(2048, r * 0.9f);

  // printLUT("boundaryLut" , "float", LUT);
  // printLUT("pressureLut", "float", spikygradientLUT);
  // printLUT("xbarLut", "float", xbarLUT);
  // printLUT("ctrLut", ctrLUT);

  writeLUT("boundaryLut", "float", LUT);
  writeLUT("pressureLut", "float", spline4gradientLUT);
  // writeLUT("pressureLut", "float", spikygradientLUT);
  writeLUT("xbarLut", "float", xbarLUT);
  writeLUT("ctrLut", ctrLUT);

  auto lookup = [&](float_u<SI::m> dist, float_u<SI::m> H, float_u<SI::volume> vol, auto LUT) {
    auto step = H / (static_cast<float>(LUT.size()) + 1) * 2.f;
    int32_t idx =
        math::clamp(static_cast<int32_t>(LUT.size()) - math::castTo<int32_t>(floor((dist.val + H.val) / step.val)), 0, static_cast<int32_t>(LUT.size()) - 1);
    return LUT[idx] / vol;
  };
  auto estimateGradient = [&](float_u<SI::m> dist, float_u<SI::m> H, float_u<SI::volume> vol, auto LUT) {
    auto step = H / (static_cast<float>(LUT.size()) + 1) * 2.f;
    int32_t idx =
        math::clamp(static_cast<int32_t>(LUT.size()) - math::castTo<int32_t>(floor((dist.val + H.val) / step.val)), 1, static_cast<int32_t>(LUT.size()) - 2);
    auto a = LUT[idx + 1] / vol;
    auto b = LUT[idx] / vol;
    auto ab = (b - a) / step;
    return ab;
  };
  auto lookupGradient = [&](float_u<SI::m> dist, float_u<SI::m> H, float_u<SI::volume> vol, auto LUT) {
    auto step = H / (static_cast<float>(LUT.size() + 1)) * 2.f;
    auto h_0 = support_from_volume(float_u<SI::volume>(1.f));
    auto h_c = support_from_volume(vol);

    auto ratio = h_0 / h_c;
    int32_t idx =
        math::clamp(static_cast<int32_t>(LUT.size()) - math::castTo<int32_t>(floor((dist.val + H.val) / step.val)), 0, static_cast<int32_t>(LUT.size()) - 1);
    return LUT[idx] / vol * ratio;
  };

  volume /= 10.f;
  h = support_from_volume(volume);
  H = h * kernelSize();
  position = float4_u<SI::m>{0.f, 2.f, 0.f, h.val};

  auto step = H / 32;
  // test kernel LUT
  if (test_spline4Kernel_LUT) {
    std::cout << std::endl
              << "Testing LUT for "
              << "spline4 kernel" << std::endl;
    for (float_u<SI::m> it = H; it >= -H - step; it -= step) {
      float_u<SI::recip_3<SI::m>> kernelSum{0.f};
      math::unit_assign<1>(position, it);
      for (const auto &var : virtual_particles(position, volume, r)) {
        if (math::unit_get<1>(var) > 0.0_m){
          continue;
		}
        kernelSum += kernel(position, var);
      }

      std::cout << std::setw(7) << std::fixed << std::setprecision(4) << it / H;
      std::cout << " -> " << std::fixed << std::setprecision(4) << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1) << math::unit_get<1>(kernelSum).val
                << " : " << std::setw(18) << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << lookup(it, H, volume, LUT).val << " : " << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << lookup(it, H, volume, LUT) / math::unit_get<1>(kernelSum) << std::endl;
    }
  }
  if (test_spline4Gradient_LUT) {
    std::cout << std::endl
              << "Testing LUT for "
              << "spline4 gradient" << std::endl;
    for (float_u<SI::m> it = H; it >= -H - step; it -= step) {
      float4_u<SI::multiply_ratios<SI::m, ratio<-4, 1>>> gradientSum{0.f, 0.f, 0.f, 0.f};
      math::unit_assign<1>(position, it);
      for (const auto &var : virtual_particles(position, volume, r)) {
        if (math::unit_get<1>(var) > 0._m){
          continue;
		}
        gradientSum += spline4_gradient(position, var);
      }

      std::cout << std::setw(7) << std::fixed << std::setprecision(4) << it / H;
      std::cout << " -> " << std::fixed << std::setprecision(4) << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << math::unit_get<1>(gradientSum).val << " : " << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << lookupGradient(it, H, volume, spline4gradientLUT).val << " : " << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << lookupGradient(it, H, volume, spline4gradientLUT) / math::unit_get<1>(gradientSum) << std::endl;
    }
  }
  if (test_estimateSpline4Gradient_LUT) {
    std::cout << std::endl
              << "Testing LUT for "
              << "estimating spline4 gradient" << std::endl;
    for (float_u<SI::m> it = H; it >= -H - step; it -= step) {
      float4_u<SI::SI_Unit<SI::m, ratio<-4, 1>>> gradientSum{0.f, 0.f, 0.f, 0.f};
      math::unit_assign<1>(position, it);
      for (const auto &var : virtual_particles(position, volume, r)) {
        if (math::unit_get<1>(var) > 0._m){
          continue;
		}
        gradientSum += spline4_gradient(position, var);
      }

      std::cout << std::setw(7) << std::fixed << std::setprecision(4) << it / H;
      std::cout << " -> " << std::fixed << std::setprecision(4) << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << math::unit_get<1>(gradientSum).val << " : " << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << estimateGradient(it, H, volume, LUT).val << " : " << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << estimateGradient(it, H, volume, LUT) / math::unit_get<1>(gradientSum) << std::endl;
    }
  }
  if (test_testSpikyGradient_LUT) {
    std::cout << std::endl
              << "Testing LUT for "
              << "spiky gradient" << std::endl;
    for (float_u<SI::m> it = H; it >= -H - step; it -= step) {
      float4_u<SI::SI_Unit<SI::m, ratio<-4, 1>>> gradientSum{0.f, 0.f, 0.f, 0.f};
      math::unit_assign<1>(position, it);
      for (const auto &var : virtual_particles(position, volume, r)) {
        if (math::unit_get<1>(var) > 0._m){
          continue;
		}
        gradientSum += PressureKernel<kernel_kind::spline4>::gradient(position, var);
      }

      std::cout << std::setw(7) << std::fixed << std::setprecision(4) << it / H;
      std::cout << " -> " << std::fixed << std::setprecision(4) << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << math::unit_get<1>(gradientSum).val << " : " << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << lookupGradient(it, H, volume, spikygradientLUT).val << " : " << std::setw(18)
                << std::setprecision(std::numeric_limits<long double>::digits10 + 1)
                << lookupGradient(it, H, volume, spikygradientLUT) / math::unit_get<1>(gradientSum) << std::endl;
    }
  }
  // getchar();
}
#else
struct mem {
  int *x, *y, *z, *w;
};

int main() {
  mem arrays;
  cache_arrays((m_01, x));
  cache_arrays((m_02, x), (m_03, y));

  float4_u<SI::m> test_x;
  float4_u<SI::velocity> test_v;
  float_u<SI::s> test_s;
  test_x + test_v *test_s;
  test_x = test_v * test_s;
  std::cout << SI::is_same_unit<decltype(test_x)::unit, decltype(test_v * test_s)::unit>::value << std::endl;

  using Tuple1 = decltype(test_x)::unit;
  using Tuple2 = decltype(test_v * test_s)::unit;

  std::cout << typeid(Tuple1).name() << std::endl;
  std::cout << typeid(Tuple2).name() << std::endl;
  std::cout << typeid(SI::flatten2<Tuple2>::type).name() << std::endl;

  // constexpr auto s1 = std::tuple_size<Tuple1>::value;
  // constexpr auto s2 = std::tuple_size<Tuple2>::value;

  return 0;
}
#endif


#define GLM_ENABLE_EXPERIMENTAL 1
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <Halide.h>
#include <fstream>
#include <glm/ext.hpp>
#include <glm/glm.hpp>
#include <halide_image_io.h>
#include <iostream>
#include <stdio.h>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <ImfRgbaFile.h>
#include <ImfStringAttribute.h>
#include <ImfMatrixAttribute.h>
#include <ImfArray.h>
#include <algorithm>
#include <ImfNamespace.h>
#include <tuple>
#include <math.h>

namespace IMF = OPENEXR_IMF_NAMESPACE;

using namespace IMF;
using namespace IMATH_NAMESPACE;

using std::string;
using std::stringstream;
using std::vector;

using namespace std;
using namespace Halide;
using namespace Halide::Tools;
using namespace Eigen;

using Vector4h = Matrix<Halide::Expr, 4, 1>;
using Matrix4h = Matrix<Halide::Expr, 4, 4>;

Var x, y, c;

struct Point
{
    float x;
    float y;
    float z;
    uint16_t r;
    uint16_t g;
    uint16_t b;
};

Matrix4h read4x4MatFromCSV(string csvFile)
{
    Matrix4h transfMat;
    ifstream in(csvFile);
    vector<float> floatVec;
    if (in)
    {
        string line;
        while (getline(in, line))
        {
            stringstream sep(line);
            string field;
            while (getline(sep, field, ','))
            {
                float val = stod(field);
                floatVec.push_back(val);
            }
        }
    }
    transfMat << floatVec[0], floatVec[1], floatVec[2], floatVec[3],
        floatVec[4], floatVec[5], floatVec[6], floatVec[7],
        floatVec[8], floatVec[9], floatVec[10], floatVec[11],
        floatVec[12], floatVec[13], floatVec[14], floatVec[15];
    return transfMat;
}

ostream &operator<<(ostream &os, const Matrix4h &m)
{
    os << m(0, 0) << " " << m(0, 1) << " " << m(0, 2) << " " << m(0, 3) << "\n"
       << m(1, 0) << " " << m(1, 1) << " " << m(1, 2) << " " << m(1, 3) << "\n"
       << m(2, 0) << " " << m(2, 1) << " " << m(2, 2) << " " << m(2, 3) << "\n"
       << m(3, 0) << " " << m(3, 1) << " " << m(3, 2) << " " << m(3, 3) << endl;
    return os;
}

void saveImageEXR(Expr result, int width, int height, const char fileName[])
{
    result = cast<float>(result);
    Target target = get_host_target();
    Func byteResult;
    byteResult(x, y) = result;
    Buffer<float> output(width, height);
    byteResult.compile_jit(target);
    byteResult.realize(output);

    Array2D<Rgba> pixels;
    pixels.resizeErase(height, width);
    for (int row{0}; row < height; row++)
    {
        for (int col{0}; col < width; col++)
        {
            pixels[row][col].r = output(col, row);
            pixels[row][col].g = output(col, row);
            pixels[row][col].b = output(col, row);
        }
    }
    RgbaOutputFile file(fileName, width, height, WRITE_RGBA);
    file.setFrameBuffer(&pixels[0][0], 1, width);
    file.writePixels(height);
}

void saveImageRaw(Expr result, size_t width, size_t height, const string &basename)
{
    Target target = get_host_target();
    Func byteResult;
    byteResult(x, y) = cast<uint8_t>(result);
    byteResult.compile_jit(target);
    Buffer<uint8_t> output(width, height);
    byteResult.realize(output);
    stringstream filename;
    filename << basename;
    save_image(output, filename.str());
}

void saveImage(Expr result, size_t width, size_t height, const string &basename)
{
    Target target = get_host_target();
    Func byteResult;
    byteResult(x, y) = cast<uint8_t>(clamp(result, 0.0f, 1.0f) * 255.0f);
    byteResult.compile_jit(target);
    Buffer<uint8_t> output(width, height);
    byteResult.realize(output);
    stringstream filename;
    filename << basename;
    save_image(output, filename.str());
}

ostream &operator<<(ostream &os, const Vector4h &v)
{
    os << v(0) << "\n"
       << v(1) << "\n"
       << v(2) << "\n"
       << v(3) << endl;
    return os;
}

Vector4h makeAxBzCLine(Vector4h a, Vector4h b)
{
    Vector4h result;
    result(0) = a(2) - b(2);
    result(1) = b(0) - a(0);
    result(2) = a(0) * b(2) - b(0) * a(2);
    result(3) = 1.0f;
    return result;
}

Vector4h cross3(Vector4h a, Vector4h b)
{
    Vector4h res;
    res(0) = a(1) * b(2) - a(2) * b(1);
    res(1) = a(2) * b(0) - a(0) * b(2);
    res(2) = a(0) * b(1) - a(1) * b(0);
    res(3) = 0;
    return res;
}

Expr dot3(Vector4h a, Vector4h b)
{
    Expr res;
    res = a(0) * b(0) + a(1) * b(1) + a(2) * b(2);
    return res;
}

tuple<Vector4h, Vector4h> pluckerLine(Vector4h p1, Vector4h p2)
{
    Vector4h l;
    Vector4h l_dash;
    l = p1(3) * p2 - p2(3) * p1;
    l(3) = 0;
    l_dash = cross3(p1, p2);
    return {l, l_dash};
}

Vector4h pluckerPlane(Vector4h l, Vector4h l_dash, Vector4h pluckPt)
{
    Vector4h u;
    u = -pluckPt(3) * l_dash + cross3(pluckPt, l);
    u(3) = dot3(pluckPt, l_dash);
    return u;
}

Vector4h intersectionLinePlane(Vector4h l, Vector4h l_dash, Vector4h plPlane)
{
    Vector4h intersection;
    intersection = -plPlane(3) * l + cross3(plPlane, l_dash);
    intersection(3) = dot3(plPlane, l);
    return intersection;
}

template <typename Function>
void writeBufferToXYZFile(Buffer<float> &buffer, Buffer<uint8_t> &color, string filename, string deliminator, Function condFunc)
{
    std::vector<Point> points;
    for (int j = 0; j < buffer.height(); j++)
    {
        for (int i = 0; i < buffer.width(); i++)
        {
            const auto x = buffer(i, j, 0);
            const auto y = buffer(i, j, 1);
            const auto z = buffer(i, j, 2);
            uint8_t r = color(i, j, 0);
            uint8_t g = color(i, j, 1);
            uint8_t b = color(i, j, 2);
            if (condFunc(x, y, z))
            {
                continue;
            }
            points.push_back({x, y, z, r, g, b});
        }
    }
    std::ofstream outFile;
    outFile.open(filename);
    outFile << points.size() << "\n";
    const auto d = deliminator;
    for (const auto &point : points)
    {
        outFile << point.x << d << point.y << d << point.z << d << point.r << d << point.g << d << point.b << "\n";
    }
}

void writeXYZ(Expr xs, Expr ys, Expr zs, int width, int height, string filename)
{
    Target target = get_host_target();
    Func result;
    result(x, y, c) = 0.0f;
    result(x, y, 0) = xs;
    result(x, y, 1) = ys;
    result(x, y, 2) = zs;

    Buffer<float> output(width, height, 3);
    result.compile_jit(target);
    result.realize(output);

    std::vector<Point> points;
    for (int j = 0; j < output.height(); j++)
    {
        for (int i = 0; i < output.width(); i++)
        {
            const auto x = output(i, j, 0);
            const auto y = output(i, j, 1);
            const auto z = output(i, j, 2);
            uint16_t zero = 0;
            points.push_back(Point{x, y, z, zero, zero, zero});
        }
    }
    std::ofstream outFile;
    outFile.open(filename);
    outFile << points.size() << "\n";
    const auto d = ";";
    for (const auto &point : points)
    {
        outFile << point.x << d << point.y << d << point.z << d << point.r << d << point.g << d << point.b << "\n";
    }
}

void debugImageEXR(Expr channel1Expr, Expr channel2Expr, Expr channel3Expr, int width, int height, const char fileName[])
{
    channel1Expr = cast<float>(channel1Expr);
    channel2Expr = cast<float>(channel2Expr);
    channel3Expr = cast<float>(channel3Expr);
    Target target = get_host_target();
    Func byteResult;
    byteResult(x, y, c) = 0.0f;
    byteResult(x, y, 0) = channel1Expr;
    byteResult(x, y, 1) = channel2Expr;
    byteResult(x, y, 2) = channel3Expr;
    Buffer<float> output(width, height, 3);
    byteResult.compile_jit(target);
    byteResult.realize(output);

    Array2D<Rgba> pixels;
    pixels.resizeErase(height, width);
    for (int row{0}; row < height; row++)
    {
        for (int col{0}; col < width; col++)
        {
            pixels[row][col].r = output(col, row, 0);
            pixels[row][col].g = output(col, row, 1);
            pixels[row][col].b = output(col, row, 2);
        }
    }
    RgbaOutputFile file(fileName, width, height, WRITE_RGBA);
    file.setFrameBuffer(&pixels[0][0], 1, width);
    file.writePixels(height);
}

int main(int argc, char **argv)
{
    const string LEFTCAM_SCAN_PNG = "images/scan/left-cam-tilt.png";
    const string RIGHTCAM_SCAN_PNG = "images/scan/right-cam-tilt.png";
    const string DEBUG_PNG = "images/debug/debug.png";
    const string CAM_MAT_PATH = "matrices/camera-matrix.csv";
    const string INV_CAM_MAT_PATH = "matrices/inv-camera-matrix.csv";
    const string TRANSF_LEFTCAM_PROJ_PATH = "matrices/transf-leftcam-proj.csv";
    const string TRANSF_RIGHTCAM_PROJ_PATH = "matrices/transf-rightcam-proj.csv";
    const string POINTCLOUD_LEFTCAM = "pointclouds/left-cam.txt";

    Buffer<uint8_t> leftCamScan = load_image(LEFTCAM_SCAN_PNG);
    Buffer<uint8_t> rightCamScan = load_image(RIGHTCAM_SCAN_PNG);
    const int HEIGHT = leftCamScan.height();
    const int WIDTH = leftCamScan.width();

    Expr leftCamScanRed = leftCamScan(x, y, 0);
    Expr leftCamScanFilterRed = leftCamScanRed * (leftCamScanRed > 80.0f);
    Expr rightCamScanRed = rightCamScan(x, y, 0);
    Expr rightCamScanFilterRed = rightCamScanRed * (rightCamScanRed > 80.0f);

    // read matrices
    Matrix4h camMat = read4x4MatFromCSV(CAM_MAT_PATH);
    Matrix4h invCamMat = read4x4MatFromCSV(INV_CAM_MAT_PATH);
    Matrix4h transfLeftCamProj = read4x4MatFromCSV(TRANSF_LEFTCAM_PROJ_PATH);
    Matrix4h transfRightCamProj = read4x4MatFromCSV(TRANSF_RIGHTCAM_PROJ_PATH);

    // define 3 points on projector to define laser scan
    Vector4h projPt1{0.0f, 0.0f, 0.0f, 1.0f};
    Vector4h projPt2{0.0f, 1.0f, 0.0f, 1.0f};
    Vector4h projPt3{0.0f, 0.0f, 1.0f, 1.0f};

    Vector4h projPt1_lcFrame = transfLeftCamProj * projPt1;
    Vector4h projPt2_lcFrame = transfLeftCamProj * projPt2;
    Vector4h projPt3_lcFrame = transfLeftCamProj * projPt3;

    Vector4h projPt1_rcFrame = transfRightCamProj * projPt1;
    Vector4h projPt2_rcFrame = transfRightCamProj * projPt2;
    Vector4h projPt3_rcFrame = transfRightCamProj * projPt3;

    const auto [projLine_lcFrame, projLineDash_lcFrame] = pluckerLine(projPt1_lcFrame, projPt2_lcFrame);
    Vector4h projPlane_lcFrame = pluckerPlane(projLine_lcFrame, projLineDash_lcFrame, projPt3_lcFrame);

    const auto [projLine_rcFrame, projLineDash_rcFrame] = pluckerLine(projPt1_rcFrame, projPt2_rcFrame);
    Vector4h projPlane_rcFrame = pluckerPlane(projLine_rcFrame, projLineDash_rcFrame, projPt3_rcFrame);

    Vector4h px_lc{x, y, 1.0f, 1.0f};
    Vector4h norm1_lc = invCamMat * px_lc;
    norm1_lc(3) = 1.0f;
    Vector4h norm2_lc = norm1_lc * 2.0f;
    norm2_lc(3) = 1.0f;

    Vector4h px_rc{x, y, 1.0f, 1.0f};
    Vector4h norm1_rc = invCamMat * px_rc;
    norm1_rc(3) = 1.0f;
    Vector4h norm2_rc = norm1_lc * 2.0f;
    norm2_rc(3) = 1.0f;

    const auto [pixelLine_lc, pixelLineDash_lc] = pluckerLine(norm1_lc, norm2_lc);
    Vector4h intersection_lc = intersectionLinePlane(pixelLine_lc, pixelLineDash_lc, projPlane_lcFrame);

    const auto [pixelLine_rc, pixelLineDash_rc] = pluckerLine(norm1_rc, norm2_rc);
    Vector4h intersection_rc = intersectionLinePlane(pixelLine_rc, pixelLineDash_rc, projPlane_rcFrame);

    intersection_lc *= leftCamScanFilterRed > 0;
    intersection_lc /= intersection_lc(3);
    debugImageEXR(intersection_lc(0), intersection_lc(1), intersection_lc(2), WIDTH, HEIGHT, "images/debug/intersection-lc.exr");
    writeXYZ(intersection_lc(0), intersection_lc(1), intersection_lc(2), WIDTH, HEIGHT, POINTCLOUD_LEFTCAM);

    intersection_rc *= rightCamScanFilterRed > 0;
    intersection_rc /= intersection_rc(3);
    writeXYZ(intersection_rc(0), intersection_rc(1), intersection_rc(2), WIDTH, HEIGHT, "pointclouds/right-cam.txt");

    //Func result;
    //result(x, y, c) = 0.0f;
    //result(x, y, 0) = intersection(0);
    //result(x, y, 1) = intersection(1);
    //result(x, y, 2) = intersection(2);

    //Buffer<float> output(WIDTH, HEIGHT, 3);
    //result.realize(output);

    //std::vector<Point> points;
    //for (int j = 0; j < output.height(); j++)
    //{
    //for (int i = 0; i < output.width(); i++)
    //{
    //const auto x = output(i, j, 0);
    //const auto y = output(i, j, 1);
    //const auto z = output(i, j, 2);
    //uint16_t zero = 0;
    //points.push_back(Point{x, y, z, zero, zero, zero});
    //}
    //}
    //std::ofstream outFile;
    //outFile.open(POINTCLOUD_LEFTCAM);
    //outFile << points.size() << "\n";
    //const auto d = ";";
    //for (const auto &point : points)
    //{
    //outFile << point.x << d << point.y << d << point.z << d << point.r << d << point.g << d << point.b << "\n";
    //}
    /*
    const float FOCAL_LEN = 36.1;
    const float PX_DIM = 10 * 10e-6;
    const bool POINT_CLOUD_GLOBAL_FRAME = true;
    const bool POINT_CLOUD_PROJ_FRAME = true;
    const string PROJECTOR_X_PNG = "images/x-val-img.png";
    const string TRANSF_MAT_WORLD_TO_PROJ_CSV = "matrices/transf-world-proj.csv";
    const string TRANSF_PROJ_CAM_CSV = "matrices/transf-proj-cam.csv";
    const string INV_CAM_MAT_CSV = "matrices/inv-cam-mat.csv";
    const string SAVE_DEPTH_IMAGE = "images/depth_img.png";
    const string SAVE_XYZ_PROJ = "pointclouds/proj.txt";
    const string SAVE_XYZ_WORLD = "pointclouds/world.txt";
    const string NO_PROJECTOR_PNG = "images/no-projector.png";

    Halide::Buffer<uint8_t> input = load_image(PROJECTOR_X_PNG);
    const int HEIGHT = input.height();
    const int WIDTH = input.width();

    Matrix4h InvK = read4x4MatFromCSV(INV_CAM_MAT_CSV);
    Matrix4h transMatProjToCam = read4x4MatFromCSV(TRANSF_PROJ_CAM_CSV);
    Matrix4h transMatWorldToProj = read4x4MatFromCSV(TRANSF_MAT_WORLD_TO_PROJ_CSV);

    cout << "Tranformation Matrix World to Projector" << endl;
    cout << transMatWorldToProj << endl;

    Vector4h px{x, y, 1.0f, 0.0f};

    Vector4h normCam1 = InvK * px;
    normCam1(3) = 1.0f;
    Vector4h normCam2 = 2 * normCam1;
    normCam2(3) = 1.0f;

    Vector4h pointCam1 = transMatProjToCam * normCam1;
    pointCam1 = pointCam1 / pointCam1(3);
    Vector4h pointCam2 = transMatProjToCam * normCam2;
    pointCam2 /= pointCam2(3);

    const auto [camLine, camLineDash] = pluckerLine(pointCam1, pointCam2);
    //Vector4h cameraLine = makeAxBzCLine(pointCam1, pointCam2);

    Vector4h pxProj{input(x, y) / 255.0f * 1920.0f, 0.0f, 1.0f, 0.0f};
    Vector4h normProj1 = InvK * pxProj;
    normProj1 = normProj1 / normProj1(2);
    normProj1(3) = 1.0f;
    Vector4h normProj2{0.0f, 0.0f, 0.0f, 1.0f};
    Vector4h normProj3{0.0f, 1.0f, 0.0f, 1.0f};

    const auto [projLine, projLineDash] = pluckerLine(normProj2, normProj3);
    Vector4h ProjPlane = pluckerPlane(projLine, projLineDash, normProj1);
    Vector4h intersection = intersectionLinePlane(camLine, camLineDash, ProjPlane);
    intersection = intersection / intersection(3);
    intersection *= (input(x, y) > 0); // does this make sense?

    Expr depth = intersection(3);

    saveImage(depth, input.width(), input.height(), SAVE_DEPTH_IMAGE);

    Halide::Buffer<uint8_t> color = load_image(NO_PROJECTOR_PNG);

    if (POINT_CLOUD_PROJ_FRAME)
    {
        Func result;
        result(x, y, c) = 0.0f;
        result(x, y, 0) = intersection(0);
        result(x, y, 1) = intersection(1);
        result(x, y, 2) = intersection(2);

        Buffer<float> output(input.width(), input.height(), 3);
        result.realize(output);
        writeBufferToXYZFile(output, color, SAVE_XYZ_PROJ, ";", conditionProjector);
    }
      */

    return 0;
}


/** \file   cmr_spirit_recon.cpp
    \brief  Implement some functionalities commonly used in cmr applications for spirit recon
    \author Hui Xue
*/

#include "cmr_spirit_recon.h"

#include "mri_core_utility.h"
#include "mri_core_spirit.h"
#include "hoNDArray_reductions.h"
#include "hoNDFFT.h"
#include "hoSPIRIT2DOperator.h"
#include "hoLsqrSolver.h"
#include "hoSPIRIT2DTDataFidelityOperator.h"
#include "hoWavelet2DTOperator.h"
#include "hoGdSolver.h"
#include <boost/make_shared.hpp>

namespace Gadgetron {

    template <typename T>
    void perform_spirit_recon_linear_2DT(const Gadgetron::hoNDArray<T>& kspace, size_t startE1, size_t endE1, const Gadgetron::hoNDArray<T>& kerIm,
                                const Gadgetron::hoNDArray<T>& kspaceInitial, Gadgetron::hoNDArray<T>& res, size_t iter_max, double iter_thres, bool print_iter)
    {
        try
        {
            size_t RO = kspace.get_size(0);
            size_t E1 = kspace.get_size(1);
            size_t CHA = kspace.get_size(2);
            size_t N = kspace.get_size(3);
            size_t S = kspace.get_size(4);

            size_t ref_N = kerIm.get_size(4);
            size_t ref_S = kerIm.get_size(5);

            if(startE1>=E1) startE1 = 0;
            if(endE1>=E1) startE1 = 0;
            if(startE1>endE1)
            {
                startE1 = 0;
                endE1 = E1 - 1;
            }

            long long num = N*S;

            hoNDArray<T> ker_Shifted(kerIm);
            Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->ifftshift2D(kerIm, ker_Shifted);

            hoNDArray<T> kspace_Shifted;
            kspace_Shifted = kspace;
            Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->ifftshift2D(kspace, kspace_Shifted);

            hoNDArray<T> kspace_initial_Shifted;
            bool hasInitial = false;
            if ( kspaceInitial.dimensions_equal(&kspace) )
            {
                kspace_initial_Shifted = kspaceInitial;
                Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->ifftshift2D(kspaceInitial, kspace_initial_Shifted);
                hasInitial = true;
            }

            res = kspace;

#ifdef USE_OMP
            int numThreads = (int)num;
            if (numThreads > omp_get_num_procs()) numThreads = omp_get_num_procs();
            GDEBUG_CONDITION_STREAM(print_iter, "numThreads : " << numThreads);
#endif // USE_OMP

            long long ii;

            std::vector<size_t> dim(3, 1);
            dim[0] = RO;
            dim[1] = E1;
            dim[2] = CHA;

#pragma omp parallel default(none) private(ii) shared(num, N, S, RO, E1, CHA, dim, startE1, endE1, ref_N, ref_S, kspace, res, kspace_Shifted, ker_Shifted, kspace_initial_Shifted, hasInitial, iter_max, iter_thres, print_iter) num_threads(numThreads) if(num>1)
            {
                boost::shared_ptr< hoSPIRIT2DOperator< T > > oper(new hoSPIRIT2DOperator< T >(&dim));
                hoSPIRIT2DOperator< T >& spirit = *oper;
                spirit.use_non_centered_fft_ = true;
                spirit.no_null_space_ = false;

                if (ref_N == 1 && ref_S == 1)
                {
                    boost::shared_ptr<hoNDArray<T> > ker(new hoNDArray< T >(RO, E1, CHA, CHA, ker_Shifted.begin()));
                    spirit.set_forward_kernel(*ker, false);
                }

                hoLsqrSolver< T > cgSolver;
                cgSolver.set_tc_tolerance((float)iter_thres);
                cgSolver.set_max_iterations( (unsigned int)iter_max);
                cgSolver.set_output_mode(print_iter ? hoLsqrSolver< T >::OUTPUT_VERBOSE : hoLsqrSolver< T >::OUTPUT_SILENT);
                cgSolver.set_encoding_operator(oper);

                hoNDArray< T > b(RO, E1, CHA);
                hoNDArray< T > unwarppedKSpace(RO, E1, CHA);

#pragma omp for
                for (ii = 0; ii < num; ii++)
                {
                    size_t s = ii / N;
                    size_t n = ii - s*N;

                    // check whether the kspace is undersampled
                    bool undersampled = false;
                    for (size_t e1 = startE1; e1 <= endE1; e1++)
                    {
                        if ((std::abs(kspace(RO / 2, e1, CHA - 1, n, s)) == 0)
                            && (std::abs(kspace(RO / 2, e1, 0, n, s)) == 0))
                        {
                            undersampled = true;
                            break;
                        }
                    }

                    T* pKpaceShifted = &(kspace_Shifted(0, 0, 0, n, s));
                    T* pRes = &(res(0, 0, 0, n, s));

                    if (!undersampled)
                    {
                        memcpy(pRes, pKpaceShifted, sizeof(T)*RO*E1*CHA);
                        continue;
                    }

                    long long kernelN = n;
                    if (kernelN >= (long long)ref_N) kernelN = (long long)ref_N - 1;

                    long long kernelS = s;
                    if (kernelS >= (long long)ref_S) kernelS = (long long)ref_S - 1;

                    boost::shared_ptr< hoNDArray< T > > acq(new hoNDArray< T >(RO, E1, CHA, pKpaceShifted));
                    spirit.set_acquired_points(*acq);

                    boost::shared_ptr< hoNDArray<T> > initialAcq;
                    if ( hasInitial )
                    {
                        initialAcq = boost::shared_ptr< hoNDArray<T> >(new hoNDArray<T>(RO, E1, CHA, &kspace_initial_Shifted(0, 0, 0, n, s)));
                        cgSolver.set_x0(initialAcq);
                    }
                    else
                    {
                        cgSolver.set_x0(acq);
                    }

                    if (ref_N == 1 && ref_S == 1)
                    {
                        spirit.compute_righ_hand_side(*acq, b);
                        cgSolver.solve(&unwarppedKSpace, &b);
                    }
                    else
                    {
                        T* pKer = &(ker_Shifted(0, 0, 0, kernelN, kernelS));
                        boost::shared_ptr<hoNDArray< T > > ker(new hoNDArray< T >(RO, E1, CHA, CHA, pKer));
                        spirit.set_forward_kernel(*ker, false);

                        spirit.compute_righ_hand_side(*acq, b);
                        cgSolver.solve(&unwarppedKSpace, &b);
                    }

                    // restore the acquired points
                    spirit.restore_acquired_kspace(*acq, unwarppedKSpace);
                    memcpy(pRes, unwarppedKSpace.begin(), unwarppedKSpace.get_number_of_bytes());
                }
            }

            Gadgetron::hoNDFFT<typename realType<T>::Type>::instance()->fftshift2D(res, kspace_Shifted);
            res = kspace_Shifted;
        }
        catch(...)
        {
            GADGET_THROW("Exceptions happened in perform_spirit_recon_linear_2DT(...) ... ");
        }
    }

    template EXPORTCMR void perform_spirit_recon_linear_2DT(const Gadgetron::hoNDArray< std::complex<float> >& kspace, size_t startE1, size_t endE1, const Gadgetron::hoNDArray< std::complex<float> >& kerIm, const Gadgetron::hoNDArray< std::complex<float> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<float> >& res, size_t iter_max, double iter_thres, bool print_iter);
    template EXPORTCMR void perform_spirit_recon_linear_2DT(const Gadgetron::hoNDArray< std::complex<double> >& kspace, size_t startE1, size_t endE1, const Gadgetron::hoNDArray< std::complex<double> >& kerIm, const Gadgetron::hoNDArray< std::complex<double> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<double> >& res, size_t iter_max, double iter_thres, bool print_iter);

    // ---------------------------------------------------------------------

    template <typename T>
    class sCB : public hoGdSolverCallBack< hoNDArray< T >, hoWavelet2DTOperator< T > >
    {
    public:
        typedef hoGdSolverCallBack< hoNDArray<T>, hoWavelet2DTOperator<T> > BaseClass;

        sCB() : BaseClass() {}
        virtual ~sCB() {}

        void execute(const hoNDArray<T>& b, hoNDArray<T>& x)
        {
            typedef hoSPIRIT2DTDataFidelityOperator<T> SpiritOperType;
            SpiritOperType* pOper = dynamic_cast<SpiritOperType*> (this->solver_->oper_system_);
            pOper->restore_acquired_kspace(x);
        }
    };

    template <typename T>
    void perform_spirit_recon_non_linear_2DT(const Gadgetron::hoNDArray<T>& kspace, const Gadgetron::hoNDArray<T>& kerIm,
                                            const Gadgetron::hoNDArray<T>& coil_map, const Gadgetron::hoNDArray<T>& kspaceLinear, Gadgetron::hoNDArray<T>& res,
                                            size_t iter_max, double iter_thres, double data_fidelity_lamda, double image_reg_lamda, double reg_N_weighting_ratio,
                                            bool reg_use_coil_sen_map, bool reg_with_approx_coeff, const std::string& wav_name, bool print_iter)
    {
        size_t RO = kspace.get_size(0);
        size_t E1 = kspace.get_size(1);
        size_t CHA = kspace.get_size(2);
        size_t N = kspace.get_size(3);
        size_t S = kspace.get_size(4);

        size_t ref_N = kerIm.get_size(4);
        size_t ref_S = kerIm.get_size(5);

        res = kspace;

        size_t s;

        for (s=0; s<S; s++)
        {
            boost::shared_ptr< hoNDArray< T > > coilMap;

            bool hasCoilMap = false;
            if (coil_map.get_size(0) == RO && coil_map.get_size(1) == E1 && coil_map.get_size(2)==CHA)
            {
                if (ref_N < N)
                {
                    coilMap = boost::shared_ptr< hoNDArray< T > >(new hoNDArray< T >(RO, E1, CHA, const_cast<T*>(coil_map.begin()) ));
                }
                else
                {
                    coilMap = boost::shared_ptr< hoNDArray< T > >(new hoNDArray< T >(RO, E1, CHA, ref_N, const_cast<T*>(coil_map.begin()) ));
                }

                hasCoilMap = true;
            }

            size_t s_used = s;
            if(s_used>=ref_S) s_used = ref_S;

            boost::shared_ptr<hoNDArray< T > > ker(new hoNDArray< T >(RO, E1, CHA, CHA, ref_N, const_cast<T*>(kerIm.begin())+s_used*RO*E1*CHA*CHA*ref_N) );
            boost::shared_ptr<hoNDArray< T > > acq(new hoNDArray< T >(RO, E1, CHA, N, const_cast<T*>(kspace.begin())+s*RO*E1*CHA*N) );
            hoNDArray< T > kspaceInitial(RO, E1, CHA, N, const_cast<T*>(kspaceLinear.begin())+s*RO*E1*CHA*N );
            hoNDArray< T > res2DT(RO, E1, CHA, N, res.begin()+s*RO*E1*CHA*N );

            if (data_fidelity_lamda > 0)
            {
                GDEBUG_STREAM("Start the NonLinear SPIRIT data fidelity iteration - regularization strength : "
                    << image_reg_lamda
                    << " - number of iteration : "                      << iter_max
                    << " - proximity across cha : "                     << false
                    << " - redundant dimension weighting ratio : "      << reg_N_weighting_ratio
                    << " - using coil sen map : "                       << reg_use_coil_sen_map
                    << " - with approx coeff : "                        << reg_with_approx_coeff
                    << " - iter thres : "                               << iter_thres
                    << " - wavelet name : "                             << wav_name
                    );

                typedef hoGdSolver< hoNDArray< T >, hoWavelet2DTOperator< T > > SolverType;
                SolverType solver;
                solver.iterations_ = iter_max;
                solver.set_output_mode(print_iter ? SolverType::OUTPUT_VERBOSE : SolverType::OUTPUT_SILENT);
                solver.grad_thres_ = iter_thres;
                solver.proximal_strength_ratio_ = image_reg_lamda;

                boost::shared_ptr< hoNDArray< T > > x0 = boost::make_shared< hoNDArray< T > >(kspaceInitial);
                solver.set_x0(x0);

                // parallel imaging term
                std::vector<size_t> dims;
                acq->get_dimensions(dims);
                hoSPIRIT2DTDataFidelityOperator< T > spirit(&dims);
                spirit.set_forward_kernel(*ker, false);
                spirit.set_acquired_points(*acq);

                // image reg term
                hoWavelet2DTOperator< T > wav3DOperator(&dims);
                wav3DOperator.set_acquired_points(*acq);
                wav3DOperator.scale_factor_first_dimension_ = 1;
                wav3DOperator.scale_factor_second_dimension_ = 1;
                wav3DOperator.scale_factor_third_dimension_ = reg_N_weighting_ratio;
                wav3DOperator.with_approx_coeff_ = reg_with_approx_coeff;
                wav3DOperator.change_coeffcients_third_dimension_boundary_ = true;
                wav3DOperator.proximity_across_cha_ = false;
                wav3DOperator.no_null_space_ = true;
                wav3DOperator.input_in_kspace_ = true;
                wav3DOperator.select_wavelet(wav_name);

                if (reg_use_coil_sen_map && hasCoilMap)
                {
                    wav3DOperator.coil_map_ = *coilMap;
                }

                // set operators

                solver.oper_system_ = &spirit;
                solver.oper_reg_ = &wav3DOperator;

                solver.solve(*acq, res2DT);
            }
            else
            {
                GDEBUG_STREAM("Start the NonLinear SPIRIT iteration with regularization strength : "
                    << image_reg_lamda
                    << " - number of iteration : " << iter_max
                    << " - proximity across cha : " << false
                    << " - redundant dimension weighting ratio : " << reg_N_weighting_ratio
                    << " - using coil sen map : " << reg_use_coil_sen_map
                    << " - with approx coeff : "  << reg_with_approx_coeff
                    << " - iter thres : " << iter_thres
                    << " - wavelet name : " << wav_name
                    );

                typedef hoGdSolver< hoNDArray< T >, hoWavelet2DTOperator< T > > SolverType;
                SolverType solver;
                solver.iterations_ = iter_max;
                solver.set_output_mode(print_iter ? SolverType::OUTPUT_VERBOSE : SolverType::OUTPUT_SILENT);
                solver.grad_thres_ = iter_thres;
                solver.proximal_strength_ratio_ = image_reg_lamda;

                boost::shared_ptr< hoNDArray< T > > x0 = boost::make_shared< hoNDArray< T > >(kspaceInitial);
                solver.set_x0(x0);

                // parallel imaging term
                std::vector<size_t> dims;
                acq->get_dimensions(dims);

                hoSPIRIT2DTOperator< T > spirit(&dims);
                spirit.set_forward_kernel(*ker, false);
                spirit.set_acquired_points(*acq);
                spirit.no_null_space_ = true;
                spirit.use_non_centered_fft_ = false;

                // image reg term
                std::vector<size_t> dim;
                acq->get_dimensions(dim);

                hoWavelet2DTOperator< T > wav3DOperator(&dim);
                wav3DOperator.set_acquired_points(*acq);
                wav3DOperator.scale_factor_first_dimension_ = 1;
                wav3DOperator.scale_factor_second_dimension_ = 1;
                wav3DOperator.scale_factor_third_dimension_ = reg_N_weighting_ratio;
                wav3DOperator.with_approx_coeff_ = reg_with_approx_coeff;
                wav3DOperator.change_coeffcients_third_dimension_boundary_ = true;
                wav3DOperator.proximity_across_cha_ = false;
                wav3DOperator.no_null_space_ = true;
                wav3DOperator.input_in_kspace_ = true;
                wav3DOperator.select_wavelet(wav_name);

                if (reg_use_coil_sen_map && hasCoilMap)
                {
                    wav3DOperator.coil_map_ = *coilMap;
                }

                // set operators
                solver.oper_system_ = &spirit;
                solver.oper_reg_ = &wav3DOperator;

                // set call back
                sCB<T> cb;
                cb.solver_ = &solver;
                solver.call_back_ = &cb;

                hoNDArray< T > b(kspaceInitial);
                Gadgetron::clear(b);

                solver.solve(b, res2DT);
                // if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "spirit_nl_2DT_res");

                spirit.restore_acquired_kspace(*acq, res2DT);

                // if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "spirit_nl_2DT_res_restored");
            }
        }
    }

    template EXPORTCMR void perform_spirit_recon_non_linear_2DT(const Gadgetron::hoNDArray< std::complex<float> >& kspace, const Gadgetron::hoNDArray< std::complex<float> >& kerIm,
        const Gadgetron::hoNDArray< std::complex<float> >& coilMap, const Gadgetron::hoNDArray< std::complex<float> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<float> >& res,
        size_t iter_max, double iter_thres, double data_fidelity_lamda, double image_reg_lamda, double reg_N_weighting_ratio,
        bool reg_use_coil_sen_map, bool reg_with_approx_coeff, const std::string& wav_name, bool print_iter);

    template EXPORTCMR void perform_spirit_recon_non_linear_2DT(const Gadgetron::hoNDArray< std::complex<double> >& kspace, const Gadgetron::hoNDArray< std::complex<double> >& kerIm,
        const Gadgetron::hoNDArray< std::complex<double> >& coilMap, const Gadgetron::hoNDArray< std::complex<double> >& kspaceInitial, Gadgetron::hoNDArray< std::complex<double> >& res,
        size_t iter_max, double iter_thres, double data_fidelity_lamda, double image_reg_lamda, double reg_N_weighting_ratio,
        bool reg_use_coil_sen_map, bool reg_with_approx_coeff, const std::string& wav_name, bool print_iter);
}