/*//////////////////////////////////////////////////////////////////////////////
// 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 >

  ------------------------------------------------------------------
  This library is free software;  you  can  redistribute  it  and/or
  modify it  under the  terms  of  the  GNU  Lesser  General  Public
  License as published  by  the  Free  Software  Foundation;  either
  version 3.0 of the License, or (at your option) any later version.

  This library is  distributed in  the hope that it will  be useful,
  but  WITHOUT ANY WARRANTY;  without  even  the implied warranty of
  MERCHANTABILITY  or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of  the GNU Lesser General  Public
  License along with this library.

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

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <vtkImageExport.h>
#include <vtkImageImport.h>
#include <vtkSmartPointer.h>

#include <vtkFloatArray.h>
#include <vtkPointData.h>

#include <vtkGenericDataObjectReader.h>
#include <vtkImageShiftScale.h>

#include <vtkColorTransferFunction.h>
#include <vtkPiecewiseFunction.h>
#include <vtkSmartVolumeMapper.h>
#include <vtkVolumeProperty.h>
#include <vtkMatrix4x4.h>

#include <vtkActor.h>
#include <vtkPolyDataMapper.h>
#include <vtkProperty.h>

#include <Math/VoxImage.h>

#include "Vtk/Math/vtkVoxImage.h"
#include "Vtk/Math/vtkDense.h"

#include <vtkAutoInit.h>
VTK_MODULE_INIT(vtkRenderingVolumeOpenGL2);
VTK_MODULE_INIT(vtkRenderingOpenGL2);
VTK_MODULE_INIT(vtkInteractionStyle);

namespace uLib {
namespace Vtk {

void VoxImage::UpdateFromContent() {
  Vector3i ev_dims = this->m_model->GetDims();
  m_Image->SetDimensions(ev_dims.data());

  Vector3f ev_spacing = this->m_model->GetSpacing();
  m_Image->SetSpacing(ev_spacing(0), ev_spacing(1), ev_spacing(2));

  Vector3f ev_pos = this->m_model->GetPosition();
  m_Image->SetOrigin(ev_pos(0), ev_pos(1), ev_pos(2));

  vtkFloatArray *array =
      vtkFloatArray::SafeDownCast(m_Image->GetPointData()->GetScalars());
  if (!array) {
    array = vtkFloatArray::New();
    m_Image->GetPointData()->SetScalars(array);
    array->Delete();
  }

  array->SetNumberOfTuples(this->m_model->GetDims().prod());
  Vector3i index(0, 0, 0);
  int i = 0;
  for (int zv = 0; zv < ev_dims(2); ++zv) {
    for (int yv = 0; yv < ev_dims(1); ++yv) {
      for (int xv = 0; xv < ev_dims(0); ++xv) {
        index << xv, yv, zv;
        array->SetValue(i++, this->m_model->GetValue(index));
      }
    }
  }
}

void VoxImage::UpdateToContent() {
  int *ext = m_Image->GetExtent();
  int dims[3] = {ext[1] - ext[0] + 1, ext[3] - ext[2] + 1, ext[5] - ext[4] + 1};
  this->m_model->SetDims(Vector3i(dims[0], dims[1], dims[2]));

  double *spacing = m_Image->GetSpacing();
  this->m_model->SetSpacing(Vector3f(spacing[0], spacing[1], spacing[2]));

  double *pos = m_Image->GetOrigin();
  this->m_model->SetPosition(Vector3f(pos[0], pos[1], pos[2]));

  vtkFloatArray *array =
      vtkFloatArray::SafeDownCast(m_Image->GetPointData()->GetScalars());
  if (array) {
    Vector3i index(0, 0, 0);
    int i = 0;
    for (int zv = 0; zv < dims[2]; ++zv) {
      for (int yv = 0; yv < dims[1]; ++yv) {
        for (int xv = 0; xv < dims[0]; ++xv) {
          index << xv, yv, zv;
          this->m_model->SetValue(index, array->GetValue(i++));
        }
      }
    }
  } else {
    std::cerr << "Error reading array Value Data\n";
  }
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
// VTK VOXIMAGE

VoxImage::VoxImage(Content *content)
    : ObjectWrapper(content), m_Actor(vtkVolume::New()),
      m_Asm(vtkAssembly::New()),
      m_Image(vtkImageData::New()), m_Outline(vtkCubeSource::New()),
      m_OutlineActor(vtkActor::New()),
      m_Reader(NULL), m_Writer(NULL), writer_factor(1.E6),
      m_Window(1.0), m_Level(0.5), m_ShadingPreset(0) {
  // Transfer functions
  m_ColorFun = vtkColorTransferFunction::New();
  m_OpacityFun = vtkPiecewiseFunction::New();
  m_UpdateConnection = Object::connect(this->m_model.get(), &uLib::Object::Updated, this, &VoxImage::Update);
  
  UpdateFromContent();
  InstallPipe();
  RescaleShaderRange();
  ULIB_ACTIVATE_DISPLAY_PROPERTIES;
}

VoxImage::~VoxImage() {
  m_Image->Delete();
  m_Actor->Delete();
  m_Asm->Delete();
  m_Outline->Delete();
  m_OutlineActor->Delete();
  m_ColorFun->Delete();
  m_OpacityFun->Delete();
}

vtkImageData *VoxImage::GetImageData() {
  UpdateFromContent();
  return m_Image;
}

void VoxImage::SaveToXMLFile(const char *fname) {
  vtkSmartPointer<vtkXMLImageDataWriter> writer =
      vtkSmartPointer<vtkXMLImageDataWriter>::New();
  writer->SetFileName(fname);
  UpdateFromContent();
  vtkSmartPointer<vtkImageShiftScale> vtkscale =
      vtkSmartPointer<vtkImageShiftScale>::New();

#if VTK_MAJOR_VERSION <= 5
  vtkscale->SetInputConnection(m_Image->GetProducerPort());
#else
  vtkscale->SetInputData(m_Image);
#endif
  vtkscale->SetScale(writer_factor);
  vtkscale->Update();
  writer->SetInputConnection(vtkscale->GetOutputPort());
  writer->Update();
  writer->Write();
}

void VoxImage::ReadFromVKTFile(const char *fname) {
  vtkSmartPointer<vtkGenericDataObjectReader> reader =
      vtkSmartPointer<vtkGenericDataObjectReader>::New();
  reader->SetFileName(fname);
  reader->Update();
  if (reader->IsFileStructuredPoints()) {
    vtkSmartPointer<vtkImageShiftScale> vtkscale =
        vtkSmartPointer<vtkImageShiftScale>::New();
    vtkscale->SetInputConnection(reader->GetOutputPort());
    vtkscale->SetScale(1 / writer_factor);
    vtkscale->Update();

    m_Image->DeepCopy(vtkscale->GetOutput());
    UpdateToContent();
    RescaleShaderRange();
  } else {
    std::cerr << "Error: file does not contain structured points\n";
  }
  m_Actor->Update();
}

void VoxImage::ReadFromXMLFile(const char *fname) {
  vtkSmartPointer<vtkXMLImageDataReader> reader =
      vtkSmartPointer<vtkXMLImageDataReader>::New();
  reader->SetFileName(fname);
  reader->Update();
  vtkSmartPointer<vtkImageShiftScale> vtkscale =
      vtkSmartPointer<vtkImageShiftScale>::New();
  vtkscale->SetInputConnection(reader->GetOutputPort());
  vtkscale->SetScale(1 / writer_factor);
  vtkscale->Update();

  m_Image->DeepCopy(vtkscale->GetOutput());
  UpdateToContent();
  RescaleShaderRange();
}

void VoxImage::setShadingPreset(int blendType) {
  m_ShadingPreset = blendType;
  vtkSmartVolumeMapper *mapper = (vtkSmartVolumeMapper *)m_Actor->GetMapper();
  if (!mapper) return;
  vtkVolumeProperty *property = m_Actor->GetProperty();

  float window = m_Window;
  float level = m_Level;

  property->SetColor(m_ColorFun);
  property->SetScalarOpacity(m_OpacityFun);
  property->SetInterpolationTypeToLinear();

  m_ColorFun->RemoveAllPoints();
  m_OpacityFun->RemoveAllPoints();

  switch (blendType) {
  case 0: // MIP
    m_ColorFun->AddRGBPoint(level - 0.5 * window, 0, 0, 0);
    m_ColorFun->AddRGBPoint(level + 0.5 * window, 1, 1, 1);
    m_OpacityFun->AddSegment(level - 0.5 * window, 0.0, level + 0.5 * window, 1.0);
    mapper->SetBlendModeToMaximumIntensity();
    break;
  case 1: // Composite
    m_ColorFun->AddRGBPoint(level - 0.5 * window, 0, 0, 0);
    m_ColorFun->AddRGBPoint(level + 0.5 * window, 1, 1, 1);
    m_OpacityFun->AddSegment(level - 0.5 * window, 0.0, level + 0.5 * window, 1.0);
    mapper->SetBlendModeToComposite();
    property->ShadeOff();
    break;
  case 2: // Composite Shaded
    m_ColorFun->AddRGBPoint(level - 0.5 * window, 0, 0, 0);
    m_ColorFun->AddRGBPoint(level + 0.5 * window, 1, 1, 1);
    m_OpacityFun->AddSegment(level - 0.5 * window, 0.0, level + 0.5 * window, 1.0);
    mapper->SetBlendModeToComposite();
    property->ShadeOn();
    break;
  case 3: // Rainbow MIP
    m_ColorFun->AddRGBPoint(level - 0.5 * window, 0, 0, 1);
    m_ColorFun->AddRGBPoint(level, 0, 1, 0);
    m_ColorFun->AddRGBPoint(level + 0.5 * window, 1, 1, 0);
    m_ColorFun->AddRGBPoint(level + window, 1, 0, 0);
    m_OpacityFun->AddSegment(level - 0.5 * window, 0.0, level + 0.5 * window, 1.0);
    mapper->SetBlendModeToMaximumIntensity();
    break;
  case 4: // Additive
    m_ColorFun->AddRGBPoint(level - 0.5 * window, 0, 0, 0);
    m_ColorFun->AddRGBPoint(level + 0.5 * window, 1, 1, 1);
    m_OpacityFun->AddSegment(level - 0.5 * window, 0.0, level + 0.5 * window, 1.0);
    mapper->SetBlendModeToAdditive();
    break;
  default:
    break;
  }
}

void VoxImage::RescaleShaderRange() {
  double range[2];
  m_Image->GetScalarRange(range);
  m_Level = (range[0] + range[1]) / 2.0;
  m_Window = range[1] - range[0];
  if (m_Window <= 1e-9)
    m_Window = 1.0;
  setShadingPreset(m_ShadingPreset);
}

void VoxImage::SetRepresentation(Representation mode) {
  Prop3D::SetRepresentation(mode); // Ensure base class data state is updated

  if (mode == Wireframe) {
    m_Actor->SetVisibility(0);
    m_OutlineActor->SetVisibility(1);
    m_OutlineActor->GetProperty()->SetRepresentationToWireframe();
  } else if (mode == Volume) {
    m_Actor->SetVisibility(1);
    m_OutlineActor->SetVisibility(1);
    m_OutlineActor->GetProperty()->SetRepresentationToWireframe();
  } else {
    // Other representations (Points, Surface, etc) are handled by basic Prop3D
    // behavior which affects the m_Asm parts.
  }
}

void VoxImage::serialize_display(uLib::Archive::display_properties_archive & ar, const unsigned int version) {
    // Call base class to show Transform and Appearance properties
    // Prop3D::serialize_display(ar, version);
    
    // Use the member variables for volume rendering parameters
    ar & boost::serialization::make_hrp("Window", m_Window);
    ar & boost::serialization::make_hrp("Level", m_Level);
    ar & boost::serialization::make_hrp_enum("Shading", m_ShadingPreset, 
      {"MIP", "Composite", "Composite Shaded", "MIP Bone", "MIP Hot", "Additive"});
}

void VoxImage::SyncFromVtk() {
  if (auto *root = this->GetProxyProp()) {
    vtkMatrix4x4 *rootMat = root->GetUserMatrix();
    if (rootMat) {
      Matrix4f vtkLocal = VtkToMatrix4f(rootMat);
      // Synchronize TRS from VTK, compensating for local volume offset
      this->m_model->FromMatrix(vtkLocal); // * this->m_model->GetLocalMatrix().inverse());
      this->m_model->Updated();
    }
  }
}

void VoxImage::Update() {
  if (auto *root = vtkProp3D::SafeDownCast(this->GetProp())) {
    vtkNew<vtkMatrix4x4> m;
    Matrix4fToVtk(this->m_model->GetMatrix(), m); // * this->m_model->GetLocalMatrix(), m);
    root->SetUserMatrix(m);
    root->Modified();
    // std::cout << "[VoxImage::Update] Set Proxy UserMatrix:" << std::endl;
    // std::cout << this->m_model->GetMatrix() << std::endl;
  }
  setShadingPreset(m_ShadingPreset);
  m_Actor->Update();
  m_Outline->SetBounds(m_Image->GetBounds());
  m_Outline->Update();

  ConnectionBlock blocker(m_UpdateConnection);
  this->Prop3D::Update();
}


void VoxImage::InstallPipe() {
  vtkSmartPointer<vtkSmartVolumeMapper> mapper =
      vtkSmartPointer<vtkSmartVolumeMapper>::New();
#if VTK_MAJOR_VERSION <= 5
  mapper->SetInputConnection(m_Image->GetProducerPort());
#else
  mapper->SetInputData(m_Image);
#endif
  mapper->Update();

  m_Actor->SetMapper(mapper);
  this->setShadingPreset(m_ShadingPreset);
  mapper->Update();

  m_Outline->SetBounds(m_Image->GetBounds());
  vtkSmartPointer<vtkPolyDataMapper> mmapper =
      vtkSmartPointer<vtkPolyDataMapper>::New();
  mmapper->SetInputConnection(m_Outline->GetOutputPort());
  
  m_OutlineActor->SetMapper(mmapper);
  m_OutlineActor->GetProperty()->SetRepresentationToWireframe();
  m_OutlineActor->GetProperty()->SetAmbient(0.7);
  
  m_Asm->AddPart(m_Actor);
  m_Asm->AddPart(m_OutlineActor);
  this->SetProp(m_Asm);
  
  // Default look
  this->SetRepresentation(Volume);
}

} // namespace Vtk
} // namespace uLib