/*! \file image.c
 *
 * \author John Reppy
 *
 * \brief code to support the loading of image data from Nrrd files.
 */

/*
 * COPYRIGHT (c) 2011 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 */

#include "Diderot/diderot.h"
#include <teem/nrrd.h>

// we use double precision representations of the transforms here to
// improve robustness, even when the results are stored as single-precision
typedef double Matrix2x2_t[4];
typedef double Matrix3x3_t[9];
typedef double Matrix4x4_t[16];

typedef double Vec2_t[2];
typedef double Vec3_t[3];

static void PrintMat2x2 (const char *title, Matrix2x2_t m)
{
#define M(i)    m[i]
    printf ("%s @ %p:\n", title, (void *)m);
    for (int i = 0;  i < 2;  i++) {
        int j = 2*i;
        printf ("  [ %6lf  %6lf ]\n", M(j+0), M(j+1));
    }
#undef M

}

static void PrintMat3x3 (const char *title, Matrix3x3_t m)
{
#define M(i)    m[i]
    printf ("%s @ %p:\n", title, (void *)m);
    for (int i = 0;  i < 3;  i++) {
        int j = 3*i;
        printf ("  [ %6lf  %6lf  %6lf ]\n", M(j+0), M(j+1), M(j+2));
    }
#undef M

}

static void PrintMat4x4 (const char *title, Matrix4x4_t m)
{
#define M(i)    m[i]
    printf ("%s @ %p:\n", title, (void *)m);
    for (int i = 0;  i < 4;  i++) {
        int j = 4*i;
        printf ("  [ %6lf  %6lf  %6lf  %6lf ]\n", M(j+0), M(j+1), M(j+2), M(j+3));
    }
#undef M

}

/* Loading transfomation information from Nrrd files */

static void LoadTransform1D (Nrrd *nin, double *s, double *t) 
{ 
  // compute the offset to the first space axis
    int base = nin->dim - nin->spaceDim;

  // Image axis Scaling and Rotation
    *s = nin->axis[base+0].spaceDirection[0];

  // Image location
    *t = nin->spaceOrigin[0];
}

static void LoadTransform2D (Nrrd *nin, Matrix2x2_t m, Vec2_t t) 
{ 
#define M(i,j)  m[(i)*2+(j)]

  // compute the offset to the first space axis
    int base = nin->dim - nin->spaceDim;

  // Image axis Scaling and Rotation
    M(0,0) = nin->axis[base+0].spaceDirection[0];
    M(1,0) = nin->axis[base+0].spaceDirection[1];
    M(0,1) = nin->axis[base+1].spaceDirection[0];
    M(1,1) = nin->axis[base+1].spaceDirection[1];

  // Image location
    t[0] = nin->spaceOrigin[0];
    t[1] = nin->spaceOrigin[1];

#undef M
}

static void LoadTransform3D (Nrrd *nin, Matrix3x3_t m, Vec3_t t) 
{ 
#define M(i,j)  m[3*(i)+(j)]

  // compute the offset to the first space axis
    int base = nin->dim - nin->spaceDim;

  // Image axis Scaling and Rotation
    M(0,0) = nin->axis[base+0].spaceDirection[0];
    M(1,0) = nin->axis[base+0].spaceDirection[1];
    M(2,0) = nin->axis[base+0].spaceDirection[2];
    M(0,1) = nin->axis[base+1].spaceDirection[0];
    M(1,1) = nin->axis[base+1].spaceDirection[1];
    M(2,1) = nin->axis[base+1].spaceDirection[2];
    M(0,2) = nin->axis[base+2].spaceDirection[0];
    M(1,2) = nin->axis[base+2].spaceDirection[1];
    M(2,2) = nin->axis[base+2].spaceDirection[2];

  // Image location
    t[0] = nin->spaceOrigin[0];
    t[1] = nin->spaceOrigin[1];
    t[2] = nin->spaceOrigin[2];

#undef M
}


/* Transformation matrix operations */

STATIC_INLINE double Det2 (
    double a, double b,
    double c, double d) 
{
    return (a*d - b*c); 
}

STATIC_INLINE double Det3 (
    double a, double b, double c,
    double d, double e, double f,
    double g, double h, double i) 
{
    return (a*e*i 
          + d*h*c 
          + g*b*f 
          - g*e*c 
          - d*b*i 
          - a*h*f); 
}

static double DetM3x3 (Matrix3x3_t m) 
{
#define M(i)    m[i]
    return Det3(M(0), M(1), M(2), 
                M(3), M(4), M(5), 
                M(6), M(7), M(8));
#undef M
}

static double DetM4x4 (Matrix4x4_t m) 
{
#define M(i)    m[i]
        return  (M( 0) * Det3(M( 5), M( 6), M( 7), 
                              M( 9), M(10), M(11), 
                              M(13), M(14), M(15))                                                 
               - M( 1) * Det3(M( 4), M( 6), M( 7), 
                              M( 8), M(10), M(11), 
                              M(12), M(14), M(15)) 
               + M( 2) * Det3(M( 4), M( 5), M( 7), 
                              M( 8), M( 9), M(11), 
                              M(12), M(13), M(15))                                                 
               - M( 3) * Det3(M( 4), M( 5), M( 6), 
                              M( 8), M( 9), M(10), 
                              M(12), M(13), M(14))); 
#undef M

}

/*! \brief compute the inverse of \arg m, storing the result in \arg i.
 *  \param m the matrix to invert
 *  \param i the inverted matrix
 */
static void InvertM2x2 (Matrix2x2_t i, Matrix2x2_t m) 
{  
#define M(i)    m[i]
#define I(j)    i[j]
        double scale = 1.0 / Det2(M(0), M(1), M(2), M(3));

        I(0) =  scale * M(3);
        I(1) = -scale * M(1);
        I(2) = -scale * M(2);
        I(3) =  scale * M(0);
#undef M
#undef I
} 

/*! \brief compute the inverse of \arg m, storing the result in \arg i.
 *  \param m the matrix to invert
 *  \param i the inverted matrix
 */
static void InvertM3x3 (Matrix3x3_t i, Matrix3x3_t m) 
{  
        double scale = 1.0 / DetM3x3(m);

#define M(i)    m[i]
#define I(j)    i[j]
        I(0) = scale * Det2 (M(4), M(5), M(7), M(8));
        I(1) = scale * Det2 (M(2), M(1), M(8), M(7));
        I(2) = scale * Det2 (M(1), M(2), M(4), M(5));
        I(3) = scale * Det2 (M(5), M(3), M(8), M(6));
        I(4) = scale * Det2 (M(0), M(2), M(6), M(8));
        I(5) = scale * Det2 (M(2), M(0), M(5), M(3));
        I(6) = scale * Det2 (M(3), M(4), M(6), M(7));
        I(7) = scale * Det2 (M(1), M(0), M(7), M(6));
        I(8) = scale * Det2 (M(0), M(1), M(3), M(4));
#undef M
#undef I
}

/*! \brief compute the inverse of \arg m, storing the result in \arg i.
 *  \param m the matrix to invert
 *  \param i the inverted matrix
 */
static void InvertM4x4 (Matrix4x4_t i, Matrix4x4_t m) 
{  
        double scale = 1.0 / DetM4x4(m);

#define M(i)    m[i]
#define I(j)    i[j]
        I( 0) =  Det3(M( 5),M( 6),M( 7),
                      M( 9),M(10),M(11),
                      M(13),M(14),M(15)) * scale;
        
        I( 1) = -Det3(M( 1),M( 2),M( 3),
                      M( 9),M(10),M(11),
                      M(13),M(14),M(15)) * scale;
                                           
        I( 2) =  Det3(M( 1),M( 2),M( 3),
                      M( 5),M( 6),M( 7),
                      M(13),M(14),M(15)) * scale;
                                           
        I( 3) = -Det3(M( 1),M( 2),M( 3),
                      M( 5),M( 6),M( 7),
                      M( 9),M(10),M(11)) * scale;
                                           
        I( 4) = -Det3(M( 4),M( 6),M( 7),
                      M( 8),M(10),M(11),
                      M(12),M(14),M(15)) * scale;
                                           
        I( 5) =  Det3(M( 0),M( 2),M( 3),
                      M( 8),M(10),M(11),
                      M(12),M(14),M(15)) * scale;
                                           
        I( 6) = -Det3(M( 0),M( 2),M( 3),
                      M( 4),M( 6),M( 7),
                      M(12),M(14),M(15)) * scale;
                                           
        I( 7) =  Det3(M( 0),M( 2),M( 3),
                      M( 4),M( 6),M( 7),
                      M( 8),M(10),M(11)) * scale;
                                           
        I( 8) =  Det3(M( 4),M( 5),M( 7),
                      M( 8),M( 9),M(11),
                      M(12),M(13),M(15)) * scale;
                                           
        I( 9) = -Det3(M( 0),M( 1),M( 3),
                      M( 8),M( 9),M(11),
                      M(12),M(13),M(15)) * scale;
                                           
        I(10) =  Det3(M( 0),M( 1),M( 3),
                      M( 4),M( 5),M( 7),
                      M(12),M(13),M(15)) * scale;
                                           
        I(11) = -Det3(M( 0),M( 1),M( 3),
                      M( 4),M( 5),M( 7),
                      M( 8),M( 9),M(11)) * scale;
                                           
        I(12) = -Det3(M( 4),M( 5),M( 6),
                      M( 8),M( 9),M(10),
                      M(12),M(13),M(14)) * scale;
                                           
        I(13) =  Det3(M( 0),M( 1),M( 2),
                      M( 8),M( 9),M(10),
                      M(12),M(13),M(14)) * scale;
                                           
        I(14) = -Det3(M( 0),M( 1),M( 2),
                      M( 4),M( 5),M( 6),
                      M(12),M(13),M(14)) * scale;
                                           
        I(15) =  Det3(M( 0),M( 1),M( 2),
                      M( 4),M( 5),M( 6),
                      M( 8),M( 9),M(10)) * scale;
#undef M
#undef I
} 

static Nrrd *LoadNrrdFile (const char *filename) 
{
  /* create a nrrd; at this point this is just an empty container */
    Nrrd *nin = nrrdNew();

  /* read in the nrrd from file */
    if (nrrdLoad(nin, filename, NULL)) {
        char *err = biffGetDone(NRRD);
        fprintf (stderr, "Diderot: trouble reading \"%s\":\n%s", filename, err);
        free (err);
        return NULL;
    }
         
    return nin; 
  
}

STATIC_INLINE bool CheckAxis (Nrrd *nin, int i)
{
    return (nin->axis[i].kind != nrrdKindSpace);
}

/* load image data from Nrrd files */
Status_t Diderot_LoadImage1D (Diderot_string_t name, Diderot_image1D_t **imgOut)
{
    Nrrd *nin = LoadNrrdFile (name);
    Diderot_image1D_t *img = (Diderot_image1D_t *)malloc(sizeof(Diderot_image1D_t));

  // compute the offset to the first space axis
    int base = nin->dim - nin->spaceDim;

// FIXME: we should also be checking the shape of the voxels and the type of the samples.
  // check that the image format is what we expect
    if (nin->spaceDim != 1) {
        fprintf(stderr, "nrrd file \"%s\" has unexpected dimension %d, expected 1\n",
            name, nin->spaceDim);
        return DIDEROT_FAIL;
    }
    else if (CheckAxis(nin, base)) {
        fprintf(stderr,
            "nrrd file \"%s\" has unexpected axis structure: %s\n",
            name,
            airEnumStr(nrrdKind, nin->axis[base].kind));
        return DIDEROT_FAIL;
    }

    img->dim = 1;
    img->size[0] = nin->axis[base+0].size;
    img->data = nin->data;
    img->dataSzb =  nrrdElementSize(nin) * nrrdElementNumber(nin);

//printf("LoadImage \"%s\": space dim = 1, axes = <%d>\n",
//name, img->size[base+0]);

  // from the Nrrd file, we load the scaling and translation
    double s, t;
    LoadTransform1D (nin, &s, &t);

    img->s = (Diderot_real_t)(1.0 / s);
    img->t = (Diderot_real_t)(-t / s);

    *imgOut = img;
    return DIDEROT_OK;
}

Status_t Diderot_LoadImage2D (Diderot_string_t name, Diderot_image2D_t **imgOut)
{
    Nrrd *nin = LoadNrrdFile (name);
    Diderot_image2D_t *img = (Diderot_image2D_t *)malloc(sizeof(Diderot_image2D_t));

  // compute the offset to the first space axis
    int base = nin->dim - nin->spaceDim;

  // check that the image format is what we expect
    if (nin->spaceDim != 2) {
        fprintf(stderr, "nrrd file \"%s\" has unexpected dimension %d, expected 2\n",
            name, nin->spaceDim);
        return DIDEROT_FAIL;
    }
    else if (CheckAxis(nin, base) || CheckAxis(nin, base+1)) {
        fprintf(stderr,
            "nrrd file \"%s\" has unexpected axis structure: %s %s\n",
            name,
            airEnumStr(nrrdKind, nin->axis[base].kind),
            airEnumStr(nrrdKind, nin->axis[base+1].kind));
        return DIDEROT_FAIL;
    }

    img->dim = 2;
    img->size[0] = nin->axis[base+0].size;
    img->size[1] = nin->axis[base+1].size;
    img->data = nin->data;
    img->dataSzb =  nrrdElementSize(nin) * nrrdElementNumber(nin);

//printf("LoadImage \"%s\": space dim = 2, axes = <%d, %d>\n",
//name, img->size[0], img->size[1]);

  // from the Nrrd file, we load the affine image-to-world transform matrix
    Matrix2x2_t m;              // rotation and scaling
    Vec2_t t;                   // translation part of the transform
    LoadTransform2D (nin, m, t);

//PrintMat2x2("m", m);
//printf("t = <%f, %f>\n", t[0], t[1]);

  // compute inverse of m, which is the transform from the world basis to the image basis
    Matrix2x2_t mInv;
    InvertM2x2 (mInv, m);
//PrintMat2x2("w2i", mInv);

  // copy results into the image data structure
    for (int i = 0;  i < 2;  i++) {
        for (int j = 0;  j < 2;  j++) {
            img->w2i[i].r[j] = mInv[2*i+j];
        }
    }
    transpose2x2 (img->w2iT, img->w2i);
    Diderot_vec2_t tVec = vec2(t[0], t[1]);
    img->tVec = -mulMat2x2Vec2(img->w2i, tVec);  // inv([M t]) = [inv(M) -inv(M)t]
//printf("tVec = <%f, %f>\n", ((union2f_t)img->tVec).r[0], ((union2f_t)img->tVec).r[1]);

    *imgOut = img;
    return DIDEROT_OK;
}

Status_t Diderot_LoadImage3D (Diderot_string_t name, Diderot_image3D_t **imgOut)
{
    Nrrd *nin = LoadNrrdFile (name);

  // compute the offset to the first space axis
    int base = nin->dim - nin->spaceDim;

  // check that the image format is what we expect
    if (nin->spaceDim != 3) {
        fprintf(stderr, "nrrd file \"%s\" has unexpected dimension %d, expected 3\n",
            name, nin->spaceDim);
        return DIDEROT_FAIL;
    }
    else if (CheckAxis(nin, base) || CheckAxis(nin, base+1) || CheckAxis(nin, base+2)) {
        fprintf(stderr,
            "nrrd file \"%s\" has unexpected axis structure: %s %s %s\n",
            name,
            airEnumStr(nrrdKind, nin->axis[base].kind),
            airEnumStr(nrrdKind, nin->axis[base+1].kind),
            airEnumStr(nrrdKind, nin->axis[base+2].kind));
        return DIDEROT_FAIL;
    }

    Diderot_image3D_t *img = (Diderot_image3D_t *)malloc(sizeof(Diderot_image3D_t));

    img->dim = 3;
    img->size[0] = nin->axis[base+0].size;
    img->size[1] = nin->axis[base+1].size;
    img->size[2] = nin->axis[base+2].size;
    img->data = nin->data;
    img->dataSzb =  nrrdElementSize(nin) * nrrdElementNumber(nin);

//printf("LoadImage \"%s\": space dim = 3, axes = <%d, %d, %d>\n",
//name, img->size[0], img->size[1], img->size[2]);

  // from the Nrrd file, we load the affine image-to-world transform matrix
    Matrix3x3_t m;              // rotation and scaling
    Vec3_t t;                   // translation part of the transform
    LoadTransform3D (nin, m, t);

//PrintMat3x3("m", m);
//printf("t = <%f, %f, %f>\n", t[0], t[1], t[2]);

  // compute inverse of m, which is the transform from the world basis to the image basis
    Matrix3x3_t mInv;
    InvertM3x3 (mInv, m);
//PrintMat3x3("w2i", mInv);

  // copy results into the image data structure
    for (int i = 0;  i < 3;  i++) {
        for (int j = 0;  j < 3;  j++) {
            img->w2i[i].r[j] = mInv[3*i+j];
        }
    }
    transpose3x3 (img->w2iT, img->w2i);
    Diderot_vec3_t tVec = vec3(t[0], t[1], t[2]);
    img->tVec = -mulMat3x3Vec3(img->w2i, tVec);  // inv([M t]) = [inv(M) -inv(M)t]
//printf("tVec = <%f, %f, %f>\n", ((Diderot_union3_t)img->tVec).r[0], ((Diderot_union3_t)img->tVec).r[1], ((Diderot_union3_t)img->tVec).r[2]);

    *imgOut = img;
    return DIDEROT_OK;
}