import sys
sys.path.insert(1,"../../")
import h2o
from tests import pyunit_utils
from h2o.estimators.gbm import H2OGradientBoostingEstimator



def metric_accessors():

    cars = h2o.import_file(path=pyunit_utils.locate("smalldata/junit/cars_20mpg.csv"))
    r = cars[0].runif()
    train = cars[r > .2]
    valid = cars[r <= .2]

    # regression
    response_col = "economy"
    distribution = "gaussian"
    predictors = ["displacement","power","weight","acceleration","year"]

    gbm = H2OGradientBoostingEstimator(nfolds=3,
                                       distribution=distribution,
                                       fold_assignment="Random")
    gbm.train(x=predictors, y=response_col, training_frame=train, validation_frame=valid)

    #   mse
    mse1 = gbm.mse(train=True,  valid=False, xval=False)
    assert isinstance(mse1, float)

    mse2 = gbm.mse(train=False, valid=True,  xval=False)
    assert isinstance(mse2, float)

    mse3 = gbm.mse(train=False, valid=False, xval=True)
    assert isinstance(mse3, float)

    mse = gbm.mse(train=True,  valid=True,  xval=False)
    assert "train" in list(mse.keys()) and "valid" in list(mse.keys()), "expected training and validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert len(mse) == 2, "expected only training and validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert isinstance(mse["train"], float) and isinstance(mse["valid"], float), "expected training and validation metrics to be floats, but got {0} and {1}".format(type(mse["train"]), type(mse["valid"]))
    assert mse["valid"] == mse2

    mse = gbm.mse(train=True,  valid=False, xval=True)
    assert "train" in list(mse.keys()) and "xval" in list(mse.keys()), "expected training and cross validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert len(mse) == 2, "expected only training and cross validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert isinstance(mse["train"], float) and isinstance(mse["xval"], float), "expected training and cross validation metrics to be floats, but got {0} and {1}".format(type(mse["train"]), type(mse["xval"]))
    assert mse["xval"] == mse3

    mse = gbm.mse(train=True,  valid=True,  xval=True)
    assert "train" in list(mse.keys()) and "valid" in list(mse.keys()) and "xval" in list(mse.keys()), "expected training, validation, and cross validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert len(mse) == 3, "expected training, validation and cross validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert isinstance(mse["train"], float) and isinstance(mse["valid"], float) and isinstance(mse["xval"], float), "expected training, validation, and cross validation metrics to be floats, but got {0}, {1}, and {2}".format(type(mse["train"]), type(mse["valid"]), type(mse["xval"]))

    mse = gbm.mse(train=False, valid=False, xval=False) # default: return training metrics
    assert isinstance(mse, float)
    assert mse == mse1

    mse = gbm.mse(train=False, valid=True,  xval=True)
    assert "valid" in list(mse.keys()) and "xval" in list(mse.keys()), "expected validation and cross validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert len(mse) == 2, "expected validation and cross validation metrics to be returned, but got {0}".format(list(mse.keys()))
    assert isinstance(mse["valid"], float) and isinstance(mse["xval"], float), "validation and cross validation metrics to be floats, but got {0} and {1}".format(type(mse["valid"]), type(mse["xval"]))

    #   mean_residual_deviance
    mean_residual_deviance1 = gbm.mean_residual_deviance(train=True,  valid=False, xval=False)
    assert isinstance(mean_residual_deviance1, float)

    mean_residual_deviance2 = gbm.mean_residual_deviance(train=False, valid=True,  xval=False)
    assert isinstance(mean_residual_deviance2, float)

    mean_residual_deviance3 = gbm.mean_residual_deviance(train=False, valid=False, xval=True)
    assert isinstance(mean_residual_deviance3, float)

    mean_residual_deviance = gbm.mean_residual_deviance(train=True,  valid=True,  xval=False)
    assert "train" in list(mean_residual_deviance.keys()) and "valid" in list(mean_residual_deviance.keys()), "expected training and validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert len(mean_residual_deviance) == 2, "expected only training and validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert isinstance(mean_residual_deviance["train"], float) and isinstance(mean_residual_deviance["valid"], float), "expected training and validation metrics to be floats, but got {0} and {1}".format(type(mean_residual_deviance["train"]), type(mean_residual_deviance["valid"]))
    assert mean_residual_deviance["valid"] == mean_residual_deviance2

    mean_residual_deviance = gbm.mean_residual_deviance(train=True,  valid=False, xval=True)
    assert "train" in list(mean_residual_deviance.keys()) and "xval" in list(mean_residual_deviance.keys()), "expected training and cross validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert len(mean_residual_deviance) == 2, "expected only training and cross validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert isinstance(mean_residual_deviance["train"], float) and isinstance(mean_residual_deviance["xval"], float), "expected training and cross validation metrics to be floats, but got {0} and {1}".format(type(mean_residual_deviance["train"]), type(mean_residual_deviance["xval"]))
    assert mean_residual_deviance["xval"] == mean_residual_deviance3

    mean_residual_deviance = gbm.mean_residual_deviance(train=True,  valid=True,  xval=True)
    assert "train" in list(mean_residual_deviance.keys()) and "valid" in list(mean_residual_deviance.keys()) and "xval" in list(mean_residual_deviance.keys()), "expected training, validation, and cross validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert len(mean_residual_deviance) == 3, "expected training, validation and cross validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert isinstance(mean_residual_deviance["train"], float) and isinstance(mean_residual_deviance["valid"], float) and isinstance(mean_residual_deviance["xval"], float), "expected training, validation, and cross validation metrics to be floats, but got {0}, {1}, and {2}".format(type(mean_residual_deviance["train"]), type(mean_residual_deviance["valid"]), type(mean_residual_deviance["xval"]))

    mean_residual_deviance = gbm.mean_residual_deviance(train=False, valid=False, xval=False) # default: return training metrics
    assert isinstance(mean_residual_deviance, float)
    assert mean_residual_deviance == mean_residual_deviance1

    mean_residual_deviance = gbm.mean_residual_deviance(train=False, valid=True,  xval=True)
    assert "valid" in list(mean_residual_deviance.keys()) and "xval" in list(mean_residual_deviance.keys()), "expected validation and cross validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert len(mean_residual_deviance) == 2, "expected validation and cross validation metrics to be returned, but got {0}".format(list(mean_residual_deviance.keys()))
    assert isinstance(mean_residual_deviance["valid"], float) and isinstance(mean_residual_deviance["xval"], float), "validation and cross validation metrics to be floats, but got {0} and {1}".format(type(mean_residual_deviance["valid"]), type(mean_residual_deviance["xval"]))


    # binomial
    cars = h2o.import_file(path=pyunit_utils.locate("smalldata/junit/cars_20mpg.csv"))
    cars["economy_20mpg"] = cars["economy_20mpg"].asfactor()
    r = cars[0].runif()
    train = cars[r > .2]
    valid = cars[r <= .2]
    response_col = "economy_20mpg"
    distribution = "bernoulli"
    predictors = ["displacement","power","weight","acceleration","year"]
    gbm = H2OGradientBoostingEstimator(nfolds=3, distribution=distribution, fold_assignment="Random")
    gbm.train(y=response_col, x=predictors, validation_frame=valid, training_frame=train)

    #   auc
    auc1 = gbm.auc(train=True,  valid=False, xval=False)
    assert isinstance(auc1, float)

    auc2 = gbm.auc(train=False, valid=True,  xval=False)
    assert isinstance(auc2, float)

    auc3 = gbm.auc(train=False, valid=False, xval=True)
    assert isinstance(auc3, float)

    auc = gbm.auc(train=True,  valid=True,  xval=False)
    assert "train" in list(auc.keys()) and "valid" in list(auc.keys()), "expected training and validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert len(auc) == 2, "expected only training and validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert isinstance(auc["train"], float) and isinstance(auc["valid"], float), "expected training and validation metrics to be floats, but got {0} and {1}".format(type(auc["train"]), type(auc["valid"]))
    assert auc["valid"] == auc2

    auc = gbm.auc(train=True,  valid=False, xval=True)
    assert "train" in list(auc.keys()) and "xval" in list(auc.keys()), "expected training and cross validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert len(auc) == 2, "expected only training and cross validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert isinstance(auc["train"], float) and isinstance(auc["xval"], float), "expected training and cross validation metrics to be floats, but got {0} and {1}".format(type(auc["train"]), type(auc["xval"]))
    assert auc["xval"] == auc3

    auc = gbm.auc(train=True,  valid=True,  xval=True)
    assert "train" in list(auc.keys()) and "valid" in list(auc.keys()) and "xval" in list(auc.keys()), "expected training, validation, and cross validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert len(auc) == 3, "expected training, validation and cross validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert isinstance(auc["train"], float) and isinstance(auc["valid"], float) and isinstance(auc["xval"], float), "expected training, validation, and cross validation metrics to be floats, but got {0}, {1}, and {2}".format(type(auc["train"]), type(auc["valid"]), type(auc["xval"]))

    auc = gbm.auc(train=False, valid=False, xval=False) # default: return training metrics
    assert isinstance(auc, float)
    assert auc == auc1

    auc = gbm.auc(train=False, valid=True,  xval=True)
    assert "valid" in list(auc.keys()) and "xval" in list(auc.keys()), "expected validation and cross validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert len(auc) == 2, "expected validation and cross validation metrics to be returned, but got {0}".format(list(auc.keys()))
    assert isinstance(auc["valid"], float) and isinstance(auc["xval"], float), "validation and cross validation metrics to be floats, but got {0} and {1}".format(type(auc["valid"]), type(auc["xval"]))

    # roc
    (fprs1, tprs1) = gbm.roc(train=True,  valid=False, xval=False)
    assert isinstance(fprs1, list)
    assert isinstance(tprs1, list)

    (fprs2, tprs2) = gbm.roc(train=False, valid=True,  xval=False)
    assert isinstance(fprs2, list)
    assert isinstance(tprs2, list)

    (fprs3, tprs3) = gbm.roc(train=False, valid=False, xval=True)
    assert isinstance(fprs3, list)
    assert isinstance(tprs3, list)

    roc = gbm.roc(train=True,  valid=True,  xval=False)
    assert "train" in list(roc.keys()) and "valid" in list(roc.keys()), "expected training and validation metrics to be returned, but got {0}".format(list(roc.keys()))
    assert len(roc) == 2, "expected only training and validation metrics to be returned, but got {0}".format(list(roc.keys()))
    assert isinstance(roc["train"], tuple) and isinstance(roc["valid"], tuple), "expected training and validation metrics to be tuples, but got {0} and {1}".format(type(roc["train"]), type(roc["valid"]))
    assert roc["valid"][0] == fprs2
    assert roc["valid"][1] == tprs2

    roc = gbm.roc(train=True,  valid=False, xval=True)
    assert "train" in list(roc.keys()) and "xval" in list(roc.keys()), "expected training and cross validation metrics to be returned, but got {0}".format(list(roc.keys()))