/*  MELODIC - Multivariate exploratory linear optimized decomposition into
    independent components

    meldata.h - data container class

    Christian F. Beckmann, FMRIB Analysis Group

    Copyright (C) 1999-2013 University of Oxford */

/*  CCOPYRIGHT */


#ifndef __MELODICDATA_h
#define __MELODICDATA_h

#include <random>

#include "CiftiLib/CiftiFile.h"
#include "armawrap/newmat.h"
#include "newimage/newimageall.h"
#include "utils/log.h"
#include "meloptions.h"
#include "melhlprfns.h"

namespace Melodic{

  class MelodicData{
  public:

    //constructor
    MelodicData(MelodicOptions &popts, Utilities::Log &plogger):
      opts(popts),logger(plogger)
    {
      after_mm = false;
      Resels = 0;
    }

    void save();

    cifti::CiftiFile inputCifti;
    NEWMAT::ReturnMatrix process_file(std::string fname, int numfiles = 1);

    inline void save4D(NEWMAT::Matrix what, std::string fname){
      if ( !opts.readCIFTI.value() ) //Process NIFTI
        {

          NEWIMAGE::volume4D<float> tempVol;
          tempVol.setmatrix(what,Mask);
          NEWIMAGE::save_volume4D(tempVol,logger.appendDir(fname));
          message("  " << logger.appendDir(fname) << std::endl);
        } else { //Process CIFTI save ICs as float
        cifti::CiftiFile outputFile;
        outputFile.setWritingFile(logger.appendDir(fname)+".nii");//sets up on-disk writing with default writing version
        cifti::CiftiXML xml(inputCifti.getCiftiXML());
        cifti::CiftiScalarsMap scalarsMap;
        std::vector<char> foo = xml.writeXMLToVector();
        scalarsMap.setLength(what.Nrows());
        xml.setMap(0, scalarsMap);
        outputFile.setCiftiXML(xml,false);
        std::vector<float> scratchRow(what.Nrows());//read/write a row at a time
        for (int64_t row=0;row<what.Ncols();row++) {
          for (int64_t col=0;col<what.Nrows();col++)
            scratchRow[col]=what(col+1,row+1);
          outputFile.setRow(scratchRow.data(),row);
        }
      }
    }

    inline void saveascii(NEWMAT::Matrix what, std::string fname){
      MISCMATHS::write_ascii_matrix(logger.appendDir(fname),what);
      message("  " << logger.appendDir(fname) << std::endl);
    }

    inline void savebinary(NEWMAT::Matrix what, std::string fname){
      MISCMATHS::write_binary_matrix(what,logger.appendDir(fname));
      message("  " << logger.appendDir(fname) << std::endl);
    }

    int  remove_components();

    void setup_classic();
    void setup_migp();
    void setup();

    void status(const std::string &txt);

    inline NEWMAT::Matrix& get_pcaE() {return pcaE;}
    inline void set_pcaE(NEWMAT::Matrix& Arg) {pcaE = Arg;}

    inline NEWMAT::RowVector& get_pcaD() {return pcaD;}
    inline void set_pcaD(NEWMAT::RowVector& Arg) {pcaD = Arg;}

    inline NEWMAT::Matrix& get_data() {return Data;}
    inline void set_data(NEWMAT::Matrix& Arg) {Data = Arg;}

    inline NEWMAT::Matrix& get_IC() {return IC;}
    inline void set_IC(NEWMAT::Matrix& Arg) {IC = Arg;}
    inline void set_IC(int ctr, NEWMAT::Matrix& Arg) {IC.Row(ctr) = Arg;}

    inline std::vector<NEWMAT::Matrix>& get_Smodes() {return Smodes;}
    inline NEWMAT::Matrix& get_Smodes(int what) {return Smodes.at(what);}
    inline void add_Smodes(NEWMAT::Matrix& Arg) {Smodes.push_back(Arg);}
    inline void save_Smodes(){
      if(Smodes.size()>0){
        NEWMAT::Matrix tmp = Smodes.at(0);
        for(unsigned int ctr = 1; ctr < Smodes.size(); ctr++)
          tmp |= Smodes.at(ctr);
        saveascii(tmp,opts.outputfname.value() + "_Smodes");
      }
    }

    inline std::vector<NEWMAT::Matrix>& get_Tmodes() {return Tmodes;}
    inline NEWMAT::Matrix& get_Tmodes(int what) {return Tmodes.at(what);}
    inline void add_Tmodes(NEWMAT::Matrix& Arg) {Tmodes.push_back(Arg);}
    inline void save_Tmodes(){
      if(Tmodes.size()>0){
        NEWMAT::Matrix tmp = Tmodes.at(0);
        outMsize("tmp",tmp);
        for(unsigned int ctr = 1; ctr < Tmodes.size(); ctr++){
          outMsize("Tmodes ",Tmodes.at(ctr));
          tmp |= Tmodes.at(ctr);
        }
        saveascii(tmp,opts.outputfname.value() + "_Tmodes");
      }
    }

    void set_TSmode_depr();
    void set_TSmode();

    inline NEWMAT::Matrix& get_param() {return param;}
    inline void set_param(NEWMAT::Matrix& Arg) {param = Arg;}

    inline NEWMAT::Matrix& get_paramS() {return paramS;}
    inline void set_paramS(NEWMAT::Matrix& Arg) {paramS = Arg;}

    inline NEWMAT::Matrix& get_white() {return whiteMatrix;}
    inline void set_white(NEWMAT::Matrix& Arg) {whiteMatrix = Arg;}

    inline NEWMAT::Matrix& get_dewhite() {return dewhiteMatrix;}
    inline void set_dewhite(NEWMAT::Matrix& Arg) {dewhiteMatrix = Arg;}

    inline NEWMAT::Matrix& get_meanC() {return meanC;}
    inline NEWMAT::Matrix& get_meanR() {return meanR;}

    inline NEWMAT::Matrix& get_stdDevi() {return stdDevi;}
    inline void set_stdDevi(NEWMAT::Matrix& Arg) {stdDevi = Arg;}

    inline NEWMAT::Matrix& get_mix() {return mixMatrix;}

    inline void set_mix(NEWMAT::Matrix& Arg) {
      mixMatrix = Arg;
      if (Tmodes.size() < 1)
        for (int ctr = 1; ctr <= Arg.Ncols(); ctr++){
          NEWMAT::Matrix tmp = Arg.Column(ctr);
          add_Tmodes(tmp);
        }
    }

    NEWMAT::Matrix expand_mix();
    NEWMAT::Matrix expand_dimred(const NEWMAT::Matrix& Mat);
    NEWMAT::Matrix reduce_dimred(const NEWMAT::Matrix& Mat);

    inline NEWMAT::Matrix& get_fmix() {return mixFFT;}
    inline void set_fmix(NEWMAT::Matrix& Arg) {mixFFT = Arg;}

    inline NEWMAT::Matrix& get_unmix() {return unmixMatrix;}
    inline void set_unmix(NEWMAT::Matrix& Arg) {unmixMatrix = Arg;}

    inline NEWIMAGE::volume<float>& get_mask() {return Mask;}
    inline void set_mask(NEWIMAGE::volume<float>& Arg) {Mask = Arg;}

    inline NEWIMAGE::volume<float>& get_mean() {return Mean;}
    inline void set_mean(NEWIMAGE::volume<float>& Arg) {Mean = Arg;}

    inline NEWIMAGE::volume<float>& get_bg() {
      if(opts.bgimage.value()>"")
        return background;
      else
        return Mean;
    }
    inline void set_bg(NEWIMAGE::volume<float>& Arg) {background = Arg;}

    inline NEWMAT::Matrix& get_Data() {return Data;}
    inline void set_Data(NEWMAT::Matrix& Arg) {Data = Arg;}

    inline NEWMAT::Matrix& get_RXweight() {return RXweight;}
    inline void set_RXweight(NEWMAT::Matrix& Arg) {RXweight = Arg;}

    inline NEWMAT::Matrix& get_ICstats() {return ICstats;}
    inline void set_ICstats(NEWMAT::Matrix& Arg) {ICstats = Arg;}

    inline NEWMAT::Matrix& get_EVP() {return EVP;}
    inline void set_EVP(NEWMAT::Matrix& Arg) {if(EVP.Storage()==0)
        EVP = Arg;}

    inline NEWMAT::Matrix& get_EV() {return EV;}
    inline void set_EV(NEWMAT::Matrix& Arg) {if(EV.Storage()==0)
        EV = Arg;}

    inline NEWMAT::Matrix& get_PPCA() {return PPCA;}
    inline void set_PPCA(NEWMAT::Matrix& Arg) {if(PPCA.Storage()==0)
        PPCA = Arg;}

    inline NEWMAT::Matrix& get_stdNoisei() {return stdNoisei;}
    inline void set_stdNoisei(NEWMAT::Matrix& Arg) {stdNoisei = Arg;}

    inline int data_dim() {return Data.Nrows();}
    inline int data_samples() {return Data.Ncols();}

    inline float get_resels() {return Resels;}
    inline void set_resels(float& Arg) {Resels = Arg;}

    inline int get_numfiles() {return numfiles;}

    inline void set_after_mm(bool val) {after_mm = val;}

    inline void flipres(int num){
      IC.Row(num) = -IC.Row(num);
      mixMatrix.Column(num) = -mixMatrix.Column(num);
      mixFFT=calc_FFT(mixMatrix);
      unmixMatrix = pinv(mixMatrix);
      if(ICstats.Storage()>0&&ICstats.Ncols()>3){
        double tmp;
        tmp = ICstats(num,3);
        ICstats(num,3) = -1.0*ICstats(num,4);
        ICstats(num,4) = -1.0*tmp;
      }
    }

    void sort();
    void dual_regression();

    std::vector<NEWMAT::Matrix> DWM, WM;
    basicGLM glmT, glmS;
    NEWMAT::Matrix Tdes, Tcon, TconF, Sdes, Scon, SconF, param, paramS;
    NEWMAT::RowVector explained_var;

  private:
    MelodicOptions &opts;
    Utilities::Log &logger;

    NEWMAT::Matrix pcaE;
    NEWMAT::RowVector pcaD;
    NEWMAT::Matrix whiteMatrix;
    NEWMAT::Matrix dewhiteMatrix;
    NEWMAT::Matrix meanC;
    NEWMAT::Matrix meanR;
    NEWMAT::Matrix stdDev;
    NEWMAT::Matrix stdDevi;
    NEWMAT::Matrix RXweight;
    NEWMAT::Matrix mixMatrix;
    NEWMAT::Matrix unmixMatrix;
    NEWMAT::Matrix mixFFT;
    NEWMAT::Matrix IC;
    NEWMAT::Matrix ICstats;
    std::vector<NEWMAT::Matrix> Tmodes;
    std::vector<NEWMAT::Matrix> Smodes;

    NEWMAT::Matrix EVP;
    NEWMAT::Matrix EV;
    NEWMAT::Matrix stdNoisei;
    NEWIMAGE::volume<float> Mask;
    NEWIMAGE::volume<float> Mean;
    NEWIMAGE::volume<float> background;


    NEWMAT::Matrix insta_mask;

    NEWMAT::Matrix Data;
    NEWMAT::Matrix PPCA;
    NEWMAT::Matrix jointCC;

    /*
     * Random number generator (seeded/used
     * alongside rand/srand for some operations).
     */
    std::mt19937_64 rng;

    bool after_mm;

    float Resels;
    int numfiles;

    char Mean_fname[1000];

    void setup_misc();
    void create_mask(NEWIMAGE::volume<float>& theMask);
    void create_RXweight();
    void est_smoothness();

    unsigned long standardise(NEWIMAGE::volume<float>& mask,
                              NEWIMAGE::volume4D<float>& R);
    float est_resels(NEWIMAGE::volume4D<float> R, NEWIMAGE::volume<float> mask);
  };
}

#endif