/* Distributed under the Apache License, Version 2.0.
   See accompanying NOTICE file for details.*/

#include "ITKUtils.h"
#include "LBMUtils.h"

namespace LBM
{
namespace ITK
{

bool Compute( Configuration const& cfg, LBM::Run &lbm)
{
  try
  {
    if (cfg.image_filename.empty())
      return false;

    lbm.Info("Processing "+ cfg.image_filename);
    itk::ImageIOBase::Pointer image_io =
      itk::ImageIOFactory::CreateImageIO(cfg.image_filename.c_str(), itk::ImageIOFactory::FileModeEnum::ReadMode);
    if (image_io == nullptr)
    {
      return false;
    }
    image_io->SetFileName(cfg.image_filename);
    image_io->ReadImageInformation();
    // Get the pixel type
    using IOPixelType = itk::ImageIOBase::IOPixelType;
    const IOPixelType pixel_type = image_io->GetPixelType();
    lbm.Info("Pixel Type is "+itk::ImageIOBase::GetPixelTypeAsString(pixel_type));
    // Get the component type
    using IOComponentType = itk::ImageIOBase::IOComponentType;
    const IOComponentType component_type = image_io->GetComponentType();
    lbm.Info("Component Type is "+image_io->GetComponentTypeAsString(component_type));
    // Print out some more information
    const unsigned int image_dimension = image_io->GetNumberOfDimensions();
    lbm.Info("Image Dimension is "+std::to_string(image_dimension));
    if (image_dimension != 3)
    {
      lbm.Error("Currently, Only 3D images are supported");
      return false;
    }

    // This filter handles all types on input, but only produces
    // signed types
    switch (component_type)
    {
    case itk::ImageIOBase::UCHAR:
    {
      using PixelType = unsigned char;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::CHAR:
    {
      using PixelType = char;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::USHORT:
    {
      using PixelType = unsigned short;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::SHORT:
    {
      using PixelType = short;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::UINT:
    {
      using PixelType = unsigned int;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::INT:
    {
      using PixelType = int;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::ULONG:
    {
      using PixelType = unsigned long;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::LONG:
    {
      using PixelType = long;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::FLOAT:
    {
      using PixelType = float;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::DOUBLE:
    {
      using PixelType = double;
      using ImageType = itk::Image<PixelType, 3>;
      LBM::ITK::Run<PixelType, ImageType> lbm_itk(cfg,lbm);
      return lbm_itk.Compute();
    }
    case itk::ImageIOBase::UNKNOWNCOMPONENTTYPE:
    default:
      lbm.Error("Unknown input image pixel component type: " + itk::ImageIOBase::GetComponentTypeAsString(component_type));
      return false;
      break;
    }
  }
  catch (itk::ExceptionObject& excep)
  {
    std::stringstream ss;
    ss << "Unable to read file : " << cfg.image_filename << " : with exception :" << excep;
    lbm.Error(ss);
  }
  return false;
}

template<LBMITK_TEMPLATE>
LBM::ITK::Run<LBMITK_TYPES>::Run(Configuration const& cfg, LBM::Run& lbm) : m_cfg(cfg), m_lbm(lbm)
{

}

template<LBMITK_TEMPLATE>
LBM::ITK::Run<LBMITK_TYPES>::~Run()
{

}

template<LBMITK_TEMPLATE>
bool LBM::ITK::Run<LBMITK_TYPES>::Compute()
{
  if (!ReadImage(m_cfg.image_filename))
    return false;
  if (!NormalizeImage())
    return false;
  if (!CropImage())
    return false;
  if (!DownsampleImage())
    return false;
  if (!OrientImageToX())
    return false;

  m_lbm.Info("Processing boundary labels...");
  auto origin = m_image->GetOrigin();
  auto spacing = m_image->GetSpacing();
  auto size = m_image->GetLargestPossibleRegion().GetSize();

  // Convert image to an LBM grid
  m_lbm.in.dimensions[0] = size[0];
  m_lbm.in.dimensions[1] = size[1];
  m_lbm.in.dimensions[2] = size[2];
  m_lbm.in.grid_spacing = spacing[0] * m_cfg.grid_spacing_conversion;
  itk::ImageRegionIteratorWithIndex<ImageType>
    m2iIt(m_image, m_image->GetLargestPossibleRegion());
  
  for (;!m2iIt.IsAtEnd();++m2iIt)
    m_lbm.in.source_labels.push_back(m2iIt.Get());

  if (!m_lbm.Compute())
    return false;
  return true;
}

template<LBMITK_TEMPLATE>
bool LBM::ITK::Run<LBMITK_TYPES>::ReadImage(std::string const& filename)
{
  if (filename.empty())
    return false;

  m_lbm.Info("Reading image...");
  m_image = ImageType::New();
  using ImageReaderType = itk::ImageFileReader<ImageType>;
  auto reader = ImageReaderType::New();
  reader->SetFileName(filename);
  try
  {
    reader->Update();
  }
  catch (itk::ExceptionObject& e)
  {
    m_ss<<"Error reading image : "<<e.what();
    m_lbm.Error(m_ss);
    return false;
  }
  m_image->Graft(reader->GetOutput());

  // Print out what we have
  m_ss << "Resulting Model properties : \n";
  m_ss << "  Origin  : " << m_image->GetOrigin() << "\n";
  m_ss << "  Size    : " << m_image->GetLargestPossibleRegion().GetSize() << "\n";
  m_ss << "  Spacing : " << m_image->GetSpacing();
  m_lbm.Info(m_ss);
  return true;
}

template<LBMITK_TEMPLATE>
bool LBM::ITK::Run<LBMITK_TYPES>::WriteImage(std::string const& filename)
{
  if (filename.empty())
    return false;
  try
  {
    using ImageWriterType = itk::ImageFileWriter<ImageType>;
    auto writer = ImageWriterType::New();
    writer->SetInput(m_image);
    writer->SetFileName(filename);
    writer->Update();
  }
  catch (itk::ExceptionObject& e)
  {
    m_ss << "Error writing image : "<< e.what();
    m_lbm.Error(m_ss);
    return false;
  }
  return true;
}

template<LBMITK_TEMPLATE>
bool LBM::ITK::Run<LBMITK_TYPES>::NormalizeImage()
{
  try
  {
    m_lbm.Info("Normalizing image spacing...");
    auto interp = InterpType::New();
    auto resample_filter = ResampleFilterType::New();

    auto origin = m_image->GetOrigin();
    auto spacing = m_image->GetSpacing();
    auto size = m_image->GetLargestPossibleRegion().GetSize();

    resample_filter->SetInput(m_image);
    resample_filter->SetInterpolator(interp);
    resample_filter->SetDefaultPixelValue(0);
    resample_filter->SetOutputOrigin(origin);
    resample_filter->SetOutputDirection(m_image->GetDirection());

    // Resize everything to the axis with the smallest spacing
    if (spacing[0] <= spacing[1] && spacing[0] <= spacing[2])
    {
      size[1] = size[1] * spacing[1] / spacing[0];
      spacing[1] = spacing[0];
      size[2] = size[2] * spacing[2] / spacing[0];
      spacing[2] = spacing[0];
    }
    else if (spacing[1] <= spacing[0] && spacing[1] <= spacing[2])
    {
      size[0] = size[0] * spacing[0] / spacing[1];
      spacing[0] = spacing[1];
      size[2] = size[2] * spacing[2] / spacing[1];
      spacing[2] = spacing[1];
    }
    else
    {
      size[0] = size[0] * spacing[0] / spacing[2];
      spacing[0] = spacing[2];
      size[1] = size[1] * spacing[1] / spacing[2];
      spacing[1] = spacing[2];
    }

    resample_filter->SetSize(size);
    resample_filter->SetOutputSpacing(spacing);
    resample_filter->Update();
    m_image = resample_filter->GetOutput();

    // Print out what we did
    m_ss << "Resulting Model properties : \n";
    m_ss << "  Origin  : " << m_image->GetOrigin() << "\n";
    m_ss << "  Size    : " << m_image->GetLargestPossibleRegion().GetSize() << "\n";
    m_ss << "  Spacing : " << m_image->GetSpacing();
    m_lbm.Info(m_ss);
  }
  catch (itk::ExceptionObject & excep)
  {
    m_ss << "Unable to normalize image with exception :" << excep;
    m_lbm.Error(m_ss);
    return false;
  }
  return true;
}

template<LBMITK_TEMPLATE>
bool LBM::ITK::Run<LBMITK_TYPES>::CropImage()
{
  try
  {
    m_lbm.Info("Cropping boundary...");
    auto i2m_filter = LabelImageToMapFilterType::New();

    i2m_filter->SetInput(m_image);
    i2m_filter->SetBackgroundValue(0);
    i2m_filter->Update();

    auto crop = CropType::New();
    crop->SetInput(i2m_filter->GetOutput());
    CropSize border;
    border[0] = 5;
    border[1] = 5;
    border[2] = 5;
    crop->SetCropBorder(border);
    crop->Update();

    auto m2i_filter = LabelMapToImageFilterType::New();
    m2i_filter->SetInput(crop->GetOutput());
    m2i_filter->Update();
    m_image = m2i_filter->GetOutput();

    // Get our new origin
    auto origin = m_image->GetOrigin();
    auto region = m_image->GetRequestedRegion();
    m_image->TransformIndexToPhysicalPoint(region.GetIndex(), origin);

    // Resample to this new origin/grid
    auto interp = InterpType::New();
    auto resample_filter = ResampleFilterType::New();
    auto spacing = m_image->GetSpacing();
    auto size = m_image->GetLargestPossibleRegion().GetSize();

    resample_filter->SetInput(m_image);
    resample_filter->SetInterpolator(interp);
    resample_filter->SetOutputOrigin(origin);
    resample_filter->SetOutputDirection(m_image->GetDirection());
    resample_filter->SetSize(size);
    resample_filter->SetOutputSpacing(spacing);
    resample_filter->Update();
    m_image = resample_filter->GetOutput();

    // Print out what we did
    m_ss << "Resulting Model properties : \n";
    m_ss << "  Origin  : " << m_image->GetOrigin() << "\n";
    m_ss << "  Size    : " << m_image->GetLargestPossibleRegion().GetSize() << "\n";
    m_ss << "  Spacing : " << m_image->GetSpacing();
    m_lbm.Info(m_ss);
  }
  catch (itk::ExceptionObject & excep)
  {
    m_ss << "Unable to crop image with exception :" << excep;
    m_lbm.Error(m_ss);
    return false;
  }
  return true;
}

template<LBMITK_TEMPLATE>
bool LBM::ITK::Run<LBMITK_TYPES>::DownsampleImage()
{
  try
  {
    auto origin = m_image->GetOrigin();
    auto spacing = m_image->GetSpacing();
    auto size = m_image->GetLargestPossibleRegion().GetSize();

    double downsample_fraction = 0;
    if (m_cfg.downsample_type == LBM::ITK::DownsampleType::Length)
    {
      if (m_cfg.downsample_value <= spacing[0])
      {
        // Nothing to do
        m_lbm.Info("Requested downsample length smaller that image spacing");
        m_lbm.Info("Running at resolution of provided image ("+std::to_string(spacing[0])+")");
        return true;
      }
      downsample_fraction = spacing[0] / m_cfg.downsample_value;
    }
    else
      downsample_fraction = m_cfg.downsample_value;
    if (downsample_fraction == 0)
    {
      m_lbm.Info("Downsampling not requested");
      return true;
    }

    m_lbm.Info("Downsampling image...");
    auto interp = InterpType::New();
    auto resample_filter = ResampleFilterType::New();

    resample_filter->SetInput(m_image);
    resample_filter->SetInterpolator(interp);
    resample_filter->SetOutputOrigin(origin);
    resample_filter->SetOutputDirection(m_image->GetDirection());

    Size temp_size;
    temp_size[0] = size[0] * downsample_fraction;
    temp_size[1] = size[1] * downsample_fraction;
    temp_size[2] = size[2] * downsample_fraction;

    Size new_size;
    new_size[0] = temp_size[0] - (temp_size[0] % 16);
    new_size[1] = temp_size[1] - (temp_size[1] % 16);
    new_size[2] = temp_size[2] - (temp_size[2] % 16);

    Spacing new_spacing;
    new_spacing[0] = spacing[0] * (static_cast<float>(size[0]) / static_cast<float>(new_size[0]));
    new_spacing[1] = spacing[1] * (static_cast<float>(size[1]) / static_cast<float>(new_size[1]));
    new_spacing[2] = spacing[2] * (static_cast<float>(size[2]) / static_cast<float>(new_size[2]));

    resample_filter->SetSize(new_size);
    resample_filter->SetOutputSpacing(new_spacing);
    resample_filter->Update();
    m_image = resample_filter->GetOutput();

    // Print out what we did
    m_ss << "Resulting Model properties : \n";
    m_ss << "  Origin  : " << m_image->GetOrigin() << "\n";
    m_ss << "  Size    : " << m_image->GetLargestPossibleRegion().GetSize() << "\n";
    m_ss << "  Spacing : " << m_image->GetSpacing();
    m_lbm.Info(m_ss);
  }
  catch (itk::ExceptionObject & excep)
  {
    m_ss << "Unable to downsample image with exception :" << excep;
    m_lbm.Error(m_ss);
    return false;
  }
  return true;
}

template<LBMITK_TEMPLATE>
bool LBM::ITK::Run<LBMITK_TYPES>::OrientImageToX()
{
  try
  {
    m_lbm.Info("Orienting image...");
    // Rotate image for long axis to be along the X-axis
    auto interp = InterpType::New();
    auto resample_filter = ResampleFilterType::New();
    auto transform = TransformType::New();

    auto origin = m_image->GetOrigin();
    auto spacing = m_image->GetSpacing();
    auto size = m_image->GetLargestPossibleRegion().GetSize();

    resample_filter->SetInput(m_image);
    resample_filter->SetInterpolator(interp);
    resample_filter->SetDefaultPixelValue(0);

    // Translate image to origin
    TransformType::OutputVectorType trans1;
    trans1[0] = -1 * (origin[0] + (spacing[0] * size[0] * 0.5));
    trans1[1] = -1 * (origin[1] + (spacing[1] * size[1] * 0.5));
    trans1[2] = -1 * (origin[2] + (spacing[2] * size[2] * 0.5));
    transform->Translate(trans1);
    // Rotate 90 degrees around Y, then 90 degrees around Z
    // Which would be the x axis after we do the Y rotate...
    TransformType::OutputVectorType yAxis, zAxis;
    yAxis[0] = 0; yAxis[1] = 1; yAxis[2] = 0;
    zAxis[0] = 1; zAxis[1] = 0; zAxis[2] = 0;
    transform->Rotate3D(yAxis, 1.5708, false);
    transform->Rotate3D(zAxis, 1.5708, false);
    // Adjust our size for this new orientation
    Size new_size;
    new_size[0] = size[1];
    new_size[1] = size[2];
    new_size[2] = size[0];
    Spacing new_spacing;
    new_spacing[0] = spacing[1];
    new_spacing[1] = spacing[2];
    new_spacing[2] = spacing[0];
    //// Now translate back
    TransformType::OutputVectorType trans2;
    trans2[0] = -1 * trans1[0];
    trans2[1] = -1 * trans1[1];
    trans2[2] = -1 * trans1[2];
    transform->Translate(trans2);
    resample_filter->SetTransform(transform);
    // Adjust our origin based on our rotation
    Origin new_origin;
    new_origin[0] = trans2[0] - (new_spacing[0] * new_size[0] * 0.5);
    new_origin[1] = trans2[1] - (new_spacing[1] * new_size[1] * 0.5);
    new_origin[2] = trans2[2] - (new_spacing[2] * new_size[2] * 0.5);

    resample_filter->SetSize(new_size);
    resample_filter->SetOutputOrigin(new_origin);
    resample_filter->SetOutputSpacing(new_spacing);
    resample_filter->SetOutputDirection(m_image->GetDirection());
    // Execute our operations
    resample_filter->Update();
    m_image = resample_filter->GetOutput();

    // Print out what we did
    m_ss << "Resulting Model properties : \n";
    m_ss << "  Origin  : " << m_image->GetOrigin() << "\n";
    m_ss << "  Size    : " << m_image->GetLargestPossibleRegion().GetSize() << "\n";
    m_ss << "  Spacing : " << m_image->GetSpacing();
    m_lbm.Info(m_ss);
  }
  catch (itk::ExceptionObject& excep)
  {
    m_ss << "Unable to rotate imate : with exception :" << excep;
    m_lbm.Error(m_ss);
    return false;
  }
  return true;
}

bool LoadConfiguration(std::string const& config_file, LBM::Run& lbm, Configuration& itk_cfg)
{
  if (!LBM::LoadConfiguration(config_file, lbm))
    return false;
  try
  {
    std::string input_dir = config_file.substr(0,config_file.find_last_of("/"));

    std::string line;
    std::istringstream iss;

    /////////////////
    // CONFIG FILE //
    /////////////////
    std::ifstream cFile(config_file, std::ios::in);

    float value;
    std::string name;
    while (std::getline(cFile, line))
    {
      if (line.empty() || line.find("#") != std::string::npos)// Skip a comment
        continue;
      if (line.find("ITK") == std::string::npos)// Only load an ITK configuration property
        continue;
      // Look for a valid ITK extension
      if (line.find(".nrrd") != std::string::npos)
      {
        itk_cfg.image_filename = input_dir+"/"+line.substr(0,line.find(".nrrd")+5);
        continue;
      }
      iss.clear();
      iss.str(line);
      if (!(iss >> value >> name))
      {
        lbm.Error("Unable to parse line from: " + config_file);
        lbm.Error("  " + line);
        return false;
      }
      if (name.compare("downsample_type") == 0)
      {
        if (value == 0)
          itk_cfg.downsample_type = DownsampleType::Fraction;
        else
          itk_cfg.downsample_type = DownsampleType::Length;
        continue;
      }
      if (name.compare("downsample") == 0)
      {
        itk_cfg.downsample_value = value; continue;
      }
      if (name.compare("grid_spacing_conversion") == 0)
      {
        itk_cfg.grid_spacing_conversion = value; continue;
      }
    }
  }
  catch (std::exception ex)
  {
    lbm.Error(ex.what());
    return false;
  }
  return true;
}


}
}
