// -------------------------------------------------------------------------
// -----                    CbmTrdSetTracksPidANN source file          -----
// -----                  Created 06/03/07 by Simeon Lebedev           -----
// -------------------------------------------------------------------------
#include "CbmTrdSetTracksPidANN.h"

#include "CbmTrdTrack.h"
#include "CbmTrdHit.h"

#include "CbmRootManager.h"

#include "TClonesArray.h"
#include "TMultiLayerPerceptron.h"
#include "TTree.h"
#include "TFile.h"
#include "TMath.h"

#include <iostream>
#include <vector>
#include <fstream>
using std::cout;
using std::endl;


// -----   Default constructor   -------------------------------------------
CbmTrdSetTracksPidANN::CbmTrdSetTracksPidANN() {

  fTrackArray  = NULL;
  fTrdHitArray = NULL;
  fNN          = NULL;
  fNofTracks   = 0;
  ANNPar1      = 1.279;
  ANNPar2      = 0.61;
  NeuralNetWeightsFile = "";
}
// -------------------------------------------------------------------------



// -----   Standard constructor   ------------------------------------------
CbmTrdSetTracksPidANN::CbmTrdSetTracksPidANN(const char* name,
				 const char* title)
  : CbmTask(name) {
  fTrackArray  = NULL;
  fTrdHitArray = NULL;
  fNN          = NULL;
  fNofTracks   = 0;
  ANNPar1      = 1.279;
  ANNPar2      = 0.61;
  NeuralNetWeightsFile = "";

}
// -------------------------------------------------------------------------



// -----   Destructor   ----------------------------------------------------
CbmTrdSetTracksPidANN::~CbmTrdSetTracksPidANN() { }
// -------------------------------------------------------------------------



// -----  SetParContainers -------------------------------------------------
void CbmTrdSetTracksPidANN::SetParContainers()
{

}
// -------------------------------------------------------------------------


// -----  RaedData -------------------------------------------------
Bool_t CbmTrdSetTracksPidANN::ReadData() {

  //
  // Read the weights for the neural network.
  //

  // Check if the input filename is valid

  const char* File = (const char*)NeuralNetWeightsFile;

  ifstream myfile(File);
  if (!myfile.is_open()) {
    cout << "-E- CbmTrdSetTracksPidANN::Init: "
  	 << "Could not open input file." << endl;
    return kFALSE;
  }
  myfile.close();
  
  cout << "-I- CbmTrdSetTracksPidANN::Init : get NeuralNet weight parameters from: " << NeuralNetWeightsFile << endl;

  Float_t inVector[12];
	
  Int_t x13;//output value

  //init TTree as input data to neural net
  TTree *simu = new TTree ("MonteCarlo","MontecarloData");
  simu->Branch("x1", &inVector[0],"x1/F");
  simu->Branch("x2", &inVector[1],"x2/F");
  simu->Branch("x3", &inVector[2],"x3/F");
  simu->Branch("x4", &inVector[3],"x4/F");
  simu->Branch("x5", &inVector[4],"x5/F");
  simu->Branch("x6", &inVector[5],"x6/F");
  simu->Branch("x7", &inVector[6],"x7/F");
  simu->Branch("x8", &inVector[7],"x8/F");
  simu->Branch("x9", &inVector[8],"x9/F");
  simu->Branch("x10", &inVector[9],"x10/F");
  simu->Branch("x11", &inVector[10],"x11/F");
  simu->Branch("x12", &inVector[11],"x12/F");
  simu->Branch("x13", &x13,"x13/I");
  fNN = new TMultiLayerPerceptron("x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12:30:x13",simu);

  fNN->LoadWeights(File);   
  fNN->DumpWeights("test.txt");  

  return kTRUE;
}

// -----   Public method Init (abstract in base class)  --------------------
InitStatus CbmTrdSetTracksPidANN::Init() {

  // Read the data from ROOT file. In case of problems return kFATAL;

  if(!NeuralNetWeightsFile) {
    cout << "-E- CbmTrdSetTracksPidANN::Init : "
	 << "No input file for weigths set." << endl;
    return kFATAL;
  }

  if (!ReadData()) return kFATAL;

  // Get and check CbmRootManager
  CbmRootManager* ioman = CbmRootManager::Instance();
  if (! ioman) {
    cout << "-E- CbmTrdSetTracksPidANN::Init: "
	 << "RootManager not instantised!" << endl;
    return kFATAL;
  }

  // Get TrdTrack array
  fTrackArray  = (TClonesArray*) ioman->GetObject("TRDTrack"); 
  if ( ! fTrackArray) {
    cout << "-E- CbmTrdSetTracksPidANN::Init: No TrdTrack array!"
	 << endl;
    return kERROR;
  }

  // Get TrdTrack array
  fTrdHitArray  = (TClonesArray*) ioman->GetObject("TRDHit");
  if ( ! fTrdHitArray) {
    cout << "-E- CbmTrdSetTracksPidANN::Init: No TrdHit array!"
	 << endl;
    return kERROR;
  }

  return kSUCCESS;

}
// -------------------------------------------------------------------------



// -----   Public method Exec   --------------------------------------------
void CbmTrdSetTracksPidANN::Exec(Option_t* opt) {

 if ( !fTrackArray ) return; 
  	
 Int_t nTracks = fTrackArray->GetEntriesFast();
 std::vector<Double_t> eLossVector;
 
   for (Int_t iTrack=0; iTrack<nTracks; iTrack++) {

     eLossVector.clear();
    	
     CbmTrdTrack* pTrack = (CbmTrdTrack*)fTrackArray->At(iTrack);
    
    // Up to now only for tracks with twelve hits the Ann can be calculated
    if (pTrack->GetNofTrdHits() < 12 ) {
      fNofTracks++;
      pTrack->SetPidANN(-2);
      continue;
    }
    
    for (Int_t iTRD=0; iTRD < pTrack->GetNofTrdHits(); iTRD++){
      Int_t TRDindex = pTrack->GetTrdHitIndex(iTRD);
      CbmTrdHit* trdHit = (CbmTrdHit*) fTrdHitArray->At(TRDindex);
      eLossVector.push_back(trdHit->GetELoss());
    }

    //------------------transform Data BEGIN--------------------------

    for (UInt_t j = 0; j<eLossVector.size();j++) {
      eLossVector[j]=eLossVector[j]*1.e6;
      eLossVector[j]=(eLossVector[j]-ANNPar1)/ANNPar2-0.225;
    }

    sort(eLossVector.begin(), eLossVector.end());

    for (UInt_t j = 0; j<eLossVector.size();j++) {
      eLossVector[j]=TMath::LandauI(eLossVector[j]);
    }

    //------------------transform Data END----------------------------

    Double_t params[12];
    for (UInt_t k = 0 ; k < eLossVector.size(); k++)
      params[k] =  eLossVector[k];
    
    Double_t nnEval = fNN->Evaluate(0,params);
    
    if (TMath::IsNaN(nnEval) == 1) {
      cout << " -W- CbmTrdSetTracksPidANN: nnEval nan " << endl;
      nnEval = -2;
    }
    
    pTrack->SetPidANN(nnEval);

  }
}
// -------------------------------------------------------------------------



// -----   Public method Finish   ------------------------------------------
void CbmTrdSetTracksPidANN::Finish() { }


ClassImp(CbmTrdSetTracksPidANN)