/*  Multimodal Image Segmentation Tool (MIST)  */
/*  Eelke Visser  */

/*  Copyright (c) 2016 University of Oxford  */

/*  CCOPYRIGHT  */

#ifndef TRANSFORMATION_H
#define TRANSFORMATION_H

#include <string>
#include <vector>
#include <boost/log/trivial.hpp>
#include "armawrap/newmat.h"
// removed because warpfns.h needs to define EXPOSE_TREACHEROUS before including this
// #include "newimage.h"
#include "warpfns/warpfns.h"

// NB: All classes in here have default operator=, copy ctor, etc!


class IdentityTransformation;
class AffineTransformation;
class NonlinearTransformation;

class Transformation
{
public:
    class TransformationException : public std::runtime_error
    {
    public:
        TransformationException(const std::string &what) : runtime_error(what) { };
    };

    Transformation();
    virtual ~Transformation() { }

    virtual void ReleaseMemory() const { }

    virtual NEWMAT::ColumnVector InverseTransformPoint(const NEWMAT::ColumnVector &p) const = 0;

    // Can't think of a better way :(
    template<class T, class U>
    NEWIMAGE::volume<T> TransformVolume(const NEWIMAGE::volume<T> &invol, const NEWIMAGE::volume<U> &refvol) const
    {
        return apply_derived_transformation(this, invol, refvol);
    }
};


class IdentityTransformation : public Transformation
{
public:
    virtual NEWMAT::ColumnVector InverseTransformPoint(const NEWMAT::ColumnVector &p) const;

    template<class T, class U>
    NEWIMAGE::volume<T> TransformVolumeImplementation(const NEWIMAGE::volume<T> &invol, const NEWIMAGE::volume<U> &refvol) const
    {
        BOOST_LOG_TRIVIAL(warning) << "IdentityTransformation::TransformVolumeImplementation() is untested!";

        NEWMAT::Matrix identity(4, 4);
        identity << 1.0 << 0.0 << 0.0 << 0.0
                 << 0.0 << 1.0 << 0.0 << 0.0
                 << 0.0 << 0.0 << 1.0 << 0.0
                 << 0.0 << 0.0 << 0.0 << 1.0;

        NEWIMAGE::volume<T> outvol(refvol.xsize(), refvol.ysize(), refvol.zsize());
        NEWIMAGE::copybasicproperties(refvol, outvol);

        // No warp
        NEWIMAGE::volume4D<float> warp;
        std::vector<int> defdir;

        // No derivatives
        std::vector<int> derivdir;
        NEWIMAGE::volume4D<T> deriv;

        NEWIMAGE::raw_general_transform(invol, identity, warp, defdir, derivdir, outvol, deriv);

        return outvol;
    }
};


class AffineTransformation : public Transformation
{
public:
    AffineTransformation(const std::string &filename);
    AffineTransformation(const NEWMAT::Matrix &matrix);

    AffineTransformation Concatenate(const AffineTransformation &second) const;
    NonlinearTransformation Concatenate(const NonlinearTransformation &second) const;

    virtual NEWMAT::ColumnVector InverseTransformPoint(const NEWMAT::ColumnVector &p) const;

    template<class T, class U>
    NEWIMAGE::volume<T> TransformVolumeImplementation(const NEWIMAGE::volume<T> &invol, const NEWIMAGE::volume<U> &refvol) const
    {
        NEWIMAGE::volume<T> outvol(refvol.xsize(), refvol.ysize(), refvol.zsize());
        NEWIMAGE::copybasicproperties(refvol, outvol);

        // No warp
        NEWIMAGE::volume4D<float> warp;
        std::vector<int> defdir;

        // No derivatives
        std::vector<int> derivdir;
        NEWIMAGE::volume4D<T> deriv;

        NEWIMAGE::raw_general_transform(invol, m_inverse.i(), warp, defdir, derivdir, outvol, deriv);

        return outvol;
    }

private:
    friend class NonlinearTransformation;

    NEWMAT::Matrix m_inverse;
};


class NonlinearTransformation : public Transformation
{
public:
    NonlinearTransformation(const std::string &filename, bool includeaffine);
    NonlinearTransformation(const NEWIMAGE::volume4D<float> &warpfield);

    NonlinearTransformation Concatenate(const NonlinearTransformation &second) const;

    template<class T>
    NonlinearTransformation Concatenate(const AffineTransformation &second, const NEWIMAGE::volume<T> &reference) const
    {
        BOOST_LOG_TRIVIAL(warning) << "NonlinearTransformation::Concatenate(AffineTransformation&) is untested!";

        NEWIMAGE::volume4D<float> secondfield;
        affine2warp(second.m_inverse.i(), secondfield, reference);

        NEWIMAGE::volume4D<float> combinedfield;
        concat_warps(AccessField(), secondfield, combinedfield);

        return NonlinearTransformation(combinedfield);
    }

    virtual NEWMAT::ColumnVector InverseTransformPoint(const NEWMAT::ColumnVector &p) const;

    template<class T, class U>
    NEWIMAGE::volume<T> TransformVolumeImplementation(const NEWIMAGE::volume<T> &invol, const NEWIMAGE::volume<U> &refvol) const
    {
        NEWIMAGE::volume<T> outvol(refvol.xsize(), refvol.ysize(), refvol.zsize());
        NEWIMAGE::copybasicproperties(refvol, outvol);

        NEWIMAGE::volume4D<float> relwarp(AccessField());
        NEWIMAGE::convertwarp_abs2rel(relwarp);
        NEWIMAGE::apply_warp(invol, outvol, relwarp);

        return outvol;
    }

    virtual void ReleaseMemory() const;

private:
    friend class AffineTransformation;

    mutable bool m_ready;
    std::string m_backingFile;
    bool m_includeAffineFromBackingFile;

    mutable NEWIMAGE::volume4D<float> m_field;

    NEWIMAGE::volume4D<float> &AccessField() const;

    void Load() const;
};


template<class T, class U>
NEWIMAGE::volume<T> apply_derived_transformation(const Transformation *xfm, const NEWIMAGE::volume<T> &invol,
                                                 const NEWIMAGE::volume<U> &refvol)
{
    if (auto ix = dynamic_cast<const IdentityTransformation *>(xfm))
        return ix->TransformVolumeImplementation(invol, refvol);
    else if (auto ax = dynamic_cast<const AffineTransformation *>(xfm))
        return ax->TransformVolumeImplementation(invol, refvol);
    else if (auto nx = dynamic_cast<const NonlinearTransformation *>(xfm))
        return nx->TransformVolumeImplementation(invol, refvol);

    throw std::logic_error("Cannot downcast Transformation object");
}


#endif // TRANSFORMATION_H
