#!/usr/bin/env python
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.
"""
Script to reprocess ZDAB files from SNO. We reprocess the files since the only
files I have access to right now were done for LETA and in those files the PMT
style banks are missing from the ZDAB file. They also might have the charge
stored in counts above pedestal instead of being normalized by the high half
point (although I'm not 100% sure about this). To reprocess a single file:

    $ ./zdab-reprocess FILENAME

and to batch reprocess:

    $ ./zdab-reprocess FILENAME FILENAME ...

You can also specify a minimum nhit value to reprocess. For example:

    $ ./zdab-reprocess --min-nhit 100 SNOCR_00000100004_000_p2.xzdab

will only reprocess events with an nhit greater than 100. One thing to note is
that this nhit value is just the total number of PMT bundles and so includes
channels with bad calibration, OWLs, and FEC/D hits.

By default, the reprocessed files will be stored in the current working
directory with _reprocessed appended to the name, but this can be changed by
passing a different suffix on the command line. For example:

    $ ./zdab-reprocess --min-nhit 100 --suffix reduced SNOCR_00000100004_000_p2.xzdab

will produce a file named SNOCR_00000100004_000_p2_reduced.xzdab.
"""

from __future__ import print_function, division
import string
import os
import re

def GetSettingsFile(run):
    """
    Returns the settings file for a given run number.

    Taken from GetSettingsFile() in autosno/lib/AutoSNO_Module.pm.
    """
    if   run < 19999: return "load_d2o_settings.cmd"
    elif run < 33907: return "load_salt_settings.cmd"
    elif run < 39999: return "load_d2o_2_settings.cmd"
    elif run < 67900: return "load_ncd_settings.cmd"
    else:             return "load_h2o_settings.cmd"

# from autosno/run_lists/RunList_clock_force_no_fix.dat
RUN_LIST_CLOCK_FORCE_NO_FIX = [
    11371,
    11377,
    11399,
    11436,
    11575,
    11681,
    11701,
    11911,
    11976,
    12187,
    12233,
    13401,
    13423,
    13431,
    13874,
    14264,
    51456,
]

def get_clock_force_no_fix_string(run):
    if run in RUN_LIST_CLOCK_FORCE_NO_FIX:
        return "set bank TUPK 1 word 3 to -1"
    else:
        return ""

# This template comes from autosno/prod/Reconstruct.pm. I also added the lines
# which autosno adds in Run() of the Reconstruct module.
SNOMAN_TEMPLATE = \
"""
* This file was automatically created by zdab-reprocess.
*
* Note: This command file was created by attempting to "reverse engineer" how
*       autosno fits neutrino runs.
*
* Author: Anthony LaTorre
* Created on August 26, 2019.
file inp 1 @input
file out 1 @output

titles @dqxx
titles @anxx_pca
titles @anxx_neutrino

@@@settings_file

$flt_set_prescale @prescale_setting
$define_test 1 $line_1 'nhits float_equals EV+$KEV_NPM:@lower_nhit @upper_nhit;'
$starting_seed @seed
$starting_seed_2 @seed2

@clock_force_no_fix

@input_nhit_thresh

* We want to process the entire file.
$num_events 0

* Titles and database checks.
set bank TMTT 1 word 0 to 1

* HCA Control.
$hca_mode $perform_correction
$hca_correction $new_correction
$hca_consecutive $yes

* QRC Control.
* We need the cross talk cut.
$apply_xtalk_cut

*Settings for muon fitters.
* Set the charge threshold for multi photon timing correction
$set_mpca_threshold 1.5
titles mpca_data.dat
* Turn on the channel dependent gain calibration
set bank TCAL 1 word 4 to 54

*Load charge probability titles file
titles charge_photoelectrons.dat

* Low nhit prescale
$flt_set_prescale_2 0.0

* Primary NHIT Cut.
$enable_test 1
$define_test 1 $line_1 'nhits float_equals EV+$KEV_NPM:18.5 9999.;'

* Primary Fitter NHIT Cut.
$enable_test 3
$define_test 3 $line_1 'nhits float_equals EV+$KEV_NPM:9.5 999.;'

* NCD Event Tag.
$enable_test 4
$define_test 4 $line_1 'ncd float_equals EV+$KEV_NCD_STATUS:0.5 3.5;'

* Muon Candidate Tag.
$enable_test 5
$define_test 5 $line_1 'nhits float_equals EV+$KEV_NPM:249.5 9999.;'

* A do nothing filter that allows autosno to filter maintenece events.
$enable_test 9
$define_test 9 $line_1 'one equals 1.0;'

define event_loop
   call INP
   call UPK
   call CAL
   call FLT($PERM_BANK_CUT)
   if_not_ok goto WRITE_OUT

EVENT_REJECT:
   call FLT(9)            * User defined cut. Default passes all.
   if_not_ok quit_event
   call FLT($JUNK_CUT)    * Junk. We still want to keep NCD events.
   if_not_ok quit_event
   call FLT(1)            * Primary NHIT Cut.
   if_ok goto FITTERS
   call FLT($PRESCALE_2_TAG)
   if_ok goto FITTERS
   goto WRITE_OUT

FITTERS:
   call FLT($ITC_CUT)
   if_not_ok goto WRITE_OUT
   call FTT
   call FTQ
   call FTP
   call RSP($KFT_FTP)

WRITE_OUT:
   call PRU
   call OUT
   quit_event
end_def

*Pruning Control
$prune_pix $drop
$prune_px $drop
$prune_pf $drop
$prune_nesg $drop
$prune_nes $drop
$prune_nemg $drop

* Run with the Database
@@run_snodb

!endfile member=fit_neutrino
""".strip()

class MyTemplate(string.Template):
    delimiter = '@'

def splitext(path):
    """
    Like os.path.splitext() except it returns the full extension if the
    filename has multiple extensions, for example:

        splitext('foo.tar.gz') -> 'foo', '.tar.gz'
    """
    full_root, full_ext = os.path.splitext(path)
    while True:
        root, ext = os.path.splitext(full_root)
        if ext:
            full_ext = ext + full_ext
            full_root = root
        else:
            break

    return full_root, full_ext

def GetRunRange(run, num=1000):
    """
    Returns the run range string for ANXX neutrino files.

    Example:

        GetRunRange(10000,1000) -> "10000-10999"

    Comes from GetRunRange() in autosno/lib/FileUtil.pm.
    """
    base_run = int(run//num)*num
    return "%i-%i" % (base_run, base_run + num - 1)

def get_anxx_neutrino_filename(run, dir):
    """
    Returns the ANXX neutrino file with the highest pass number for a given
    run. The files should all be located in dir.

    FIXME: Is it OK to just get the file with the highest pass number?
    Technically this info should come from autosno, but it's not clear exactly
    how it's getting the pass number in the function GetPassNum() in
    AutoSNO_Module.pm.
    """
    filename_re = "anxx_nu_(\d+)_p(\d+).dat"
    p = re.compile(filename_re)

    filenames = []
    for filename in os.listdir(dir):
        match = p.match(filename)

        if int(match.group(1)) != run:
            continue

        filenames.append((match.group(2),filename))

    return sorted(filenames)[-1][1]

def get_anxx_pca_next(run,dir):
    """
    Returns the next ANXX PCA file after `run`.

    See Get_ANxx_PCA_Next() in autosno/lib/GetTitles.pm.
    """
    filename_re = "anxx_pca_(\d+)_p(\d+).dat"
    p = re.compile(filename_re)

    filenames = []
    for filename in os.listdir(dir):
        match = p.match(filename)

        if match is None:
            continue

        if int(match.group(1)) < run:
            continue

        filenames.append((match.group(1),match.group(2),filename))

    return sorted(filenames)[-1][2]


if __name__ == '__main__':
    import argparse
    import sys
    import os
    from subprocess import check_call

    parser = argparse.ArgumentParser("reprocess zdabs")
    parser.add_argument("filenames", nargs="+", help="filenames of zdabs")
    parser.add_argument("--dir", default=None, help="directory to store reprocessed zdab files")
    parser.add_argument("--min-nhit", type=int, default=None, help="minimum nhit value")
    parser.add_argument("--lower-nhit", type=float, default=18.5, help="lower value for primary nhit cut")
    parser.add_argument("--upper-nhit", type=float, default=9999, help="upper value for primary nhit cut")
    parser.add_argument("--suffix", default="reprocessed", help="suffix appended to output filename")
    args = parser.parse_args()

    filename_re = "(?:SNOCR_|SNO)(\d+)_\w+.x?zdab"

    p = re.compile(filename_re)

    template = MyTemplate(SNOMAN_TEMPLATE)

    for filename in args.filenames:
        head, tail = os.path.split(filename)
        match = p.match(tail)

        if not match:
            print("Unable to parse filename '%s'" % tail)
            sys.exit(1)

        try:
            run = int(match.group(1))
        except Exception as e:
            print("Unable to convert run number to int: %s" % str(e))
            sys.exit(1)

        root, ext = splitext(tail)

        output = "%s_%s%s" % (root,args.suffix,ext)

        if args.dir:
            output = os.path.join(args.dir,output)

        cmd_filename = "%s.cmd" % root

        dqxx_filename = "/sno/mcprod/dqxx/DQXX_%010i.dat" % run
        anxx_neutrino_directory = "/sno/mcprod/anxx/titles/neutrino/%s" % GetRunRange(run)
        anxx_neutrino_filename = get_anxx_neutrino_filename(run, anxx_neutrino_directory)
        anxx_neutrino_filename = os.path.join(anxx_neutrino_directory,anxx_neutrino_filename)
        anxx_pca_directory = "/sno/mcprod/anxx/titles/pca"
        anxx_pca_filename = get_anxx_pca_next(run, anxx_pca_directory)
        anxx_pca_filename = os.path.join(anxx_pca_directory,anxx_pca_filename)

        if args.min_nhit:
            input_nhit_thresh = "$input_nhit_thresh %i" % args.min_nhit
        else:
            input_nhit_thresh = ''

        with open(cmd_filename, "w") as f:
            f.write(template.safe_substitute(input=filename,
                                             output=output,
                                             dqxx=dqxx_filename,
                                             anxx_neutrino=anxx_neutrino_filename,
                                             anxx_pca=anxx_pca_filename,
                                             settings_file=GetSettingsFile(run),
                                             prescale_setting=1.0, # no idea what this is
                                             lower_nhit = args.lower_nhit,
                                             upper_nhit = args.upper_nhit,
                                             seed=0,
                                             seed2=0,
                                             clock_force_no_fix=get_clock_force_no_fix_string(run),
                                             input_nhit_thresh=input_nhit_thresh))

        check_call(["snoman.exe","-c",cmd_filename])