from util import BaseCheckList, invalidate_records, annotate_all, update_variable_qc
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
    detector_temp_too_high = (frame['detectorTemp'] > 90)
    frame['detector_check'] = detector_temp_too_high * 1
    frame['qc_detectorTemp'] = detector_temp_too_high * 1
    annotate_all(frame, detector_temp_too_high, 'detector temperature too high')
    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
    hbbb_too_low = (frame['HBBbottomTemp'] - 333) < -2
    hbba_too_low = (frame['HBBapexTemp'] - 333) < -2
    hbbt_too_low = (frame['HBBtopTemp'] - 333) < -2
    hbbb_too_high = (frame['HBBbottomTemp'] - 333) > 2
    hbba_too_high = (frame['HBBapexTemp'] - 333) > 2
    hbbt_too_high = (frame['HBBtopTemp'] - 333) > 2
    hbbb_problem = hbbb_too_low | hbbb_too_high
    hbba_problem = hbba_too_low | hbba_too_high
    hbbt_problem = hbbt_too_low | hbbt_too_high

    # Record qc for each thermistor
    # qc variables are probabilites between 0 and 1
    variable_qcs = pd.DataFrame({
        'qc_HBBbottomTemp':hbbb_problem * 1,
        'qc_HBBapexTemp' : hbba_problem * 1,
        'qc_HBBtopTemp' : hbbt_problem * 1
    }, index=frame.index)
    frame = update_variable_qc(frame, variable_qcs)

    # Compute overall BB quality
    frame['hbb_thermistor_check'] = (hbbb_problem | hbba_problem | hbbt_problem) * 1
    annotate_all(frame, hbbb_too_low, 'HBB bottom temperature too low')
    annotate_all(frame, hbbt_too_low, 'HBB top temperature too low')
    annotate_all(frame, hbba_too_low, 'HBB apex temperature too low')
    annotate_all(frame, hbbb_too_high, 'HBB bottom temperature too high')
    annotate_all(frame, hbbt_too_high, 'HBB top temperature too high')
    annotate_all(frame, hbba_too_high, 'HBB apex temperature too high')
    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')],
        'qc_notes':''
    }), {})
    assert all(frame['hbb_thermistor_check'] == [1,0,1])
    assert all('HBB {} temperature too low'.format(x) in frame.iloc[0].qc_notes.split(',') for x in ['bottom','apex','top'])
    assert all('HBB {} temperature too high'.format(x) in frame.iloc[2].qc_notes.split(',') for x in ['bottom','apex','top'])

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')],
        'qc_notes':''
    }), {})
    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')],
        'qc_notes':''
    }), {})
    assert all(frame['detector_check'] == [0,1])
    assert frame.iloc[1].qc_notes == 'detector temperature too high'