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

  Copyright (c) Kitware, Inc.
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/VolViewCopyright.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 "vtkVVSaveVolume.h"

#include "vtkErrorCode.h"
#include "vtkImageData.h"
#include "vtkImageReader2.h"
#include "vtkObjectFactory.h"
#include "vtkImageClip.h"
#include "vtkImageShiftScale.h"
#include "vtkVolumeProperty.h"

#include "vtkKWApplication.h"
#include "vtkKWMessageDialog.h"
#include "vtkKWProgressCommand.h"
#include "vtkKWProgressGauge.h"
#include "vtkKWOpenWizard.h"
#include "vtkKWOpenFileProperties.h"
#include "vtkKWOpenFileHelper.h"
#include "vtkKWWindow.h"

#include "vtkVVDataItemVolume.h"

// Writers

#include "vtkBMPWriter.h"
#include "vtkImageData.h"
#include "vtkImageWriter.h"
#include "vtkJPEGWriter.h"
#include "vtkPNGWriter.h"
#include "vtkPNMWriter.h"
#include "vtkTIFFWriter.h"
#include "vtkWriter.h"
#include "vtkXMLImageDataWriter.h"
#include "vtkMetaImageWriter.h"

#include <vtksys/SystemTools.hxx>

//----------------------------------------------------------------------------
vtkStandardNewMacro(vtkVVSaveVolume);

vtkSetObjectImplementationMacro(vtkVVSaveVolume, DataItemVolume, vtkVVDataItemVolume);
vtkSetObjectImplementationMacro(vtkVVSaveVolume, Window, vtkKWWindow);

//----------------------------------------------------------------------------
vtkVVSaveVolume::vtkVVSaveVolume()
{
  this->DataItemVolume = NULL;
  this->FileName       = NULL;
  this->Writer         = NULL;
  this->ImageWriter    = NULL;
  this->Window         = NULL;
  this->IsImage        = 0;
}

//----------------------------------------------------------------------------
vtkVVSaveVolume::~vtkVVSaveVolume()
{
  this->SetDataItemVolume(NULL);
  this->SetFileName(NULL);
  if (this->Writer)
    {
    this->Writer->SetInputConnection(0,NULL);
    this->Writer->Delete();
    }
  if (this->ImageWriter)
    {
    this->ImageWriter->SetInput(NULL);
    this->ImageWriter->Delete();
    }
  if (this->Window)
    {
    this->SetWindow(NULL);
    }
}

//----------------------------------------------------------------------------
int vtkVVSaveVolume::Write()
{
  int success = 0;

  if (!this->DataItemVolume || !this->FileName)
    {
    vtkErrorMacro("Input or filename not set. Don't know how to continue");
    return success;
    }

  // Write file

  if (!this->InstantiateWriter(this->FileName))
    {
    vtkErrorMacro("Can not instantiate writer");
    return success;
    }

  char *filepattern = new char [strlen(this->FileName) + 10];
  *filepattern = '\0';
  
  // ImageWriter = a 2D format writer, series

  if (this->IsImage)
    {
    success = this->WriteImages(this->FileName, filepattern);
    }
  else
    {
    vtkKWProgressCommand *cb = vtkKWProgressCommand::New();
    cb->SetWindow(this->Window);
    cb->SetStartMessage("Saving volume to disk");

    // Writer = vtkXMLImageDataWriter

    if (vtkXMLImageDataWriter *xmliw = vtkXMLImageDataWriter::SafeDownCast(this->Writer))
      {
      xmliw->SetFileName(this->FileName);
      xmliw->SetInput(this->DataItemVolume->GetImageData());
      xmliw->EncodeAppendedDataOn();
      xmliw->AddObserver(vtkCommand::StartEvent, cb);
      xmliw->AddObserver(vtkCommand::ProgressEvent, cb);
      xmliw->AddObserver(vtkCommand::EndEvent, cb);
      xmliw->Write();
      xmliw->RemoveObserver(cb);
      xmliw->SetInput(0);

      success = (xmliw->GetErrorCode() == 
                 vtkErrorCode::OutOfDiskSpaceError) ? 0 : 1;
      }

    else if (vtkMetaImageWriter *mw = vtkMetaImageWriter::SafeDownCast(this->Writer))
      {
      mw->SetFileName(this->FileName);
      mw->SetInput(this->DataItemVolume->GetImageData());
      mw->AddObserver(vtkCommand::StartEvent, cb);
      mw->AddObserver(vtkCommand::ProgressEvent, cb);
      mw->AddObserver(vtkCommand::EndEvent, cb);
      mw->Write();
      mw->RemoveObserver(cb);
      mw->SetInput(0);

      success = (mw->GetErrorCode() == 
                 vtkErrorCode::OutOfDiskSpaceError) ? 0 : 1;
      }

    // ImageWriter = vtkImageWriter (RAW)

    else
      {
      this->ImageWriter->SetFileName(this->FileName);
      this->ImageWriter->SetInput(this->DataItemVolume->GetImageData());
      this->ImageWriter->SetFileDimensionality(3);
      this->ImageWriter->AddObserver(vtkCommand::StartEvent, cb);
      this->ImageWriter->AddObserver(vtkCommand::ProgressEvent, cb);
      this->ImageWriter->AddObserver(vtkCommand::EndEvent, cb);
      this->ImageWriter->Write();
      this->ImageWriter->RemoveObserver(cb);
      this->ImageWriter->SetInput(0);
        
      success = (this->ImageWriter->GetErrorCode() == 
                 vtkErrorCode::OutOfDiskSpaceError) ? 0 : 1;

      strcpy(filepattern, this->FileName);
      }

    cb->Delete();
    }

  // Write the .vvi file

  if (success)
    {
    vtkKWOpenWizard *wiz = vtkKWOpenWizard::New();
    vtkKWOpenFileProperties *file_prop = wiz->GetOpenFileProperties();
    vtkImageData *img_data = this->DataItemVolume->GetImageData();
    file_prop->CopyFromImageData(img_data);

    file_prop->SetIndependentComponents(
      this->DataItemVolume->GetVolumeProperty()->GetIndependentComponents());
    file_prop->SetDistanceUnits(this->DataItemVolume->GetDistanceUnits());
    int i;
    for (i = 0; i < img_data->GetNumberOfScalarComponents(); i++)
      {
      file_prop->SetScalarUnits(i, this->DataItemVolume->GetScalarUnits(i));
      }
    file_prop->SetScope(this->DataItemVolume->GetScope());
    
    // Specify info regarding pattern/dimensionality
    // We have to create a fake image reader here

    //vtkImageReader2 *rdr = NULL;
    if (this->ImageWriter && *filepattern)
      {
      //rdr = vtkImageReader2::New();
      //rdr->SetFilePattern(filepattern);
      file_prop->SetFilePattern(filepattern);
      //rdr->SetDataExtent(valid_info->GetWholeExtent());
      //rdr->SetFileDimensionality(this->ImageWriter->GetFileDimensionality());
      file_prop->SetFileDimensionality(this->ImageWriter->GetFileDimensionality());
      //wiz->SetLastReader(rdr);
      if(vtkBMPWriter::SafeDownCast(this->ImageWriter))
        {
        file_prop->SetNumberOfScalarComponents(3);
        if (img_data->GetNumberOfScalarComponents() == 1)
          {
          file_prop->SetIndependentComponents(0);
          }
        }
      }
    else
      {
      file_prop->SetFilePattern(NULL);
      }

    // Write the VVI

    wiz->WriteVVIForFile(this->FileName);

    wiz->Delete();
    }

  delete [] filepattern;

  return success;
}

//----------------------------------------------------------------------------
int vtkVVSaveVolume::InstantiateWriter(const char* file)
{  
  int success = 0;

  // Check for extention

  const char *ext3 = NULL;
  const char *ext4 = NULL;
  if (strlen(file) > 4)
    {
    ext3 = file + strlen(file) - 4;
    }
  if (strlen(file) > 5)
    {
    ext4 = file + strlen(file) - 5;
    }

  if (this->Writer)
    {
    this->Writer->Delete();
    this->Writer = NULL;
    }

  if (this->ImageWriter)
    {
    this->ImageWriter->Delete();
    this->ImageWriter = NULL;
    }

  this->IsImage = 0;

  if (!ext3 && !ext4)
    {
    vtkErrorMacro("Filename too short, don't know how to write it.");
    }
  else if (!strcmp(ext3, ".tif") || !strcmp(ext4, ".tiff"))
    {
    this->ImageWriter = vtkTIFFWriter::New();
    this->IsImage = 1;
    success = 1;
    }
  else if (!strcmp(ext3, ".jpg") || !strcmp(ext4, ".jpeg"))
    {
    this->ImageWriter = vtkJPEGWriter::New();
    this->IsImage = 1;
    success = 1;
    }
  else if (!strcmp(ext3, ".bmp"))
    {
    this->ImageWriter = vtkBMPWriter::New();
    this->IsImage = 1;
    success = 1;
    }
  else if (!strcmp(ext3, ".png"))
    {
    this->ImageWriter = vtkPNGWriter::New();
    this->IsImage = 1;
    success = 1;
    }
  else if (!strcmp(ext3, ".ppm") || !strcmp(ext3, ".pgm"))
    {
    this->ImageWriter = vtkPNMWriter::New();
    this->IsImage = 1;
    success = 1;
    }
  else if (!strcmp(ext3, ".vti"))
    {
    this->Writer = vtkXMLImageDataWriter::New();
    success = 1;
    }
  else if (!strcmp(ext3, ".mha"))
    {
    // Metaimage file.. Set compression true on this one
    this->Writer = vtkMetaImageWriter::New();
    static_cast< vtkMetaImageWriter *>(this->Writer)->SetCompression(true);
    success = 1;
    }
  else if (!strcmp(ext3, ".mhd"))
    {
    // Metaimage file.. Set compression to false on this one. Most people tend
    // preserve the header and then have an accompanying raw version that they
    // can import into other tools. Hence we'd like the raw part un-compressed.
    this->Writer = vtkMetaImageWriter::New();
    static_cast< vtkMetaImageWriter *>(this->Writer)->SetCompression(false);
    success = 1;
    }
  else if (!strcmp(ext3, ".raw"))
    {
    this->ImageWriter = vtkImageWriter::New();
    success = 1;
    }
  
  return success;
}

//----------------------------------------------------------------------------
int vtkVVSaveVolume::WriteImages(const char* filename, char *filepattern)
{
  if (!this->DataItemVolume || !filename)
    {
    vtkErrorMacro("Input or filename not set. Don't know how to continue");
    return 0;
    }

  int cc, i;

  // Get file extension and base name

  const char *ext3 = NULL;
  const char *ext4 = NULL;
  const char *ext = NULL;
  char basename[1024];

  if (strlen(filename) > 4)
    {
    ext3 = filename + strlen(filename) - 4;
    if (ext3[0] == '.')
      {
      strcpy(basename, filename);
      basename[strlen(filename) - 4] = '\0';
      ext3++;
      ext = ext3;
      }    
    }

  if (strlen(filename) > 5)
    {
    ext4 = filename + strlen(filename) - 5;
    if (ext4[0] == '.')
      {
      strcpy(basename, filename);
      basename[strlen(filename) - 5] = '\0';
      ext4++;
      ext = ext4;
      }
    }

  if (!ext)
    {
    return 0;
    }

  // Check if it is a series, and eventually delete the old series first

  int notpattern = 0;
  char *pattern = new char [strlen(filename) + 20];
  int zmin = 0, zmax = 0;
  
  vtkKWOpenFileHelper::FindSeriesPattern(filename, pattern, &zmin, &zmax);
  if (zmin == zmax)
    {
    notpattern = 1;    
    }
  else
    {
    if (!this->AskAndDeleteSeries(filename, pattern, zmin, zmax))
      {
      delete [] pattern;
      return 0;
      }
    }

  // Make new filename

  char *newfile = new char[strlen(filename) + 10];

  // Setup pipeline

  vtkImageClip *imc = vtkImageClip::New();
  vtkImageShiftScale *iss = NULL;
  vtkImageData *img_data = this->DataItemVolume->GetImageData();

  if(vtkPNGWriter::SafeDownCast(this->ImageWriter) &&
     (img_data->GetScalarType() == VTK_UNSIGNED_SHORT ||
      img_data->GetScalarType() == VTK_UNSIGNED_CHAR))
    {
    imc->SetInput(img_data);
    }
  else
    {
    iss = vtkImageShiftScale::New();
    iss->SetInput(img_data);
    if (vtkPNGWriter::SafeDownCast(this->ImageWriter))
      {
      iss->SetOutputScalarTypeToUnsignedShort();
      }
    else
      {
      iss->SetOutputScalarTypeToUnsignedChar();
      }

    imc->SetInput(iss->GetOutput());
    }

  this->ImageWriter->SetInput(imc->GetOutput());
  img_data->Update();

  // We need to scale image down!!

  if(iss)
    {
    double* range = img_data->GetScalarRange();
    double trange = (double)range[1] - (double)range[0];
    if (trange > 255)
      {
      iss->SetShift(-range[0]);
      iss->SetScale(255.0 / trange);
      }
    }
  // imc->GetOutput()->Update(); // this would duplicate the volume in mem !
  imc->GetOutput()->UpdateInformation();

  // Make a copy of the extent as it can change in the for loop

  int *twext = imc->GetOutput()->GetWholeExtent();

  int wext[6];
  for(i =0; i < 6; i++)
    {
    wext[i] = twext[i];
    }
  
  // compute the format string
  // %0Xd where X is the number of 0's required to pad the digits

  char formatstring[1024];
  int depth = 0;
  int val;
  
  for(val = wext[5]; val > 0; val /= 10)
    {
    ++depth;
    }
  if (notpattern)
    {
    sprintf(formatstring, "%s.%%0%dd.%s", basename, depth, ext);
    }
  else
    {
    strcpy(formatstring, pattern);
    }

  int success = 1;
  
  for (cc = wext[4]; cc <= wext[5]; cc++)
    {
    sprintf(newfile, formatstring, cc);
    if (!vtksys::SystemTools::FileExists(newfile))
      {
      vtkKWOpenFileHelper::FindSeriesPattern(newfile, pattern, &zmin, &zmax);
      if (zmin != zmax && 
          !this->AskAndDeleteSeries(newfile, pattern, zmin, zmax))
        {
        success = 0;
        }
      }
    }

  // Write each slice

  if (success)
    {
    this->Window->SetStatusText("Saving volume to disk as series");

    for (cc = wext[4]; cc <= wext[5]; cc++)
      {
      this->Window->GetProgressGauge()->SetValue(
        (int)(100.0 * (float)(cc - wext[4]) / (float)(wext[5] - wext[4])));

      imc->SetOutputWholeExtent(wext[0], wext[1], wext[2], wext[3], cc, cc);
      imc->Update();
      sprintf(newfile, formatstring, cc);
      this->ImageWriter->SetFileName(newfile);
      this->ImageWriter->Write();

      // Did not work: remove all files

      if (this->ImageWriter->GetErrorCode() == 
          vtkErrorCode::OutOfDiskSpaceError)
        {
        success = 0;
        for (i = wext[4]; i < cc; i++)
          {
          sprintf(newfile, formatstring, i);
          vtksys::SystemTools::RemoveFile(newfile);
          }
        break;
        }
      }
  
    this->Window->GetProgressGauge()->SetValue(0);
    this->Window->SetStatusText("Saving volume to disk as series -- Done");

    if (success && filepattern)
      {
      strcpy(filepattern, formatstring);
      }
    }

  // Cleanup
  
  if (iss)
    {
    iss->Delete();
    }
  imc->Delete();  
  
  delete [] newfile;  
  delete [] pattern;
  
  return success;
}

//----------------------------------------------------------------------------
int vtkVVSaveVolume::AskAndDeleteSeries(
  const char* filename, const char* pattern, int zmin, int zmax)
{
  ostrstream str;
  char *nfilename = new char [strlen(filename) + 20];    

  str << "The series of images you selected already exists. "
    "VolView will overwrite the old series. In the process files:\n";
  sprintf(nfilename, pattern, zmin);
  str << nfilename << "\n ... \n";
  sprintf(nfilename, pattern, zmax);
  str << nfilename << "\nwill be deleted.\n"
    "Are you sure you want to do that?" << ends;

  int res = vtkKWMessageDialog::PopupYesNo(
    this->GetApplication(), this->Window, "Save Series", 
    str.str(), vtkKWMessageDialog::WarningIcon);
  str.rdbuf()->freeze(0);

  if (res)
    {
    int cc;
    for (cc = zmin; cc <= zmax; cc++)
      {
      sprintf(nfilename, pattern, cc);
      vtksys::SystemTools::RemoveFile(nfilename);
      }
    }

  delete [] nfilename;

  return res;
}

//----------------------------------------------------------------------------
void vtkVVSaveVolume::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  
  os << indent << "DataItemVolume: " << this->DataItemVolume << endl;
  os << indent << "FileName: " << this->FileName << endl;
}

