//____________________________________________________________________________
/*!

\class   KM3NeTWrite

\brief   A class to write gSeaGen output in a KM3NeT file

\author  Alfonso Andres Garcia Soto <alfonsog@nikhef.nl>
         Nikhef, Netherlands

\created 2019

\cpright  Copyright (c) 2019, The KM3NeT Collaboration
          For the full text of the license see $GSEAGEN/LICENSE

*/
//____________________________________________________________________________


#ifdef _KM3NET_ENABLED__

#ifndef _xxKM3NETWRITE_H__
#define _xxKM3NETWRITE_H__

#include <cstdio>
#include <fstream>
#include <ctime>
#include <string>
#include <map>

#include <TSystem.h>

#include "PropaMuon/PropaTracks.h"
#include "SeaEvent/GBinParam.h"
#include "SeaEvent/GSeaEvent.h"

#include <ostream>
#include "Head.hh"
#include "Evt.hh"
#include "../definitions/w2list_gseagen.hh"
#include "../definitions/trkmembers.hh"
#include "TFile.h"
#include "TTree.h"

using namespace std;
using namespace genie;

class KM3NeTWrite {

public :

  GenParam* fGenPar;
  TFile* fOutFile;
  TTree* fTree;
  Evt*   fEvt;

  KM3NeTWrite(GenParam* GenPar ) {

    fGenPar = GenPar;

    ostringstream fnstr;
    fnstr << fGenPar->EvFilePrefix  << "." << fGenPar->RunNu << ".aa.root";

    string OutFileName=fnstr.str();

    // setup km3net output tree
    fEvt     = new Evt();
    fOutFile = new TFile( OutFileName.c_str() , "RECREATE");
    fTree    = new TTree( "E", "KM3NeT Evt Tree");
    fTree->Branch("Evt", &fEvt, 32000, 4 );
  }

  ~KM3NeTWrite() {
    fOutFile->Write();
    delete fOutFile;
  }

  void WriteHeader() {

    string prefix = "";

    fOutFile->cd();

    Head h;
    tm *date;
    date = localtime (&fGenPar->RunTime);

    char ch[1024];

    unsigned iPhysics=0;  // physics tag counter
    unsigned iSimul=0;    // simul tag counter
    unsigned iSeed=0;     // seed tag counter

    sprintf(ch,"%d",fGenPar->RunNu);
    h.set_line(prefix+"start_run",ch);

    if ( (fGenPar->GenMode.compare("BIN")==0) || (fGenPar->GenMode.compare("EVT")==0) ){
      sprintf(ch,"CORSIKA %5.3f %2.2d  %2.2d%2.2d ", fGenPar->CORSIKAVer,fGenPar->CORSIKADate,0,0);
      if (iSimul==0) h.set_line(prefix+"simul",ch);
      else           h.set_line(prefix+"simul_"+to_string(iSimul),ch);
      iSimul++;
    }

    sprintf(ch,"gSeaGen %s %2.2d%2.2d%2.2d  %2.2d%2.2d ", fGenPar->gSeaGenVer.c_str(),
      date->tm_year-100,date->tm_mon+1,date->tm_mday,date->tm_hour,date->tm_min);
    if (iSimul==0) h.set_line(prefix+"simul",ch);
    else           h.set_line(prefix+"simul_"+to_string(iSimul),ch);
    iSimul++;


    if ( fGenPar->GenMode.compare("BIN")==0 ){
      sprintf(ch,"%10.2E %10.2E %10.2E %10.2E",fGenPar->EPrimMin,fGenPar->EPrimMax,fGenPar->PrimCtMin,fGenPar->PrimCtMax);
      h.set_line(prefix+"cut_primary",ch);
    }

    if ( fGenPar->GenMode.compare("BIN")==0 ){
      sprintf(ch,"%10.2E %d %d %d",fGenPar->EmuMin,0,0,0);
      h.set_line(prefix+"cut_seamuon",ch);
    }

    if ( fGenPar->GenMode.compare("BIN")==0 ) sprintf(ch,"%10.2E %d %d %d",fGenPar->EvMin,0,0,0);
    else sprintf(ch,"%10.2E %10.2E %10.2E %10.2E",fGenPar->EvMin,fGenPar->EvMax,fGenPar->CtMin,fGenPar->CtMax);
    h.set_line(prefix+"cut_nu",ch);

    if ( (fGenPar->GenMode.compare("BIN")==0) || (fGenPar->GenMode.compare("EVT")==0) ) h.set_line(prefix+"source_mode","FILE");
    else h.set_line(prefix+"source_mode",fGenPar->GenMode);

    sprintf(ch,"%+4.2f", -1*fGenPar->Alpha);
    h.set_line(prefix+"spectrum",ch);

    if ( fGenPar->GenMode.compare("POINT")==0 || fGenPar->GenMode.compare("DIFFUSE")==0 ) {

//      sprintf(ch, "GENIE %s %2.2d%2.2d%2.2d  %2.2d%2.2d%2.2d",fGenPar->GenieVer.c_str(),
//        date->tm_year-100,date->tm_mon+1,date->tm_mday,date->tm_hour,date->tm_min,date->tm_sec);

      sprintf(ch, "GENIE %s",fGenPar->GenieVer.c_str());

      if (iPhysics==0) h.set_line(prefix+"physics",ch);
      else             h.set_line(prefix+"physics_"+to_string(iPhysics),ch);
      iPhysics++;


#ifdef _GENIEVERSIONGTEQ3__
      string ReWgtVer;
      string VerFile = gSystem->Getenv("GENIE_REWEIGHT");
      VerFile.append("/VERSION");
      ifstream file_stream(VerFile.c_str(), ios::in);
      file_stream>>ReWgtVer;
      file_stream.close();

//      sprintf(ch, "GENIE_REWEIGHT %s %2.2d%2.2d%2.2d  %2.2d%2.2d%2.2d",ReWgtVer.c_str(),
//        date->tm_year-100,date->tm_mon+1,date->tm_mday,date->tm_hour,date->tm_min,date->tm_sec);

      sprintf(ch, "GENIE_REWEIGHT %s",ReWgtVer.c_str());

      if (iPhysics==0) h.set_line(prefix+"physics",ch);
      else             h.set_line(prefix+"physics_"+to_string(iPhysics),ch);
      iPhysics++;

      sprintf(ch, "GENIE_TUNE %s %s", fGenPar->GenTune.c_str(), fGenPar->EventGeneratorList.c_str());
      if (iPhysics==0) h.set_line(prefix+"physics",ch);
      else             h.set_line(prefix+"physics_"+to_string(iPhysics),ch);
      iPhysics++;

#endif

      h.set_line(prefix+"XsecFile",fGenPar->InpXSecFile);
      
      if ( fGenPar->GenMode == "POINT") {
        if ( fGenPar->SourceName != "" ){
          ostringstream name;
          name<<"Dec:";
          name.precision(3);
          name<<fixed;    
          name << fGenPar->Declination*180./M_PI;
          cout << name.str();
        }
        sprintf(ch,"%f %f %f %s",fGenPar->RightAscension,fGenPar->Declination,fGenPar->SourceRadius*180./M_PI,fGenPar->SourceName.c_str());
        h.set_line(prefix+"source_info",ch);
      }         

    }
    else if ( (fGenPar->GenMode.compare("BIN")==0) || (fGenPar->GenMode.compare("EVT")==0) ) {
      sprintf(ch, "CORSIKA %s %s",fGenPar->CORSIKAHadronicModelLE.c_str(),fGenPar->CORSIKAHadronicModelHE.c_str());

      if (iPhysics==0) h.set_line(prefix+"physics",ch);
      else             h.set_line(prefix+"physics_"+to_string(iPhysics),ch);
      iPhysics++;
    }

    if(fGenPar->GenMJD){
      sprintf(ch,"%d %d", fGenPar->TimeStampStart,fGenPar->TimeStampStop);
      h.set_line(prefix+"time_interval",ch);
    }

    h.set_line(prefix+"drawing",fGenPar->Drawing);

    if(fGenPar->detector.size()>0){
      sprintf(ch, "gSeaGen %s", fGenPar->detector.c_str());
      h.set_line(prefix+"detector",ch);
    }

    sprintf(ch, "%10.3f %10.3f %10.3f %10.3f %10.3f",
      fGenPar->OriginX,fGenPar->OriginY,fGenPar->Can1-fGenPar->OriginZ,fGenPar->Can2-fGenPar->OriginZ,fGenPar->Can3);
    h.set_line(prefix+"fixedcan",ch);

    sprintf(ch,"%10.6f %10.6f",fGenPar->SiteLatitude,fGenPar->SiteLongitude);
    h.set_line(prefix+"coord_site",ch);

    sprintf(ch,"%0.3f", fGenPar->SiteDepth);
    h.set_line(prefix+"depth",ch);

    if(fGenPar->HBedRock>0){
      sprintf(ch,"%10.3f", fGenPar->HBedRock );
      h.set_line(prefix+"bedrock",ch);
    } 

    double IntVol=M_PI*pow(fGenPar->RVol,2);
    IntVol *= fGenPar->HSeaWater*fGenPar->RhoSW+fGenPar->HRock*fGenPar->RhoSR;

    double Zmin=fGenPar->Can1-fGenPar->HRock+fGenPar->OriginZ;
    double Zmax=fGenPar->Can1+fGenPar->HSeaWater+fGenPar->OriginZ;

    if ( (fGenPar->GenMode.compare("BIN")==0) || (fGenPar->GenMode.compare("EVT")==0) ){
      sprintf(ch,"0 0 0 0 %12.2E",fGenPar->NTot);
      h.set_line(prefix+"genvol",ch);
    }
    else {
      sprintf(ch,"%10.3f %10.3f %10.3f %12.2E %12.2E",Zmin,Zmax,fGenPar->RVol,IntVol,fGenPar->NTot);
      h.set_line(prefix+"genvol",ch);
    }

    if ( (fGenPar->GenMode.compare("BIN")==0) || (fGenPar->GenMode.compare("EVT")==0) ){
      sprintf(ch,"CORSIKA %d %d", fGenPar->CORSIKASeed[0], fGenPar->CORSIKASeed[1] );
      if (iSeed==0) h.set_line(prefix+"seed",ch);
      else          h.set_line(prefix+"seed_"+to_string(iSeed),ch);
      iSeed++;
    }

    sprintf(ch,"gSeaGen %d", fGenPar->RanSeed );
    if (iSeed==0) h.set_line(prefix+"seed",ch);
    else          h.set_line(prefix+"seed_"+to_string(iSeed),ch);
    iSeed++;
    

    if ( fGenPar->GenMode.compare("BIN")==0 ){
      sprintf(ch,"%10d", fGenPar->Primary );
      h.set_line(prefix+"primary",ch);
    }
    else{
      ostringstream nulist;
      for(unsigned iNuType=0; iNuType<fGenPar->NuList.size(); iNuType++)nulist << fGenPar->NuList[iNuType]  << " ";
      string snulist=nulist.str();
      h.set_line(prefix+"primary",snulist);  
    }


    if(gSystem->Getenv("GXMLPATH")){
      string gxmlpath=gSystem->Getenv("GXMLPATH");
      sprintf(ch, "%s",gxmlpath.c_str());
      h.set_line(prefix+"GXMLPATH",ch);
    }

    for(unsigned iFlux=0; iFlux<fGenPar->FluxFiles.size(); iFlux++){
      if ( (fGenPar->GenMode.compare("BIN")==0) || (fGenPar->GenMode.compare("EVT")==0) ){
        sprintf(ch,"%d %s", fGenPar->Primary, fGenPar->CRModel.c_str() );
      }
      else sprintf(ch,"%d %s", fGenPar->FluxFiles[iFlux].NuType, fGenPar->FluxFiles[iFlux].FluxSimul.c_str() );

      for(unsigned iFile=0; iFile<fGenPar->FluxFiles[iFlux].FileNames.size(); iFile++) sprintf(ch,"%s %s", ch, fGenPar->FluxFiles[iFlux].FileNames[iFile].c_str());
      if (iFlux==0) h.set_line(prefix+"flux",ch);
      else          h.set_line(prefix+"flux_"+to_string(iFlux),ch);
    }

    if(fGenPar->TrackEvts>0){
      sprintf(ch, "gSeaGen %s %s", fGenPar->PropCode.c_str(),fGenPar->PropCodeVer.c_str());
      if (iPhysics==0) h.set_line(prefix+"physics",ch);
      else             h.set_line(prefix+"physics_"+to_string(iPhysics),ch);
      iPhysics++;
    }

    h.Write("Head");

  }


  void write( GSeaTrack& gst, Trk& t ) {
    t.id        = gst.Id;
    t.status    = gst.Status;
    t.type      = gst.PdgCode;
    t.mother_id = gst.MotherId ;
    t.E         = gst.E;
    t.pos       = Vec( gst.V1, gst.V2, gst.V3 );
    t.dir       = Vec( gst.D1, gst.D2, gst.D3 );
    t.len       = gst.Length;
    t.t         = gst.T;
    t.usr.push_back( gst.GenCounter );
    t.usr_names.push_back( "generation_counter" );
  }

  Trk to_trk( GSeaTrack& gst ) {
    Trk t;
    write( gst, t );
    return t;
  }

  void WriteEvent(GSeaEvent * SeaEvt) {

    fEvt -> id = SeaEvt->iEvt;
    fEvt -> mc_run_id = fGenPar->RunNu;
    fEvt -> mc_event_time.SetSec( SeaEvt->EvtTime[0] );
    fEvt -> mc_event_time.SetNanoSec( SeaEvt->EvtTime[1] * 16 );
    fEvt -> mc_trks.clear();
    for( auto gst : SeaEvt->Tracks ) {
      if ( gst.Status==0 ) {
        int apdg = TMath::Abs(gst.PdgCode);
        if ( apdg==12 || apdg==14 || apdg==16 ) gst.Status = TRK_ST_PRIMARYNEUTRINO;
        else                                    gst.Status = TRK_ST_PRIMARYCOSMIC;
      }
      fEvt->mc_trks.push_back ( to_trk( gst ));
    }
    // weight information
    fEvt -> w = { SeaEvt->Agen, SeaEvt->GenWeight, SeaEvt->EvtWeight };

    int lepincan = SeaEvt->LepInCan ? 1 : 0;
    int verincan = SeaEvt->VerInCan;

    //see https://git.km3net.de/common/km3net-dataformat/-/blob/master/definitions/w2list_gseagen.csv
    fEvt -> w2list.resize(23);
    fEvt -> w2list[W2LIST_GSEAGEN_PS]            = fGenPar->GlobalGenWeight;
    fEvt -> w2list[W2LIST_GSEAGEN_EG]            = pow(SeaEvt->Neutrino.E,fGenPar->Alpha);
    fEvt -> w2list[W2LIST_GSEAGEN_XSEC_MEAN]     = SeaEvt->XSecMean;
    fEvt -> w2list[W2LIST_GSEAGEN_COLUMN_DEPTH]  = SeaEvt->ColumnDepth;
    fEvt -> w2list[W2LIST_GSEAGEN_P_EARTH]       = SeaEvt->PEarth;
    fEvt -> w2list[W2LIST_GSEAGEN_WATER_INT_LEN] = SeaEvt->WaterIntLen;
    fEvt -> w2list[W2LIST_GSEAGEN_P_SCALE]       = SeaEvt->PScale;
    fEvt -> w2list[W2LIST_GSEAGEN_BX]            = SeaEvt->Bx;
    fEvt -> w2list[W2LIST_GSEAGEN_BY]            = SeaEvt->By;
    fEvt -> w2list[W2LIST_GSEAGEN_ICHAN]         = double( SeaEvt->ScattId );
    fEvt -> w2list[W2LIST_GSEAGEN_CC]            = double( SeaEvt->InterId );
    fEvt -> w2list[W2LIST_GSEAGEN_DISTAMAX]      = SeaEvt->DistaMax;
    fEvt -> w2list[W2LIST_GSEAGEN_WATERXSEC]     = SeaEvt->SWXSec;
    fEvt -> w2list[W2LIST_GSEAGEN_XSEC]          = SeaEvt->XSec;
    fEvt -> w2list[W2LIST_GSEAGEN_DXSEC]         = SeaEvt->DXSec;
    fEvt -> w2list[W2LIST_GSEAGEN_TARGETA]       = double( SeaEvt->TargetA );
    fEvt -> w2list[W2LIST_GSEAGEN_TARGETZ]       = double( SeaEvt->TargetZ );
    fEvt -> w2list[W2LIST_GSEAGEN_VERINCAN]      = double( verincan );
    fEvt -> w2list[W2LIST_GSEAGEN_LEPINCAN]      = double( lepincan );
    fEvt -> w2list[W2LIST_GSEAGEN_N_RETRIES]     = SeaEvt->NRetries;            // save the number of retries to account for it in the weights
    fEvt -> w2list[W2LIST_GSEAGEN_CUSTOM_YAW]    = fGenPar->CustomRotation[0];  // user-specified rotation of CORSIKA showers (around z-axis)
    fEvt -> w2list[W2LIST_GSEAGEN_CUSTOM_PITCH]  = fGenPar->CustomRotation[1];  // user-specified rotation of CORSIKA showers (around y-axis)
    fEvt -> w2list[W2LIST_GSEAGEN_CUSTOM_ROLL]   = fGenPar->CustomRotation[2];  // user-specified rotation of CORSIKA showers (around x-axis)
    fTree->Fill();
  }


};

#endif // _GKM3NETWRITE_H__

#endif