#ifndef COVARIANCEMATRIX_H_107I2BH1
#define COVARIANCEMATRIX_H_107I2BH1

#include "TMatrixD.h"
#include "vector"
#include "TH2D.h"
#include "TVectorD.h"
#include "TString.h"
#include <math.h>
#include "TCanvas.h"
using namespace std;

class CovarianceMatrix { // consider inheriting from TMatrix...

public:

  CovarianceMatrix(int _dim){
     LateConstructor(_dim);
  }
  
  CovarianceMatrix(){
     // Construct at first fill
     dim = 0;
  }
  
  void LateConstructor(int _dim){
     dim = _dim;
     covMat = new TMatrixD(dim, dim);
     corrMat = new TMatrixD(dim, dim);
     covNu = new TMatrixD(dim, dim);
     covMu = new TVectorD(dim);
     N = new TMatrixD(dim, dim);
     NMu = new TVectorD(dim);
  }
  
  ~CovarianceMatrix(){
     delete covMat;
     delete corrMat;
     delete covMu;
     delete covNu;
     delete N;
     delete NMu;
  }
  
  void Fill(std::vector<Double_t> input){
     if (dim == 0) {  // Create structure based on first input vector (useful to avoid keeping track of the number of variables doing the analysis)
        LateConstructor(input.size());
     }


    for (unsigned int i = 0; i < input.size(); ++i)
      {
        if (!isnan(input.at(i)) and !isinf(input.at(i))){
          (*covMu)[i] += input.at(i);
          (*NMu)[i] += 1;
        
          for (unsigned int j = 0; j < input.size(); ++j)
            {
              if (!isnan(input.at(j)) and !isinf(input.at(j))){
                (*covNu)[i][j] += input.at(i) * input.at(j);
                (*N)[i][j] += 1;
              }
            }
        }
      }

  }
  
  void Finalize() const {
    for (unsigned int i = 0; i < dim; ++i)
      {
        for (unsigned int j = 0; j < dim; ++j)
          {
            (*covMat)[i][j] =  ((*N)[i][j] != 0.0) ? (*covNu)[i][j] / (*N)[i][j] - (*covMu)[i]/(*NMu)[i] * (*covMu)[j] / (*NMu)[j] : 0.0;
          }
      }

      for(size_t i = 0; i < dim; ++i)
      {
         for(size_t j = 0; j < dim; ++j)
         {
            (*corrMat)[i][j] = ((*covMat)[i][j] != 0.0) ? (*covMat)[i][j]/(sqrt((*covMat)[i][i]) * sqrt((*covMat)[j][j])) : 0.0;
         }
      }
  }


  TMatrixD GetCovarianceMatrix(){
    Finalize();
    return (*covMat); // Return copy
  }
  TMatrixD Get(){
    return (*covMat);// Return Copy
  }

  void SetLabels(std::vector<TString> labels) {
     _labels = labels;
  }

  TMatrixD GetCorrelationMatrix(){
    Finalize();
    return (*corrMat); // Return copy
  }

  void Print(){
    Finalize();
    covMat->Print();
    corrMat->Print();
  }
  
  
  TH2D* PlotCorrelation() const {
     Finalize();
     
     TH2D* h2 = new TH2D( *corrMat );


     
     h2->SetNameTitle( "Correlations", "Correlations" );
     // 
     
     // From TMVA
     if (_labels.size() > 0) {
        for (UInt_t ivar=0; ivar< dim; ivar++) {
           h2->GetXaxis()->SetBinLabel( ivar+1, _labels.at(ivar) );
           h2->GetYaxis()->SetBinLabel( ivar+1, _labels.at(ivar) );
        }
     }
     
     h2->Scale( 100.0  ); 
     for (UInt_t ibin=1; ibin<=dim; ibin++) {
        for (UInt_t jbin=1; jbin<=dim; jbin++) {
           h2->SetBinContent( ibin, jbin, Int_t(h2->GetBinContent( ibin, jbin )) );
        }
     }

     // style settings
     const Float_t labelSize = 0.02;
     h2->SetStats( 0 );
     h2->GetXaxis()->SetLabelSize( labelSize );
     h2->GetYaxis()->SetLabelSize( labelSize );
     h2->SetMarkerSize( 0.8 );
     h2->SetMarkerColor( 1 );
     h2->LabelsOption( "d" ); // diagonal labels on x axis
     h2->SetLabelOffset( 0.008 );// label offset on x axis
     h2->SetMinimum( -100.0 );
     h2->SetMaximum( +100.0 );

     
     return h2;
     
  }
  
  
CovarianceMatrix operator+(CovarianceMatrix &other) 
  { 
    CovarianceMatrix temp(dim); 

    (*temp.N) += (*N);
    (*temp.N) += (*other.N);
    
    (*temp.covMu) += (*covMu);
    (*temp.covMu) += (*other.covMu);

    (*temp.covNu) += (*covNu);
    (*temp.covNu) += (*other.covNu);
    
    return temp; 
  }

  void operator+=(CovarianceMatrix &other) 
    { 
       if (dim == other.dim) {
         (*N) += (*other.N);
         (*covMu) += (*other.covMu);
         (*covNu) += (*other.covNu);
      }

    }



public:

  TMatrixD *N;
  TVectorD *NMu;

  TMatrixD *covMat;
  TMatrixD *corrMat;

  TVectorD *covMu;
  TMatrixD *covNu;

  unsigned int dim;
  
  std::vector<TString> _labels;


};


#endif /* end of include guard: COVARIANCEMATRIX_H_107I2BH1 */