//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2015 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
//
#include "ITKSource.hpp"

#include <QFile>
#include <QDialog>
#include <QFileDialog>

#include <itkExtractImageFilter.h>

#include "Geometry.hpp"
#include "Progress.hpp"
#include "Misc.hpp"

#ifdef WIN32
#  define FileDialogOptions QFileDialog::DontUseNativeDialog
#else
#  define FileDialogOptions 0
#endif

namespace mgx
{
  bool ITKImageReader::initialize(QWidget* parent)
  {
    bool is_work_store = stringToWorkStore(parm("Store"));
    if(!checkState().store((is_work_store ? STORE_WORK : STORE_MAIN), parm("Stack number").toInt()))
      return false;
    filename = parm("Filename");

    // Get the file name
    if(parent)
      filename = QFileDialog::getOpenFileName(parent, QString("Select stack file"), filename,
                                            QString("All Stack files (*.*)"), 0, FileDialogOptions);
    if(filename.isEmpty())
      return false;
    setParm("Filename", filename);

    // Check if file can be opened for reading
    QFile file(filename);
    if(!file.open(QIODevice::ReadOnly))
      throw(QString("ITKImageReader::Cannot open input file: %1").arg(filename));
    file.close();

    // Create reader
    QByteArray infile = filename.toLocal8Bit();
    imageIO = itk::ImageIOFactory::CreateImageIO(infile.data(), itk::ImageIOFactory::ReadMode);

    // Set up dialog
    QDialog dlg(parent);
    Ui_ImportStackDialog ui;
    this->ui = &ui;
    ui.setupUi(&dlg);
    QObject::connect(ui.NextSeries, SIGNAL(clicked()), this, SLOT(NextSeriesSlot()));
    QObject::connect(ui.PrevSeries, SIGNAL(clicked()), this, SLOT(PrevSeriesSlot()));

    selectSeries = 1;
    LoadSeries(false);
    if(selectSeries != 0)
      throw(QString("ITKImageReader::Cannot open input file: %1").arg(filename));

    if(dlg.exec() != QDialog::Accepted)
      return false;

    selectComponent = ui.LoadComponent->value() - 1;
    selectChannel = ui.LoadChannel->value() - 1;
    selectTimePoint = ui.LoadTimePoint->value() - 1;
    std::cout << "Image Size:" << size << std::endl;
    std::cout << "Image Step:" << step << std::endl;

    return true;
  }

  void ITKImageReader::LoadSeries(bool next)
  {
    QByteArray infile;
    int series = selectSeries + (next ? 1 : -1);
    if(series < 0)
      return;
    if(series == 0)
      infile = filename.toLocal8Bit();
    else
      infile = (QString("!%1!").arg(series) + filename).toLocal8Bit();

    QApplication::setOverrideCursor(QCursor(Qt::WaitCursor));

    try {
      imageIO = itk::ImageIOFactory::CreateImageIO(infile.data(), itk::ImageIOFactory::ReadMode);
      if(!imageIO)
        throw("Bad imageIO");

      imageIO->SetFileName(infile.data());
      imageIO->ReadImageInformation();
    }
    catch(...) {
      QApplication::restoreOverrideCursor();
      return;
    }
    QApplication::restoreOverrideCursor();

    pixelType = QString::fromStdString(imageIO->GetComponentTypeAsString(imageIO->GetComponentType()));
    componentSize = imageIO->GetComponentSize();
    components = imageIO->GetNumberOfComponents();

    // Find the size and spacing of the image
    size = Point5i(1, 1, 1, 1, 1);
    step = Point5f(.0f, .0f, .0f, .0f, .0f);
    uint dims = imageIO->GetNumberOfDimensions();
    for(uint i = 0; i < dims; i++) {
      size[i] = imageIO->GetDimensions(i);
      step[i] = imageIO->GetSpacing(i) * 1000.0;
    }
    // If we got this far the series is good
    selectSeries = series;
    ui->Series->setText(QString("%1").arg(series + 1));

    ui->Components->setText(QString("%1").arg(components));
    ui->LoadComponent->setRange(1, components);

    ui->Channels->setText(QString("%1").arg(size[4]));
    ui->LoadChannel->setRange(1, size[4]);

    ui->TimePoints->setText(QString("%1").arg(size[3]));
    ui->LoadTimePoint->setRange(1, size[3]);

    ui->StackSize->setText(QString("%1 x %2 x %3").arg(size[0]).arg(size[1]).arg(size[2]));
    ui->StackStep->setText(
      QString("%1%2 x %3%4 x %5%6").arg(step[0]).arg(UM).arg(step[1]).arg(UM).arg(step[2]).arg(UM));
  }

  bool ITKImageReader::run(Stack* stack, Store* store, QString filename)
  {
    if(filename.isEmpty())
      return setErrorMessage("ITKImageReader::Error trying to load a stack from an empty filename.");
    typedef itk::Image<unsigned short int, 5> U5ImageType;
    typedef itk::ImageFileReader<U5ImageType> ReaderType;

    // Start the progress bar
    progressStart(QString("Loading %1 Stack %2")
              .arg((store == stack->main()) ? "main" : "work").arg(stack->userId()), 0, false);

    // Define the reader
    QByteArray infile;
    if(selectSeries == 0)
      infile = filename.toLocal8Bit();
    else
      infile = (QString("!%1!").arg(selectSeries) + filename).toLocal8Bit();
    ReaderType::Pointer reader = ReaderType::New();
    reader->SetFileName(infile.data());
    reader->Update();

    // Get the entire region, and reduce to a single channel
    U5ImageType::RegionType inputRegion = reader->GetOutput()->GetLargestPossibleRegion();
    U5ImageType::SizeType size = inputRegion.GetSize();
    //U5ImageType::SpacingType step = reader->GetOutput()->GetSpacing();

    size[3] = 0;
    size[4] = 0;
    U5ImageType::IndexType start = inputRegion.GetIndex();

    start[3] = selectTimePoint;
    start[4] = selectChannel;
    U5ImageType::RegionType desiredRegion(start, size);

    // Define the filter to extract a channel
    typedef itk::ExtractImageFilter<U5ImageType, UImageType> ExtractorType;
    ExtractorType::Pointer extractor = ExtractorType::New();
    extractor->SetInput(reader->GetOutput());
    extractor->SetExtractionRegion(desiredRegion);
    extractor->SetDirectionCollapseToIdentity();
    extractor->Update();

    UImageType* outImage = extractor->GetOutput();
    // Change spacing, default for ITK is millimeters
    outImage->SetSpacing(outImage->GetSpacing() * 1000.0);

    bool result = UImageConverter::TransferImage(stack, store, outImage);
    if(!result)
      return (setErrorMessage("ITKImageReader::Unable to open stack:" + filename));

    // Check if 8bit and convert to 16
    if(componentSize < 2) {
      ushort maxVal = 0;
      HVecUS& data = store->data();
      for(uint idx = 0; idx < data.size(); idx++)
        if(maxVal < data[idx])
          maxVal = data[idx];
      if(maxVal < 256)
        for(uint idx = 0; idx < data.size(); idx++)
          data[idx] *= 256;
    }

    stack->center();
    store->changed();
    SETSTATUS("Loaded stack " << stack->userId() << ", file:" << filename << " size:" << stack->size()
                              << " step:" << stack->step());
    return true;
  }

  bool ITKImageReader::run()
  {
    int stack_id = parm("Stack number").toInt();
    if(!checkState().stack(STACK_ANY, stack_id))
      return false;
    Stack* stack = this->stack(stack_id);
    Store* store = (stringToWorkStore(parm("Store")) ? stack->work() : stack->main());
    bool res = run(stack, store, parm("Filename"));
    if(res)
      store->show();
    return res;
  }
  REGISTER_PROCESS(ITKImageReader);
}
