#!/usr/bin/env python

import sys
from Detector.UtilsLogging import STR_LEVEL_NAMES, DICT_NAME_TO_LEVEL, logging
logger = logging.getLogger(__name__)

SCRNAME = sys.argv[0].rsplit('/')[-1]

USAGE = '%s -d <dataset-name> -s<detector-name> -n <number-of-events> -m 0 <number-events-to-skip> -c <calib-dir> -L <log-level-str>'%SCRNAME\
      + '\n  Examples:'\
      + '\n  %s -d exp=xpptut15:run=450 -s pnccd -n 100'%SCRNAME\
      + '\n  %s -d exp=cxilv9518:run=8 -s CxiDs1.0:Jungfrau.0 -n 10 -m 0 -L DEBUG'%SCRNAME\
      + '\n  %s -d exp=cxilv9518:run=8 -s CxiDs1.0:Jungfrau.0 -n 10'%SCRNAME\
      + '\n  %s -d exp=mfxp17518:run=108 -s MfxEndstation.0:Epix10ka2M.0 -n 10'%SCRNAME\
      + '\n  %s -d exp=meclv2518:run=99:dir=/reg/d/ffb/mec/meclv2518/xtc/:smd -s MecTargetChamber.0:Epix10kaQuad.0 -n 10 -m 0 -c /reg/neh/operator/mecopr/calib -L DEBUG'%SCRNAME\
      + '\n  %s -d exp=xpplw3319:run=293 -s XppGon.0:Epix100a.3 -n 1000 -S 0:180,620: -M p'%SCRNAME

if len(sys.argv)<3: sys.exit('%s\n EXIT - MISSING ARGUMENTS\n' % USAGE)

def argument_parser():
    import argparse

    d_dsname   = 'exp=cxilv9518:run=8'
    d_detname  = 'CxiDs1.0:Jungfrau.0'
    d_events   = 100
    d_evskip   = 0
    d_logmode  = 'INFO'
    d_calibdir = None #/reg/neh/operator/mecopr/calib/
    d_prefix   = 'img'
    d_aslice   = ':'
    d_mode     = 'p'

    h_dsname   = 'dataset name, default = %s' % d_dsname
    h_detname  = 'detector name, default = %s' % d_detname
    h_events   = 'number of events to collect, default = %s' % d_events
    h_evskip   = 'number of events to skip, default = %s' % d_evskip
    h_logmode  = 'logging mode, one of %s, default = %s' % (STR_LEVEL_NAMES, d_logmode)
    h_calibdir = 'calibration directory, default = %s' % str(d_calibdir)
    h_prefix   = 'output file name prefix, default = %s' % str(d_prefix)
    h_aslice   = 'array slice, e.g. 0:180,620:, default = %s' % str(d_aslice)
    h_mode     = 'mode of data: r/c/i/p : raw/calib/image/raw-peds default = %s' % str(d_aslice)

    parser = argparse.ArgumentParser(usage=USAGE, description='Accumulates det.calib array average and max')
    parser.add_argument('-d', '--dsname',   default=d_dsname,   type=str,   help=h_dsname)
    parser.add_argument('-s', '--detname',  default=d_detname,  type=str,   help=h_detname)
    parser.add_argument('-n', '--events',   default=d_events,   type=int,   help=h_events)
    parser.add_argument('-m', '--evskip',   default=d_evskip,   type=int,   help=h_evskip)
    parser.add_argument('-L', '--logmode',  default=d_logmode,  type=str,   help=h_logmode)
    parser.add_argument('-c', '--calibdir', default=d_calibdir, type=str,   help=h_calibdir)
    parser.add_argument('-f', '--prefix',   default=d_prefix,   type=str,   help=h_prefix)
    parser.add_argument('-S', '--aslice',   default=d_aslice,   type=str,   help=h_aslice)
    parser.add_argument('-M', '--mode',     default=d_mode,     type=str,   help=h_mode)

    return parser

parser = argument_parser()
args = parser.parse_args()
print('parser.parse_args: %s' % str(args))

dsname   = args.dsname
detname  = args.detname
events   = args.events
evskip   = args.evskip
calibdir = args.calibdir
logmode  = args.logmode
prefix   = args.prefix
aslice   = args.aslice
mode     = args.mode

logging.basicConfig(format='[%(levelname).1s] %(name)s L%(lineno)04d: %(message)s', level=DICT_NAME_TO_LEVEL[logmode])

import numpy as np
from psana import DataSource, Detector, setOption
from Detector.GlobalUtils import info_ndarr, selected_record
import PSCalib.GlobalUtils as gu

_slice = eval('np.s_[%s]' % aslice)

nev_sum = 0
nda_sum = None
nda_max = None

peds = None

if True:

    if calibdir is not None:
        setOption('psana.calib-dir', calibdir)
        logging.info('set non-default psana.calib-dir:' + str(calibdir))

    ds = DataSource(dsname)
    det = Detector(detname)

    logger.info('calibDir: %s' % ds.env().calibDir())
    logger.info('det.shape:' + str(det.shape()))

    expname = ds.env().experiment()
    print('experiment: %s' % expname)
    dettype = gu.det_type_from_source(det.source)
    print('dettype: %s' % dettype)
    runnum = None

    do_break = False

    for orun in ds.runs():
      runnum = orun.run()
      print('==== run %04d ====' % runnum)

      for i, evt in enumerate(orun.events()):
        if i<evskip: continue
        if i>events:
           do_break = True
           break
        if i<5:
            peds = det.pedestals(evt)
            logging.debug(info_ndarr(peds, 'pedestals'))
            gains = det.gain(evt)
            logging.debug(info_ndarr(gains, 'gains'))
            assert peds is not None

        raw = det.raw(evt)
        logging.debug(info_ndarr(raw, 'raw'))

        if raw is None:
            logging.warning('Event %4d raw is None' % i)
            continue

        arr = raw if mode in ('r','p') else\
              det.calib(evt) if mode == 'c' else\
              det.image(evt) if mode == 'i' else\
              raw

        arr = arr.astype(peds.dtype) if aslice in (':', None) else\
              arr[_slice].astype(peds.dtype)

        logging.debug(info_ndarr(arr, 'arr'))

        if selected_record(i):
            #logging.info('Event %4d' % i)
            logging.info(info_ndarr(arr, 'Event %4d arr[%s]:' % (i, aslice)))

        if nda_max is None:
            nev_sum = 1
            nda_sum = np.array(arr).astype(peds.dtype)
            nda_max = np.array(arr).astype(peds.dtype)
        else:
            np.maximum(arr, nda_max, out=nda_max)
            nda_sum += arr
            nev_sum += 1
      if do_break: break # runs

    if nev_sum: nda_sum /= nev_sum

    if mode == 'p':
        nda_max -= peds
        nda_sum -= peds
    if nda_sum.ndim>2: nda_sum.shape = nda_max.shape

    logging.info('\nStatistics of events nevt:%d nev_sum:%d'%(i,nev_sum))

    _detname = detname.replace('.','-').replace(':','-')
    _prefix = '%s-%s-%s-r%04d-e%06d-mode-%s' % (prefix, _detname, expname, runnum, nev_sum, mode)

    fname = _prefix + '-max.npy'
    np.save(fname, nda_max)
    logging.info('saved file %s' % fname)

#    if nda_max.ndim>2: nda_max.shape = (int(nda_max.size/nda_max.shape[-1]), nda_max.shape[-1])
#    fname = _prefix + '-max.txt'
#    np.savetxt(fname, nda_max, fmt='%.3f')
#    logging.info('saved files %s and *.npy'%fname)

    fname = _prefix + '-ave.npy'
    np.save(fname, nda_sum)
    #np.savetxt(fname, nda_sum, fmt='%.3f')
    logging.info('saved file %s'%fname)

    sys.exit('END OF %s' % SCRNAME)

# EOF
