from util import BaseCheckList
import sklearn.covariance
import scipy.stats
import numpy as np

thermal_checks = BaseCheckList()

def thermal_correlation_check(frame, parameters):
    return frame

@thermal_checks.check(depends=['HBBbottomTemp','HBBtopTemp','HBBapexTemp'])
def hbb_covariance_check(frame, parameters):
    hbbb_mean = parameters.get('HBBbottomTemp_mean', 333.101)
    hbbt_mean = parameters.get('HBBtopTemp_mean', 332.993)
    hbba_mean = parameters.get('HBBapexTemp_mean', 332.818)

    hbb_contamination = parameters.get('hbb_contamination', .01)
    envelope = sklearn.covariance.EllipticEnvelope(store_precision=False, contamination=hbb_contamination)
    envelope.location_ = np.array([hbbb_mean, hbbt_mean, hbba_mean])
    envelope.covariance_ = np.array([[ 0.0179306 ,  0.0164418 ,  0.01817402],
                                     [ 0.0164418 ,  0.01898899,  0.02326341],
                                     [ 0.01817402,  0.02326341,  0.02968953]])
    distances = envelope.mahalanobis(frame[['HBBbottomTemp','HBBtopTemp','HBBapexTemp']]) ** .5
    frame['hbb_covariance_check'] = (distances > 6) * 1
    return frame