#/ Type: DRS
#/ Name: S Isotopes
#/ Authors: Bence Paul
#/ Description: A sulfur isotopes DRS
#/ References: None
#/ Version: 0.1
#/ Contact: N/A

from iolite import QtGui
from iolite import QtCore

import numpy as np


def runDRS():

    drs.message("Starting S isotopes DRS...")
    drs.progress(0)

    # Get settings
    settings = drs.settings()
    print(settings)

    if settings["IndexChannel"] not in data.timeSeriesNames(data.Input):
        IoLog.errorWithOrigin("Index channel not found. DRS cannot continue", "Sr Isotopes CaAr REE DRS")
        drs.message("DRS did not finish. Please check Messages")
        drs.progress(100)
        drs.finished()
        return

    indexChannel = data.timeSeries(settings["IndexChannel"])
    rmName = settings["ReferenceMaterial"]
    maskOption = settings["Mask"]
    maskChannel = data.timeSeries(settings["MaskChannel"])
    cutoff = settings["MaskCutoff"]
    trim = settings["MaskTrim"]

    # Create debug messages for the settings being used
    IoLog.debug("indexChannelName = %s" % indexChannel.name)
    IoLog.debug("Masking data  = True" if maskOption else "Masking data  = False")
    IoLog.debug("maskChannelName = %s" % maskChannel.name)
    IoLog.debug("maskCutoff = %f" % cutoff)
    IoLog.debug("maskTrim = %f" % trim)

    # Setup index time
    drs.message("Setting up index time...")
    drs.progress(5)
    drs.setIndexChannel(indexChannel)

    # Setup the mask
    if maskOption:
        drs.message("Making mask...")
        drs.progress(10)
        mask = drs.createMaskFromCutoff(maskChannel, cutoff, trim)
        data.createTimeSeries('mask', data.Intermediate, indexChannel.time(), mask)
    else:
        mask = np.ones_like(indexChannel.data())
        data.createTimeSeries('mask', data.Intermediate, indexChannel.time(), mask)

    # Interp onto index time and baseline subtract
    drs.message("Interpolating onto index time and baseline subtracting...")

    allInputChannels = data.timeSeriesList(data.Input)
    blGrp = None

    if len(data.selectionGroupList(data.Baseline)) < 1:
        IoLog.error("Could not find a baseline group. Combined Sr DRS cannot proceed...")
        drs.message("Please select some baseline selections")
        drs.progress(100)
        drs.finished()
        return

    if len(data.selectionGroupList(data.Baseline)) > 1:
        IoLog.error("There are more than one baseline groups. Combined Sr DRS cannot proceed...")
        drs.message("DRS did not finish. Please check Messages")
        drs.progress(100)
        drs.finished()
        return
    elif len(data.selectionGroupList(data.Baseline)) < 1:
        IoLog.error("No baseline groups found. Have you selected baselines yet?")
        drs.message("DRS did not finish. Please check Messages")
        drs.progress(100)
        drs.finished()
        return
    else:
        blGrp = data.selectionGroupList(data.Baseline)[0]

    drs.baselineSubtract(blGrp, allInputChannels, mask, 15, 35)

    drs.message("Calculating ratios...")
    drs.progress(40)

    # Get baselines subtracted channels here
    S32_CPS = data.timeSeries("S32_CPS").data()
    S33_CPS = data.timeSeries("S33_CPS").data()
    S34_CPS = data.timeSeries("S32_CPS").data()

    S33_32_raw = S33_CPS / S32_CPS
    S34_32_raw = S34_CPS / S32_CPS

    data.createTimeSeries('S33_32_raw', data.Intermediate, indexChannel.time(), S33_32_raw)
    data.createTimeSeries('S34_32_raw', data.Intermediate, indexChannel.time(), S34_32_raw)

    TotalSBeam = S32_CPS + S33_CPS + S34_CPS
    data.createTimeSeries('Total S Beam', data.Intermediate, indexChannel.time(), TotalSBeam)

    drs.message('Calculating reference material corrected results...')
    drs.progress(70)
    
    data.updateResults()
    
    # Now check if there are selections for the reference standard, and if so, generate standard-normalised ratios
    try:
        StdSpline_S33_32 = data.spline(rmName, 'S33_32_raw').data()
    except:
        IoLog.errorWithOrigin('The S Isotopes DRS requires Ref Material selections to proceed.', 'S isotope DRS')
        drs.message('DRS did not finish. Please check Messages')
        drs.progress(100)
        drs.finished()
        return

    try:
        StdSpline_S34_32 = data.spline(rmName, 'S34_32_raw').data()
    except:
        IoLog.errorWithOrigin('Could not get the S34_32_raw spline', 'S isotope DRS')
        drs.message('DRS did not finish. Please check Messages')
        drs.progress(100)
        drs.finished()
        return

    D33S =  1000. * (S33_32_raw - StdSpline_S33_32) / StdSpline_S33_32
    D34S =  1000. * (S34_32_raw - StdSpline_S34_32) / StdSpline_S34_32

    try:
        StdSpline_S32 = data.spline(rmName, 'S32_CPS').data()
    except:
        IoLog.errorWithOrigin('Could not get the S32_CPS spline', 'S isotope DRS')
        drs.message('DRS did not finish. Please check Messages')
        drs.progress(100)
        drs.finished()
        return

    R32_int = S32_CPS / StdSpline_S32


    # Try to get the standard values from the RM file
    try:
        StdValue_D33S = data.referenceMaterialData(rmName)["D33S"].value()
        StdValue_D34S = data.referenceMaterialData(rmName)["D34S"].value()
    except:
        IoLog.errorWithOrigin('Could not get one of the RM D values', 'S isotope DRS')
        drs.message('DRS did not finish. Please check Messages')
        drs.progress(100)
        drs.finished()
        return

    D33S_VCDT = D33S + StdValue_D33S
    D34S_VCDT = D34S + StdValue_D34S

    data.createTimeSeries('StdSpline_S32', data.Output, indexChannel.time(), StdSpline_S32)
    data.createTimeSeries('D33S', data.Output, indexChannel.time(), D33S)
    data.createTimeSeries('D34S', data.Output, indexChannel.time(), D34S)
    data.createTimeSeries('R32_int', data.Output, indexChannel.time(), R32_int)
    data.createTimeSeries('D33S_VCDT', data.Output, indexChannel.time(), D33S_VCDT)
    data.createTimeSeries('D34S_VCDT', data.Output, indexChannel.time(), D34S_VCDT)

    drs.message("Finished!")
    drs.progress(100)
    drs.finished()


def settingsWidget():
    """
    This function puts together a user interface to configure the DRS.

    It is important to have the last line of this function call:
    drs.setSettingsWidget(widget)
    """

    widget = QtGui.QWidget()
    formLayout = QtGui.QFormLayout()
    widget.setLayout(formLayout)

    timeSeriesNames = data.timeSeriesNames(data.Input)
    defaultChannelName = ""
    if 'TotalBeam' in timeSeriesNames:
        defaultChannelName = 'TotalBeam'
    elif timeSeriesNames:
        defaultChannelName = timeSeriesNames[0]

    drs.setSetting("IndexChannel", "S34")
    drs.setSetting("ReferenceMaterial", "STD_Py")
    drs.setSetting("Mask", True)
    drs.setSetting("MaskChannel", defaultChannelName)
    drs.setSetting("MaskCutoff", 0.05)
    drs.setSetting("MaskTrim", 0.0)

    settings = drs.settings()

    indexComboBox = QtGui.QComboBox(widget)
    indexComboBox.addItems(timeSeriesNames)
    indexComboBox.setCurrentText(settings["IndexChannel"])
    indexComboBox.textActivated.connect(lambda t: drs.setSetting("IndexChannel", t))
    formLayout.addRow("Index channel", indexComboBox)

    def updateIndexCombo():
        timeSeriesNames = data.timeSeriesNames(data.Input)
        indexComboBox.clear()
        indexComboBox.addItems(timeSeriesNames)

    data.dataChanged.connect(updateIndexCombo)

    rmComboBox = QtGui.QComboBox(widget)
    rmNames = data.selectionGroupNames(data.ReferenceMaterial)
    rmComboBox.addItems(rmNames)
    if settings["ReferenceMaterial"] in rmNames:
        rmComboBox.setCurrentText(settings["ReferenceMaterial"])
        drs.setSetting("ReferenceMaterial", settings["ReferenceMaterial"])

    rmComboBox.textActivated.connect(lambda t: drs.setSetting("ReferenceMaterial", t))
    formLayout.addRow("Reference material", rmComboBox)

    def updateRMCombo():
        rmNames = data.selectionGroupNames(data.ReferenceMaterial)
        rmComboBox.clear()
        rmComboBox.addItems(rmNames)

    data.selectionGroupsChanged.connect(updateRMCombo)

    verticalSpacer = QtGui.QSpacerItem(20, 40, QtGui.QSizePolicy.Minimum, QtGui.QSizePolicy.Minimum)
    formLayout.addItem(verticalSpacer)

    maskCheckBox = QtGui.QCheckBox(widget)
    maskCheckBox.setChecked(settings["Mask"])
    maskCheckBox.toggled.connect(lambda t: drs.setSetting("Mask", bool(t)))
    formLayout.addRow("Mask", maskCheckBox)

    maskComboBox = QtGui.QComboBox(widget)
    maskComboBox.addItems(data.timeSeriesNames(data.Input))
    maskComboBox.setCurrentText(settings["MaskChannel"])
    maskComboBox.currentTextChanged.connect(lambda t: drs.setSetting("MaskChannel", t))
    formLayout.addRow("Mask channel", maskComboBox)

    maskLineEdit = QtGui.QLineEdit(widget)
    maskLineEdit.setText(settings["MaskCutoff"])
    maskLineEdit.textChanged.connect(lambda t: drs.setSetting("MaskCutoff", float(t)))
    formLayout.addRow("Mask cutoff", maskLineEdit)

    maskTrimLineEdit = QtGui.QLineEdit(widget)
    maskTrimLineEdit.setText(settings["MaskTrim"])
    maskTrimLineEdit.textChanged.connect(lambda t: drs.setSetting("MaskTrim", float(t)))
    formLayout.addRow("Mask trim", maskTrimLineEdit)

    # Restore settings
    try:
        settings = drs.settings()
        print('Restoring settings...')
        print(settings)
        rmComboBox.setCurrentText(settings["ReferenceMaterial"])
    except KeyError:
        pass

    drs.setSettingsWidget(widget)