/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkOpenQubeElectronicData.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/

#include "vtkOpenQubeElectronicData.h"

#include "vtkDataSetCollection.h"
#include "vtkImageData.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"

#include <openqube/basisset.h>
#include <openqube/cube.h>

// Internal class to store queue/qube information along with the image
class OQEDImageData : public vtkImageData
{
public:
  vtkTypeMacro(OQEDImageData, vtkImageData);
  static OQEDImageData* New();

  vtkSetMacro(OrbitalNumber, vtkIdType);
  vtkGetMacro(OrbitalNumber, vtkIdType);

  vtkSetMacro(ImageType, OpenQube::Cube::Type);
  vtkGetMacro(ImageType, OpenQube::Cube::Type);

  vtkSetMacro(MetaSpacing, float);
  vtkGetMacro(MetaSpacing, float);

  vtkSetMacro(MetaPadding, float);
  vtkGetMacro(MetaPadding, float);

  virtual void DeepCopy(vtkDataObject* src)
  {
    this->Superclass::DeepCopy(src);
    OQEDImageData* o = OQEDImageData::SafeDownCast(src);
    if (!o)
    {
      // Just return without copying metadata
      return;
    }
    this->OrbitalNumber = o->OrbitalNumber;
    this->ImageType = o->ImageType;
    this->MetaSpacing = o->MetaSpacing;
    this->MetaPadding = o->MetaPadding;
  }

protected:
  OQEDImageData() {}
  ~OQEDImageData() {}

  vtkIdType OrbitalNumber;
  OpenQube::Cube::Type ImageType;

  float MetaSpacing;
  float MetaPadding;

private:
  OQEDImageData(const OQEDImageData&) = delete;
  void operator=(const OQEDImageData&) = delete;
};
vtkStandardNewMacro(OQEDImageData);

//------------------------------------------------------------------------------
vtkStandardNewMacro(vtkOpenQubeElectronicData);

//------------------------------------------------------------------------------
vtkOpenQubeElectronicData::vtkOpenQubeElectronicData()
  : BasisSet(nullptr)
  , Spacing(0.1)
{
  this->Padding = 2.0;
}

//------------------------------------------------------------------------------
vtkOpenQubeElectronicData::~vtkOpenQubeElectronicData() {}

//------------------------------------------------------------------------------
void vtkOpenQubeElectronicData::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "BasisSet: @" << this->BasisSet << "\n";

  // Dump Images
  os << indent << "Images: @" << this->Images << "\n";
  vtkCollectionSimpleIterator cookie;
  this->Images->InitTraversal(cookie);
  while (vtkDataSet* dataset = this->Images->GetNextDataSet(cookie))
  {
    vtkImageData* data = vtkImageData::SafeDownCast(dataset);
    // If this is somehow not a vtkImageData, print a warning and skip it
    if (!data)
    {
      vtkWarningMacro(<< "vtkDataSet in this->Images is not a vtkImageData"
                         " object. This should not happen...");
      continue;
    }

    OQEDImageData* oqeddata = OQEDImageData::SafeDownCast(dataset);
    // Identify the type of image
    if (oqeddata)
    {
      switch (oqeddata->GetImageType())
      {
        case OpenQube::Cube::MO:
          os << indent << "this->Images has molecular orbital #" << oqeddata->GetOrbitalNumber()
             << " imagedata @" << oqeddata << ":\n";
          break;
        case OpenQube::Cube::ElectronDensity:
          os << indent
             << "this->Images has electron density imagedata "
                "@"
             << oqeddata << ":\n";
          break;
        case OpenQube::Cube::VdW:
          os << indent
             << "this->Images has van der Waals imagedata"
                " @"
             << oqeddata << ":\n";
          break;
        case OpenQube::Cube::ESP:
          os << indent
             << "this->Images has electrostatic potential imagedata"
                " @"
             << oqeddata << ":\n";
          break;
        case OpenQube::Cube::FromFile:
          os << indent
             << "this->Images has file-loaded imagedata"
                " @"
             << oqeddata << ":\n";
          break;
        default:
        case OpenQube::Cube::None:
          os << indent
             << "this->Images has imagedata from an unknown source"
                " @"
             << oqeddata << ":\n";
          break;
      }
    }
    else // oqeddata null
    {
      os << indent
         << "this->Images has imagedata that was externally "
            "generated @"
         << data << ":\n";
    }
    data->PrintSelf(os, indent.GetNextIndent());
  }
}

//------------------------------------------------------------------------------
vtkIdType vtkOpenQubeElectronicData::GetNumberOfMOs()
{
  if (!this->BasisSet || !this->BasisSet->isValid())
  {
    return 0;
  }

  return this->BasisSet->numMOs();
}

//------------------------------------------------------------------------------
unsigned int vtkOpenQubeElectronicData::GetNumberOfElectrons()
{
  if (!this->BasisSet || !this->BasisSet->isValid())
  {
    return 0;
  }

  return this->BasisSet->numElectrons();
}

//------------------------------------------------------------------------------
vtkImageData* vtkOpenQubeElectronicData::GetMO(vtkIdType orbitalNumber)
{
  vtkDebugMacro(<< "Searching for MO " << orbitalNumber);
  // First check if there is an existing image for this orbital
  vtkCollectionSimpleIterator cookie;
  this->Images->InitTraversal(cookie);
  while (vtkDataSet* dataset = this->Images->GetNextDataSet(cookie))
  {
    OQEDImageData* data = OQEDImageData::SafeDownCast(dataset);
    if (!data)
    {
      continue;
    }

    if (data->GetImageType() != OpenQube::Cube::MO)
    {
      continue;
    }

    if (data->GetOrbitalNumber() != orbitalNumber)
    {
      continue;
    }

    if (data->GetMetaSpacing() != this->Spacing || data->GetMetaPadding() != this->Padding)
    {
      continue;
    }

    vtkDebugMacro(<< "Found MO " << orbitalNumber);
    return static_cast<vtkImageData*>(data);
  }

  // If there is no existing image, calculate it
  vtkDebugMacro(<< "MO " << orbitalNumber << " not found. Calculating...");
  return this->CalculateMO(orbitalNumber);
}

//------------------------------------------------------------------------------
vtkImageData* vtkOpenQubeElectronicData::GetElectronDensity()
{
  // First check if there is an existing image for this orbital
  vtkCollectionSimpleIterator cookie;
  this->Images->InitTraversal(cookie);
  while (vtkDataSet* dataset = this->Images->GetNextDataSet(cookie))
  {
    OQEDImageData* data = OQEDImageData::SafeDownCast(dataset);
    if (!data)
    {
      continue;
    }

    if (data->GetImageType() != OpenQube::Cube::ElectronDensity)
    {
      continue;
    }

    if (data->GetMetaSpacing() != this->Spacing || data->GetMetaPadding() != this->Padding)
    {
      continue;
    }

    return static_cast<vtkImageData*>(data);
  }

  // If there is no existing image, calculate it
  return this->CalculateElectronDensity();
}

//------------------------------------------------------------------------------
void vtkOpenQubeElectronicData::DeepCopy(vtkDataObject* obj)
{
  vtkOpenQubeElectronicData* oqed = vtkOpenQubeElectronicData::SafeDownCast(obj);
  if (!oqed)
  {
    vtkErrorMacro("Can only deep copy from vtkOpenQubeElectronicData "
                  "or subclass.");
    return;
  }

  // Call superclass
  this->Superclass::DeepCopy(oqed);

  // Copy all images over by hand to get OQEDImageData metadata right
  vtkCollectionSimpleIterator cookie;
  oqed->Images->InitTraversal(cookie);
  while (vtkDataSet* dataset = oqed->Images->GetNextDataSet(cookie))
  {
    // Only copy valid OQEDImageData objects
    if (OQEDImageData* oqedImage = OQEDImageData::SafeDownCast(dataset))
    {
      vtkNew<OQEDImageData> thisImage;
      thisImage->DeepCopy(oqedImage);
      this->Images->AddItem(thisImage);
    }
  }

  // Copy other ivars
  if (oqed->BasisSet)
  {
    this->BasisSet = oqed->BasisSet->clone();
  }
  this->Spacing = oqed->Spacing;
}

//------------------------------------------------------------------------------
vtkImageData* vtkOpenQubeElectronicData::CalculateMO(vtkIdType orbitalNumber)
{
  vtkDebugMacro(<< "Calculating MO " << orbitalNumber);
  if (!this->BasisSet)
  {
    vtkWarningMacro(<< "No OpenQube::BasisSet set.");
    return 0;
  }
  else if (!this->BasisSet->isValid())
  {
    vtkWarningMacro(<< "Invalid OpenQube::BasisSet set.");
    return 0;
  }

  // Create and calculate cube
  OpenQube::Cube* cube = new OpenQube::Cube();
  cube->setLimits(this->BasisSet->moleculeRef(), this->Spacing, this->Padding);

  vtkDebugMacro(<< "Calculating OpenQube::Cube for MO " << orbitalNumber);
  if (!this->BasisSet->blockingCalculateCubeMO(cube, orbitalNumber))
  {
    vtkWarningMacro(<< "Unable to calculateMO for orbital " << orbitalNumber << " in OpenQube.");
    return 0;
  }

  // Create image and set metadata
  vtkNew<OQEDImageData> image;
  image->SetMetaSpacing(this->Spacing);
  image->SetMetaPadding(this->Padding);
  vtkDebugMacro(<< "Converting OpenQube::Cube to vtkImageData for MO " << orbitalNumber);

  // Copy cube --> image
  this->FillImageDataFromQube(cube, image);
  image->SetOrbitalNumber(orbitalNumber);

  vtkDebugMacro(<< "Adding vtkImageData to this->Images for MO " << orbitalNumber);
  this->Images->AddItem(image);

  return image;
}

//------------------------------------------------------------------------------
vtkImageData* vtkOpenQubeElectronicData::CalculateElectronDensity()
{
  vtkDebugMacro(<< "Calculating electron density...");
  if (!this->BasisSet)
  {
    vtkErrorMacro(<< "No OpenQube::BasisSet set.");
    return 0;
  }
  else if (!this->BasisSet->isValid())
  {
    vtkErrorMacro(<< "Invalid OpenQube::BasisSet set.");
    return 0;
  }

  // Create and calculate cube
  OpenQube::Cube* cube = new OpenQube::Cube();
  cube->setLimits(this->BasisSet->moleculeRef(), this->Spacing, this->Padding);

  vtkDebugMacro(<< "Calculating OpenQube::Cube...");
  if (!this->BasisSet->blockingCalculateCubeDensity(cube))
  {
    vtkWarningMacro(<< "Unable to calculate density in OpenQube.");
    return 0;
  }

  // Create image and set metadata
  vtkNew<OQEDImageData> image;
  image->SetMetaSpacing(this->Spacing);
  image->SetMetaPadding(this->Padding);
  vtkDebugMacro(<< "Converting OpenQube::Cube to vtkImageData.");

  // Copy cube --> image
  this->FillImageDataFromQube(cube, image);

  vtkDebugMacro(<< "Adding vtkImageData to this->Images");
  this->Images->AddItem(image);

  return image;
}

//------------------------------------------------------------------------------
void vtkOpenQubeElectronicData::FillImageDataFromQube(OpenQube::Cube* qube, vtkImageData* image)
{
  Eigen::Vector3i dim = qube->dimensions();
  Eigen::Vector3d min = qube->min();
  Eigen::Vector3d max = qube->max();
  Eigen::Vector3d spacing = qube->spacing();

  vtkDebugMacro(<< "Converting OpenQube::Cube to vtkImageData:"
                << "\n\tDimensions: " << dim[0] << "  " << dim[1] << " " << dim[2]
                << "\n\tMinimum: " << min[0] << "  " << min[1] << " " << min[2]
                << "\n\tMaximum: " << max[0] << "  " << max[1] << " " << max[2]
                << "\n\tSpacing: " << spacing[0] << "  " << spacing[1] << " " << spacing[2]);

  if (OQEDImageData* oqedImage = OQEDImageData::SafeDownCast(image))
  {
    vtkDebugMacro("Setting cube type to " << qube->cubeType());
    oqedImage->SetImageType(qube->cubeType());
  }
  else
  {
    vtkWarningMacro(<< "Cannot cast input image to internal OQEDImageData type.");
  }

  image->SetExtent(0, dim[0] - 1, 0, dim[1] - 1, 0, dim[2] - 1);
  image->SetOrigin(min.data());
  image->SetSpacing(spacing.data());
  image->AllocateScalars(VTK_DOUBLE, 1);

  double* dataPtr = static_cast<double*>(image->GetScalarPointer());
  std::vector<double>* qubeVec = qube->data();

  const size_t qubeSize = qubeVec->size();

  if (qubeSize != static_cast<size_t>(dim[0]) * dim[1] * dim[2])
  {
    vtkWarningMacro(<< "Size of qube (" << qubeSize << ") does not equal"
                    << " product of dimensions (" << dim[0] * dim[1] * dim[2]
                    << "). Image may not be accurate.");
  }

  long int qubeInd = -1; // Incremented to 0 on first use.
  for (int i = 0; i < dim[0]; ++i)
  {
    for (int j = 0; j < dim[1]; ++j)
    {
      for (int k = 0; k < dim[2]; ++k)
      {
        dataPtr[(k * dim[1] + j) * dim[0] + i] = (*qubeVec)[++qubeInd];
      }
    }
  }

  vtkDebugMacro(<< "Copied " << qubeSize << " (actual: " << qubeInd + 1
                << ") points from qube to vtkImageData.");
}