//____________________________________________________________________
//
// Atl Task for MemTk Analysis
// 
//  
// Author: stamm <mailto: stamm@physik.hu-berlin.de>
// Update: $Id: AtlMemTkAnalysisTask.cxx,v 1.13 2015/08/26 12:40:46 mergelm Exp $
// Copyright: 2013 (C) stamm
//
#ifndef ATL_AtlMemTkAnalysisTask
#include <AtlMemTkAnalysisTask.h>
#endif
#include <TChain.h>
#include <TString.h>
#include <TObjString.h>
#include <TNamed.h>
#include <TSystem.h>
#include <iostream>
#include <fstream>
#include <TDatime.h>

using namespace std;

#ifndef __CINT__
ClassImp(AtlMemTkAnalysisTask);
#endif

//____________________________________________________________________

AtlMemTkAnalysisTask::AtlMemTkAnalysisTask(const char* name, const char* title) :
    AtlTask(name, title) {
    //
    // Default constructor
    //
    
    SetBatchJob(kTRUE);
    SetBatchNodeAutomatic(kFALSE);
    fRootScript = 0;
    fListOfUserEnvs = new TList;

    fNSubjobs = 1;
    fJobid = 1;

    fMemTkRunScript = 0;
    
    // Set environment variables for common libraries
    AddUserEnv("LIBATLASANALYSIS");
    AddUserEnv("LIBMEMTK");
}

//____________________________________________________________________

AtlMemTkAnalysisTask::~AtlMemTkAnalysisTask() {
    //
    // Default destructor
    //
    if ( fRootScript != 0 ) delete fRootScript;
    fListOfUserEnvs->Delete(); delete fListOfUserEnvs;
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::ExecInteractiveJob(Option_t *option) {
    //
    // Interactive execution
    //
    // The interactive execution is done via an intermediate step of
    // an analysis run .C script in the same manner as for batch
    // jobs. Thus only a single way of job execution has to be
    // defined. In particular all selection cuts are set in the same
    // way.
    //
    TString opt = option;
    
    // Start analysis process

    // Check input chain
    // ToDo: CheckInputChain();

    Error("ExecInteractiveJob", "Currently not supported! Abort!\n");
    gSystem->Abort(0);
}    

//____________________________________________________________________

Bool_t AtlMemTkAnalysisTask::ExecBatchJob(Option_t *option) {
    //
    // Batch job submission
    //
    // Returns kFALSE in case of error.
    //
    TString opt = option;
    
    // Create submit scripts
    CreateRunScript();

    // Submit batch job
    return SubmitBatchJob();
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::ExecGridJob(Option_t *option) {
    //
    // Grid job submission
    //
    TString opt = option;

    // Create GRID_HOME directory
    if ( gSystem->Getenv("GRID_HOME") == 0 ) {    
	Error("ExecGridJob", "Env. variable GRID_HOME is not defined, Abort!");
	gSystem->Abort(0);
    }
    if ( gSystem->Exec(Form("mkdir -p %s",
			    gSystem->Getenv("GRID_HOME"))) != 0 ) {
	Error("ExecGridJob", "Failed to create GRID_HOME directory. Abort!");
	gSystem->Abort(0);
    }
    
    // Create submit scripts
    CreateGridRunScript();

    // Submit grid job
    SubmitGridJob();
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::CreateRunScript(Option_t *option) {
    //
    // Create run script for job submission
    //
    ofstream out;
    out.open(fRunScript->Data());
    out << "#!/bin/sh" << endl
	<< "# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" << endl
	<< "# !!! This is an automatically generated file !!!" << endl
	<< "# !!! D O   N O T   E D I T                   !!!" << endl
	<< "# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" << endl
	<< "#" << endl
	<< "# Run script for A++ analysis job submission" << endl
	<< "#" << endl << endl
	<< "# Set environment" << endl;

    // Platform
    if ( gSystem->Getenv("ARCH_TYPE") != 0 ) {
	out << "export ARCH_TYPE=" << gSystem->Getenv("ARCH_TYPE") << endl;
    } else {
	Error("CreateRunScript", "Environment variable ARCH_TYPE not set. Abort!");
	gSystem->Abort(0);
    }

    if (!WriteEnvSetup(out)) {

    // ROOT
    out << "export ROOTSYS=" << gSystem->Getenv("ROOTSYS") << endl
	<< "export PATH=${ROOTSYS}/bin:${PATH}" << endl
	<< "export LD_LIBRARY_PATH=${ROOTSYS}/lib:${LD_LIBRARY_PATH}" << endl;

    // A++ and user-defined libraries
    TIter next_lib(fListOfUserEnvs);
    TNamed *env = 0;
    while ( (env = (TNamed*)next_lib()) ) {
	if ( gSystem->Getenv(env->GetName()) != 0 ) {
	    out << "export " << env->GetName() << "="
		<< gSystem->Getenv(env->GetName()) << endl
		<< "export LD_LIBRARY_PATH=${"
		<< env->GetName() << "}/lib/${ARCH_TYPE}:${LD_LIBRARY_PATH}"
		<< endl;
	} else {
	    Error("CreateRunScript", "Environment variable %s not set. Abort!",
		  env->GetName());
	    gSystem->Abort(0);
	}
    }

    }
    out	<< endl << endl
	<< "# Job execution" << endl;

    out << "JOBHOME=" << fJobHome->Data() << endl
	<< "cd $JOBHOME;" << endl	
	<< "# Remove old Logfiles " << endl
	<< "rm " << fLogFilePath->Data() << endl;

    // Create Root-Script from master file
    if ( !(fInputFiles->GetEntries() == 1) ) {
	Error("CreateRunScript",
	      "Number of inputfiles needs to be exactly one! This is not the case (n=%d)! Abort!",
	      fInputFiles->GetEntries());
	gSystem->Abort(0);
    }

    // use "|" for sed since path name contains "/"
    if ( fMemTkRunScript == 0 ) {
	Error("CreateRunScript", "No MemTk run script defined. Abort!");
	gSystem->Abort();
    }
    out << "cp $LIBMEMTK/test/" << fMemTkRunScript->Data() << " analysis_run.C"
	<< endl
	<< "sed -i 's|NSUBJOBS|"
	<< fNSubjobs
	<< "|g' analysis_run.C"
	<< endl
	<< "sed -i 's|INPUTFILE|"
	<< ((TObjString*) fInputFiles->At(0))->GetString().Data()
	<< "|g' analysis_run.C"
	<< endl
	<< "sed -i 's|OUTPUTFILE|"
	<< fOutputFileName->Data()
	<< "|g' analysis_run.C"
	<< endl
	<< "sed -i 's|SUBJOBID|"
	<< fJobid
	<< "|g' analysis_run.C"
	<< endl;
    
    out << "root -q -l -b analysis_run.C > " << fLogFilePath->Data()
	<< " 2>&1" << endl;
    out.close();
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::CreateGridRunScript() {
    //
    // Create run script for job submission
    //
    ofstream out;
    out.open(fRunScript->Data());
    out << "#!/bin/sh" << endl
	<< "# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" << endl
	<< "# !!! This is an automatically generated file !!!" << endl
	<< "# !!! D O   N O T   E D I T                   !!!" << endl
	<< "# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" << endl
	<< "#" << endl
	<< "# Run script for A++ analysis job submission" << endl
	<< "#" << endl << endl
	<< "# Set environment" << endl;

    // Platform
    out << "export ARCH_TYPE=`uname -i`-`uname -s`-gcc`gcc -dumpversion`-root`root-config --version | sed 's|/|.|'`" << endl
	<< endl;

    // Cuba
    out << "# Cuba config" << endl
	<< "export CUBACORES=0" << endl
	<< "export CUBALIB=$PWD/Cuba-3.0" << endl
	<< "export LD_LIBRARY_PATH=${CUBALIB}:${LD_LIBRARY_PATH}" << endl
	<< endl;

    // LHAPDF
    out << "# LHAPDF config" << endl
	// << "export LHAPDF=$PWD/lhapdf" << endl
	// the line below has to match with the gcc config...
	<< "export LHAPDF=/cvmfs/atlas.cern.ch/repo/sw/software/x86_64-slc6-gcc47-opt/17.8.0/sw/lcg/external/MCGenerators_lcgcmt65/lhapdf/5.8.8/x86_64-slc6-gcc47-opt/" << endl
	<< "export LD_LIBRARY_PATH=${LHAPDF}/lib/:${LD_LIBRARY_PATH}" << endl
	<< "export LHAPATH=$PWD/" << endl
	<< endl;

    // MemTk
    out << "# MemTk" << endl
	<< "export LIBMEMTK=$PWD/libMEM" << endl
	<< "export LD_LIBRARY_PATH=${LIBMEMTK}/lib/${ARCH_TYPE}:${LD_LIBRARY_PATH}" << endl
	<< endl;
    
    out	<< endl << endl
	<< "# Job execution" << endl;

    // Create Root-Script from master file
    if ( !(fInputFiles->GetEntries() == 1) ) {
	Error("CreateRunScript",
	      "Number of inputfiles needs to be exactly one! This is not the case (n=%d)! Abort!",
	      fInputFiles->GetEntries());
	gSystem->Abort(0);
    }

    // Copy config run_master from MemTk tarball (default file) to current directory
    // Modify it as it is done in CreateRunScript()
    if ( fMemTkRunScript == 0 ) {
	Error("CreateRunScript", "No MemTk run script defined. Abort!");
	gSystem->Abort();
    }
    out << "cp $LIBMEMTK/test/" << fMemTkRunScript->Data() << " analysis_run.C"
	<< endl
	<< "sed -i 's|NSUBJOBS|"
	<< fNSubjobs
	<< "|g' analysis_run.C"
	<< endl
	<< "# Move input file"
	<< endl
	<< "mv " << gSystem->BaseName(((TObjString*) fInputFiles->At(0))->GetString().Data()) << " inputfile.root"
	<< endl
	<< "sed -i 's|INPUTFILE|"
	<< "inputfile.root"
	<< "|g' analysis_run.C"
	<< endl
	<< "sed -i 's|OUTPUTFILE|"
	<< gSystem->BaseName(fOutputFileName->Data())
	<< "|g' analysis_run.C"
	<< endl
	<< "sed -i 's|SUBJOBID|"
	<< fJobid
	<< "|g' analysis_run.C"
	<< endl;
    
    out << "root -q -l -b analysis_run.C" << endl;
    // Logfile is always available on grid machines!
    out.close();

    // Copy run script to GRID_HOME (replacing existing one)
    // Keep it for now (might be useful to check quickly offline
    // what has been produced on the grid)
    gSystem->Exec(Form("cp %s %s",
		       fRunScript->Data(),
		       gSystem->Getenv("GRID_HOME")));
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::AddUserEnv(const char* env) {
    //
    // Add user-defined environment variable for shared library search
    // paths
    //
    fListOfUserEnvs->Add(new TNamed(env, ""));
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::Print(Option_t *option) const {
    //
    // Print config
    //
    cout << "=============================" << endl;
    cout << " Task: " << fName << endl;
    cout << " ----------------------------" << endl;
    cout << " JobHome : " << fJobHome->Data() << endl;
    cout << " OutputFileName: " << fOutputFileName->Data() << endl;
    cout << " " << endl;
    cout << " " << endl;
    cout << " " << endl;
    cout << "=============================" << endl;
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::SubmitGridJob() {
    //
    // Submit MemTk grid job
    //

    // ToDo: Specify the location of...
    // ... lhapdf-tarball
    // ... pdf-set
    // ... cuba tarball
    // ... MemTk tarball
    // and copy them if needed to the GRID_HOME directory
    
    // Check if cteq66.LHgrid and lhapdf-tarball is present
    TString gridhome = gSystem->ExpandPathName("$GRID_HOME");
    gSystem->ChangeDirectory(gridhome.Data());

    gSystem->Exec("pwd");

    // TString lhapdf = "lhapdf-5.8.8.tar.gz";
    // TString lhapdf_tmp = lhapdf;
    // if ( gSystem->FindFile(gridhome.Data(), lhapdf_tmp) == 0) {
    // 	Error("SubmitGridJob", "File not found! %s/%s",
    // 	      gridhome.Data(),
    // 	      lhapdf.Data());
    // 	gSystem->Abort(0);
    // }

    TString cteq66 = "cteq66.LHgrid";
    TString cteq66_tmp = cteq66;
    if ( gSystem->FindFile(gridhome.Data(), cteq66_tmp) == 0) {
	Error("SubmitGridJob", "File not found! %s/%s",
	      gridhome.Data(),
	      cteq66.Data());
	gSystem->Abort(0);
    }

    TString cuba30 = "Cuba-3.0.tar.gz";
    TString cuba30_tmp = cuba30;
    if ( gSystem->FindFile(gridhome.Data(), cuba30_tmp) == 0) {
	Error("SubmitGridJob", "File not found! %s/%s",
	      gridhome.Data(),
	      cuba30.Data());
	gSystem->Abort(0);
    }

    TString libMEM = "libMEM.tgz";
    TString libMEM_tmp = libMEM;
    if ( gSystem->FindFile(gridhome.Data(), libMEM_tmp) == 0) {
	Error("SubmitGridJob", "File not found! %s/%s",
	      gridhome.Data(),
	      libMEM.Data());
	gSystem->Abort(0);
    }

    TString jobsub_cmd("prun ");
    jobsub_cmd.Append("--destSE=DESY-ZN_LOCALGROUPDISK ");
    jobsub_cmd.Append("--long ");
    jobsub_cmd.Append("--maxCpuCount=172800 ");
    jobsub_cmd.Append("--bexec='./bexec.sh' ");
    jobsub_cmd.Append(Form("--exec='./%s' ",
			   gSystem->BaseName(fRunScript->Data())));

    // Use filename for dataset since it contains all information
    TString dsname = gSystem->BaseName(((TObjString*) fInputFiles->At(0))->GetString().Data());
    dsname.ReplaceAll(".root", "");
    TString runscript = gSystem->BaseName(fRunScript->Data());
    runscript.ReplaceAll(".run", "");
    dsname.ReplaceAll(Form("_%s", runscript.Data()), "");
    // create unique outDS
    TDatime time;
    jobsub_cmd.Append(Form("--outDS='user.%s.MemTk.%010d.%s' ",
			   fGridUserName->Data(),
			   time.Get(),
			   fGridSuffix->Data()));
    jobsub_cmd.Append(Form("--rootVer=%s ", fGridRootVersion->Data()));
    jobsub_cmd.Append(Form("--cmtConfig=%s ", fGridCmtVersion->Data()));

    // Has to be synchronised with bexec.sh
    // jobsub_cmd.Append(Form("--extFile=%s,%s,%s,%s ",
    // 			   lhapdf.Data(),
    jobsub_cmd.Append(Form("--extFile=%s,%s,%s ",
			   cteq66.Data(),
			   gSystem->BaseName(
			       ((TObjString*) fInputFiles->At(0))->GetString().Data()),
			   libMEM.Data()));
    jobsub_cmd.Append(Form("--outputs=%s ",
			   gSystem->BaseName(fOutputFileName->Data())));
    
    // Add input file to tarball
    TString fullfilename = ((TObjString*) fInputFiles->At(0))->GetString();
    // TString addfile = Form("tar -f %s -C %s -r %s",
    // 			   tarball.Data(),
    // 			   gSystem->DirName(fullfilename.Data()),
    // 			   gSystem->BaseName(fullfilename.Data()));
    TString addfile = Form("cp %s %s",
     			   fullfilename.Data(),
			   gSystem->Getenv("GRID_HOME"));
    Info("SubmitGridJob", "Add file to grid_home:\n%s", addfile.Data());
    gSystem->Exec(addfile.Data());
    // Submit prun cmd
    Info("SubmitGridJob", "Execute %s", jobsub_cmd.Data());
    gSystem->Exec(jobsub_cmd.Data());

    // Remove input file from tarball
    // TString removefile = Form("tar -f %s --delete %s",
    // 			      tarball.Data(),
    // 			      gSystem->BaseName(fullfilename.Data()));
    TString removefile = Form("rm %s/%s",
			      gSystem->Getenv("GRID_HOME"),
			      gSystem->BaseName(fullfilename.Data()));
    Info("SubmitGridJob", "Remove file from grid_home:\n%s", removefile.Data());
    gSystem->Exec(removefile.Data());
    // Remove runscript from GRID_HOME
    TString removeRunscript = Form("rm %s/%s",
				   gSystem->Getenv("GRID_HOME"),
				   gSystem->BaseName(fRunScript->Data()));
    Info("SubmitGridJob","%s", removeRunscript.Data());
    gSystem->Exec(removeRunscript.Data());


    // Create Logfile
    CreateLogFilePath();
    // Remove old logfile if it exist
    gSystem->Exec(Form("rm -f %s", fLogFilePath->Data()));

    // Add the string that is used in AtlTopLevelAnalysis::GetTaskStatus
    // to check if the job was successful

    // Get some information on the day, month, year.
    TString weekday = "";
    switch ( time.GetDayOfWeek() ) {
	case 1: weekday.Append("Mon."); break;
	case 2: weekday.Append("Tue."); break;
	case 3: weekday.Append("Wed."); break;
	case 4: weekday.Append("Thu."); break;
	case 5: weekday.Append("Fri."); break;
	case 6: weekday.Append("Sat."); break;
	case 7: weekday.Append("Sun."); break;
	default:
	    Error("SubmitGridJob", "Unkown day of week. Abort!");
	    gSystem->Abort();
    }	    
    
    
    ofstream out;
    out.open(fLogFilePath->Data());
    out << "Job Summary" << endl
	<< "Grid job was submitted on "
	<< weekday.Data() << ", "
	<< time.GetDay() << "."
	<< time.GetMonth() << "."
	<< time.GetYear() << " at "
	<< time.GetHour() << ":"
	<< time.GetMinute() << "."
	<< time.GetSecond() << endl;
    out.close();
}

//____________________________________________________________________

void AtlMemTkAnalysisTask::SetMemTkRunScript(const char* script) {
    //
    //
    //
    if ( fMemTkRunScript != 0 ) delete fMemTkRunScript;
    fMemTkRunScript = new TString(script);
}