//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class TpcdEdxTaskAlice
//      see TpcdEdxTaskAlice2.h for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Felix Boehmer, Christian Hoeppner     TUM         
//      Philipp Gadow                         TUM
//
//-----------------------------------------------------------

// Panda Headers ----------------------

// This Class' Header ------------------
#include "TpcdEdxTaskAlice2.h"

// C/C++ Headers ----------------------
#include <cmath>
#include <iostream>
#include <map>
#include <assert.h> 
#include <stdexcept>
#include <algorithm>

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "GFTrackCand.h"
#include "TpcCluster.h"
#include "TpcDigiMapper.h"
#include "TpcFrontend.h"
#include "TpcSPHit.h"
#include "PndMCTrack.h"
#include "TpcPoint.h"
#include "TpcAlignmentManager.h"
#include "TpcDigiPadID.h"


#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "TpcDigiPar.h"

#include "GeaneTrackRep.h"
#include "GFAbsTrackRep.h"
#include "GFRecoHitFactory.h"
#include "GFKalman.h"
#include "GFException.h"
#include "TH1D.h"
#include "TFile.h"
#include "TGeoTrack.h"

#include "GFDetPlane.h"
#include "TpcdEdx.h"

#include "GFAbsRecoHit.h"
#include "TVector3.h"

#include <signal.h>
#include <stdlib.h>

#include "PndDetectorList.h"

// Class Member definitions -----------


TpcdEdxTaskAlice2::TpcdEdxTaskAlice2()
    : FairTask("dE/dx Task"), _persistence(kFALSE),
    _DX(1.5), _idealdEdx(kFALSE)
{
    _trackBranchName = "TpcTrackPostFit";
    _clusterBranchName = "TpcCluster";
    _digiBranchName = "TpcDigi";
    fVerbose = 0;
}


TpcdEdxTaskAlice2::~TpcdEdxTaskAlice2()
{
}

InitStatus
        TpcdEdxTaskAlice2::Init()
{
    //Get ROOT Manager
    FairRootManager* ioman= FairRootManager::Instance();

    if(ioman==0)
    {
        Error("TpcdEdxTaskAlice2::Init","RootManager not instantiated!");
        return kERROR;
    }

    // Get input collection
    _trackArray=(TClonesArray*) ioman->GetObject(_trackBranchName);
    if(_trackArray==0)
    {
        Error("TpcdEdxTaskAlice2::Init","track-array not found!");
        return kERROR;
    }

    if(_idealdEdx) {
        _mcTrackArray=(TClonesArray*) ioman->GetObject("MCTrack");
        if(_mcTrackArray==0)
            Fatal("TpcdEdxTaskAlice2::Init","MCTrack-array not found! Cannot calculate MC dEdx");

        _pointArray=(TClonesArray*) ioman->GetObject("TpcPoint");
        if(_pointArray==0)
            Fatal("TpcdEdxTaskAlice2::Init","Point-array not found! Cannot calculate MC dEdx");
    }

    _clusterArray=(TClonesArray*) ioman->GetObject(_clusterBranchName);
    if(_clusterArray==0)
    {
        Error("TpcdEdxTaskAlice2::Init","Cluster/Hit - array not found!");
        return kERROR;
    }

    _digiArray=(TClonesArray*) ioman->GetObject(_digiBranchName);
    if(_digiArray==0)
    {
        Error("TpcdEdxTaskAlice2::Init","Digi-array not found!");
        return kERROR;
    }

    _dEdxOutArray = new TClonesArray("TpcdEdx");
    ioman->Register("dEdx","Tpc",_dEdxOutArray,_persistence);

    _dEdxMCOutArray = new TClonesArray("TpcdEdx");
    ioman->Register("dEdx_MC","Tpc",_dEdxMCOutArray,_persistence);

    // setup histograms
    _distHist=new TH1D("distHist","Stepping lengths",500,-3.,3.);
    _dirHist=new TH1D("dirHist","direction abundances",10,-2.,2.);

    _hitBranchID = FairRootManager::Instance()->GetBranchId(_clusterBranchName);

    return kSUCCESS;
}


void
        TpcdEdxTaskAlice2::SetParContainers() {

    std::cout<<"TpcClusterFinderTask::SetParContainers"<<std::endl;
    std::cout.flush();

    // Get run and runtime database
    FairRun* run = FairRun::Instance();
    if ( ! run ) Fatal("SetParContainers", "No analysis run");

    FairRuntimeDb* db = run->GetRuntimeDb();
    if ( ! db ) Fatal("SetParContainers", "No runtime database");

    // Get Tpc digitisation parameter container
    _par= (TpcDigiPar*) db->getContainer("TpcDigiPar");
    if (! _par ) Fatal("SetParContainers", "TpcDigiPar not found");

}



void
        TpcdEdxTaskAlice2::Exec(Option_t* opt)
{
    FairRootManager* ioman= FairRootManager::Instance();

    static unsigned int nTotalDigis(0);
    static unsigned int nDoubleDigis(0);
    int evts(0);

    _dEdxOutArray->Delete();
    _dEdxMCOutArray->Delete();

    bool DEBUG = false;

    Int_t ntracks=_trackArray->GetEntriesFast();

    if(ntracks>20000){
        std::cout<<"ntracks="<<ntracks<<" Evil Event! skipping"<<std::endl;
        return;
    }

    for(Int_t itr=0;itr<ntracks;++itr){
        if (fVerbose) std::cout<<"TpcdEdxTaskAlice2::Exec(): starting track "<<itr<<std::endl;
        GFTrack* trk=(GFTrack*)_trackArray->At(itr);
        if (fVerbose) std::cout<<"*** Number of hits in track: "<<trk->getNumHits()<<" ***"<<std::endl;

        TpcdEdx dedx;
        TpcdEdx dedx_MC;

        GFAbsTrackRep* theRep = trk->getCardinalRep(); //track representation, used for extrapolation
        if (theRep->getStatusFlag() != 0) {
            if (DEBUG){
                std::cout<<" Track Representation was flagged, creating empty dedx, continue with next track"<<std::endl;
            }
            // write empty dEdx and continue with next trk
            new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);

            if(_idealdEdx) {
                new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
            }

            continue;
        }
        else{
            if (DEBUG){
                std::cout<<" Track Representation was not flagged"<<std::endl;
            }
        }

        //check for GEANE trackrep
        if(dynamic_cast<GeaneTrackRep*>(theRep) != NULL) {
            ((GeaneTrackRep*)theRep)->setPropDir(0); // not needed for RKTrackRep
        }

        GFTrackCand cand = trk->getCand();    //Track candidate -- a list of cluster indices

        bool useSPHits = false;
        if(TString(_clusterArray->GetClass()->GetName()) == "TpcSPHit"){
            useSPHits = true;
            if (DEBUG){
                std::cout<<" use SPHits = true"<<std::endl;
            }
        }
        else{
            if (DEBUG){
                std::cout<<" use SPHits = false"<<std::endl;
            }
        }
        
        std::vector<unsigned int> hitIDs = cand.getHitIDs(_hitBranchID); //gets IDs of clusters in track
        
        //GET MC INFORMATION FOR CROSS-CHECK
        //only works with IdealTracking and pure trackIDs
        std::vector<TpcPoint*> pointlist;
        if(_idealdEdx) {
            if (fVerbose) std::cerr<<"collect MC points"<<std::endl;
            //loop over MC points and collect
            for(int p=0; p<_pointArray->GetEntriesFast(); ++p) {
                int id = ((TpcPoint*)_pointArray->At(p))->GetTrackID();
                if(id==0) //only primary tracks
                    pointlist.push_back((TpcPoint*)_pointArray->At(p));
            }
        }

        if (fVerbose) std::cerr<<"collect digis"<<std::endl;
        // get all digis in the track out of the
        std::vector<const TpcDigi*> digis;           //vector with digis of the track
        unsigned int count = 0;
        TClonesArray* actClArr;
        for(int i=0; i<hitIDs.size(); i++){
            TpcCluster* cl;
            if(useSPHits) {
                TpcSPHit* theHit = (TpcSPHit*) _clusterArray->At(hitIDs[i]);
                int clBrID = theHit->getClusterBranchID();
                int clID = theHit->getClusterID();
                if(count==0) {
                    TString clBrName = ioman->GetBranchName(clBrID);
                    actClArr = (TClonesArray*) ioman->GetObject(clBrName);
                }

                cl = (TpcCluster*)(actClArr->At(clID));
                count++;
            }
            else
                cl = (TpcCluster*) _clusterArray->At(hitIDs[i]);
            
            const std::map<unsigned int,double>* digiMap = cl->getDigiMap(); //all the digis in the cluster

            for(std::map<unsigned int, double>::const_iterator it=digiMap->begin();it!=digiMap->end();++it){
                TpcDigi* aDigi=new TpcDigi(*((TpcDigi*)_digiArray->At((*it).first)));
                if (DEBUG){
                    std::cout<<"processing digi with amp "<< aDigi->amp()<<std::endl;
                }
                aDigi->amp(aDigi->amp()*(*it).second); //if the digi is shared assign it a fractional amplitude (fraction is in the 2nd entry in map, usually .5)
                if (DEBUG){
                    std::cout<<" digi amp is now "<< aDigi->amp()<<std::endl;
                }
                digis.push_back(aDigi);
            }
        }

        //######################################################################################
        //           list of digis was obtained, dedx algorithm starts now
        //######################################################################################
        TpcDigiPadID tpcdigipadid;

        // sort digis with respect to their pad ids
        std::sort (digis.begin(), digis.end(), tpcdigipadid);

        TVector3 pos,mom; //get position and momentum from track representation
        if (fVerbose) std::cerr<<"get pos and mom"<<std::endl;
        try{
            pos = theRep->getPos();
            if(DEBUG){
                std::cout<<"pos x:"<<pos.x()<<", y:"<<pos.y()<<", z"<< pos.z()<<std::endl;
            }
            mom = theRep->getMom();
        }

        catch(GFException& e){
            e.what();

            // if Exception: write empty dEdx and continue with next trk
            new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);

            if (DEBUG){
                std::cout<<"exception: write empty dedx and continue with next track"<<std::endl;
            }

            if(_idealdEdx) {
                new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
            }

            continue;
        }

        dedx.setMom(mom);

        //define padrows as planes perpendicular to x direction
        GFDetPlane padrow, nextpadrow;
        TVector3 distbetweenpadrows(0.75,0,0); //distance between two padrows is 0.75 cm
        unsigned int numberofpadrows = 63;
        //set coordinates for first padrow
        double xfirstpadrow = 85.225;
        padrow.setO(xfirstpadrow,0,0); //origin of plane of first pad row is middle of first pad row
        padrow.setNormal(1,0,0); //normal of plane of first padrow shows in x-direction
        nextpadrow.setO(padrow.getO()+distbetweenpadrows);
        nextpadrow.setNormal(1,0,0);

        //get track position on the first padrow

        try{
            theRep->extrapolate(padrow);
        }
        catch(GFException& e){
            e.what();

            // if Exception: write empty dEdx and continue with next trk
            new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);

            if (DEBUG){
                std::cout<<"exception: write empty dedx and continue with next track"<<std::endl;
            }

            if(_idealdEdx) {
                new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
            }

            continue;
        }
        TVector3 startPoint = theRep->getPos();
        TVector3 nextPoint(0,0,0);

        if (DEBUG){
            std::cout<<"track was extrapolated to first pad, position of track is now"<<std::endl;
            std::cout<<"pos x:"<<theRep->getPos().x()<<", y:"<<theRep->getPos().y()<<", z"<< theRep->getPos().z()<<std::endl;
        }

        bool exc(false); //exception
        for (unsigned int iPadrow = 0; iPadrow < numberofpadrows; ++iPadrow){
            //if not the first time: move on a padrow and define new startpoint
            if(iPadrow>0){
                padrow.setO(nextpadrow.getO());
                nextpadrow.setO(padrow.getO()+distbetweenpadrows);
                startPoint = nextPoint;

                if (DEBUG){
                    std::cout<<"padrow #: "<< iPadrow<<"  position in x:"<<padrow.getO().x()<<", nextpadrow position in x:"<< nextpadrow.getO().x() <<std::endl;
                }
            }


            try{
                theRep->extrapolate(nextpadrow);
            }
            catch(GFException& e){
                e.what();

                // if Exception: write empty dEdx and continue with next trk
                new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);

                if (DEBUG){
                    std::cout<<"exception: write empty dedx and continue with next track"<<std::endl;
                }

                if(_idealdEdx) {
                    new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
                }

                continue;
            }
            nextPoint = theRep->getPos();


            if(DEBUG){
                std::cout<<"padrow: "<<iPadrow<<std::endl;
            }

            //            deviation of the digi to the vector between the planes is not implemented
            //            dedx.setDistX( (nextPoint - startPoint).X());
            //            dedx.setDistY( (nextPoint - startPoint).Y());
            //            dedx.setDistZ( (nextPoint - startPoint).Z());

            double dist = (nextPoint - startPoint).Mag();

            double dE = 0.;


            std::set<const TpcDigi*> usedDigis;

            for(unsigned int i=0;i<digis.size();++i){
                TVector3 digiPos;
                try {
                    TpcDigiMapper::getInstance()->map(digis[i],digiPos);
                }
                catch(const std::exception& ex) {
                    std::cout<<ex.what()<<std::endl;
                }

                bool isinpadrow = false;
                //calculate, which digis lie between the two rows and sum up their amplitudes
                if (digiPos.x()<nextpadrow.getO().x() && digiPos.x()>= padrow.getO().x()){
                    dE += digis[i]->amp();
                    isinpadrow = true;
                    if (usedDigis.count(digis[i])>0){
                        std::cerr << "TpcdEdxTask::Exec --> Digi used twice" << std::endl;
                        nDoubleDigis += 1;
                    }
                    usedDigis.insert(digis[i]);
                }
                if (DEBUG){
                    std::cout<<"digiPos x:"<<digiPos.x()<<", y:"<<digiPos.y()<<", z"<< digiPos.z()<<", is in padrow?:"<<isinpadrow<<std::endl;
                }
            }

            if (DEBUG){
                std::cout<<"check if add to dedx: dE>0? "<<dE<<", exception = false?"<<exc<<std::endl;
            }

            if(dE>0. && (!exc)){
                dedx.add(dE,dist);
                if (DEBUG){
                    std::cout<<"dedx successflully added to list."<<std::endl;
                }
            }

            dE = 0.;

        }

        if (DEBUG){
            std::cout<<" successfully created new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);"<<std::endl;
        }

        new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);

        if(_idealdEdx) {
            new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
        }

        for(unsigned int i=0;i<digis.size();++i){
            delete digis[i];
        }

    }
}


void 
        TpcdEdxTaskAlice2::WriteHistograms(const TString& filename){
    TFile* file = new TFile(filename,"RECREATE");
    file->mkdir("DEDX");
    file->cd("DEDX");

    _distHist->Write();
    delete _distHist;

    _dirHist->Write();
    delete _dirHist;

    file->Close();
    delete file;
}



ClassImp(TpcdEdxTaskAlice2)