#!/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

***Required for the RAA
*Return end of file to initiate analyse/restore phases
$return_eof $enabled
*raa.dat must be loaded through titles as TRAA is not stored in the database.
titles raa.dat
* Memory is tight and we are not using the reduced buffer.
titles raa_request_ncdcorr.dat
*RAA buffer file location. Warning this file can be very big.
FILE RAA 1 raa.tmp 
*Delete PMT style banks before storing data in the buffer. Reduces buffer size.
$raa_delete_pmts $yes

* 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

*FTI
* FTI titles files are not in the database yet
*Load FTI angular PDF titles file
titles fti_ang_pdfs.dat

*Load charge probability titles file
titles charge_photoelectrons.dat

titles fitter_fti.dat

* Low nhit prescale
$flt_set_prescale_2 0.01

* 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
   if_eof goto EOF
   call UPK
   call CAL
   call RAA($raa_hca,$raa_qrc,$raa_ncd_corr_tag)
   quit_event
EOF:
   call RAA
   call FLT($PERM_BANK_CUT)
   if_not_ok goto WRITE_OUT
   call UPK

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 goto JUNK_SELECT:
   call FLT(1)            * Primary NHIT Cut.
   if_ok goto FITTERS
   call FLT(4)            * NCD Cut.
   if_ok goto FITTERS
   call FLT($PRESCALE_2_TAG)
   if_ok goto FITTERS
   quit_event

JUNK_SELECT:
   call FLT(4)            * NCD Cut.
   if_ok goto WRITE_OUT
   quit_event

FITTERS:
   call CAL
   call FLT($ITC_CUT)
   if_not_ok goto WRITE_OUT
   call FLT(21)
   if_not_ok goto XFTU
   call FTN
XFTU:
   call FTU
   call FTT
   call FTQ
   call FTP
   call FTG

SKIP_FITTERS:
   call FLT(5)
   if_not_ok goto WRITE_OUT
   call FTR
   call FTM
   call FTI
   call FTY
   call FTZ

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("--suffix", default="reprocessed", help="suffix appended to output filename")
    args = parser.parse_args()

    filename_re = "SNOCR_(\d+)_\d+_p\d+.xzdab"
    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.xzdab" % (root,args.suffix)

        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,
                                             lower_nhit = 18.5,
                                             upper_nhit = 9999.0,
                                             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])