from util import BaseCheckList, annotate_all, _compute_robust_zscore, invalidate_records, update_variable_qc
import pandas as pd
import numpy as np
import scipy.stats

radiometric_checks = BaseCheckList()

@radiometric_checks.check(depends=['skyViewImaginaryRadiance2510_2515'])
def imaginary_radiance_check(frame, parameters):
    threshold = parameters.get('imaginary_radiance_threshold', 1)
    imaginary_radiance_problem = abs(frame.skyViewImaginaryRadiance2510_2515) > threshold
    frame['imaginary_radiance_check'] = imaginary_radiance_problem * 1
    annotate_all(frame, imaginary_radiance_problem, 'sky view imaginary radiance out of range')
    frame = invalidate_records(frame, 'imaginary_radiance_check')
    return frame

@radiometric_checks.check(depends=['HBBviewStdDevRadiance985_990','LW_HBB_NEN','SW_HBB_NEN'])
def hbb_radiance_check(frame, parameters):
    # Std dev, nen, lw, sw
    hbb_std_dist = scipy.stats.chi2.fit(frame.HBBviewStdDevRadiance985_990)
    _, hbb_std_dist_hi = scipy.stats.chi2.interval(.995, *hbb_std_dist)
    hbb_std_dist_problem = frame.HBBviewStdDevRadiance985_990 > hbb_std_dist_hi

    lw_hbb_nen_dist = scipy.stats.chi2.fit(frame.LW_HBB_NEN)
    _, lw_hbb_nen_hi = scipy.stats.chi2.interval(.995, *lw_hbb_nen_dist)
    lw_hbb_nen_problem = frame.LW_HBB_NEN > lw_hbb_nen_hi

    sw_hbb_nen_dist = scipy.stats.chi2.fit(frame.SW_HBB_NEN)
    _, sw_hbb_nen_hi = scipy.stats.chi2.interval(.995, *sw_hbb_nen_dist)
    sw_hbb_nen_problem = frame.SW_HBB_NEN > sw_hbb_nen_hi

    frame['hbb_radiance_check'] = (hbb_std_dist_problem | lw_hbb_nen_problem | sw_hbb_nen_problem) * 1
    annotate_all(frame, hbb_std_dist_problem, 'HBB radiance Std.Dev. too high')
    annotate_all(frame, lw_hbb_nen_problem, 'LW HBB NEN too high')
    annotate_all(frame, sw_hbb_nen_problem, 'SW HBB NEN too high')
    frame = invalidate_records(frame, 'hbb_radiance_check')
    return frame

@radiometric_checks.check(depends=['LWresponsivity','SWresponsivity'], updates=['LWresponsivity', 'SWresponsivity'])
def responsivity_check(frame, parameters):
    # lw, sw
    lw_zscore = _compute_robust_zscore(frame['LWresponsivity'], 50)
    sw_zscore = _compute_robust_zscore(frame['SWresponsivity'], 50)
    lw_problem = abs(lw_zscore) > 6
    sw_problem = abs(sw_zscore) > 6
    variable_qcs = pd.DataFrame({
        'qc_LWresponsivity':lw_problem * 1,
        'qc_SWresponsivity':sw_problem * 1
    })
    frame['responsivity_check'] = (lw_problem | sw_problem) * 1
    frame = update_variable_qc(frame, variable_qcs)
    annotate_all(frame, lw_problem, 'LW responsivity outlier')
    annotate_all(frame, sw_problem, 'SW responsivity outlier')
    frame = invalidate_records(frame, 'responsivity_check')
    return frame