/*****************************************************************************
 * $CAMITK_LICENCE_BEGIN$
 *
 * CamiTK - Computer Assisted Medical Intervention ToolKit
 * (c) 2001-2018 Univ. Grenoble Alpes, CNRS, TIMC-IMAG UMR 5525 (GMCAO)
 *
 * Visit http://camitk.imag.fr for more information
 *
 * This file is part of CamiTK.
 *
 * CamiTK is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * CamiTK 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 version 3 for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with CamiTK.  If not, see <http://www.gnu.org/licenses/>.
 *
 * $CAMITK_LICENCE_END$
 ****************************************************************************/

#include "VtkImageComponentExtension.h"
#include "VtkImageComponent.h"
#include "RawImageComponent.h"

#include <Log.h>

using namespace camitk;

// -- QT stuff
#include <QFile>
#include <QTextStream>
#include <QFileInfo>

// vtk image writers
#include <vtkImageWriter.h>
#include <vtkJPEGWriter.h>
#include <vtkPNGWriter.h>
#include <vtkTIFFWriter.h>
#include <vtkBMPWriter.h>
#include <vtkPNMWriter.h>
#include <vtkMetaImageWriter.h>
#include <vtkImageShiftScale.h>
#include <vtkType.h>

// --------------- getName -------------------
QString VtkImageComponentExtension::getName() const {
    return "vtkImages Component";
}

// --------------- getDescription -------------------
QString VtkImageComponentExtension::getDescription() const {
    return tr("Manage any file type supported by vtk in <b>CamiTK</b>");
}

// --------------- getFileExtensions -------------------
QStringList VtkImageComponentExtension::getFileExtensions() const {
    QStringList ext;
    ext << "jpg" << "png" << "tiff" << "tif" << "bmp" << "pbm" << "pgm" << "ppm";
    ext << "mhd" << "mha";
    ext << "raw";
    return ext;
}

// --------------- open -------------------
Component* VtkImageComponentExtension::open(const QString& fileName) {
    try {
        if (fileName.endsWith(".raw")) {
            return new RawImageComponent(fileName);
        }
        else {
            return new VtkImageComponent(fileName);
        }
    }
    catch (const AbortException& e) {
        throw e;
    }
}

// --------------- save -------------------
bool VtkImageComponentExtension::save(Component* component) const {
    ImageComponent* img = dynamic_cast<ImageComponent*>(component);
    if (!img) {
        return false;
    }
    else {

        QFileInfo fileInfo(component->getFileName());
        vtkSmartPointer<vtkImageData> image; //img->getImageDataWithFrameTransform();

        if (fileInfo.completeSuffix() == "raw") {
            image = img->getImageDataWithFrameTransform();
            //-- save as raw
            // vtk writer
            vtkSmartPointer<vtkImageWriter> writer;

            // write the image as raw data
            writer = vtkSmartPointer<vtkImageWriter>::New();
            writer->SetFileDimensionality(3);
            writer->SetFileName(fileInfo.absoluteFilePath().toStdString().c_str());
            writer->SetInputData(image);

            try {
                writer->Write();
            }
            catch (...) {
                return false;
            }

            // Write info in a separated file
            QFile infoFile(fileInfo.absoluteDir().absolutePath() + fileInfo.baseName() + ".info");
            infoFile.open(QIODevice::ReadWrite | QIODevice::Text);
            QTextStream out(&infoFile);
            out << "Raw data info (Written by CamiTK VtkImageComponentExtension)" << endl;
            out << "Data Filename:\t" << fileInfo.absoluteFilePath() << endl;
            int* dims = image->GetDimensions();
            out << "Image dimensions:\t" << dims[0] << "\t" << dims[1] << "\t" << dims[2] << endl;
            out << "Voxel storage type:\t" << image->GetScalarTypeAsString() << endl;
            out << "Number of scalar components:\t" << image->GetNumberOfScalarComponents() << endl;
            double* voxelSpace = image->GetSpacing();
            out << "Voxel spacing:\t" << voxelSpace[0] << "\t" << voxelSpace[1] << "\t" << voxelSpace[2] << endl;
            double* imageOrigin = image->GetOrigin();
            out << "Image origin:\t" << imageOrigin[0] << "\t" << imageOrigin[1] << "\t" << imageOrigin[2] << endl;
            infoFile.flush();
            infoFile.close();
            return true;
        }
        else {
            image = img->getImageData(); //img->getImageDataWithFrameTransform();

            //-- save as vtk image
            // vtk writer
            vtkSmartPointer<vtkImageWriter> writer;

            // filename extension
            QString fileExt = fileInfo.suffix();

            // filename prefix (for 2D image series)
            QString filePattern = fileInfo.absoluteDir().absolutePath() + "/" + fileInfo.baseName();
            filePattern.append("_%04u.").append(fileExt);

            // for any file format, image should be unsigned char casted (see vtk image writers doc).
            // to cast an image into another scalar type, use vtkImageShiftScale instead of vtkImageCast (see doc).
            vtkSmartPointer<vtkImageShiftScale> castFilter = vtkSmartPointer<vtkImageShiftScale>::New();
            // get image scalar range
            double* imgRange = image->GetScalarRange();
            double scalarTypeMin = imgRange[0];
            double scalarTypeMax = imgRange[1];

            // shift according to the (un)signed type of the scalar
            int scalarType = image->GetScalarType();
            switch (scalarType) {
                case VTK_CHAR:
                case VTK_SIGNED_CHAR:
                case VTK_SHORT:
                case VTK_INT:
                case VTK_LONG:
                case VTK_FLOAT:
                case VTK_DOUBLE: {
                    // shift only for signed scalar types
                    double shift = (scalarTypeMax - scalarTypeMin) / 2;
                    castFilter->SetShift(shift);
                }
                break;
                case VTK_UNSIGNED_CHAR:
                case VTK_UNSIGNED_SHORT:
                case VTK_UNSIGNED_INT:
                case VTK_UNSIGNED_LONG:
                default:
                    // do not shift, only scale
                    break;
            }

            // scale
            double scale = (double) 255.0 / (image->GetScalarTypeMax() - image->GetScalarTypeMin());
            castFilter->SetScale(scale);

            // convert to unsigned char
            castFilter->SetOutputScalarTypeToUnsignedChar();
            castFilter->SetInputData(image);
            castFilter->Update();

            // Writer initialization, depending on file extension
            if (QString::compare(fileExt, "jpg", Qt::CaseInsensitive) == 0) {
                writer = vtkSmartPointer<vtkJPEGWriter>::New();
                writer->SetFileDimensionality(2);
            }
            else if (QString::compare(fileExt, "png", Qt::CaseInsensitive) == 0) {
                writer = vtkSmartPointer<vtkPNGWriter>::New();
                writer->SetFileDimensionality(2);
            }
            else if ((QString::compare(fileExt, "tiff", Qt::CaseInsensitive) == 0) ||
                     (QString::compare(fileExt, "tif",  Qt::CaseInsensitive) == 0)) {
                writer = vtkSmartPointer<vtkTIFFWriter>::New();
                writer->SetFileDimensionality(2);
            }
            else if (QString::compare(fileExt, "bmp", Qt::CaseInsensitive) == 0) {
                writer = vtkSmartPointer<vtkBMPWriter>::New();
                writer->SetFileDimensionality(2);
            }
            else if ((QString::compare(fileExt, "pbm", Qt::CaseInsensitive) == 0) ||
                     (QString::compare(fileExt, "pgm", Qt::CaseInsensitive) == 0) ||
                     (QString::compare(fileExt, "ppm", Qt::CaseInsensitive) == 0)) {
                writer = vtkSmartPointer<vtkPNMWriter>::New();
                writer->SetFileDimensionality(2);
            }
            else if ((QString::compare(fileExt, "mhd", Qt::CaseInsensitive) == 0) ||
                     (QString::compare(fileExt, "mha", Qt::CaseInsensitive) == 0)) {
                vtkSmartPointer<vtkMetaImageWriter> metaImgWriter = vtkSmartPointer<vtkMetaImageWriter>::New();
                metaImgWriter->SetCompression(false);
                writer = metaImgWriter;
                writer->SetFileDimensionality(3);
            }
            else {
                CAMITK_WARNING(tr("Saving Error: cannot save file: unrecognized extension \".%1\"").arg(fileExt))
                return false;
            }

            if (writer) {
                if ((image->GetDataDimension() == 2) ||
                        (QString::compare(fileExt, "mhd", Qt::CaseInsensitive) == 0) ||
                        (QString::compare(fileExt, "mha", Qt::CaseInsensitive) == 0)) {
                    writer->SetFileName(fileInfo.absoluteFilePath().toStdString().c_str());
                }
                else {
                    writer->SetFilePattern(filePattern.toStdString().c_str());
                }

                // we save the unsigned char casted image
                writer->SetInputConnection(castFilter->GetOutputPort());
                writer->Write();

                // Save image orientation for compatible header (as it is not possible to store it either in
                // vtkImageData or vtkMetaImageWriter classes
                // we are obliged to open the whole image, as a text, replace the anatomical orientation string
                // store it back in ASCII (to preserve data format)
                if ((QString::compare(fileExt, "mha", Qt::CaseInsensitive) == 0) || (QString::compare(fileExt, "mhd", Qt::CaseInsensitive) == 0)) {
                    QByteArray fileData;
                    QFile file(component->getFileName());
                    file.open(QIODevice::ReadWrite);
                    fileData = file.readAll();

                    QString text(fileData);
                    // Store anatomical orientation in header
                    text.replace(QString("AnatomicalOrientation = ???"), QString("AnatomicalOrientation = ") + ImageOrientationHelper::getOrientationAsQString(img->getInitialOrientation()));

                    file.seek(0); // go to the beginning of the file
                    file.write(text.toLatin1()); // write back the file, in ASCII to preserve data format

                    file.close(); // close the file handle.
                }
            }

            return true;
        }
    }
}
