from util import BaseCheckList, invalidate_records
import pandas as pd
import numpy as np

def detector_check(frame, parameters):
    """
    Check that the detector temp is in range
    """
    if 'detectorTemp' not in frame.columns:
        return frame
    frame['detector_check'] = (frame['detectorTemp'] > 90) * 1
    frame = invalidate_records(frame, 'detector_check')
    return frame

def hbb_thermistor_check(frame, parameters):
    """
    Check that all HBB thermistor temps are in range
    """
    if not np.in1d(['HBBbottomTemp','HBBapexTemp','HBBtopTemp'], frame.columns).all():
        return frame
    frame['hbb_thermistor_check'] = (
        (abs(frame['HBBbottomTemp'] - 333) > 2) |
        (abs(frame['HBBapexTemp'] - 333) > 2) |
        (abs(frame['HBBtopTemp'] - 333) > 2)
    ) * 1
    frame = invalidate_records(frame, 'hbb_thermistor_check')
    return frame

def hbb_stable_check(frame, parameters):
    return frame

class CheckList(BaseCheckList):
    checks = [detector_check, hbb_thermistor_check, hbb_stable_check]


#### TESTS ####

def test_hbb_thermistor_check():
    frame = hbb_thermistor_check(pd.DataFrame({
        'HBBbottomTemp':[300,333,336],
        'HBBapexTemp':[300,333,336],
        'HBBtopTemp':[300,333,336],
        'sceneMirrorPosition':[ord('H'), ord('A'), ord('S')]
    }), {})
    assert all(frame['hbb_thermistor_check'] == [1,0,1])

def test_hbb_thermistor_check2():
    frame = hbb_thermistor_check(pd.DataFrame({
        'HBBbottomTemp':[300,333,333],
        'HBBapexTemp':[300,333,333],
        'HBBtopTemp':[300,333,333],
        'sceneMirrorPosition':[ord('H'), ord('A'), ord('S')]
    }), {})
    assert all(frame['hbb_thermistor_check'] == [1,0,1])

def test_detector_check():
    frame = detector_check(pd.DataFrame({
        'detectorTemp':[50,100],
        'sceneMirrorPosition':[ord('H'), ord('A')]
    }), {})
    assert all(frame['detector_check'] == [0,1])