/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 * @(#)root/roofitcore:$Id$
 * Authors:                                                                  *
 *   WV, Wouter Verkerke, UC Santa Barbara, verkerke@slac.stanford.edu       *
 *   DK, David Kirkby,    UC Irvine,         dkirkby@uci.edu                 *
 *                                                                           *
 * Copyright (c) 2000-2005, Regents of the University of California          *
 *                          and Stanford University. All rights reserved.    *
 *                                                                           *
 * Redistribution and use in source and binary forms,                        *
 * with or without modification, are permitted according to the terms        *
 * listed in LICENSE (http://roofit.sourceforge.net/license.txt)             *
 *****************************************************************************/

//////////////////////////////////////////////////////////////////////////////
//
// BEGIN_HTML
// RooStudyManager is a utility class to manage studies that consist of
// repeated applications of generate-and-fit operations on a workspace
//
// END_HTML
//



#include "RooFit.h"
#include "Riostream.h"

#include "RooStudyManager.h"
#include "RooWorkspace.h"
#include "RooAbsStudy.h"
#include "RooDataSet.h"
#include "RooMsgService.h"
#include "RooStudyPackage.h"
#include "TTree.h"
#include "TFile.h"
#include "TRegexp.h"
#include "TKey.h"
#include <string>
#include "TROOT.h"
#include "TSystem.h"

using namespace std ;

ClassImp(RooStudyManager)
  ;



//_____________________________________________________________________________
RooStudyManager::RooStudyManager(RooWorkspace& w)
{  
  _pkg = new RooStudyPackage(w) ;
}



//_____________________________________________________________________________
RooStudyManager::RooStudyManager(RooWorkspace& w, RooAbsStudy& study)
{  
  _pkg = new RooStudyPackage(w) ;
  _pkg->addStudy(study) ;
}


//_____________________________________________________________________________
RooStudyManager::RooStudyManager(const char* studyPackFileName) 
{
  string pwd = gDirectory->GetName() ;
  TFile *f = new TFile(studyPackFileName) ;
  _pkg = dynamic_cast<RooStudyPackage*>(f->Get("studypack")) ;
  gDirectory->cd(Form("%s:",pwd.c_str())) ;
}



//_____________________________________________________________________________
void RooStudyManager::addStudy(RooAbsStudy& study) 
{
  _pkg->addStudy(study) ;
}




//_____________________________________________________________________________
void RooStudyManager::run(Int_t nExperiments) 
{
  _pkg->driver(nExperiments) ;
}



//_____________________________________________________________________________
void RooStudyManager::runProof(Int_t nExperiments, const char* proofHost, Bool_t showGui) 
{
  // Open PROOF-Lite session
  coutP(Generation) << "RooStudyManager::runProof(" << GetName() << ") opening PROOF session" << endl ;
  void* p = (void*) gROOT->ProcessLineFast(Form("TProof::Open(\"%s\")",proofHost)) ;

  // Check that PROOF initialization actually succeeeded
  if (p==0) {
    coutE(Generation) << "RooStudyManager::runProof(" << GetName() << ") ERROR initializing proof, aborting" << endl ;
    return ;
  }

  // Suppress GUI if so requested
  if (!showGui) {
    gROOT->ProcessLineFast(Form("((TProof*)0x%lx)->SetProgressDialog(0) ;",(ULong_t)p)) ;
  }

  // Propagate workspace to proof nodes
  coutP(Generation) << "RooStudyManager::runProof(" << GetName() << ") sending work package to PROOF servers" << endl ;
  gROOT->ProcessLineFast(Form("((TProof*)0x%lx)->AddInput((TObject*)0x%lx) ;",(ULong_t)p,(ULong_t)_pkg) ) ;

  // Run selector in parallel
  coutP(Generation) << "RooStudyManager::runProof(" << GetName() << ") starting PROOF processing of " << nExperiments << " experiments" << endl ;
			 
  gROOT->ProcessLineFast(Form("((TProof*)0x%lx)->Process(\"RooProofDriverSelector\",%d) ;",(ULong_t)p,nExperiments)) ;

  // Aggregate results data
  coutP(Generation) << "RooStudyManager::runProof(" << GetName() << ") aggregating results data" << endl ;
  TList* olist = (TList*) gROOT->ProcessLineFast(Form("((TProof*)0x%lx)->GetOutputList()",(ULong_t)p)) ;
  aggregateData(olist) ;

  // cleaning up                                                                                                                                           
  coutP(Generation) << "RooStudyManager::runProof(" << GetName() << ") cleaning up input list" << endl ;                                                   
  gROOT->ProcessLineFast(Form("((TProof*)0x%lx)->GetInputList()->Remove((TObject*)0x%lx) ;",(ULong_t)p,(ULong_t)_pkg) ) ;                                                   
  
}


//_____________________________________________________________________________
void RooStudyManager::closeProof(Option_t *option)
{
  // "Option_t *option" takes the parameters forwarded to gProof->Close(option).
  //
  // This function is intended for scripts that run in loops
  // where it is essential to properly close all connections and delete
  // the TProof instance (frees ports).

  if (gROOT->GetListOfProofs()->LastIndex() != -1  &&  gROOT->ProcessLineFast("gProof;"))
  {
    gROOT->ProcessLineFast(Form("gProof->Close(\"%s\") ;",option)) ;
    gROOT->ProcessLineFast("gProof->CloseProgressDialog() ;") ;

    // CloseProgressDialog does not do anything when run without GUI. This detects
    // whether the proof instance is still there and deletes it if that is the case.
    if (gROOT->GetListOfProofs()->LastIndex() != -1  &&  gROOT->ProcessLineFast("gProof;")) {
      gROOT->ProcessLineFast("delete gProof ;") ;
    }
  } else {
    ooccoutI((TObject*)NULL,Generation) << "RooStudyManager: No global Proof objects. No connections closed." << endl ;
  }
}



//_____________________________________________________________________________
void RooStudyManager::prepareBatchInput(const char* studyName, Int_t nExpPerJob, Bool_t unifiedInput=kFALSE) 
{
  TFile f(Form("study_data_%s.root",studyName),"RECREATE") ;
  _pkg->Write("studypack") ;
  f.Close() ;

  if (unifiedInput) {

    // Write header of driver script
    ofstream bdr(Form("study_driver_%s.sh",studyName)) ;
    bdr << "#!/bin/sh" << endl 
        << Form("if [ ! -f study_data_%s.root ] ; then",studyName) << endl 
        << "uudecode <<EOR" << endl ;
    bdr.close() ;

    // Write uuencoded ROOT file (base64) in driver script
    gSystem->Exec(Form("cat study_data_%s.root | uuencode -m study_data_%s.root >> study_driver_%s.sh",studyName,studyName,studyName)) ;

    // Write remainder of deriver script
    ofstream bdr2 (Form("study_driver_%s.sh",studyName),ios::app) ;
    bdr2 << "EOR" << endl 
	 << "fi" << endl 
	 << "root -l -b <<EOR" << endl 
	 << Form("RooStudyPackage::processFile(\"%s\",%d) ;",studyName,nExpPerJob) << endl 
	 << ".q" << endl 
	 << "EOR" << endl ;  
    // Remove binary input file
    gSystem->Unlink(Form("study_data_%s.root",studyName)) ;    

    coutI(DataHandling) << "RooStudyManager::prepareBatchInput batch driver file is '" << Form("study_driver_%s.sh",studyName) << "," << endl
			<< "     input data files is embedded in driver script" << endl ;

  } else {

    ofstream bdr(Form("study_driver_%s.sh",studyName)) ;
    bdr << "#!/bin/sh" << endl 
	<< "root -l -b <<EOR" << endl 
	<< Form("RooStudyPackage::processFile(\"%s\",%d) ;",studyName,nExpPerJob) << endl 
	<< ".q" << endl 
	<< "EOR" << endl ;  

    coutI(DataHandling) << "RooStudyManager::prepareBatchInput batch driver file is '" << Form("study_driver_%s.sh",studyName) << "," << endl
			<< "     input data file is " << Form("study_data_%s.root",studyName) << endl ;

  }
}




//_____________________________________________________________________________
void RooStudyManager::processBatchOutput(const char* filePat) 
{
  list<string> flist ;
  expandWildCardSpec(filePat,flist) ;

  TList olist ;

  for (list<string>::iterator iter = flist.begin() ; iter!=flist.end() ; ++iter) {
    coutP(DataHandling) << "RooStudyManager::processBatchOutput() now reading file " << *iter << endl ;
    TFile f(iter->c_str()) ;

    TList* list = f.GetListOfKeys() ;
    TIterator* kiter = list->MakeIterator();
    
    TObject* obj ;
    TKey* key ;
    while((key=(TKey*)kiter->Next())) {      
      obj = f.Get(key->GetName()) ;
      TObject* clone = obj->Clone(obj->GetName()) ;
      olist.Add(clone) ;
    }
    delete kiter ;
  }
  aggregateData(&olist) ;
  olist.Delete() ;
}


//_____________________________________________________________________________
void RooStudyManager::aggregateData(TList* olist) 
{
  for (list<RooAbsStudy*>::iterator iter=_pkg->studies().begin() ; iter!=_pkg->studies().end() ; iter++) {
    (*iter)->aggregateSummaryOutput(olist) ;
  }  
}




//_____________________________________________________________________________
void RooStudyManager::expandWildCardSpec(const char* name, list<string>& result)
{
  // case with one single file
  if (!TString(name).MaybeWildcard()) {
    result.push_back(name) ;
    return ;
  }

   // wildcarding used in name
   TString basename(name);

   Int_t dotslashpos = -1;
   {
      Int_t next_dot = basename.Index(".root");
      while(next_dot>=0) {
         dotslashpos = next_dot;
         next_dot = basename.Index(".root",dotslashpos+1);
      }
      if (basename[dotslashpos+5]!='/') {
         // We found the 'last' .root in the name and it is not followed by
         // a '/', so the tree name is _not_ specified in the name.
         dotslashpos = -1;
      }
   }
   //Int_t dotslashpos = basename.Index(".root/");
   TString behind_dot_root;
   if (dotslashpos>=0) {
      // Copy the tree name specification
      behind_dot_root = basename(dotslashpos+6,basename.Length()-dotslashpos+6);
      // and remove it from basename
      basename.Remove(dotslashpos+5);
   }

   Int_t slashpos = basename.Last('/');
   TString directory;
   if (slashpos>=0) {
      directory = basename(0,slashpos); // Copy the directory name
      basename.Remove(0,slashpos+1);      // and remove it from basename
   } else {
      directory = gSystem->UnixPathName(gSystem->WorkingDirectory());
   }

   const char *file;
   void *dir = gSystem->OpenDirectory(gSystem->ExpandPathName(directory.Data()));

   if (dir) {
      //create a TList to store the file names (not yet sorted)
      TList l;
      TRegexp re(basename,kTRUE);
      while ((file = gSystem->GetDirEntry(dir))) {
         if (!strcmp(file,".") || !strcmp(file,"..")) continue;
         TString s = file;
         if ( (basename!=file) && s.Index(re) == kNPOS) continue;
         l.Add(new TObjString(file));
      }
      gSystem->FreeDirectory(dir);
      //sort the files in alphanumeric order
      l.Sort();
      TIter next(&l);
      TObjString *obj;
      while ((obj = (TObjString*)next())) {
         file = obj->GetName();
         if (behind_dot_root.Length() != 0)
	   result.push_back(Form("%s/%s/%s",directory.Data(),file,behind_dot_root.Data())) ;
         else
	   result.push_back(Form("%s/%s",directory.Data(),file)) ;
      }
      l.Delete();
   }
}
 RooStudyManager.cxx:1
 RooStudyManager.cxx:2
 RooStudyManager.cxx:3
 RooStudyManager.cxx:4
 RooStudyManager.cxx:5
 RooStudyManager.cxx:6
 RooStudyManager.cxx:7
 RooStudyManager.cxx:8
 RooStudyManager.cxx:9
 RooStudyManager.cxx:10
 RooStudyManager.cxx:11
 RooStudyManager.cxx:12
 RooStudyManager.cxx:13
 RooStudyManager.cxx:14
 RooStudyManager.cxx:15
 RooStudyManager.cxx:16
 RooStudyManager.cxx:17
 RooStudyManager.cxx:18
 RooStudyManager.cxx:19
 RooStudyManager.cxx:20
 RooStudyManager.cxx:21
 RooStudyManager.cxx:22
 RooStudyManager.cxx:23
 RooStudyManager.cxx:24
 RooStudyManager.cxx:25
 RooStudyManager.cxx:26
 RooStudyManager.cxx:27
 RooStudyManager.cxx:28
 RooStudyManager.cxx:29
 RooStudyManager.cxx:30
 RooStudyManager.cxx:31
 RooStudyManager.cxx:32
 RooStudyManager.cxx:33
 RooStudyManager.cxx:34
 RooStudyManager.cxx:35
 RooStudyManager.cxx:36
 RooStudyManager.cxx:37
 RooStudyManager.cxx:38
 RooStudyManager.cxx:39
 RooStudyManager.cxx:40
 RooStudyManager.cxx:41
 RooStudyManager.cxx:42
 RooStudyManager.cxx:43
 RooStudyManager.cxx:44
 RooStudyManager.cxx:45
 RooStudyManager.cxx:46
 RooStudyManager.cxx:47
 RooStudyManager.cxx:48
 RooStudyManager.cxx:49
 RooStudyManager.cxx:50
 RooStudyManager.cxx:51
 RooStudyManager.cxx:52
 RooStudyManager.cxx:53
 RooStudyManager.cxx:54
 RooStudyManager.cxx:55
 RooStudyManager.cxx:56
 RooStudyManager.cxx:57
 RooStudyManager.cxx:58
 RooStudyManager.cxx:59
 RooStudyManager.cxx:60
 RooStudyManager.cxx:61
 RooStudyManager.cxx:62
 RooStudyManager.cxx:63
 RooStudyManager.cxx:64
 RooStudyManager.cxx:65
 RooStudyManager.cxx:66
 RooStudyManager.cxx:67
 RooStudyManager.cxx:68
 RooStudyManager.cxx:69
 RooStudyManager.cxx:70
 RooStudyManager.cxx:71
 RooStudyManager.cxx:72
 RooStudyManager.cxx:73
 RooStudyManager.cxx:74
 RooStudyManager.cxx:75
 RooStudyManager.cxx:76
 RooStudyManager.cxx:77
 RooStudyManager.cxx:78
 RooStudyManager.cxx:79
 RooStudyManager.cxx:80
 RooStudyManager.cxx:81
 RooStudyManager.cxx:82
 RooStudyManager.cxx:83
 RooStudyManager.cxx:84
 RooStudyManager.cxx:85
 RooStudyManager.cxx:86
 RooStudyManager.cxx:87
 RooStudyManager.cxx:88
 RooStudyManager.cxx:89
 RooStudyManager.cxx:90
 RooStudyManager.cxx:91
 RooStudyManager.cxx:92
 RooStudyManager.cxx:93
 RooStudyManager.cxx:94
 RooStudyManager.cxx:95
 RooStudyManager.cxx:96
 RooStudyManager.cxx:97
 RooStudyManager.cxx:98
 RooStudyManager.cxx:99
 RooStudyManager.cxx:100
 RooStudyManager.cxx:101
 RooStudyManager.cxx:102
 RooStudyManager.cxx:103
 RooStudyManager.cxx:104
 RooStudyManager.cxx:105
 RooStudyManager.cxx:106
 RooStudyManager.cxx:107
 RooStudyManager.cxx:108
 RooStudyManager.cxx:109
 RooStudyManager.cxx:110
 RooStudyManager.cxx:111
 RooStudyManager.cxx:112
 RooStudyManager.cxx:113
 RooStudyManager.cxx:114
 RooStudyManager.cxx:115
 RooStudyManager.cxx:116
 RooStudyManager.cxx:117
 RooStudyManager.cxx:118
 RooStudyManager.cxx:119
 RooStudyManager.cxx:120
 RooStudyManager.cxx:121
 RooStudyManager.cxx:122
 RooStudyManager.cxx:123
 RooStudyManager.cxx:124
 RooStudyManager.cxx:125
 RooStudyManager.cxx:126
 RooStudyManager.cxx:127
 RooStudyManager.cxx:128
 RooStudyManager.cxx:129
 RooStudyManager.cxx:130
 RooStudyManager.cxx:131
 RooStudyManager.cxx:132
 RooStudyManager.cxx:133
 RooStudyManager.cxx:134
 RooStudyManager.cxx:135
 RooStudyManager.cxx:136
 RooStudyManager.cxx:137
 RooStudyManager.cxx:138
 RooStudyManager.cxx:139
 RooStudyManager.cxx:140
 RooStudyManager.cxx:141
 RooStudyManager.cxx:142
 RooStudyManager.cxx:143
 RooStudyManager.cxx:144
 RooStudyManager.cxx:145
 RooStudyManager.cxx:146
 RooStudyManager.cxx:147
 RooStudyManager.cxx:148
 RooStudyManager.cxx:149
 RooStudyManager.cxx:150
 RooStudyManager.cxx:151
 RooStudyManager.cxx:152
 RooStudyManager.cxx:153
 RooStudyManager.cxx:154
 RooStudyManager.cxx:155
 RooStudyManager.cxx:156
 RooStudyManager.cxx:157
 RooStudyManager.cxx:158
 RooStudyManager.cxx:159
 RooStudyManager.cxx:160
 RooStudyManager.cxx:161
 RooStudyManager.cxx:162
 RooStudyManager.cxx:163
 RooStudyManager.cxx:164
 RooStudyManager.cxx:165
 RooStudyManager.cxx:166
 RooStudyManager.cxx:167
 RooStudyManager.cxx:168
 RooStudyManager.cxx:169
 RooStudyManager.cxx:170
 RooStudyManager.cxx:171
 RooStudyManager.cxx:172
 RooStudyManager.cxx:173
 RooStudyManager.cxx:174
 RooStudyManager.cxx:175
 RooStudyManager.cxx:176
 RooStudyManager.cxx:177
 RooStudyManager.cxx:178
 RooStudyManager.cxx:179
 RooStudyManager.cxx:180
 RooStudyManager.cxx:181
 RooStudyManager.cxx:182
 RooStudyManager.cxx:183
 RooStudyManager.cxx:184
 RooStudyManager.cxx:185
 RooStudyManager.cxx:186
 RooStudyManager.cxx:187
 RooStudyManager.cxx:188
 RooStudyManager.cxx:189
 RooStudyManager.cxx:190
 RooStudyManager.cxx:191
 RooStudyManager.cxx:192
 RooStudyManager.cxx:193
 RooStudyManager.cxx:194
 RooStudyManager.cxx:195
 RooStudyManager.cxx:196
 RooStudyManager.cxx:197
 RooStudyManager.cxx:198
 RooStudyManager.cxx:199
 RooStudyManager.cxx:200
 RooStudyManager.cxx:201
 RooStudyManager.cxx:202
 RooStudyManager.cxx:203
 RooStudyManager.cxx:204
 RooStudyManager.cxx:205
 RooStudyManager.cxx:206
 RooStudyManager.cxx:207
 RooStudyManager.cxx:208
 RooStudyManager.cxx:209
 RooStudyManager.cxx:210
 RooStudyManager.cxx:211
 RooStudyManager.cxx:212
 RooStudyManager.cxx:213
 RooStudyManager.cxx:214
 RooStudyManager.cxx:215
 RooStudyManager.cxx:216
 RooStudyManager.cxx:217
 RooStudyManager.cxx:218
 RooStudyManager.cxx:219
 RooStudyManager.cxx:220
 RooStudyManager.cxx:221
 RooStudyManager.cxx:222
 RooStudyManager.cxx:223
 RooStudyManager.cxx:224
 RooStudyManager.cxx:225
 RooStudyManager.cxx:226
 RooStudyManager.cxx:227
 RooStudyManager.cxx:228
 RooStudyManager.cxx:229
 RooStudyManager.cxx:230
 RooStudyManager.cxx:231
 RooStudyManager.cxx:232
 RooStudyManager.cxx:233
 RooStudyManager.cxx:234
 RooStudyManager.cxx:235
 RooStudyManager.cxx:236
 RooStudyManager.cxx:237
 RooStudyManager.cxx:238
 RooStudyManager.cxx:239
 RooStudyManager.cxx:240
 RooStudyManager.cxx:241
 RooStudyManager.cxx:242
 RooStudyManager.cxx:243
 RooStudyManager.cxx:244
 RooStudyManager.cxx:245
 RooStudyManager.cxx:246
 RooStudyManager.cxx:247
 RooStudyManager.cxx:248
 RooStudyManager.cxx:249
 RooStudyManager.cxx:250
 RooStudyManager.cxx:251
 RooStudyManager.cxx:252
 RooStudyManager.cxx:253
 RooStudyManager.cxx:254
 RooStudyManager.cxx:255
 RooStudyManager.cxx:256
 RooStudyManager.cxx:257
 RooStudyManager.cxx:258
 RooStudyManager.cxx:259
 RooStudyManager.cxx:260
 RooStudyManager.cxx:261
 RooStudyManager.cxx:262
 RooStudyManager.cxx:263
 RooStudyManager.cxx:264
 RooStudyManager.cxx:265
 RooStudyManager.cxx:266
 RooStudyManager.cxx:267
 RooStudyManager.cxx:268
 RooStudyManager.cxx:269
 RooStudyManager.cxx:270
 RooStudyManager.cxx:271
 RooStudyManager.cxx:272
 RooStudyManager.cxx:273
 RooStudyManager.cxx:274
 RooStudyManager.cxx:275
 RooStudyManager.cxx:276
 RooStudyManager.cxx:277
 RooStudyManager.cxx:278
 RooStudyManager.cxx:279
 RooStudyManager.cxx:280
 RooStudyManager.cxx:281
 RooStudyManager.cxx:282
 RooStudyManager.cxx:283
 RooStudyManager.cxx:284
 RooStudyManager.cxx:285
 RooStudyManager.cxx:286
 RooStudyManager.cxx:287
 RooStudyManager.cxx:288
 RooStudyManager.cxx:289
 RooStudyManager.cxx:290
 RooStudyManager.cxx:291
 RooStudyManager.cxx:292
 RooStudyManager.cxx:293
 RooStudyManager.cxx:294
 RooStudyManager.cxx:295
 RooStudyManager.cxx:296
 RooStudyManager.cxx:297
 RooStudyManager.cxx:298
 RooStudyManager.cxx:299
 RooStudyManager.cxx:300
 RooStudyManager.cxx:301
 RooStudyManager.cxx:302
 RooStudyManager.cxx:303
 RooStudyManager.cxx:304
 RooStudyManager.cxx:305
 RooStudyManager.cxx:306
 RooStudyManager.cxx:307
 RooStudyManager.cxx:308
 RooStudyManager.cxx:309
 RooStudyManager.cxx:310
 RooStudyManager.cxx:311
 RooStudyManager.cxx:312
 RooStudyManager.cxx:313
 RooStudyManager.cxx:314
 RooStudyManager.cxx:315
 RooStudyManager.cxx:316
 RooStudyManager.cxx:317
 RooStudyManager.cxx:318
 RooStudyManager.cxx:319
 RooStudyManager.cxx:320
 RooStudyManager.cxx:321
 RooStudyManager.cxx:322