/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <addons/nifti/niftiimage.hh>
#include <nifti1_io.h>

#include <mia/core/attribute_names.hh>

using namespace mia;
using namespace std;

NS_BEGIN(niftiimage)

static const char *AttrID_nifti_sform = "nifti-sform";
static const char *AttrID_nifti_sform_code = "nifti-sform-code";
static const char *AttrID_nifti_qform_code = "nifti-qform-code";
static const char *AttrID_nifti_intent_code = "nifti-intent-code";
static const char *AttrID_nifti_intent_name = "nifti-intent-name";
static const char *AttrID_nifti_intent_p1 = "nifti-intent_p1";
static const char *AttrID_nifti_intent_p2 = "nifti-intent_p2";
static const char *AttrID_nifti_intent_p3 = "nifti-intent_p3";



static const char *AttrID_nifti_toffset = "nifti-toffset";       // float
static const char *AttrID_nifti_xyz_units = "nifti-xyz_units";   // int
static const char *AttrID_nifti_time_units = "nifti-time_units"; // int

/*
  Within MIA the orientation of the images is stored like given in DICOM.
  To convert to a nifti-like orientation the following formula is used:

  input:
     SIZE image seize in pixels
     SCALE pixeldim in mm
     ORIG_IN location of first pixel in pyhsical space
     R_IN    rotation orientation matrix
  outut:
     ORIG_OUT location of first pixel in pyhsical space (nifti)
     R_OUT    rotation orientation matrix (nifti)

                 / -1, -1,  1 \
  R_OUT = R_IN * | -1, -1,  1 |
                 \  1,  1, -1 /

    dz = (SIZE.z - 1) * scale.z;
  ORIG_OUT.x = - R_IN.z.x * dz - ORIG_IN.x;
  ORIG_OUT.y = - R_IN.z.y * dz - ORIG_IN.y;
  ORIG_OUT.z =   R_IN.z.z * dz + ORIG_IN.z;

  When reading a nifti image the inverse operation is applied.

*/


CNifti3DImageIOPlugin::CNifti3DImageIOPlugin():
       C3DImageIOPlugin("nifti")
{
       add_supported_type(it_sbyte);
       add_supported_type(it_ubyte);
       add_supported_type(it_sshort);
       add_supported_type(it_ushort);
       add_supported_type(it_sint);
       add_supported_type(it_uint);
       add_supported_type(it_slong);
       add_supported_type(it_ulong);
       add_supported_type(it_float);
       add_supported_type(it_double);
       add_suffix(".nii");
       add_suffix(".NII");
       CVFloatTranslator::register_for(AttrID_nifti_sform);
       CSITranslator::register_for(AttrID_nifti_sform_code);
       CSITranslator::register_for(AttrID_nifti_qform_code);
       CSITranslator::register_for(AttrID_nifti_intent_code);
       CFloatTranslator::register_for("nifti-intent_p1");
       CFloatTranslator::register_for("nifti-intent_p2");
       CFloatTranslator::register_for("nifti-intent_p3");
       CFloatTranslator::register_for(IDRescaleSlope);
       CFloatTranslator::register_for(IDRescaleIntercept);
       CFloatTranslator::register_for(AttrID_nifti_toffset);
       CSITranslator::register_for(AttrID_nifti_xyz_units);
       CSITranslator::register_for(AttrID_nifti_time_units);
}

struct NiftiDeallocator {
       void operator () (nifti_image *image)
       {
              nifti_image_free(image);
       }
};

void copy_attributes(C3DImage& image, const nifti_image& ni)
{
       // set position and orientation
       image.set_attribute(AttrID_nifti_qform_code, ni.qform_code);

       if (ni.qform_code == 0) { // method 1
              image.set_orientation(ior_default);
              image.set_voxel_size(C3DFVector(ni.dx,  ni.dy, ni.dz));
       } else { // Method 2
              image.set_orientation(ior_xyz);
              image.set_voxel_size(C3DFVector(ni.dx,  ni.dy, ni.dz));
              mat44 mat = nifti_quatern_to_mat44( ni.quatern_b, ni.quatern_c, ni.quatern_d,
                                                  0.0f, 0.0f, 0.0f,
                                                  1.0f, 1.0f, 1.0f, ni.qfac);
              C3DDMatrix rot(C3DDVector(-mat.m[0][0], -mat.m[1][0], mat.m[2][0]),
                             C3DDVector(-mat.m[0][1], -mat.m[1][1], mat.m[2][1]),
                             C3DDVector( mat.m[0][2], mat.m[1][2], -mat.m[2][2]));
              image.set_rotation(rot);
              const float dz = (ni.nz - 1) * ni.dz;
              C3DFVector org(- rot.z.x * dz - ni.qoffset_x,
                             - rot.z.y * dz - ni.qoffset_y,
                             - rot.z.z * dz + ni.qoffset_z);
              image.set_origin(org);
       }

       if (ni.sform_code > 0) { // method 3
              vector<float> am = {ni.sto_xyz.m[0][0], ni.sto_xyz.m[0][1], ni.sto_xyz.m[0][2], ni.sto_xyz.m[0][3],
                                  ni.sto_xyz.m[1][0], ni.sto_xyz.m[1][1], ni.sto_xyz.m[1][2], ni.sto_xyz.m[1][3],
                                  ni.sto_xyz.m[2][0], ni.sto_xyz.m[2][1], ni.sto_xyz.m[2][2], ni.sto_xyz.m[2][3],
                                  ni.sto_xyz.m[3][0], ni.sto_xyz.m[3][1], ni.sto_xyz.m[3][2], ni.sto_xyz.m[3][3]
                                 };
              image.set_attribute(AttrID_nifti_sform, am);
              image.set_attribute(AttrID_nifti_sform_code, ni.sform_code);

              // in this case use s-from for orientation
              if (ni.qform_code == 0) {
                     C3DDMatrix rot(C3DDVector(-ni.sto_xyz.m[0][0] / ni.dx,
                                               -ni.sto_xyz.m[1][0] / ni.dy,
                                               ni.sto_xyz.m[2][0] / ni.dz),
                                    C3DDVector(-ni.sto_xyz.m[0][1] / ni.dx,
                                               -ni.sto_xyz.m[1][1] / ni.dy,
                                               ni.sto_xyz.m[2][1] / ni.dz),
                                    C3DDVector( ni.sto_xyz.m[0][2] / ni.dx,
                                                ni.sto_xyz.m[1][2] / ni.dy,
                                                -ni.sto_xyz.m[2][2] / ni.dz));
                     image.set_rotation(rot);
                     const float dz = (ni.nz - 1) * ni.dz;
                     C3DFVector org(- rot.z.x * dz - ni.qoffset_x,
                                    - rot.z.y * dz - ni.qoffset_y,
                                    - rot.z.z * dz + ni.qoffset_z);
                     image.set_origin(org);
              }
       }

       image.set_attribute(AttrID_nifti_toffset, ni.toffset);
       image.set_attribute(AttrID_nifti_xyz_units, ni.xyz_units);
       image.set_attribute(AttrID_nifti_time_units, ni.time_units);
       image.set_attribute(AttrID_nifti_intent_code, ni.intent_code);
       image.set_attribute(AttrID_nifti_intent_p1, ni.intent_p1);
       image.set_attribute(AttrID_nifti_intent_p2, ni.intent_p2);
       image.set_attribute(AttrID_nifti_intent_p3, ni.intent_p3);

       if (strlen(ni.intent_name) > 0)
              image.set_attribute(AttrID_nifti_intent_name, string(ni.intent_name));

       image.set_attribute(IDRescaleSlope, ni.scl_slope);
       image.set_attribute(IDRescaleIntercept, ni.scl_inter);
}


template <typename Image>
CNifti3DImageIOPlugin::PData read_images(const C3DBounds& size, const nifti_image& ni)
{
       int timesteps = 1;

       for (int extra_dims = 4; extra_dims <= ni.ndim; ++extra_dims)
              timesteps *= ni.dim[extra_dims];

       cvdebug() << "Loading " << timesteps << " images\n";
       CNifti3DImageIOPlugin::PData result(new C3DImageVector);
       result->reserve(timesteps);
       typedef typename Image::value_type TInPixel;
       const TInPixel *in_data = reinterpret_cast<const TInPixel *>(ni.data);
       size_t stride = size.product();

       for (int i = 0; i < timesteps; ++i, in_data += stride) {
              Image *img = new Image(size);
              copy_attributes(*img, ni);
              copy(in_data, in_data + img->size(), img->begin());
              result->push_back(P3DImage(img));
       }

       return result;
}


CNifti3DImageIOPlugin::PData CNifti3DImageIOPlugin::do_load(const std::string&  filename) const
{
       if (!is_nifti_file(filename.c_str()))
              return PData();

       unique_ptr<nifti_image, NiftiDeallocator> image(nifti_image_read(filename.c_str(), 1));

       if (!image) {
              cvdebug() << "File '" <<  filename << "' not a NIFTI file";
              return PData();
       }

       switch (image->intent_code) {
       case 0: // no specific intent, assume image data
       case NIFTI_INTENT_DIMLESS:
              break;

       default:
              throw create_exception<invalid_argument>("Found a NIFTI-1 file with intent ",
                            image->intent_code, " that is not supported by MIA.");
       }

       if (image->ndim < 3 || image->ndim > 5)
              throw create_exception<invalid_argument>("Nifti: 3D(+t) image expected but ",
                            image->ndim, " dimensions available in '", filename);

       C3DBounds size(image->nx, image->ny, image->nz);

       switch (image->datatype) {
       case DT_UINT8:
              return read_images<C3DUBImage>(size, *image);

       case DT_INT16:
              return read_images<C3DSSImage>(size, *image);

       case DT_INT32:
              return read_images<C3DSIImage>(size, *image);

       case DT_FLOAT32:
              return read_images<C3DFImage>(size,  *image);

       case DT_FLOAT64:
              return read_images<C3DDImage>(size,  *image);

       case DT_INT8:
              return read_images<C3DSBImage>(size, *image);

       case DT_UINT16:
              return read_images<C3DUSImage>(size, *image);

       case DT_UINT32:
              return read_images<C3DUIImage>(size, *image);

       case DT_INT64:
              return read_images<C3DSLImage>(size, *image);

       case DT_UINT64:
              return read_images<C3DULImage>(size, *image);

       default:
              throw create_exception<invalid_argument>("NIFTI: input format ", image->datatype, " not supported");
       }
}

static int datatype_to_nifti(EPixelType ptype)
{
       switch (ptype) {
       case it_sbyte:
              return DT_INT8;

       case it_ubyte:
              return DT_UINT8;

       case it_sshort:
              return DT_INT16;

       case it_ushort:
              return DT_UINT16;

       case it_sint:
              return DT_INT32;

       case it_uint:
              return DT_UINT32;

       case it_slong:
              return DT_INT64;

       case it_ulong:
              return DT_UINT64;

       case it_float:
              return DT_FLOAT32;

       case it_double:
              return DT_FLOAT64;

       default:
              throw create_exception<invalid_argument>("NIFTI-1 ImageIO: Trying to save pixel type '",
                            CPixelTypeDict.get_name(ptype),
                            "' that is not supported by nifti");
       }
}


class FCopyImageDataToNifti : public TFilter <void>
{
public:
       FCopyImageDataToNifti(nifti_image& ni, EPixelType pt):
              m_ni(ni),
              m_pt(pt),
              m_npixels_written(0)
       {
       }

       template <typename T>
       void operator () (const T3DImage<T>& image)
       {
              // sanity type check for multi-images
              if (m_pt != image.get_pixel_type()) {
                     throw create_exception<invalid_argument>("NIFTI-IO: image series containing images of varying pixel types is not supported");
              }

              if (static_cast<unsigned>(m_ni.nx) != image.get_size().x ||
                  static_cast<unsigned>(m_ni.ny) != image.get_size().y ||
                  static_cast<unsigned>(m_ni.nz) != image.get_size().z) {
                     throw create_exception<invalid_argument>("NIFTI-IO: image series containing images of varying sizes is not supported");
              }

              T *out_data = reinterpret_cast<T *>(m_ni.data);
              copy(image.begin(), image.end(), out_data + m_npixels_written);
              m_npixels_written += image.size();
       }
private:
       nifti_image& m_ni;
       EPixelType m_pt;
       long m_npixels_written;
};



bool CNifti3DImageIOPlugin::do_save(const std::string& fname, const Data& data) const
{
       assert(!data.empty());
       const C3DImage& image = *data[0];
       int dims[8] = {3, 1, 1, 1, 1, 1, 1, 1};
       C3DBounds size = image.get_size();
       dims[1] = size.x;
       dims[2] = size.y;
       dims[3] = size.z;
       auto pixel_type = image.get_pixel_type();
       int datatype = datatype_to_nifti(pixel_type);

       if (data.size() > 1) {
              cvdebug() << "Saving " << data.size() << " images\n";
              dims[0] = 5;
              dims[5] = data.size();
       }

       // create the output image
       unique_ptr<nifti_image, NiftiDeallocator> output(nifti_make_new_nim(dims, datatype, 1));

       if (!output) {
              throw create_exception<runtime_error>("NIFTI-imageio: Unable to create output image of size [",
                                                    size, "]x[", data.size(), "]");
       }

       output->pixdim[0] = dims[0];
       auto scale = image.get_voxel_size();
       output->dx = output->pixdim[1] = scale.x;
       output->dy = output->pixdim[2] = scale.y;
       output->dz = output->pixdim[3] = scale.z;
       // it would make sense to add this next one as an attribute
       output->dt = output->pixdim[4] = 1;
       // not supported
       output->du = output->pixdim[5] = 1;
       output->dv = output->pixdim[6] = 1;
       output->dw = output->pixdim[7] = 1;
       output->scl_slope = image.get_attribute_as<float>(IDRescaleSlope, 0.0f);
       output->scl_inter = image.get_attribute_as<float>(IDRescaleIntercept, 0.0f);
       // need to see whether to support this
       output->cal_min = 0.0f;
       output->cal_max = 0.0f ;

       if (image.has_attribute(AttrID_nifti_qform_code)) {
              output->qform_code = image.get_attribute_as<int>(AttrID_nifti_qform_code);
       } else {
              output->qform_code = 1; // default to method 2
       }

       auto org = image.get_origin();

       // Analyze 7.5 like data
       if (output->qform_code == 0) {
              // here starts the orientation insanity
              // todo: re-check how this qfac is actually used
              output->qoffset_x = org.x;
              output->qoffset_y = org.y;
              output->qoffset_z = org.z;
              output->quatern_b = 0.0;
              output->quatern_c = 0.0;
              output->quatern_d = 0.0;

              switch (image.get_orientation()) {
              case ior_undefined:
              case ior_xyz:
                     break;

              default:
                     cvwarn() << __FUNCTION__ << "FIXME:manual set orientation detected, but ignored\n";
              }
       } else {
              output->qform_code = NIFTI_XFORM_SCANNER_ANAT;
              auto rot = image.get_rotation().as_matrix_3x3();
              cvdebug() << "rot-quat= " << image.get_rotation().as_quaternion() << "\n";
              cvdebug() << "org= " << org << "\n";
              cvdebug() << "rot= " << rot << "\n";
              const float dz = (size.z - 1) * scale.z;
              output->qoffset_x = - rot.z.x * dz - org.x;
              output->qoffset_y = - rot.z.y * dz - org.y;
              output->qoffset_z = rot.z.z * dz + org.z;
              // this seems to be the proper transaltion
              output->qto_xyz = nifti_make_orthog_mat44(
                                       -rot.x.x, -rot.y.x, rot.z.x,
                                       -rot.x.y, -rot.y.y, rot.z.y,
                                       rot.x.z, rot.y.z, -rot.z.z);
              nifti_mat44_to_quatern(output->qto_xyz,
                                     &output->quatern_b, &output->quatern_c, &output->quatern_d,
                                     NULL, NULL, NULL, NULL, NULL, NULL, &output->qfac);
              output->pixdim[0] = output->qfac;
       }

       // copy s-form if available
       if (image.has_attribute(AttrID_nifti_sform_code)) {
              output->sform_code = image.get_attribute_as<int>(AttrID_nifti_sform_code);
              auto am = image.get_attribute_as<vector<float>>(AttrID_nifti_sform);
              output->sto_xyz.m[0][0] = am[0];
              output->sto_xyz.m[0][1] = am[1];
              output->sto_xyz.m[0][2] = am[2];
              output->sto_xyz.m[0][3] = am[3];
              output->sto_xyz.m[1][0] = am[4];
              output->sto_xyz.m[1][1] = am[5];
              output->sto_xyz.m[1][2] = am[6];
              output->sto_xyz.m[1][3] = am[7];
              output->sto_xyz.m[2][0] = am[8];
              output->sto_xyz.m[2][1] = am[9];
              output->sto_xyz.m[2][2] = am[10];
              output->sto_xyz.m[2][3] = am[11];
              output->sto_xyz.m[3][0] = am[12];
              output->sto_xyz.m[3][1] = am[13];
              output->sto_xyz.m[3][2] = am[14];
              output->sto_xyz.m[3][3] = am[15];
       } else { // without a sform code leave it empty (for now)
              output->sform_code = NIFTI_XFORM_SCANNER_ANAT;
              auto rot = image.get_rotation().as_matrix_3x3();
              output->sto_xyz.m[0][0] = -rot.x.x * scale.x;
              output->sto_xyz.m[0][1] = -rot.y.x * scale.y;
              output->sto_xyz.m[0][2] =  rot.z.x * scale.z;
              output->sto_xyz.m[0][3] = output->qoffset_x;
              output->sto_xyz.m[1][0] = -rot.x.y * scale.x;
              output->sto_xyz.m[1][1] = -rot.y.y * scale.y;
              output->sto_xyz.m[1][2] =  rot.z.y * scale.z;
              output->sto_xyz.m[1][3] = output->qoffset_y;
              output->sto_xyz.m[2][0] = rot.x.z * scale.y;
              output->sto_xyz.m[2][1] = rot.y.z * scale.y;
              output->sto_xyz.m[2][2] = -rot.z.z * scale.z;
              output->sto_xyz.m[2][3] = output->qoffset_z;
              output->sto_xyz.m[3][0] = 0;
              output->sto_xyz.m[3][1] = 0;
              output->sto_xyz.m[3][2] = 0;
              output->sto_xyz.m[3][3] = 1;
       }

       output->toffset = image.get_attribute_as<float>(AttrID_nifti_toffset, 0.0f);
       output->xyz_units  = image.get_attribute_as<int>(AttrID_nifti_xyz_units, NIFTI_UNITS_MM);
       output->time_units = image.get_attribute_as<int>(AttrID_nifti_time_units, NIFTI_UNITS_SEC);
       output->intent_code = image.get_attribute_as<int>(AttrID_nifti_intent_code, 0);
       output->intent_p1 = image.get_attribute_as<float>(AttrID_nifti_intent_p1, 0);
       output->intent_p2 = image.get_attribute_as<float>(AttrID_nifti_intent_p2, 0);
       output->intent_p3 = image.get_attribute_as<float>(AttrID_nifti_intent_p3, 0);
       string intent_name = image.get_attribute_as<string>(AttrID_nifti_intent_name, "");

       if (!intent_name.empty()) {
              strncpy(output->intent_name, intent_name.c_str(),  14);
       }

       output->num_ext = 0;
       // set the data
       FCopyImageDataToNifti nsave(*output, image.get_pixel_type());

       for (auto i = data.begin(); i != data.end(); ++i) {
              ::mia::accumulate(nsave, **i);
       }

       // it could be that this needs the file name without extension
       if (nifti_set_filenames(output.get(), fname.c_str(), 1, 0) != 0) {
              throw create_exception<runtime_error>("NIFTI-imageio: Something went wrong with setting the filename to '",
                                                    fname, "'");
       }

       auto resultfile = nifti_image_write_hdr_img(output.get(), 1, "wb");
       return resultfile ? false : true;
}

const std::string CNifti3DImageIOPlugin::do_get_descr() const
{
       return "NIFTI-1 3D image IO. The orientation is transformed in the same way like "
              "it is done with 'dicomtonifti --no-reorder' from the vtk-dicom package.";
}

const std::string CNifti3DImageIOPlugin::do_get_preferred_suffix() const
{
       return "nii";
}

extern "C" EXPORT CPluginBase *get_plugin_interface()
{
       return new CNifti3DImageIOPlugin();
}

NS_END
