#
# Filter_combined_ZYP.py
#
# ---------------------------------------------------------------------------
# This is a combined script for event filter by either sample log or time.
# The two ways of conducting event filter is not compatible with each other
# and users can choose which one to use by setting the `filter_by_time`
# variable in the 'Input section' below to either `True` (to enable filter by
# time) or `False` (to enable filter by sample log).
# ---------------------------------------------------------------------------
#
# ----------------------------
# Yuanpeng Zhang @ Nov-08-2024
# ----------------------------
#
from mantid.simpleapi import *
from datetime import datetime, timedelta

###############################################################################
# ------------------------------ Input section --------------------------------
###############################################################################

# Inputs for runs information.
title = 'output_th'
IPTS = 32368
run = 1532451
background = None
BackgroundScale = 1
vanadium = 1530874  # Run number or `None`
vanadium_IPTS = 23858
out_dir = "/SNS/users/y8z/Temp"

# Inputs for runs control.
normaliseBy = 'Monitor'  # One on (None, Monitor, Time)
units = 'Theta'  # One of (Theta, ElasticQ, ElasticDSpacing)
Binning = '5,125,1200'  # Min,Max,Number_of_bins
log_binning = False

# Inputs concerning event filter. First, choose to filter by time or sample log,
# depending on the value of `filter_by_time`.
#
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# N.B. The option to filter by time and sample log is NOT compatible with each other.
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
#
filter_by_time = True

# Inputs concerning event filter by time.
if filter_by_time:
    start_time = '0'
    stop_time = '900'
    time_interval = '500'
    unit_of_time = 'Seconds'
    time_tolerance = 0

# Inputs concerning event filter by sample log.
LogName = 'HB2C:SE:SampleTemp'
LogValueInterval = 5
MinimumLogValue = None
MaximumLogValue = None
log_val_tolerance = LogValueInterval / 2.0
log_change_direction = 'Both'

###############################################################################
# ---------------------------- Execution section ------------------------------
###############################################################################

# Load vandium if not already loaded.
vanadium_ws = 'HB2C_{}'.format(vanadium)

if vanadium_ws not in mtd:
    LoadWAND(IPTS=vanadium_IPTS,
             RunNumbers=vanadium,
             OutputWorkspace=vanadium_ws)

# Load background if needed and if not already loaded.
if background is not None:
    background_ws = 'HB2C_{}'.format(background)
    if background_ws not in mtd:
        LoadEventNexus(Filename='/HFIR/HB2C/IPTS-{}/nexus/HB2C_{}.nxs.h5'.format(IPTS, background),
                       OutputWorkspace=background_ws,
                       LoadMonitors=True)
else:
    background_ws = None

# Load data if not already loaded
ws = 'HB2C_{}'.format(run)

if ws not in mtd:
    LoadEventNexus(Filename='/HFIR/HB2C/IPTS-{}/nexus/HB2C_{}.nxs.h5'.format(IPTS, run),
                   OutputWorkspace=ws,
                   LoadMonitors=True)
    # Mask detectors to be the same as vanadium
    MaskDetectors(ws, MaskedWorkspace=vanadium_ws)

# Filter events
if filter_by_time:
    GenerateEventsFilter(InputWorkspace=ws,
                         OutputWorkspace='filter',
                         InformationWorkspace='info',
                         StartTime=start_time,
                         StopTime=stop_time,
                         TimeInterval=time_interval,
                         TimeTolerance=time_tolerance,
                         UnitOfTime=unit_of_time)
    FilterEvents(InputWorkspace=ws,
                 SplitterWorkspace='filter',
                 OutputWorkspaceBaseName=ws + '_filtered',
                 InformationWorkspace='info',
                 GroupWorkspaces=True,
                 FilterByPulseTime=True,
                 OutputWorkspaceIndexedFrom1=True)
    FilterEvents(InputWorkspace=ws + '_monitors',
                 SplitterWorkspace='filter',
                 OutputWorkspaceBaseName=ws + '_filtered_mon',
                 InformationWorkspace='info',
                 GroupWorkspaces=True,
                 FilterByPulseTime=True,
                 SpectrumWithoutDetector='Skip only if TOF correction',
                 OutputWorkspaceIndexedFrom1=True)
else:
    GenerateEventsFilter(InputWorkspace=ws,
                         OutputWorkspace='filter',
                         InformationWorkspace='info',
                         LogName=LogName,
                         LogValueInterval=LogValueInterval,
                         MinimumLogValue=MinimumLogValue,
                         MaximumLogValue=MaximumLogValue,
                         LogValueTolerance=log_val_tolerance,
                         FilterLogValueByChangingDirection=log_change_direction
                         )
    FilterEvents(InputWorkspace=ws,
                 SplitterWorkspace='filter',
                 OutputWorkspaceBaseName=ws + '_filtered',
                 InformationWorkspace='info',
                 GroupWorkspaces=True,
                 FilterByPulseTime=True,
                 OutputWorkspaceIndexedFrom1=True)
    FilterEvents(InputWorkspace=ws + '_monitors',
                 SplitterWorkspace='filter',
                 OutputWorkspaceBaseName=ws + '_filtered_mon',
                 InformationWorkspace='info',
                 GroupWorkspaces=True,
                 FilterByPulseTime=True,
                 SpectrumWithoutDetector='Skip only if TOF correction',
                 OutputWorkspaceIndexedFrom1=True)

# Set the monitor count on filtered WS
for n in range(mtd[ws + '_filtered'].getNumberOfEntries()):
    AddSampleLog(mtd[ws + '_filtered'].getItem(n),
                 LogName="gd_prtn_chrg",
                 LogType='Number',
                 NumberType='Double',
                 LogText=str(mtd[ws + '_filtered_mon'].getItem(n).getNumberEvents()))

if background is not None:
    AddSampleLog(background_ws, LogName="gd_prtn_chrg",
                 LogType='Number',
                 NumberType='Double',
                 LogText=str(mtd[background_ws + '_monitors'].getNumberEvents()))

# Run powder diffraction reduction
xmin, xmax, bins = Binning.split(',')
WANDPowderReduction(ws + '_filtered',
                    CalibrationWorkspace=vanadium_ws,
                    BackgroundWorkspace=background_ws,
                    BackgroundScale=BackgroundScale,
                    XMin=xmin,
                    XMax=xmax,
                    NumberBins=bins,
                    Target=units,
                    LogBinning=log_binning,
                    NormaliseBy=normaliseBy,
                    OutputWorkspace=title)

run_start = mtd[ws].getRun().getLogData("run_start").value
run_start = run_start.split(".")[0] + "." + run_start.split(".")[1][:6]
starting_time = datetime.strptime(run_start, '%Y-%m-%dT%H:%M:%S.%f')
for count in range(mtd[ws + '_filtered'].getNumberOfEntries()):
    out_wksp = mtd[title].getItem(count)
    if filter_by_time:
        tname_part = "{:d}".format(int(start_time) + count * int(time_interval))
        out_wksp_new = f"{str(out_wksp)}_{tname_part}s"
        RenameWorkspace(
            InputWorkspace=out_wksp,
            OutputWorkspace=out_wksp_new
        )
        fn_tmp = f"{title}_{tname_part}s_"
        duration = count * int(time_interval)
        stop_time = starting_time + timedelta(seconds=duration)
        stop_time = stop_time.strftime('%Y-%m-%dT%H:%M:%S.%f')
        fn_tmp += f"{stop_time}.dat"
        SaveAscii(
            InputWorkspace=out_wksp_new,
            Filename=os.path.join(
                out_dir,
                fn_tmp
            )
        )
    else:
        SaveAscii(
            InputWorkspace=out_wksp,
            Filename=os.path.join(
                out_dir,
                f"{title}_{item}.dat"
            )
        )