#! /usr/bin/env python

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

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

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

import sys
if sys.version_info[:3] < (2,4,0):
    print "rivet scripts require Python version >= 2.4.0... exiting"
    sys.exit(1)

import logging, os


## Handle command line
from optparse import OptionParser
parser = OptionParser(usage=__doc__)
parser.add_option("--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_option("-q", "--quiet", dest="LOGLEVEL", default=logging.INFO,
                  action="store_const", const=logging.WARNING, help="only write out warning and error messages")
parser.add_option("-v", "--verbose", dest="LOGLEVEL", default=logging.INFO,
                  action="store_const", const=logging.DEBUG, help="provide extra debugging messages")
parser.add_option("-i", "--inline-info", dest="INLINE", action="store_true",
                  default=False, help="Put analysis info into source file instead of separate data file.")
opts, args = parser.parse_args()
logging.basicConfig(format="%(msg)s", level=opts.LOGLEVEL)
ANANAMES = args

## Work out installation paths
ANAROOT = os.path.abspath(opts.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 opts.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.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>"
    ANAYEAR = "<Insert year of publication>"
    INSPIRE_SPIRES = None
    ANAINSPIREID = "<Insert the Inspire ID>"
    if match:
        STDANA = True
        ANAEXPT = match.group(1)
        ANAYEAR = match.group(2)
        INSPIRE_SPIRES = match.group(3)
        ANAINSPIREID = match.group(4)
    if INSPIRE_SPIRES == "S":
        ANAREFREPO = "Spires"
    else:
        ANAREFREPO = "Inspire"
    KEYWORDS = {
        "ANANAME" : ANANAME,
        "ANAEXPT" : ANAEXPT,
        "ANAYEAR" : ANAYEAR,
        "ANAREFREPO" : ANAREFREPO,
        "ANAINSPIREID" : ANAINSPIREID
        }

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


    ## Try do download AIDA data file from hepdata
    if STDANA:
        try:
            import urllib
            hdurl = "http://hepdata.cedar.ac.uk/view/irn%s/aida" % ANAINSPIREID
            logging.debug("Getting data file from HepData at %s" % hdurl)
            httpstream = urllib.urlopen(hdurl)
            aidastr = httpstream.read()
            if "<html>" in aidastr:
                logging.warning("Problem encountered when getting data from HepData (%s). No reference data file written." % hdurl)
            else:
                f = open("%s.aida" % ANANAME, "w")
                f.write(aidastr)
                f.close()
            httpstream.close()
        except Exception, e:
            logging.error("HepData oops: %s" % e)


    if opts.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 spiresId()           const { return "%(ANASPIRESID)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"]
    else:
        INLINEMETHODS=""

    KEYWORDS["INLINEMETHODS"] = INLINEMETHODS


    ANASRCFILE = os.path.join(ANASRCDIR, ANANAME+".cc")
    logging.debug("Writing implementation template to %s" % ANASRCFILE)
    f = open(ANASRCFILE, "w")
    src = """\
// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/RivetAIDA.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Projections/FinalState.hh"
/// @todo Include more projections as required, e.g. ChargedFinalState, FastJets, ZFinder...

namespace Rivet {


  class %(ANANAME)s : public Analysis {
  public:

    /// @name Constructors etc.
    //@{

    /// Constructor
    %(ANANAME)s()
      : Analysis("%(ANANAME)s")
    {    }

    //@}


  public:

    /// @name Analysis methods
    //@{

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

      /// @todo Initialise and register projections here

      /// @todo Book histograms here, e.g.:
      // _h_XXXX = bookProfile1D(1, 1, 1);
      // _h_YYYY = bookHistogram1D(2, 1, 1);

    }


    /// Perform the per-event analysis
    void analyze(const Event& event) {
      const double weight = event.weight();

      /// @todo Do the event by event analysis here

    }


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

      /// @todo Normalise, scale and otherwise manipulate histograms here

      // scale(_h_YYYY, crossSection()/sumOfWeights()); # norm to cross section
      // normalize(_h_YYYY); # normalize to unity

    }

    //@}


  private:

    // Data members like post-cuts event weight counters go here


  private:

    /// @name Histograms
    //@{
    AIDA::IProfile1D *_h_XXXX;
    AIDA::IHistogram1D *_h_YYYY;
    //@}
%(INLINEMETHODS)s

  };



  // The hook for the plugin system
  DECLARE_RIVET_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=[Uncomment and insert title for histogram d01-x01-y01 here]
#XLabel=[Uncomment and insert x-axis label for histogram d01-x01-y01 here]
#YLabel=[Uncomment and 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 opts.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: <Insert short %(ANANAME)s description>
Experiment: %(ANAEXPT)s
Collider: <Insert %(ANANAME)s collider>
%(ANAREFREPO)sID: %(ANAINSPIREID)s
Status: UNVALIDATED
Authors:
 - Your Name <your@email.address>
#References:
# - <Example: Phys.Lett.B639:151-158,2006, Erratum-ibid.B658:285-289,2008>
# - <Example: doi:10.1016/j.physletb.2006.04.048>
# - <Example: arXiv:hep-ex/0511054 (plus erratum)>
RunInfo:
  <Insert event types (not gen-specific), energy, any kinematic
  efficiency cut(s) that may be needed; essentially any details needed to set
  up a generator to reproduce the data.>
NumEvents: 1000000
NeedCrossSection: no
#Beams: <Insert beam pair(s), e.g. [p-, p+] or [[p-, e-], [p-, e+]]>
#Energies: <Insert list of run energies or beam energy pairs in GeV,
#           e.g. [1960] or [[8.0, 3.5]] or [630, 1800]. Order pairs to match "Beams">
#PtCuts: <Insert list of kinematic pT cuts in GeV, e.g. [0, 20]>
#NeedCrossSection: True
Description:
  '<Insert a fairly long description, including what is measured
  and if possible what it is useful for in terms of MC validation
  and tuning. Use \LaTeX for maths like $\pT > \unit{50}{\GeV}$.
  Use single quotes around the block if required (see YAML 1.2 manual)>'
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-buildplugin Rivet%s.so %s.cc' to compile the plugin" % (ANANAME, ANANAME))
