#! /usr/bin/env python3

import dactylos as dact
import hepbasestack as hep

import concurrent.futures as fut
import os
import sys
import time
import tqdm
import pylab as p
import hjson

p.ion()
from os.path import join as _j
ROOT_FOUND = True
try:
    import ROOT
except ImportError as e:
    ROOT_FOUND = False

#shaping_spree  = True
runtime = 240
outdir = 'run12-Wed03052020'

# some global variables
NCHANNELS = 8
NBINS = 16834 # 14-bit digitizer

MPI = False
if MPI:
    EXECUTOR = fut.ProcessPoolExecutor(max_workers=2)


def setup_root_histos():
    """
    Prepare energy histograms in ROOT.TH1D style.
    """
    if not ROOT_FOUND:
        print ("Failed to import root.... Some functionality might not be available")
        print ("Make sure $ROOTSYS/lib is in your $PYTHONPATH")
        raise ImportError("This requires CERN's root python library to run!")


    histos = []
    for i in range(NCHANNELS):
        h = ROOT.TH1D()
        h.SetName("ehistch" + str(i))
        h.SetTitle("DigitizerChannel Ch" + str(i))
        h.SetBins(NBINS, 0, NBINS-1)
        histos.append(h)
    return histos

def histogram_energy(input_file):
    """
    Create root TH1D histograms for the energy
    """
    if not ROOT_FOUND:
        print ("Failed to import root.... Some functionality might not be available")
        print ("Make sure $ROOTSYS/lib is in your $PYTHONPATH")
        raise ImportError("This requires CERN's root python library to run!")

    f = ROOT.TFile(input_file,"UPDATE")        
    f.cd()
    histos = setup_root_histos()
    for ch in range(7):
        for event in f.__getattr__(f'ch{ch}'):
            histos[ch].Fill(event.energy)  
        histos[ch].Write()  
        f.Write()
    f.Close()




if __name__ == '__main__':
    
    import argparse
    parser = argparse.ArgumentParser(description='Take data with CAEN N6725 digitizer. If desired, operate SUNEC13 chamber.')
    #parser.add_argument('--shaping-spree', action='store_true', default=False, help='Take data for different shaping times from a pre-defined (currently hardcoded) list')
    parser.add_argument('--test', action='store_true', default=False)
    parser.add_argument('--scope', action='store_true', default=False)
    parser.add_argument('--read-waveforms', action='store_true', default=False, help='Save the waveforms as well.')
    parser.add_argument('-o','--outdir', type=str, default='run-digitizer-data')
    parser.add_argument('-i','--infile', type=str, default='infile', help='Only to use in combination with --create-histograms-only')
    parser.add_argument('--create-histograms-only', action='store_true', default=False, help='add histograms to the outfile after the run. Basically, histogram the "energy" field in the root file. Requires -i infile to be given.')
    parser.add_argument('--loglevel', type=int, default='20', help='loglevel ')
    parser.add_argument('-r','--runtime', type=int, default=20, help='runtime in seconds, default 20')
    #parser.add_argument('--detector-name', type=str, default='', help='add the id/name of the detector for identification. This is a MUST')
    #parser.add_argument('--run-config', type=str, default='', help='path to a run configuration file')
    parser.add_argument('--digitizer-config', type=str, default='config_x.json', help='path to a digitizer configuration file')
    parser.add_argument('--wf-readout-fw-test', action='store_true', default=False, help='test the functionality for the wf fw. DEVELOPMENT ONLY!')
    parser.add_argument('--energy-histo-test',\
                        action='store_true',\
                        default=False,\
                        help='test the energy histogram + event live view for dpp-pha. DEVELOPMENT ONLY!')


    args = parser.parse_args()
    logger = hep.logger.get_logger(args.loglevel)
    config = args.digitizer_config
    #if not args.run_config:
    #    shaping_spree = args.shaping_spree
    #    #shaping_times = [100, 500, 1000, 1500, 2000, 3000, 5000, 6000, 7000, 10000]
    #    #shaping_times = [4000]
    #    shaping_times = [500,1000, 2000, 4000, 5000, 8000, 12000, 16000, 20000, 30000]
    #    detector_sequence = False
    #else:
    #    run_config = hjson.load(open(args.run_config)) 
    #    args.read_waveforms = run_config['save-waveforms']
    #    shaping_spree = len(run_config['peaking-times']) > 0
    #    if shaping_spree:
    #        shaping_times = run_config['peaking-times'] 
    #    detector_sequence = run_config['detector-sequence']        

    outdir        = args.outdir

    if args.create_histograms_only:
        histogram_energy(args.infile)
        logger.info('done!')
        sys.exit()

    # for anything else we need a detector name
    #if not args.detector_name:
    #    raise ValueError('No detector name/id given. abort!')

    outfilename = _j(outdir, f"data-")

    try:
        os.mkdir(outdir)
    except Exception as e:
        logger.warning(f'Can not create {outdir} - exception {e}!')
    
    if args.wf_readout_fw_test:
        logger.warning('Just for testing...')
        digi = dact.CaenN6725(dact.CaenN6725.parse_configfile(config), logger=logger, has_dpp_pha_firmware=False)
        logger.info(f"Got handle {digi.get_handle()}")
        if args.scope:
            for wf in digi.live_view(20):
                print (wf)
        else:
            runtime = args.runtime
            logger.info(f"Will run digitizer for {runtime} seconds and store to file 'example.root'")
            digi.run_digitizer(runtime, rootfilename=_j(outdir,'example.root'))
            logger.info("...done")
        del digi
        sys.exit(0)

    if args.energy_histo_test:
        logger.warn('TEST -> energy histogram..')
        digi = dact.CaenN6725(dact.CaenN6725.parse_configfile(config), logger=logger)
        digi.digitizer.clear_energy_histogram()
        for wf in digi.live_view(20, fill_histogram=True):
            pass
        print(wf)
        #t_elapsed = 0
        #while t_elapsed < 20:
        #    events = digi.live_view(fill_histogram=True)
        #    time.sleep(1)
        #    t_elapsed += 1
        for ch in range(8):
            histo = digi.digitizer.get_energy_histogram(ch)
            p.plot(histo)
            p.savefig(f'testhist{ch}.png')
        del digi
        sys.exit(0)

    if args.test:
        logger.info('Starting test session...')
        digi = dact.CaenN6725(dact.CaenN6725.parse_configfile(config), logger=logger)
        handle = digi.get_handle()
        logger.info(f"Got handle {handle}")
        #digi2 = dact.CaenN6725({})
        #digi2.set_handle(handle)
        #temps = digi2.get_temperatures()
        #print (type(temps))
        #print (temps)
        #logger.info(f"temps : {temps}")
        # show the current firmware version
        logger.info(f"We are connected to the digitizer. Boardinfo : {digi.get_board_info()}")
        logger.info(f"Will run digitizer for 20 seconds...")
        digi.run_digitizer(20, rootfilename=_j(outdir,'example.root'),\
                           read_waveforms=args.read_waveforms)
        del digi
        sys.exit(0)
    #del digi # closes the connection to the digitzer
    
    #if shaping_spree:
    #    for stime in tqdm.tqdm(shaping_times):
    #        print (f"Taking data for shaping time {stime} ..")
    #        outfile = outfilename + f'_stime{stime}.root'
    #        print (f"Writing data to file {outfile}")
    #        digi = dact.CaenN6725(dact.CaenN6725.parse_configfile(config), shaping_time=stime)
    #        #ex = hep.timed_progressbar(runtime) # show a progressbar for runtime seconds
    #        digi.run_digitizer(runtime, rootfilename=outfile,\
    #                           read_waveforms=args.read_waveforms)
    #        del digi  # close the digitizer. This is important, since only one at the time
    #                  # can be open.
    #        #ex.shutdown()

    
    if args.scope:
        if args.wf_readout_fw_test:
            print ('Not implemented yet')
            sys.exit(1)
        try:
            os.mkdir(outdir)
        except Exception as e:
            print (e)
            pass
    
        digi = dact.CaenN6725(dact.CaenN6725.parse_configfile(config))
        digi.init_scope()
        while True:
            logger.info("Getting shaping...")
            digi.oscilloscope(filename = _j(outdir,"shaping.png"))
            logger.info("Getting trigger...")
            digi.oscilloscope(filename = _j(outdir, "trigger.png"),
                              trace1 = dact.DPPVirtualProbe1.Delta2,\
                              trace2 = dact.DPPVirtualProbe2.Input,\
                              dtrace1 = dact.DPPDigitalProbe1.TRGWin)
            time.sleep(5)
        #digi.oscilloscope(filename = _j(outdir, "test3.png"),
        #                  trace1 = dact.DPPVirtualProbe1.Trapezoid,\
        #                  trace2 = dact.DPPVirtualProbe2.Input,\
        #                  dtrace1 = dact.DPPDigitalProbe1.Armed)
        #digi.oscilloscope(filename = _j(outdir,"test4.png"),
        #                  trace1 = dact.DPPVirtualProbe1.Input,\
        #                  trace2 = dact.DPPVirtualProbe2.TrapezoidReduced,\
        #                  dtrace1 = dact.DPPDigitalProbe1.PkRun)
        #digi.oscilloscope(filename = _j(outdir,"test5.png"),
        #                  trace1 = dact.DPPVirtualProbe1.Input,\
        #                  trace2 = dact.DPPVirtualProbe2.Baseline,\
        #                  dtrace1 = dact.DPPDigitalProbe1.Busy)
        #digi.oscilloscope(filename = _j(outdir,"test6.png"),
        #                  trace1 = dact.DPPVirtualProbe1.Input,\
        #                  trace2 = dact.DPPVirtualProbe2.Threshold,\
        #                  dtrace1 = dact.DPPDigitalProbe1.PileUp)
        #digi.oscilloscope(filename = _j(outdir,"test7.png"),
        #                  trace1 = dact.DPPVirtualProbe1.Input,\
        #                  trace2 = dact.DPPVirtualProbe2.Threshold,\
        #                  dtrace1 = dact.DPPDigitalProbe1.TRGHoldoff)
        
    
    
