#!/usr/bin/env python
"""
  Merges panel gain and offset constants for specified experiment and jungfrau detector

  Created on 2018-04-30 by Mikhail Dubrovin
"""
from __future__ import print_function

import logging
logger = logging.getLogger()

import os
import sys
import psana

import Detector.UtilsJungfrau as ju
from Detector.jungfrau_datasets import DSNAMES, DETNAMES
from pyimgalgos.GlobalUtils import print_command_line_parameters
from time import time
import PSCalib.GlobalUtils as gu # str_tstamp, deploy_file, log_rec_on_start, replace
import Detector.UtilsLogging as ul
from Detector.UtilsJungfrauCalib import uc, DIR_REPO, CALIB_REPO_JUNGFRAU, DIR_LOG_AT_START, fname_prefix_merge

#DIR_REPO = os.path.join(DIR_ROOT, 'detector/gains/jungfrau/')

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

def usage():
    return  'Usage: jungfrau_gain_constants -d <DSNAME> [-s <DETNAME>] [-t <TSTAMP>] [-D]\n'\
           +'       where <DSNAME> is like one of the list:\n         %s\n' % '\n         '.join(DSNAMES[:3])\
           +'       <DETNAME> is optional parameterr like one of the list:\n         %s\n' % '\n         '.join(DETNAMES[:3])\
           +'       <TSTAMP> - timestamp in format YYYYMMDDHHMMSS - beginning of the validity range'\
           +'\n\nTest 4M: jungfrau_gain_constants -d exp=xpptut15:run=610 -s jungfrau4M -D -c calib -t 20220908160000'\
           +'\n     1M: jungfrau_gain_constants -d exp=xpptut15:run=430 -s Jungfrau1M -D -c calib -t 20220908160000'\
           +'\n   0.5M: jungfrau_gain_constants -d exp=xpptut15:run=410 -s Jungfrau512k -D -c calib -t 20220908160000'\
           +'\n\nTry:  jungfrau_gain_constants -h\n'


def str_exp_run(dsname):
    """Returns str like xppx1004421-r1 from dsname
       exp=xppx1004421:run=1:dir=/cds/data/drpsrcf/xpp/xppx1004421/xtc/:stream=0-79
    """
    #return '-'.join(dsname.split(':',2)[:2]).replace('exp=','').replace('run=','r')
    return '-'.join([p for p in dsname.split(':') if len(p)>3 and p[:3] in ('exp','run')]).replace('exp=','').replace('run=','r')


def jungfrau_gain_constants():
    """Main script for CL
    """
    print(80*'_')

    parser = option_parser()

    if len(sys.argv)==1:
        parser.print_help()
        print('\n%s' % usage())
        sys.exit('\nWARNING: MISSING ARGUMENTS\n')

    (popts, pargs) = parser.parse_args()
    nargs = len(pargs)
    DSNAME = pargs[0] if nargs>0 else popts.dsname
    DETNAME= pargs[1] if nargs>1 else popts.source
    TSTAMP = pargs[2] if nargs>2 else popts.tstamp
    clbdir = popts.clbdir
    deploy = popts.deploy
    upload = popts.upload
    loglev = popts.loglev
    repdir = popts.repdir

    defs = vars(parser.get_default_values())
    DSNAME_DEF  = defs['dsname']
    DETNAME_DEF = defs['source']
    TSTAMP_DEF  = defs['tstamp']

    verbose = loglev != 'INFO'
    if verbose: print_command_line_parameters(parser)

    repoman = uc.RepoManager(dirrepo=repdir, dir_log_at_start=DIR_LOG_AT_START,\
                             dirmode=popts.dirmode, filemode=popts.filemode, group=popts.group)
    logname = repoman.logname('%s_%s_%s' % (SCRNAME, str_exp_run(DSNAME), gu.get_login()))
    ul.init_logger(loglevel=loglev, logfname=logname, group=popts.group, fmt='[%(levelname).1s] %(filename)s L%(lineno)04d %(message)s')
    logger.info('log file: %s' % logname)
    repoman.save_record_at_start(SCRNAME, adddict={'logfile':logname})
    logger.info('Generator of gain correction constants for Jungfrau')

    ds = psana.DataSource(DSNAME)
    env = ds.env()
    source = ju.psana_source(env, DETNAME)
    ssource = ju.string_from_source(source)
    strsrc  = ssource.replace(':','-')
    npanels = ju.number_of_modules_in_jungfrau(env, source)
    logger.info('Dataset "%s", name pattern "%s"' % (DSNAME, DETNAME)\
               +'\n      found source "%s" with number of panels: %s' %(ssource, str(npanels)))
    jfid = ju.id_jungfrau(env, source)

    co = ju.jungfrau_config_object(env, source)
    if co.Version < 3: # for old versions of config jfid="MecTargetChamber.0-Jungfrau.0"
        msg = "WARNING: Jungfrau id is not available for config V%d. Constants can't be auto-generated." % co.Version
        logger.warning(msg)
        sys.exit()

    cdirs = ju.find_panel_calib_dirs(jfid, dname=repdir, tstamp=TSTAMP)  # dname=ju.DIRNAME
    if None in cdirs:
        msg = "Calibdirs can't be found for all panels"
        sys.exit(msg)

    dirs = ['%s/%s' % (repdir, d) for d in cdirs]

    msg = 'Selected panel repository directories:'
    for d in dirs: msg += '\n      %s' % d
    logger.info(msg)

    ctype_gain   = 'pixel_gain'
    ctype_offset = 'pixel_offset'

    exp      = env.experiment()
    runnum   = env.configStore().get(psana.EventId).run()
    calibdir = env.calibDir().replace('//','/') if clbdir is None else clbdir
    #ctypedir = '/some-path/calib/Jungfrau::CalibV1/CxiEndstation.0:Jungfrau.0/'
    ctypedir = '%s/Jungfrau::CalibV1/%s' % (calibdir, ssource)
    ofname   = '%d-end.data' % runnum
    lfname   = None
    verbos   = False
    dirmode  = 0o2775
    filemode = 0o664
    group    = 'ps-users'

    tstamp_run, tstamp_now = uc.tstamps_run_and_now(env)
    tstamp = tstamp_run

#    tstamp = tstamp_run if TSTAMP is None else\
#             TSTAMP if int(TSTAMP)>9999 else\
#             tstamp_run

    dmerge = repoman.makedir_merge()
    fmerge_prefix = fname_prefix_merge(dmerge, ssource, tstamp, exp, runnum)
    #fmerge_prefix  = '%s-r%04d_%s'%(exp, runnum, ssource)
    fprefix_g = '%s_%s'%(fmerge_prefix, ctype_gain)
    fprefix_o = '%s_%s'%(fmerge_prefix, ctype_offset)

    logger.info('Begin merging panel arrays')
    ju.merge_panel_constants(dirs, ifname='%s/g%d_gain.npy',   ofname=fprefix_g, ofmt='%.4f', filemode=filemode, group=group)
    ju.merge_panel_constants(dirs, ifname='%s/g%d_offset.npy', ofname=fprefix_o, ofmt='%.1f', filemode=filemode, group=group)
    #else: logger.warning('Add option -M to merge panel arrays')

    logger.debug('calibdir: %s'%calibdir)
    logger.debug('destination file name: %s'%ofname)

    if deploy:
        logger.info('Deploy calib files under %s' % ctypedir)
        gu.deploy_file('%s.txt'%fprefix_g, ctypedir, ctype_gain,   ofname, lfname, verbose,\
                       filemode=filemode, dirmode=dirmode, group=group)
        gu.deploy_file('%s.txt'%fprefix_o, ctypedir, ctype_offset, ofname, lfname, verbose,\
                       filemode=filemode, dirmode=dirmode, group=group)
    else:
        logger.warning('Add option -D to deploy calib files under directory %s' % ctypedir)

    if upload:
        #dcs add <fname>.h5 -t pixel_gain -f <ifname>.txt -T 0 -m "comment"
        msg = 'Commands to upload calibration constants in repository %s:' % repdir
        #cmt = 'Id found in %s for %s' % (DSNAME, ssource)
        cmt = 'See command: jungfrau_gain_constants -d %s -U' % (DSNAME)
        fmt = 'dcs add %s/jungfrau-%s.h5 -t %s -f %s.txt -T 1500000000 -m "%s" -e %s -r %s -d %s'
        msg += '\n==== upload command for %s:\n' % ctype_gain
        msg += fmt % (repdir, jfid, ctype_gain,   fprefix_g, cmt, exp, runnum, ssource)
        msg += '\n==== upload command for %s:\n' % ctype_offset
        msg += fmt % (repdir, jfid, ctype_offset, fprefix_o, cmt, exp, runnum, ssource)
        logger.info(msg)
    #else:
    #    logger.warning('Add option -R to learn about uploading constants in repository %s' % repdir)

    os.chmod(logname, popts.filemode)
    uc.cgu.change_file_ownership(logname, user=None, group=popts.group)

    logger.info('Log file: %s' % logname)
    #fhandler.close()


def option_parser():
    from optparse import OptionParser

    d_dsname = 'exp=mecls3216:run=226' # 'exp=xpptut15:run=430'
    d_source = 'Jungfrau'
    d_tstamp = '20150101000000' # gu.str_tstamp(fmt='%Y%m%d%H%M%S', time_sec=None)
    d_deploy = False
    d_upload = False
    d_clbdir = None
    d_loglev = 'INFO'
    d_lfpath = 'DEPRECATED: /reg/g/psdm/logs/calibman/jungfrau_gain - see .../detector/gains/jungfrau/logs'
    d_repdir = DIR_REPO  # '/reg/g/psdm/detector/gains/jungfrau' # CALIB_REPO_JUNGFRAU
    d_dirmode  = 0o2775
    d_filemode = 0o664
    d_group    = 'ps-users'

    h_dsname='dataset name, default = %s' % d_dsname
    h_source='input ndarray source/detector name, default = %s' % d_source
    h_tstamp='time stamp for beginning of the validity range in format YYYYMMDDHHMMSS, default (now) = %s' % d_tstamp
    h_deploy='deploy files in experimental calib directory, default = %s' % d_deploy
    h_upload='add to repository resulting files with gains and offsets for merged panesls, default = %s' % d_upload
    h_clbdir='non-default calib directory for deployment of calibration files, default = %s' % d_clbdir
    h_loglev='dataset name, default = %s' % d_loglev
    h_lfpath='path to log file, default = %s' % d_lfpath
    h_repdir='repository of gain files, default = %s' % d_repdir
    h_dirmode = 'mode for all mkdir, default = %s' % oct(d_dirmode)
    h_filemode= 'mode for all saved files, default = %s' % oct(d_filemode)
    h_group   = 'group ownership for all files, default = %s' % d_group

    parser = OptionParser(description='Generator of gain correction constants for Jungfrau', usage = usage())
    parser.add_option('-d', '--dsname', dest='dsname', default=d_dsname, action='store', type='string', help=h_dsname)
    parser.add_option('-s', '--source', dest='source', default=d_source, action='store', type='string', help=h_source)
    parser.add_option('-t', '--tstamp', dest='tstamp', default=d_tstamp, action='store', type='string', help=h_tstamp)
    parser.add_option('-D', '--deploy', dest='deploy', default=d_deploy, action='store_true',           help=h_deploy)
    parser.add_option('-U', '--upload', dest='upload', default=d_upload, action='store_true',           help=h_upload)
    parser.add_option('-c', '--clbdir', dest='clbdir', default=d_clbdir, action='store', type='string', help=h_clbdir)
    parser.add_option('-L', '--loglev', dest='loglev', default=d_loglev, action='store', type='string', help=h_loglev)
    parser.add_option('-P', '--lfpath', dest='lfpath', default=d_lfpath, action='store', type='string', help=h_lfpath)
    parser.add_option('-R', '--repdir', dest='repdir', default=d_repdir, action='store', type='string', help=h_repdir)
    parser.add_option('--dirmode',       default=d_dirmode, action='store', type='int',    help=h_dirmode)
    parser.add_option('--filemode',      default=d_filemode,action='store', type='int',    help=h_filemode)
    parser.add_option('--group',         default=d_group,   action='store', type='string', help=h_group)

    return parser


def test_str_exp_run():
    dsname = 'exp=xppx1004421:run=1:dir=/cds/data/drpsrcf/xpp/xppx1004421/xtc/:stream=0-79'
    #dsname = 'exp=xppx1004421:run=1'
    print('str_exp_run(%s) = %s' % (dsname, str_exp_run(dsname)))


if __name__ == "__main__":
    #test_str_exp_run()
    jungfrau_gain_constants()
    sys.exit(0)

# EOF
