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):
    interval_seconds = frame['datetime'].diff().astype(np.int64) / 1e9
    hbbb_diff = frame['HBBbottomTemp'].diff() / interval_seconds
    hbba_diff = frame['HBBapexTemp'].diff() / interval_seconds
    hbbt_diff = frame['HBBtopTemp'].diff() / interval_seconds
    hsr = parameters.get('hbb_stable_rate', .002)
    hbbb_diff_problem = abs(hbbb_diff.fillna(0)) > hsr
    hbba_diff_problem = abs(hbba_diff.fillna(0)) > hsr
    hbbt_diff_problem = abs(hbbt_diff.fillna(0)) > hsr

    variable_qcs = pd.DataFrame({
        'qc_HBBbottomTemp':hbbb_diff_problem * 1,
        'qc_HBBapexTemp' : hbba_diff_problem * 1,
        'qc_HBBtopTemp' : hbbt_diff_problem * 1
    }, index=frame.index)
    frame = update_variable_qc(frame, variable_qcs)

    frame['hbb_stable_check'] = (hbbb_diff_problem | hbbt_diff_problem | hbbt_diff_problem) * 1
    annotate_all(frame, hbbb_diff_problem, 'HBB bottom temperature not stable')
    annotate_all(frame, hbba_diff_problem, 'HBB apex temperature not stable')
    annotate_all(frame, hbbt_diff_problem, 'HBB top temperature not stable')
    frame = invalidate_records(frame, 'hbb_stable_check')
    return frame

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


#### TESTS ####

def test_hbb_stable_check():
    # Check with rates of .0016
    dummy_data = pd.DataFrame({
            'HBBbottomTemp':np.arange(0,2,.1),
            'HBBapexTemp':np.arange(0,2,.1),
            'HBBtopTemp':np.arange(0,2,.1),
            'datetime':pd.date_range('1/1/2000', periods=20, freq='60s'),
            'qc_notes':'',
            'qc_percent':0,
            'sceneMirrorPosition':ord('S')
        })
    frame = hbb_stable_check(
        dummy_data, {'hbb_stable_rate':.002}
    )
    assert all(frame['hbb_stable_check'] == 0)
    
    dummy_data.ix[10:,'HBBbottomTemp'] = np.arange(1,3,.2)
    dummy_data.ix[10:,'HBBtopTemp'] = np.arange(1,3,.2)
    dummy_data.ix[10:,'HBBbottomTemp'] = np.arange(1,3,.2)

    frame = hbb_stable_check(
        dummy_data, {'hbb_stable_rate':.002}
    )
    assert all(frame.ix[:10, 'hbb_stable_check'] == 0)
    assert all(frame.ix[11:, 'hbb_stable_check'] == 1)

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'