#
# Filter_by_time_stamp_ZYP.py
#
# ---------------------------------------------------------------------------
# Script for filtering events based on time stamp
# ---------------------------------------------------------------------------
#
# ----------------------------
# Yuanpeng Zhang @ Nov-10-2024
# ----------------------------
#
from mantid.simpleapi import *

###############################################################################
# ------------------------------ 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 for time stamp for filtering
# 'time_interval = 100' here means all data falling in between the start
# and stop time will be filtered into a single dataset.
start_time = '2024-10-23T03:44:06'
stop_time = '2024-10-23T03:52:26'
time_interval = '100'
unit_of_time = 'Percent'
time_tolerance = 0

###############################################################################
# ---------------------------- 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
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)

# 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)

out_wksp = mtd[title]
fn_tmp = f"{title}_"
fn_tmp += f"{start_time}_{stop_time}.dat"
SaveAscii(
    InputWorkspace=out_wksp,
    Filename=os.path.join(
        out_dir,
        fn_tmp
    )
)