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

  Program:   Visualization Toolkit
  Module:    vtkPMultiResolutionGenericIOReader.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 "vtkPMultiResolutionGenericIOReader.h"

#include "vtkCallbackCommand.h"
#include "vtkCompositeDataPipeline.h"
#include "vtkDataArraySelection.h"
#include "vtkFileSeriesReader.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPGenericIOMultiBlockReader.h"
#include "vtkSmartPointer.h"
#include "vtkStringArray.h"
#include "vtkUnstructuredGrid.h"

#include "vtk_jsoncpp.h"

#include <algorithm>
#include <cassert>
#include <map>
#include <set>
#include <string>
#include <sstream>

struct resolution_t
{
  std::string FileName;
  vtkSmartPointer< vtkFileSeriesReader > Reader;
};

class vtkPMultiResolutionGenericIOReader::vtkInternal
{
public:
  std::vector< resolution_t > Resolutions;
  int NumberOfBlocksPerLevel;
  Json::Reader JsonReader;

  void AddLevel(int level)
    {
    resolution_t newLevel;
    newLevel.Reader = vtkSmartPointer< vtkFileSeriesReader >::New();
    vtkNew< vtkPGenericIOMultiBlockReader > internalReader;
    newLevel.Reader->SetReader(internalReader.GetPointer());
    newLevel.Reader->SetFileNameMethod("SetFileName");
    this->Resolutions.insert(this->Resolutions.begin() + level, newLevel);
    }
#define JSON_READ_ERROR() std::cerr << "Error reading in JSON" << std::endl; \
      return false

  vtkPGenericIOMultiBlockReader* GetReaderForLevel(int level)
    {
    return vtkPGenericIOMultiBlockReader::SafeDownCast(
      this->Resolutions[level].Reader->GetReader());
    }

  void Clear()
    {
    this->Resolutions.clear();
    this->NumberOfBlocksPerLevel = -1;
    }

  bool ParseJson(const std::string& parentDir, std::istream& jsonIn)
    {
    Json::Value root;
    if (!JsonReader.parse(jsonIn,root))
      {
      JSON_READ_ERROR();
      }
    if (!root.isMember("levels"))
      {
      JSON_READ_ERROR();
      }
    const Json::Value& levelsArray = root.get("levels",Json::Value::nullRef);
    if (!levelsArray.isArray())
      {
      JSON_READ_ERROR();
      }
    for (Json::Value::ArrayIndex i = 0; i < levelsArray.size(); ++i)
      {
      const Json::Value& level = levelsArray[i];
      if (!level.isMember("timesteps"))
        {
        JSON_READ_ERROR();
        }
      const Json::Value& timestepsArray = level.get("timesteps",Json::Value::nullRef);
      if (!timestepsArray.isArray())
        {
        JSON_READ_ERROR();
        }
      this->AddLevel(i);
      std::map< double, std::string > timeFiles;
      for (Json::Value::ArrayIndex j = 0; j < timestepsArray.size(); ++j)
        {
        const Json::Value& timestep = timestepsArray[j];
        if (!(timestep.isMember("time") && timestep.isMember("file")))
          {
          JSON_READ_ERROR();
          }
        const Json::Value& time = timestep.get("time",Json::Value::nullRef);
        const Json::Value& file = timestep.get("file",Json::Value::nullRef);
        if (!time.isNumeric())
          {
          JSON_READ_ERROR();
          }
        double t = time.asDouble();
        std::string f = file.asString();
        // allow paths relative to the config file
        if (f[0] != '/')
          {
          f = parentDir + f;
          }
        timeFiles.insert(std::pair< double, std::string >(t,f));
        }
      for (std::map< double, std::string >::iterator itr = timeFiles.begin();
        itr != timeFiles.end(); ++itr)
        {
        // vtkFileSeriesReader ignores time on inputs unless they already have
        // a time range, but iteration over the map ensures they are in order...
        // Eventually we may need to take into account the time in the JSON,
        // which is why it is there.
        this->Resolutions[i].Reader->AddFileName(itr->second.c_str());
        }
      }
    return true;
    }
};

vtkStandardNewMacro(vtkPMultiResolutionGenericIOReader);
//----------------------------------------------------------------------------
vtkPMultiResolutionGenericIOReader::vtkPMultiResolutionGenericIOReader()
{
  this->Internal = new vtkInternal();
  this->FileName = NULL;
  this->Internal->NumberOfBlocksPerLevel = -1;
  this->SetNumberOfInputPorts(0);
  this->SetNumberOfOutputPorts(1);

  this->XAxisVariableName = NULL;
  this->YAxisVariableName = NULL;
  this->ZAxisVariableName = NULL;

  this->SetXAxisVariableName("x");
  this->SetYAxisVariableName("y");
  this->SetZAxisVariableName("z");

  this->PointDataArraySelection = vtkDataArraySelection::New();
  this->SelectionObserver = vtkCallbackCommand::New();
  this->SelectionObserver->SetCallback(
    &vtkPMultiResolutionGenericIOReader::SelectionModifiedCallback);
  this->SelectionObserver->SetClientData(this);
  this->PointDataArraySelection->AddObserver(
    vtkCommand::ModifiedEvent,this->SelectionObserver);
}

//----------------------------------------------------------------------------
vtkPMultiResolutionGenericIOReader::~vtkPMultiResolutionGenericIOReader()
{
  this->SetFileName(NULL);
  delete this->Internal;
  this->PointDataArraySelection->RemoveObserver(this->SelectionObserver);
  this->SelectionObserver->Delete();
  this->PointDataArraySelection->Delete();
}

//----------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
}

//----------------------------------------------------------------------------
bool vtkPMultiResolutionGenericIOReader::CanReadFile(const char* fName)
{
  std::string filename(fName);
  std::string::size_type it = filename.find_last_of(".");
  std::string extension(filename, it+1, filename.length());
  return extension == "gios";
}

//----------------------------------------------------------------------------
static inline bool SetStringProperty(char* & destStr, const char* newStr)
{
  if (destStr == NULL && newStr == NULL) return false;
  if (destStr && newStr && strcmp(destStr,newStr) == 0) return false;
  delete [] destStr;
  if (newStr)
    {
    size_t n = strlen(newStr) +1;
    char* cp1 = new char[n];
    const char* cp2 = newStr;
    destStr = cp1;
    do { *cp1++ = *cp2++; } while (--n);
    }
  else
    {
    destStr = NULL;
    }
  return true;
}

//----------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::SetFileName(const char* fname)
{
  this->Internal->Clear();
  if (!SetStringProperty(this->FileName,fname)) return;
  if (fname)
    {
    std::ifstream fis(this->FileName);
    std::string filename(this->FileName);
    std::string::size_type idx = filename.find_last_of("/");
    std::string parentDir(filename,0,idx+1);
    if (!this->Internal->ParseJson(parentDir,fis))
      {
      // clear any partially read data if reading failed for some reason
      this->Internal->Resolutions.clear();
      }
    else
      {
      // sync selected data arrays
      this->PointDataArraySelection->CopySelections(
        this->Internal->GetReaderForLevel(0)->GetPointDataArraySelection());
      // sync x,y,z axes
      for (int i = 0; i < this->GetNumberOfLevels(); ++i)
        {
        vtkPGenericIOMultiBlockReader* r = this->Internal->GetReaderForLevel(i);
        r->SetXAxisVariableName(this->XAxisVariableName);
        r->SetYAxisVariableName(this->YAxisVariableName);
        r->SetZAxisVariableName(this->ZAxisVariableName);
        }
      }
    }
  this->Modified();
}

//----------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::SetXAxisVariableName(const char* arg)
{
  if (SetStringProperty(this->XAxisVariableName,arg))
    {
    this->Modified();
    for (unsigned i = 0; i < this->Internal->Resolutions.size(); ++i)
      {
      this->Internal->GetReaderForLevel(i)->SetXAxisVariableName(arg);
      }
    }
}

//----------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::SetYAxisVariableName(const char* arg)
{
  if (SetStringProperty(this->YAxisVariableName,arg))
    {
    this->Modified();
    for (unsigned i = 0; i < this->Internal->Resolutions.size(); ++i)
      {
      this->Internal->GetReaderForLevel(i)->SetYAxisVariableName(arg);
      }
    }
}

//----------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::SetZAxisVariableName(const char* arg)
{
  if (SetStringProperty(this->ZAxisVariableName,arg))
    {
    this->Modified();
    for (unsigned i = 0; i < this->Internal->Resolutions.size(); ++i)
      {
      this->Internal->GetReaderForLevel(i)->SetZAxisVariableName(arg);
      }
    }
}

//----------------------------------------------------------------------------
vtkStringArray* vtkPMultiResolutionGenericIOReader::GetArrayList()
{
  if (this->GetNumberOfLevels() > 0)
    {
    return this->Internal->GetReaderForLevel(0)->GetArrayList();
    }
  else
    {
    return NULL;
    }
}

//----------------------------------------------------------------------------
bool vtkPMultiResolutionGenericIOReader::InsertLevel(const char* fileName, int level)
{
  if (fileName == NULL)
    {
    return false;
    }
  assert(this->Internal != NULL);
  if (level < 0 || static_cast<unsigned>(level) > this->Internal->Resolutions.size())
    {
    vtkErrorMacro(<< "Level is out of range.");
    return false;
    }
  //this->Internal->AddLevel(level,fileName);

  if (this->GetNumberOfLevels() != 0)
    {
    this->Internal->GetReaderForLevel(level)->GetPointDataArraySelection()
      ->CopySelections(this->PointDataArraySelection);
    }
  else
    {
    this->PointDataArraySelection->CopySelections(
      this->Internal->GetReaderForLevel(level)->GetPointDataArraySelection());
    }
  this->Modified();
  return true;
}

//----------------------------------------------------------------------------
int vtkPMultiResolutionGenericIOReader::GetNumberOfLevels() const
{
  assert(this->Internal != NULL);
  return this->Internal->Resolutions.size();
}

//----------------------------------------------------------------------------
const char* vtkPMultiResolutionGenericIOReader::GetFileNameForLevel(int level) const
{
  assert(this->Internal != NULL);
  if (level < 0 || static_cast<unsigned>(level) >= this->Internal->Resolutions.size())
    {
    return (const char*) NULL;
    }
  return this->Internal->Resolutions[level].FileName.c_str();
}

//----------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::RemoveAllLevels()
{
  this->Internal->Resolutions.clear();
  this->Modified();
}

//------------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::SelectionModifiedCallback(
  vtkObject* vtkNotUsed(caller), unsigned long vtkNotUsed(eid), void* clientdata, void* vtkNotUsed(calldata))
{
  assert(clientdata != NULL);
  vtkPMultiResolutionGenericIOReader* reader =
    static_cast<vtkPMultiResolutionGenericIOReader*>(clientdata);
  for (int i = 0; i < reader->GetNumberOfLevels(); ++i)
    {
    reader->Internal->GetReaderForLevel(i)->GetPointDataArraySelection()
      ->CopySelections(reader->PointDataArraySelection);
    }
  reader->Modified();
}

//------------------------------------------------------------------------------
int vtkPMultiResolutionGenericIOReader::GetNumberOfPointArrays()
{
  return this->PointDataArraySelection->GetNumberOfArrays();
}

//------------------------------------------------------------------------------
const char* vtkPMultiResolutionGenericIOReader::GetPointArrayName(int i)
{
  assert("pre: array index is out-of-bounds!" &&
    (i >= 0) && (i < this->GetNumberOfPointArrays()));
  return this->PointDataArraySelection->GetArrayName(i);
}

//------------------------------------------------------------------------------
int vtkPMultiResolutionGenericIOReader::GetPointArrayStatus(const char* name)
{
  assert(this->PointDataArraySelection->ArrayExists(name));
  return this->PointDataArraySelection->ArrayIsEnabled(name);
}

//------------------------------------------------------------------------------
void vtkPMultiResolutionGenericIOReader::SetPointArrayStatus(
  const char* name, int status)
{
  if (!this->PointDataArraySelection->ArrayExists(name))
    {
    assert(this->Internal->Resolutions.size() > 1);
    this->PointDataArraySelection->CopySelections(
      this->Internal->GetReaderForLevel(0)->GetPointDataArraySelection());
    vtkPMultiResolutionGenericIOReader::SelectionModifiedCallback(NULL,0l,(void*)this,NULL);
    assert(this->PointDataArraySelection->ArrayExists(name));
    }
  if (status)
    {
    this->PointDataArraySelection->EnableArray(name);
    }
  else
    {
    this->PointDataArraySelection->DisableArray(name);
    }
  for (int i = 0; i < this->GetNumberOfLevels(); ++i)
    {
    if (this->Internal->GetReaderForLevel(i)->GetArrayList()->LookupValue(name) >=0)
      {
      this->Internal->GetReaderForLevel(i)->SetPointArrayStatus(name,status);
      }
    }
}

void printBoundsOfBlockInfo(vtkMultiBlockDataSet* ds,int offset)
{
  for (unsigned i = 0; i < ds->GetNumberOfBlocks(); ++i)
    {
    vtkInformation* info = ds->GetMetaData(i);
    std::cerr << "Bounds for block " << (i + offset) << ": ";
    if (info->Has(vtkCompositeDataPipeline::BOUNDS()))
      {
      double bb[6];
      info->Get(vtkCompositeDataPipeline::BOUNDS(),bb);
      std::cerr << "[ ";
      for (int j = 0; j < 6; ++j)
        {
        std::cerr << bb[j] << ", ";
        }
      std::cerr << "]" << std::endl;
      }
    }
}

//----------------------------------------------------------------------------
int vtkPMultiResolutionGenericIOReader::RequestInformation(vtkInformation *request,
        vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  // call the superclass's method
  this->Superclass::RequestInformation(request,inputVector,outputVector);

  // get the output info object
  vtkInformation* outInfo = outputVector->GetInformationObject(0);
  assert(outInfo != NULL);

  // this filter does handle streaming pieces of the dataset
  outInfo->Set(CAN_HANDLE_PIECE_REQUEST(),1);

  // create a dummy dataset that will have the right structure
  vtkSmartPointer< vtkMultiBlockDataSet > infoSet =
    vtkSmartPointer< vtkMultiBlockDataSet >::New();
  infoSet->SetNumberOfBlocks(this->GetNumberOfLevels());
  this->Internal->NumberOfBlocksPerLevel = -1;


  // request the information from each internal reader and put them
  // into the information object for this reader as blocks
  for (unsigned i = 0; i < this->Internal->Resolutions.size(); ++i)
    {
    vtkMultiBlockDataSet* dataSet;
    this->Internal->Resolutions[i].Reader->ProcessRequest(
      request,inputVector,outputVector);
    if (i == 0)
      {
      this->PointDataArraySelection->CopySelections(
        this->Internal->GetReaderForLevel(0)->GetPointDataArraySelection());
      }
    dataSet = vtkMultiBlockDataSet::SafeDownCast(
      outInfo->Get(
        vtkCompositeDataPipeline::COMPOSITE_DATA_META_DATA()));
    // also make sure that they have the same number of blocks
    // (important for this kind of data)
    if (this->Internal->NumberOfBlocksPerLevel == -1)
      {
      this->Internal->NumberOfBlocksPerLevel =
        dataSet->GetNumberOfBlocks();
      }
    else if ((unsigned)(this->Internal->NumberOfBlocksPerLevel) !=
        dataSet->GetNumberOfBlocks())
      {
      vtkErrorMacro(<< "Wrong number of blocks in level " << i);
      return 0;
      }

//    printBoundsOfBlockInfo(dataSet,i * this->Internal->NumberOfBlocksPerLevel);
    infoSet->SetBlock(i,dataSet);
    outInfo->Remove(
      vtkCompositeDataPipeline::COMPOSITE_DATA_META_DATA());
    }
  // We assume all datasets have blocks with the same bounds.  The rest of the pipeline
  // does not know this, so this loop first finds which resolution has bounds and copies
  // those bounds to all the other resolutions for each block
  for (int i = 0; i < this->Internal->NumberOfBlocksPerLevel; ++i) // for each block
    {
    double bounds[6] = {0,0,0,0,0,0};
    // find the resolution with nonzero bounds (zero is default from GenericIO reader
    for (unsigned j = 0; j < this->Internal->Resolutions.size(); ++j)
      {
      if (bounds[0] == 0 && bounds[1] == 0)
        {
        static_cast<vtkMultiBlockDataSet*>(
              infoSet->GetBlock(j))
            ->GetMetaData(i)->Get(vtkStreamingDemandDrivenPipeline::BOUNDS(),bounds);
        }
      }
    // copy bounds to all resolutions
    for (unsigned j = 0; j < this->Internal->Resolutions.size(); ++j)
      {
      static_cast<vtkMultiBlockDataSet*>(
            infoSet->GetBlock(j))
          ->GetMetaData(i)->Set(vtkStreamingDemandDrivenPipeline::BOUNDS(),bounds,6);
      }
    }

  std::cerr.flush();
  // return the computed object as the composite data metadata
  outInfo->Set(vtkCompositeDataPipeline::COMPOSITE_DATA_META_DATA(),
    infoSet.GetPointer());
  return 1;
}

//----------------------------------------------------------------------------
// This function is a binary search that returns the index of where the item
// should be inserted if it is not already present.  More precisely, it
// returns the minimum index i such that vec[i] <= value && value < vec[i+1].
// If the item is greater than every item in the array, it returns the size of
// the array and if the item is less than every item in the array it returns 0.
// NOTE: as per binary search, vec should be a sorted array of ints
int binSearch(int* vec,int size, int value)
{
  // handle boundary conditions
  if (vec[0] > value)
    {
    return 0;
    }
  if (vec[size-1] < value)
    {
    return size;
    }
  // loop invariants:
  // vec[min] is always <= value
  // vec[max] is alwasy > value
  int min = 0, max = size;
  while (min < max -1)
    {
    int mid = (min+max)/2;
    if (vec[mid] <= value)
      {
      min = mid;
      }
    else
      {
      max = mid;
      }
    }
  return min;
}

int vtkPMultiResolutionGenericIOReader::RequestUpdateExtent(vtkInformation* request,
                                                            vtkInformationVector** inputVector,
                                                            vtkInformationVector* outputVector)
{
  // The vtkFileSeriesReader needs this to be called to update its timestep
  for (unsigned i = 0; i < this->Internal->Resolutions.size(); ++i)
    {
    this->Internal->Resolutions[i].Reader->ProcessRequest(request,inputVector,outputVector);
    this->Internal->GetReaderForLevel(i)->GetPointDataArraySelection()
      ->CopySelections(this->PointDataArraySelection);
    }
  return 1;
}

//----------------------------------------------------------------------------
int vtkPMultiResolutionGenericIOReader::RequestData(vtkInformation *request,
        vtkInformationVector** vtkNotUsed(inputVector),
        vtkInformationVector* outputVector)
{
  // Get the output dataset
  vtkInformation* outInfo = outputVector->GetInformationObject(0);
  assert(outInfo != NULL);
  vtkMultiBlockDataSet* output = vtkMultiBlockDataSet::GetData(outputVector,0);
  assert(output != NULL);
  // set the number of blocks
  output->SetNumberOfBlocks(this->GetNumberOfLevels());

  int size;
  std::vector<int> idVector;
  // find out which parts of the dataset we need to load
  if (outInfo->Has(vtkCompositeDataPipeline::LOAD_REQUESTED_BLOCKS()))
    {
    int* ids;
    size = outInfo->Length(vtkCompositeDataPipeline::UPDATE_COMPOSITE_INDICES());
    ids = outInfo->Get(vtkCompositeDataPipeline::UPDATE_COMPOSITE_INDICES());
    idVector.resize(size);
    std::copy(ids,ids+size,idVector.begin());
    }
  // default to loading all of the lowest level of detail
  else
    {
    for (int j = 0; j < this->Internal->NumberOfBlocksPerLevel; ++j)
      {
      idVector.push_back(j);
      }
    size = idVector.size();
    }

  // sort the requested blocks
  std::sort(idVector.begin(),idVector.end());
  // compute the block ids relative to the file (the internal reader needs these)
  std::vector<int> localIdVector;
  for (unsigned i = 0; i < idVector.size(); ++i)
    {
    localIdVector.push_back(idVector[i] % this->Internal->NumberOfBlocksPerLevel);
    }
  // these mark the beginning and end of the local ids for the internal reader
  // for the current itertion of the loop
  int lBound = 0, uBound;
  for (int i = 0; i < this->GetNumberOfLevels(); ++i)
    {
    // compute new bounds for current reader's blocks
    uBound = binSearch(&idVector[0],idVector.size(),this->Internal->NumberOfBlocksPerLevel*(i+1));
    int levelSize = uBound - lBound;

    vtkNew< vtkMultiBlockDataSet > dataset;
    if (levelSize > 0)
      {
      // create a new request
      vtkNew< vtkInformation > internalRequestData;
      internalRequestData->Set(vtkDemandDrivenPipeline::REQUEST_DATA());
      internalRequestData->Copy(request);

      // create a new output info
      vtkNew< vtkInformation > internalOutInfo;
      internalOutInfo->CopyEntry(outInfo,vtkStreamingDemandDrivenPipeline::TIME_STEPS());
      internalOutInfo->CopyEntry(outInfo,vtkStreamingDemandDrivenPipeline::TIME_RANGE());
      internalOutInfo->Set(vtkStreamingDemandDrivenPipeline::UPDATE_TIME_STEP(),
                           outInfo->Get(vtkStreamingDemandDrivenPipeline::UPDATE_TIME_STEP()));
      internalOutInfo->Set(vtkDataObject::DATA_OBJECT(),dataset.GetPointer());
      internalOutInfo->Set(vtkCompositeDataPipeline::UPDATE_COMPOSITE_INDICES(),
                           &localIdVector[lBound],levelSize);
      internalOutInfo->Set(vtkCompositeDataPipeline::LOAD_REQUESTED_BLOCKS(),1);

      // put the output info in a vector
      vtkNew< vtkInformationVector > internalOutVector;
      internalOutVector->Append(internalOutInfo.GetPointer());

      // ask the internal reader for its data
      this->Internal->Resolutions[i].Reader->ProcessRequest(
            internalRequestData.GetPointer(),(vtkInformationVector**)NULL,
            internalOutVector.GetPointer());
      // set the block in our output to the output of the internal reader
      }
    else
      {
      dataset->SetNumberOfBlocks(this->Internal->NumberOfBlocksPerLevel);
      }
    output->SetBlock(i,dataset.GetPointer());
    // the next reader's blocks will start where this one's blocks ended
    lBound = uBound;
    if (i == 0)
      {
      output->GetInformation()->Set(vtkDataObject::DATA_TIME_STEP(),
                                    dataset->GetInformation()->Get(vtkDataObject::DATA_TIME_STEP()));
      }
    }
  return 1;
}