/*****************************************************************************
 * $CAMITK_LICENCE_BEGIN$
 *
 * CamiTK - Computer Assisted Medical Intervention ToolKit
 * (c) 2001-2025 Univ. Grenoble Alpes, CNRS, Grenoble INP - UGA, TIMC, 38000 Grenoble, France
 *
 * 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 "Transformation.h"

#include <Log.h>

using namespace camitk;

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

// vtk image writers
// disable warning generated by clang about the surrounded headers
#include <CamiTKDisableWarnings>
// 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 <CamiTKReEnableWarnings>

#include <vtkType.h>

// --------------- getName -------------------
QString VtkImageComponentExtension::getName() const {
    return "VTK Image";
}

// --------------- getDescription -------------------
QString VtkImageComponentExtension::getDescription() const {
    return tr("Manage image file type directly supported by vtk. For MetaImage documentation, check <a href=\"https://itk.org/Wiki/ITK/MetaIO\">the official documentation (https://itk.org/Wiki/ITK/MetaIO)</a>.");
}

// --------------- getFileExtensions -------------------
QStringList VtkImageComponentExtension::getFileExtensions() const {
    QStringList ext;
    ext << "mha" << "mhd";
    ext << "png" << "tiff" << "tif" << "bmp" << "pbm" << "pgm" << "ppm" << "jpg";
    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;

        if (fileInfo.completeSuffix() == "raw") {
            image = img->getImageData(); // do not use img->getImageDataWithFrameTransform(); as image was already converted to RAI at reading time
            //-- 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 (...) {
                CAMITK_WARNING(tr("Saving Error: cannot save file: exception during writing \".%1\"").arg(fileInfo.absoluteFilePath()))
                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)" << Qt::endl;
            out << "Data Filename:\t" << fileInfo.absoluteFilePath() << Qt::endl;
            int* dims = image->GetDimensions();
            out << "Image dimensions:\t" << dims[0] << "\t" << dims[1] << "\t" << dims[2] << Qt::endl;
            out << "Voxel storage type:\t" << image->GetScalarTypeAsString() << Qt::endl;
            out << "Number of scalar components:\t" << image->GetNumberOfScalarComponents() << Qt::endl;
            double* voxelSpace = image->GetSpacing();
            out << "Voxel spacing:\t" << voxelSpace[0] << "\t" << voxelSpace[1] << "\t" << voxelSpace[2] << Qt::endl;
            double* imageOrigin = image->GetOrigin();
            out << "Image origin:\t" << imageOrigin[0] << "\t" << imageOrigin[1] << "\t" << imageOrigin[2] << Qt::endl;
            infoFile.flush();
            infoFile.close();
            return true;
        }
        else {
            image = img->getImageData();

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

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

            // Writer initialization, depending on file extension
            if ((QString::compare(fileExt, "mhd", Qt::CaseInsensitive) == 0) ||
                    (QString::compare(fileExt, "mha", Qt::CaseInsensitive) == 0)) {
                // saving meta image is simple
                vtkSmartPointer<vtkMetaImageWriter> metaImgWriter = vtkSmartPointer<vtkMetaImageWriter>::New();
                metaImgWriter->SetCompression(true);
                writer = metaImgWriter;
                writer->SetFileDimensionality(3);
                writer->SetFileName(fileInfo.absoluteFilePath().toStdString().c_str());
                writer->SetInputData(image);
            }
            else {
                // Saving as other classic 2D images

                // 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 {
                                    CAMITK_WARNING(tr("Saving Error: cannot save file: unrecognized extension \".%1\"").arg(fileExt))
                                    return false;
                                }
                            }
                        }
                    }
                }

                // for any other classic 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();

                if (image->GetDataDimension() == 2)  {
                    // saving 2D to 3D
                    writer->SetFileName(fileInfo.absoluteFilePath().toStdString().c_str());
                }
                else {
                    // saving 3D to 2D
                    // filename prefix (for 2D image series)
                    QString filePattern = fileInfo.absoluteDir().absolutePath() + "/" + fileInfo.baseName();
                    filePattern.append("_%04u.").append(fileExt);

                    writer->SetFilePattern(filePattern.toStdString().c_str());
                }

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

            // write to the file
            try {
                writer->Write();
            }
            catch (...) {
                CAMITK_WARNING(tr("Saving Error: cannot save file: exception during writing \".%1\"").arg(fileInfo.absoluteFilePath()))
                return false;
            }

            // Save image position and rotation when exporting to MHA
            // As it is not used by vtkMetaImageWriter,
            // we need to open the saved file as text and replace the corresponding fields TransformMatrix
            // and Offset
            //
            // For AnatomicalOrientation, see https://itk.org/Wiki/ITK/MetaIO, where it states that
            // the default seems to be LPS (athough this is not very clear see the corresponding
            // section in https://itk.org/Wiki/ITK/MetaIO/Documentation#Spatial_Objects)
            // In CamiTK, as soon as the image is read, it is transformed to RAI.
            // Therefore the writen image anotomical orientation is set to RAI.
            if ((QString::compare(fileExt, "mha", Qt::CaseInsensitive) == 0) || (QString::compare(fileExt, "mhd", Qt::CaseInsensitive) == 0)) {
                // prepare string replacement
                vtkSmartPointer<vtkMatrix4x4> currentMatrix = img->getMainTransformation()->getMatrix();
                QString transformMatrix = "TransformMatrix =";

                // According to https://itk.org/Wiki/MetaIO/Documentation#Associated_transformations
                // "matrix that is serialized in a column-major format" -> save column by column
                for (int i = 0; i < 3; i++) {//column
                    for (int j = 0; j < 3; j++) {//line
                        if (fabs(currentMatrix->GetElement(j, i)) < 1e-6) {
                            currentMatrix->SetElement(j, i, 0.0);
                        }

                        transformMatrix += " " + QString::number(currentMatrix->GetElement(j, i));
                    }
                }

                transformMatrix += "\n";
                QRegularExpression transformRegExp = QRegularExpression("TransformMatrix = ([-+]?[0-9]*\\.?[0-9]+([eE][-+]?[0-9]+)?[ |\\n])+");

                auto* pos = new double[3];
                img->getMainTransformation()->getTransform()->GetPosition(pos);

                for (int i = 0; i < 3; i++) {
                    if (fabs(pos[i]) < 1e-6) {
                        pos[i] = 0.0;
                    }
                }

                QString offset = "Offset = "
                                 + QString::number(pos[0], 'g', 3) + " "
                                 + QString::number(pos[1], 'g', 3) + " "
                                 + QString::number(pos[2], 'g', 3) + "\n";

                QRegularExpression offsetRegExp = QRegularExpression("Offset = ([-+]?[0-9]*\\.?[0-9]+([eE][-+]?[0-9]+)?[ |\\n])+");

                // anotomical orientation is replaced by RAI as CamiTK has transformed all orientation to RAI internally when
                // the ImageComponent was created.
                QRegularExpression anatomicalRegExp = QRegularExpression("AnatomicalOrientation = \\\?\\\?\\\?");

                // a MetaImage header line is always a capitalized word followed by " = " and a value
                QRegularExpression headerRegExp = QRegularExpression("[A-Z][A-z]* = \\.*");

                // read the file and replace the strings
                QFile file(fileInfo.absoluteFilePath());
                file.open(QIODevice::ReadOnly);
                // the new meta image with tweaked field (that the MetaImagWriter did not write properly)
                QBuffer tweakedMetaImage;
                tweakedMetaImage.open(QBuffer::WriteOnly);

                // assume that the text header lines are less than 256 char
                char buffer[256];
                qint64 readSize = file.readLine(buffer, sizeof(buffer));
                int nbOfReplace = 0;
                QString line(buffer);

                // read the header
                while (readSize > 0 && line.contains(headerRegExp) && nbOfReplace < 3) {
                    // check header
                    if (line.contains(transformRegExp)) {
                        line.replace(transformRegExp, transformMatrix);
                        nbOfReplace++;
                    }
                    else {
                        if (line.contains(offsetRegExp)) {
                            line.replace(offsetRegExp, offset);
                            nbOfReplace++;
                        }
                        else {
                            if (line.contains(anatomicalRegExp)) {
                                line.replace(anatomicalRegExp, QString("AnatomicalOrientation = RAI"));
                                nbOfReplace++;
                            }
                        }
                    }

                    tweakedMetaImage.write(line.toUtf8());

                    readSize = file.readLine(buffer, sizeof(buffer));
                    line = buffer;
                }

                if (readSize > 0) {
                    // header is finished, not the file, this is a mha
                    // flush the end of the file to tweakedBuffer
                    tweakedMetaImage.write(buffer, readSize);
                    while (readSize > 0) {
                        readSize = file.read(buffer, sizeof(buffer));
                        tweakedMetaImage.write(buffer, readSize);
                    }
                }

                // now overwrite the file
                tweakedMetaImage.close();
                file.close();
                file.open(QIODevice::WriteOnly);
                file.write(tweakedMetaImage.buffer());
                file.close();
            }
        }

        return true;
    }
}