#! /usr/bin/env python

"""\
%(prog)s: make templates of analysis source files for Rivet

Usage: %(prog)s [--help|-h] [--srcroot=<srcrootdir>] <analysisname>

Without the --srcroot flag, the analysis files will be created in the current
directory.
"""

import rivet, sys, os
rivet.util.check_python_version()
rivet.util.set_process_name(os.path.basename(__file__))
import logging


## Handle command line
import argparse
parser = argparse.ArgumentParser(usage=__doc__)
parser.add_argument("ANANAMES", nargs="+", help="names of analyses to make")
parser.add_argument("--srcroot", metavar="DIR", dest="SRCROOT", default=None,
                    help="install the templates into the Rivet source tree (rooted " +
                    "at directory DIR) rather than just creating all in the current dir")
parser.add_argument("-i", "--inline-info", dest="INLINE", action="store_true",
                    default=False, help="put analysis info into the source file instead of a separate data-file")
parser.add_argument("--no-downloads", dest="DOWNLOAD", action="store_false",
                    default=True, help="don't try to get metadata or reference data from online sources")
parser.add_argument("-m","--minimal", dest="MINIMAL", action="store_true",
                    default=False, help="fill implementation with a minimal template (empty methods)")
parser.add_argument("-q", "--quiet", dest="LOGLEVEL", default=logging.INFO,
                    action="store_const", const=logging.WARNING, help="only write out warning and error messages")
parser.add_argument("-v", "--verbose", dest="LOGLEVEL", default=logging.INFO,
                    action="store_const", const=logging.DEBUG, help="provide extra debugging messages")
args = parser.parse_args()
logging.basicConfig(format="%(msg)s", level=args.LOGLEVEL)
ANANAMES = args.ANANAMES


## Work out installation paths
ANAROOT = os.path.abspath(args.SRCROOT or os.getcwd())
if not os.access(ANAROOT, os.W_OK):
    logging.error("Can't write to source root directory %s" % ANAROOT)
    sys.exit(1)
ANASRCDIR = os.getcwd()
ANAINFODIR = os.getcwd()
ANAPLOTDIR = os.getcwd()
if args.SRCROOT:
    ANASRCDIR = os.path.join(ANAROOT, "src/Analyses")
    ANAINFODIR = os.path.join(ANAROOT, "data/anainfo")
    ANAPLOTDIR = os.path.join(ANAROOT, "data/plotinfo")
    if not (os.path.exists(ANASRCDIR) and os.path.exists(ANAINFODIR) and os.path.exists(ANAPLOTDIR)):
        logging.error("Rivet analysis dirs do not exist under %s" % ANAROOT)
        sys.exit(1)
if not (os.access(ANASRCDIR, os.W_OK) and os.access(ANAINFODIR, os.W_OK) and os.access(ANAPLOTDIR, os.W_OK)):
    logging.error("Can't write to Rivet analysis dirs under %s" % ANAROOT)
    sys.exit(1)


## Check for disallowed characters in analysis names
import string
allowedchars = string.ascii_letters + string.digits + "_"
all_ok = True
for ananame in ANANAMES:
    for c in ananame:
        if c not in allowedchars:
            logging.error("Analysis name '%s' contains disallowed character '%s'!" % (ananame, c))
            all_ok = False
            break
if not all_ok:
    logging.error("Exiting... please ensure that all analysis names are valid")
    sys.exit(1)


## Now make each analysis
for ANANAME in ANANAMES:
    logging.info("Writing templates for %s to %s" % (ANANAME, ANAROOT))

    ## Extract some metadata from the name if it matches the standard pattern
    import re
    re_stdana = re.compile(r"^(\w+)_(\d{4})_(I|S)(\d+)$")
    match = re_stdana.match(ANANAME)
    STDANA = False
    ANAEXPT = "<Insert the experiment name>"
    ANACOLLIDER = "<Insert the collider name>"
    ANAYEAR = "<Insert year of publication>"
    INSPIRE_SPIRES = 'I'
    ANAINSPIREID = "<Insert the Inspire ID>"
    if match:
        STDANA = True
        ANAEXPT = match.group(1)
        if ANAEXPT.upper() in ("ALICE", "ATLAS", "CMS", "LHCB"):
            ANACOLLIDER = "LHC"
        elif ANAEXPT.upper() in ("CDF", "D0"):
            ANACOLLIDER = "Tevatron"
        elif ANAEXPT.upper() == "BABAR":
            ANACOLLIDER = "PEP-II"
        elif ANAEXPT.upper() == "BELLE":
            ANACOLLIDER = "KEKB"
        ANAYEAR = match.group(2)
        INSPIRE_SPIRES = match.group(3)
        ANAINSPIREID = match.group(4)
    if INSPIRE_SPIRES == "I":
        ANAREFREPO = "Inspire"
    else:
        ANAREFREPO = "Spires"
    KEYWORDS = {
        "ANANAME" : ANANAME,
        "ANAEXPT" : ANAEXPT,
        "ANACOLLIDER" : ANACOLLIDER,
        "ANAYEAR" : ANAYEAR,
        "ANAREFREPO" : ANAREFREPO,
        "ANAINSPIREID" : ANAINSPIREID
        }

    ## Try to get bib info from SPIRES
    ANABIBKEY = ""
    ANABIBTEX = ""
    bibkey, bibtex = None, None
    if STDANA:
        try:
            logging.info("Getting Inspire/SPIRES biblio data for '%s'" % ANANAME)
            bibkey, bibtex = rivet.spiresbib.get_bibtex_from_repo(INSPIRE_SPIRES, ANAINSPIREID)
        except Exception as e:
            logging.error("Inspire/SPIRES oops: %s" % e)
        if bibkey and bibtex:
            ANABIBKEY = bibkey
            ANABIBTEX = bibtex
    KEYWORDS["ANABIBKEY"] = ANABIBKEY
    KEYWORDS["ANABIBTEX"] = ANABIBTEX


    ## Try to download YODA data file from HepData
    if STDANA and args.DOWNLOAD:
        try:
            try:
                import requests, gzip
            except:
                msg = "Automatic reference-data retrieval from HepData requires "
                msg += "'requests' module to be installed (e.g. using pip)."
                logging.warning(msg)
            #
            hdurl = None
            if INSPIRE_SPIRES == "I":
                hdurl = "http://www.hepdata.net/record/ins%s?format=yoda&rivet=%s" % (ANAINSPIREID, ANANAME)
            if not hdurl:
                raise Exception("Couldn't identify a URL for getting reference data from HepData")
            logging.info("Getting data file from HepData at %s" % hdurl)
            tmpfile = requests.get(hdurl)
            #
            import tarfile, io, os, shutil
            with tarfile.open(mode='r:gz', fileobj=io.BytesIO(tmpfile.content)) as tar:
                with gzip.open('%s.yoda.gz' % ANANAME, 'wb') as fout:
                    fout.write(tar.extractfile(tar.members[0].name).read())
            os.chmod('%s.yoda.gz' % ANANAME, 0o644)
        except Exception as e:
            logging.warning("Problem encountered retrieving from HepData: %s" % hdurl)
            logging.warning("No reference data file written")
            logging.debug("HepData oops: %s: %s" % (str(type(e)), e))


    INLINEMETHODS = ""
    if args.INLINE:
        KEYWORDS["ANAREFREPO_LOWER"] = KEYWORDS["ANAREFREPO"].lower()
        INLINEMETHODS = """
  public:
    string experiment()         const { return "%(ANAEXPT)s"; }
    string year()               const { return "%(ANAYEAR)s"; }
    string %(ANAREFREPO_LOWER)sId()          const { return "%(ANAINSPIREID)s"; }
    string collider()           const { return ""; }
    string summary()            const { return ""; }
    string description()        const { return ""; }
    string runInfo()            const { return ""; }
    string bibKey()             const { return "%(ANABIBKEY)s"; }
    string bibTeX()             const { return "%(ANABIBTEX)s"; }
    string status()             const { return "UNVALIDATED"; }
    vector<string> authors()    const { return vector<string>(); }
    vector<string> references() const { return vector<string>(); }
    vector<std::string> todos() const { return vector<string>(); }
    """ % KEYWORDS
        del KEYWORDS["ANAREFREPO_LOWER"]
    KEYWORDS["INLINEMETHODS"] = INLINEMETHODS


    if ANANAME.startswith("MC_"):
        HISTOBOOKING = """\
      book(_h["XXXX"], "myh1", 20, 0.0, 100.0);
      book(_h["YYYY"], "myh2", logspace(20, 1e-2, 1e3));
      book(_h["ZZZZ"], "myh3", {0.0, 1.0, 2.0, 4.0, 8.0, 16.0});
      book(_p["AAAA"], "myp", 20, 0.0, 100.0);
      book(_c["BBBB"], "myc");""" % KEYWORDS
    else:
        HISTOBOOKING = """\
      // specify custom binning
      book(_h["XXXX"], "myh1", 20, 0.0, 100.0);
      book(_h["YYYY"], "myh2", logspace(20, 1e-2, 1e3));
      book(_h["ZZZZ"], "myh3", {0.0, 1.0, 2.0, 4.0, 8.0, 16.0});
      // take binning from reference data using HEPData ID (digits in "d01-x01-y01" etc.)
      book(_h["AAAA"], 1, 1, 1);
      book(_p["BBBB"], 2, 1, 1);
      book(_c["CCCC"], 3, 1, 1);""" % KEYWORDS
    KEYWORDS["HISTOBOOKING"] = HISTOBOOKING


    ANASRCFILE = os.path.join(ANASRCDIR, ANANAME+".cc")
    logging.debug("Writing implementation template to %s" % ANASRCFILE)
    f = open(ANASRCFILE, "w")
    if args.MINIMAL:
        src = '''\
// -*- C++ -*-
#include "Rivet/Analysis.hh"

namespace Rivet {


  /// @brief Add a short analysis description here
  class %(ANANAME)s : public Analysis {
  public:

    /// Constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(%(ANANAME)s);

    /// @name Analysis methods
    /// @{

    /// Book histograms and initialise projections before the run
    void init() {

    }


    /// Perform the per-event analysis
    void analyze(const Event& event) {

    }


    /// Normalise histograms etc., after the run
    void finalize() {

    }

    /// @}

    /// @name Histograms
    /// @{
    /// @}

  };


  RIVET_DECLARE_PLUGIN(%(ANANAME)s);

}
''' % KEYWORDS
    else:
        src = '''\
// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Projections/DressedLeptons.hh"
#include "Rivet/Projections/MissingMomentum.hh"
#include "Rivet/Projections/DirectFinalState.hh"

namespace Rivet {


  /// @brief Add a short analysis description here
  class %(ANANAME)s : public Analysis {
  public:

    /// Constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(%(ANANAME)s);


    /// @name Analysis methods
    /// @{

    /// Book histograms and initialise projections before the run
    void init() {

      // Initialise and register projections

      // The basic final-state projection:
      // all final-state particles within
      // the given eta acceptance
      const FinalState fs(Cuts::abseta < 4.9);

      // The final-state particles declared above are clustered using FastJet with
      // the anti-kT algorithm and a jet-radius parameter 0.4
      // muons and neutrinos are excluded from the clustering
      FastJets jetfs(fs, FastJets::ANTIKT, 0.4, JetAlg::Muons::NONE, JetAlg::Invisibles::NONE);
      declare(jetfs, "jets");

      // FinalState of direct photons and bare muons and electrons in the event
      DirectFinalState photons(Cuts::abspid == PID::PHOTON);
      DirectFinalState bare_leps(Cuts::abspid == PID::MUON || Cuts::abspid == PID::ELECTRON);

      // Dress the bare direct leptons with direct photons within dR < 0.1,
      // and apply some fiducial cuts on the dressed leptons
      Cut lepton_cuts = Cuts::abseta < 2.5 && Cuts::pT > 20*GeV;
      DressedLeptons dressed_leps(photons, bare_leps, 0.1, lepton_cuts);
      declare(dressed_leps, "leptons");

      // Missing momentum
      declare(MissingMomentum(fs), "MET");

      // Book histograms
%(HISTOBOOKING)s

    }


    /// Perform the per-event analysis
    void analyze(const Event& event) {

      // Retrieve dressed leptons, sorted by pT
      Particles leptons = apply<FinalState>(event, "leptons").particles();

      // Retrieve clustered jets, sorted by pT, with a minimum pT cut
      Jets jets = apply<FastJets>(event, "jets").jetsByPt(Cuts::pT > 30*GeV);

      // Remove all jets within dR < 0.2 of a dressed lepton
      idiscardIfAnyDeltaRLess(jets, leptons, 0.2);

      // Select jets ghost-associated to B-hadrons with a certain fiducial selection
      Jets bjets = filter_select(jets, hasBTag(Cuts::pT > 5*GeV && Cuts::abseta < 2.5));

      // Veto event if there are no b-jets
      if (bjets.empty()) vetoEvent;

      // Apply a missing-momentum cut
      if (apply<MissingMomentum>(event, "MET").missingPt() < 30*GeV)  vetoEvent;

      // Fill histogram with leading b-jet pT
      _h["XXXX"]->fill(bjets[0].pT()/GeV);

    }


    /// Normalise histograms etc., after the run
    void finalize() {

      normalize(_h["XXXX"]); // normalize to unity
      normalize(_h["YYYY"], crossSection()/picobarn); // normalize to generated cross-section in pb (no cuts)
      scale(_h["ZZZZ"], crossSection()/picobarn/sumW()); // norm to generated cross-section in pb (after cuts)

    }

    /// @}


    /// @name Histograms
    /// @{
    map<string, Histo1DPtr> _h;
    map<string, Profile1DPtr> _p;
    map<string, CounterPtr> _c;
    /// @}
%(INLINEMETHODS)s

  };


  RIVET_DECLARE_PLUGIN(%(ANANAME)s);

}
''' % KEYWORDS
    f.write(src)
    f.close()

    ANAPLOTFILE = os.path.join(ANAPLOTDIR, ANANAME+".plot")
    logging.debug("Writing plot template to %s" % ANAPLOTFILE)
    f = open(ANAPLOTFILE, "w")
    src = '''\
BEGIN PLOT /%(ANANAME)s/d01-x01-y01
Title=[Insert title for histogram d01-x01-y01 here]
XLabel=[Insert $x$-axis label for histogram d01-x01-y01 here]
YLabel=[Insert $y$-axis label for histogram d01-x01-y01 here]
# + any additional plot settings you might like, see make-plots documentation
END PLOT

# ... add more histograms as you need them ...
''' % KEYWORDS
    f.write(src)
    f.close()

    if args.INLINE:
        sys.exit(0)
    ANAINFOFILE = os.path.join(ANAINFODIR, ANANAME+".info")
    logging.debug("Writing info template to %s" % ANAINFOFILE)
    f = open(ANAINFOFILE, "w")
    src = """\
Name: %(ANANAME)s
Year: %(ANAYEAR)s
Summary: <A one-line informative %(ANANAME)s headline>
Experiment: %(ANAEXPT)s
Collider: %(ANACOLLIDER)s
%(ANAREFREPO)sID: %(ANAINSPIREID)s
Status: UNVALIDATED NOHEPDATA SINGLEWEIGHT UNPHYSICAL
Reentrant: false
Authors:
 - Your Name <your@email.address>
References:
 - 'Example: Eur.Phys.J. C76 (2016) no.7, 392'
 - 'Example: DOI:10.1140/epjc/s10052-016-4184-8'
 - 'Example: arXiv:1605.03814'
#RunInfo: <Describe event types, cuts, and other general generator config tips.>
#Beams: [p+, p+]
#Energies: [[6500,6500]]
#Luminosity_fb: 139.0
Description:
  'A brief description of what is measured and what it is useful for.
   Use LaTeX for maths, e.g. $\pT > 50\;\GeV$.'
ValidationInfo:
  'A description of the process used to validate the Rivet code against
  the original experimental analysis. Cut-flow tables and similar information
  are welcome'
#ReleaseTests:
# - $A my-hepmc-prefix :MODE=some_rivet_flag
Keywords: []
BibKey: %(ANABIBKEY)s
BibTeX: '%(ANABIBTEX)s'
ToDo:
 - Implement the analysis, test it, remove this ToDo, and mark as VALIDATED :-)

""" % KEYWORDS
    f.write(src)
    f.close()

    logging.info("Use e.g. 'rivet-build Rivet%s.so %s.cc' to compile the plugin" % (ANANAME, ANANAME))
