"""
Tests for isotonic regression, bayesian binning with quantiles, and
reliably calibrated isotonic regression.
"""


import isotonic
import numpy as np
# BBQ-code is Matlab/Octave code:
from oct2py import octave
from sklearn.isotonic import IsotonicRegression
from sklearn.metrics import roc_auc_score
# ENIR-code is R:
import rpy2.robjects as robjects
from rpy2.robjects.packages import importr
enir = importr('enir')
r = robjects.r
# Automatic conversion or numpy arrays to R-vectors
import rpy2.robjects.numpy2ri
rpy2.robjects.numpy2ri.activate()

# Enable octave to find Naeini's functions!
# (download from https://github.com/pakdaman/calibration)
octave.eval("addpath('./calibration/BBQ/')", verbose=False)

# Set test parameters:
dataset = int(input("Select dataset to run experiment on (1, 2, or 3): "))
n_iterations = int(input("Set number of iterations (30 used in paper): "))
# metric = input("Select metric ('mse' or 'auc_roc'): ")  # Perhaps allow only auc-roc and exclude mse?
metric = 'auc_roc'  # Can also be set to 'mse' to run the algorithm with mse-based bin merges.
# reshuffle = input("Shuffle data? (y/n): ")  # Should be shuffled at least once. Perhaps remove option.
# It seems that there is no convenient way of estimating the credible intervals, not to speak of the maximum
# credible intervals for the Naeini procedure. Hence the 'Naeini vs. RCIR with d set by Naeini vs. IR'
# is pointless.
# model_comparison = input("Naeini vs. RCIR-CV ('Naeini vs. RCIR-CV') or Naeini vs. RCIR? ")

# Load dataset:
if dataset == 1:
    # Read data for test 1:
    test_description = "test1"
    data_class = isotonic.load_pickle("./data/dataset_1_class.pickle")
    data_scores = isotonic.load_pickle("./data/dataset_1_scores.pickle")
elif dataset == 2:
    # Read data for test 2:
    test_description = "test2"
    data_class = isotonic.load_pickle("./data/dataset_2_class.pickle")
    data_scores = isotonic.load_pickle("./data/dataset_2_scores.pickle")
elif dataset == 3:
    test_description = "test3"
    data_class = isotonic.load_pickle("./data/dataset_3_class.pickle")
    data_scores = isotonic.load_pickle("./data/dataset_3_scores.pickle")
elif dataset == 4:  # Hidden experiment. Does not really provide anything interesting.
    test_description = "test4"
    data_class = np.random.binomial(1, .5, 30000)
    data_scores = np.random.uniform(low=0, high=1, size=30000)
elif dataset == 5:  # Test mode
    tmp_class = isotonic.load_pickle("./data/dataset_1_class.pickle")
    tmp_scores = isotonic.load_pickle("./data/dataset_1_scores.pickle")
    test_description = "test1"
    data_class = tmp_class[:1000]
    data_scores = tmp_scores[:1000]
    print("Dataset loaded.")

else:
    print("Not a valid dataset selection. ")
    import sys
    sys.exit()
print("Dataset with " + str(data_class.shape[0]) + " samples loaded.")

if(metric == 'mse' or metric == 'auc_roc'):
    pass  # All in order.
else:
    print("Not a valid metric selection.")
    import sys
    sys.exit()

k = 100
print("Expected calibration error (ECE) and MCE calculated with k = " + str(k))

ir_top_ten_mse = []
enir_top_ten_mse = []
bir_top_ten_mse = []
rcir_top_ten_mse = []
bbq_top_ten_mse = []
ir_top_ten_auc_roc = []
enir_top_ten_auc_roc = []
bir_top_ten_auc_roc = []
rcir_top_ten_auc_roc = []
bbq_top_ten_auc_roc = []
ir_top_ten_ece = []
enir_top_ten_ece = []
bir_top_ten_ece = []
rcir_top_ten_ece = []
bbq_top_ten_ece = []
ir_top_ten_mce = []
enir_top_ten_mce = []
bir_top_ten_mce = []
rcir_top_ten_mce = []
bbq_top_ten_mce = []

for i in range(n_iterations):
    # Shuffle samples:
    # This is super important as the data is generated by a non-stationary process!
    n_rows = data_scores.shape[0]
    idx = np.random.permutation(range(n_rows))
    data_class = data_class[idx]
    data_scores = data_scores[idx]
    test_class = data_class[:n_rows * 1 // 4]
    training_class = data_class[n_rows // 4: 3 * n_rows // 4]
    validation_class = data_class[3 * n_rows // 4:]  # NOT USED FOR ANYTHING. SKIP?
    test_scores = data_scores[:n_rows * 1 // 4]
    training_scores = data_scores[n_rows // 4: 3 * n_rows // 4]
    validation_scores = data_scores[3 * n_rows // 4:]

    # Create ENIR model using R:
    enir_model = enir.enir_build(robjects.FloatVector(training_scores.tolist()),
                                 robjects.BoolVector(training_class.tolist()))
    enir_prob = enir.enir_predict(enir_model, robjects.FloatVector(test_scores.tolist()))
    # Convert to numpy.array:
    enir_prob = np.array(enir_prob)
    # Kill session?

    # Create isotonic regression model
    ir_model = IsotonicRegression(y_min=0, y_max=1, out_of_bounds='clip')
    ir_model.fit(X=training_scores, y=training_class)

    # Create BBQ-model
    octave.push('training_scores', training_scores, verbose=False)
    octave.push('training_class', training_class, verbose=False)
    octave.eval('options.N0 = 2', verbose=False)
    octave.eval("bbq_model = build(training_scores', training_class', options)", verbose=False)
    octave.push('test_scores', test_scores)
    octave.eval("test_prob = predict(bbq_model, test_scores, 1)", verbose=False)
    bbq_test_prob = octave.pull('test_prob', verbose=False)
    bbq_test_prob = [item[0] for item in bbq_test_prob]

    # Create RCIR model with d = .20 (?)
    # QUICK HACK TO TEST RCIR-CV INSTEAD OF RCIR WITH SOME d:
    # NOTE THAT RCIR-CV GETS _MORE_ _DATA_ THAN
    # rcir_model = isotonic.train_rcir_cv(training_scores, training_class , d=.1)
    rcir_model = isotonic.train_rcir_cv(training_class, training_scores,
                                        validation_class, validation_scores)

    # Create bootstrap IR model
    bir_model = isotonic.bootstrap_isotonic_regression(training_class, training_scores,
                                                       sampling_rate=.95, n_models=1000)

    # Separate the top-N% best scoring test samples (currently 100%, i.e. all):
    top_idx = np.argsort(test_scores)
    top_ten_idx = top_idx[int(len(top_idx) * .9):]
    top_ten_scores = test_scores[top_ten_idx]
    top_ten_class = test_class[top_ten_idx]
    # Estimate the probabilities for these samples for different models:
    top_ten_ir_prob = ir_model.predict(top_ten_scores)
    # ENIR
    top_ten_enir_prob = enir.enir_predict(enir_model, robjects.FloatVector(top_ten_scores.tolist()))
    top_ten_enir_prob = np.array(top_ten_enir_prob)
    # top_ten_rcir_prob = isotonic.predict_rcir(rcir_model, top_ten_scores)
    top_ten_rcir_prob = isotonic.predict(rcir_model, top_ten_scores)
    octave.push('top_ten_scores', top_ten_scores, verbose=False)
    octave.eval("top_ten_bbq_prob = predict(bbq_model, top_ten_scores, 1)", verbose=False)
    top_ten_bbq_prob = octave.pull('top_ten_bbq_prob', verbose=False)
    top_ten_bbq_prob = np.array([item[0] for item in top_ten_bbq_prob])
    top_ten_bir_prob = isotonic.bootstrap_isotonic_regression_predict(bir_model, top_ten_scores)
    # Estimate metrics for the results:
    ir_top_ten_mse.append(sum((top_ten_ir_prob - top_ten_class)**2) / len(top_ten_class))
    enir_top_ten_mse.append(sum((top_ten_enir_prob - top_ten_class)**2) / len(top_ten_class))
    bir_top_ten_mse.append(sum((top_ten_bir_prob - top_ten_class)**2) / len(top_ten_class))
    rcir_top_ten_mse.append(sum((top_ten_rcir_prob - top_ten_class)**2) / len(top_ten_class))
    bbq_top_ten_mse.append(sum((top_ten_bbq_prob - top_ten_class)**2) / len(top_ten_class))
    ir_top_ten_auc_roc.append(roc_auc_score(top_ten_class, top_ten_ir_prob))
    enir_top_ten_auc_roc.append(roc_auc_score(top_ten_class, top_ten_enir_prob))
    bir_top_ten_auc_roc.append(roc_auc_score(top_ten_class, top_ten_bir_prob))
    rcir_top_ten_auc_roc.append(roc_auc_score(top_ten_class, top_ten_rcir_prob))
    bbq_top_ten_auc_roc.append(roc_auc_score(top_ten_class, top_ten_bbq_prob))
    ir_top_ten_ece.append(isotonic.expected_calibration_error(top_ten_class, top_ten_ir_prob, k))
    enir_top_ten_ece.append(isotonic.expected_calibration_error(top_ten_class, top_ten_enir_prob, k))
    bir_top_ten_ece.append(isotonic.expected_calibration_error(top_ten_class, top_ten_bir_prob, k))
    rcir_top_ten_ece.append(isotonic.expected_calibration_error(top_ten_class, top_ten_rcir_prob, k))
    bbq_top_ten_ece.append(isotonic.expected_calibration_error(top_ten_class, top_ten_bbq_prob, k))
    ir_top_ten_mce.append(isotonic.maximum_calibration_error(top_ten_class, top_ten_ir_prob, k))
    enir_top_ten_mce.append(isotonic.maximum_calibration_error(top_ten_class, top_ten_enir_prob, k))
    bir_top_ten_mce.append(isotonic.maximum_calibration_error(top_ten_class, top_ten_bir_prob, k))
    rcir_top_ten_mce.append(isotonic.maximum_calibration_error(top_ten_class, top_ten_rcir_prob, k))
    bbq_top_ten_mce.append(isotonic.maximum_calibration_error(top_ten_class, top_ten_bbq_prob, k))
# Print results:
print("Metrics for 10% of samples with highest score.")
print("Metric \t IR \t\t ENIR \t\t BIR \t\t RCIR-CV \t\t BBQ")
print("MSE\t" + str(np.mean(ir_top_ten_mse)) + "\t" + str(np.mean(enir_top_ten_mse)) + "\t" + str(np.mean(bir_top_ten_mse)) + "\t" + str(np.mean(rcir_top_ten_mse)) + "\t" + str(np.mean(bbq_top_ten_mse)))
print("AUC-ROC\t" + str(np.mean(ir_top_ten_auc_roc)) + "\t" + str(np.mean(enir_top_ten_auc_roc)) + "\t" + str(np.mean(bir_top_ten_auc_roc)) + "\t" + str(np.mean(rcir_top_ten_auc_roc)) + "\t" + str(np.mean(bbq_top_ten_auc_roc)))
print("ECE\t" + str(np.mean(ir_top_ten_ece)) + "\t" + str(np.mean(enir_top_ten_ece)) + "\t" + str(np.mean(bir_top_ten_ece)) + "\t" + str(np.mean(rcir_top_ten_ece)) + "\t" + str(np.mean(bbq_top_ten_ece)))
print("MCE\t" + str(np.mean(ir_top_ten_mce)) + "\t" + str(np.mean(enir_top_ten_mce)) + "\t" + str(np.mean(bir_top_ten_mce)) + "\t" + str(np.mean(rcir_top_ten_mce)) + "\t" + str(np.mean(bbq_top_ten_mce)))


"""
    BBQ did not outperform IR and BIR by order or magnitude. By testing on the three
    datasets, we see that on a separate testing set BBQ outperforms IR and BIR clearly
    in _one_ case (datset 3), whereas IR and BIR outperforms BBQ on dataset 2 on all
    metrics. The results on dataset 1 are split, i.e. IR and BIR outperforms BBQ on
    all other metrics except ECE (which was the authors primary metric).
    BBQ also takes very long to train (how many models do they use?!?), hence it is not
    at all clear that BBQ is clearly better than the others.
    I interpret the result on dataset 3 as an example of where the monotonicity
    assumption is not met which would then explain why BBQ is clearly better.

    This is also interesting because the scores to be calibrated are all from similar
    tasks (conversion probability prediction based on click-stream data) and use the
    same base-learner.

    RCIR did not produce comparable results for any of the datasets with d=.2.

    ECE intuitively describes the quality guarantee we are hoping for, although it is
    by no means an unambiguous metric (e.g. think about a classifier randomizing samples
    *perfectly* and then predicting the natural rate for all samples -> ECE = 0).

    The choise of k=10 for estimating ECE is rather poor (although used by two authors
    already). It indicates that we only require that samples are "well-calibrated" in
    10% chunks of the entire dataset. We used k=100 instead (we also tested k=10, the
    results do not significantly change). K=10 would cause the last bin to be really
    wide in terms of mapping probabilities, which could be a problem with high class-
    imbalance.

    We could estimate the effect of the monotonicity assumption vs. the lack of it
    proposed by BBQ by estimating the metrics for different "bins" (i.e. chunks of
    samples related to k). If two neighboring bins for BBQ have better metrics than
    the corresponding ones for IR and BIR, that might indicate that the probabilities
    are mapped in a non-monotonic fashion by BBQ. This could of course be verified,
    and would then provide conclusive evidence for the weakness of the monotonicity-
    assumption.

    A hypothesis was that some of these methods might perform better or worse for
    some fraction of the top-scoring samples as it is highly class-imbalanced data.
    Theoretically this can be showed with beta-distributions where a low number of
    samples results in wide credible intervals.

    Metrics used were AUC-ROC, MSE, ECE, and MCE.

    *AUC-ROC as equivalent to Wilcoxon rank-sum test and Mann-Whitney U, i.e. the
    probability that a positive sample is ranked higher than a negative one (point
    estimate, ML).
    *MSE is a measure that nicely takes into account the class-imbalance.

    ******* Expected Calibration Error, ECE *********
    *How should a suitable 'k' be chosen? Does it matter?
    **ECE is just absolute mean difference between the average of probabilities in
    one group vs. the empirical frequency in that group. I.e. some sort of mean
    absolute error.
    ***Breaking the monotonicity assumption should lead to better AUC-ROC if it is
    a good thing to do. BBQ did not do that for 2/3 datasets. Hence it seems that
    the improved ECE
    is partly due to shuffling samples and predicting more average.
    ****If all samples were mapped to the same probability, and the 'calibration'
    was done randomly enough, the ECE would be 0 (i.e. the 'best'). Hence, it is
    not a particularly good metric in that sense.

    *MCE is even worse for that bin. If we perfectly separate all the samples to 0
    and 1, then the maximum calibration error for this bin will be .9!
    *Hence, at the very minimum, we need to adjust the k according to the data, and
    preferably use AUC-ROC and MSE as primary metrics.

    ****** IR vs. BIR *******
    BIR seems to produce better AUC-ROC than IR on a separate testing set. This is
    interesting as IR per default maximizes the AUC-ROC on the training set. By
    allowing BIR to split these bins while still keeping monotonicity, we end up
    with better AUC-ROC on a separate testing set!!! This indicates, that allowing
    small increasing probabilities with values within bins could produce better
    AUC-ROC for IR, i.e. that the monotonicity could even be replaced with strict
    monotonicity (g(a) > g(b) for a > b).

    Turns out RCIR is worst for data where BBQ performs best, i.e. the splitting
    up of bins enabled by RCIR (compared to IR) results in worse AUC-ROC!

    Estimated results for RCIR with different 'd' (perhaps CV?) to see whether it
    was good for something. Tried .4, .2, and .1. None of the results were better
    than IR, BIR and/or BBQ.

    Next up:
    -Read up on BBQ-paper.
    -Read up on Naeini's newer papers. There was something about a version based
    on isotonic regression there (?).
    -Does the probem with the higher mapping bins persist? Can we improve on it?

    Also, it seems that BIR produces better results on AUC-ROC and MSE for pretty
    much any sampling rates between 0.1 and .99 as long as the number of models is
    above 100.

    ???
    BBQ typically maps the highest scoring samples to probabilities that are way lower
    than those provided by IR. Analyzing the last bin using a beta-distribution, we have
    extremely good reason to believe that some samples should map to higher probabilities.
    -What's up with that? How can BBQ still score equally well or better on most metrics?

    If BBQ has a tendency to map high scores to lower probabilities than we consider
    reasonable (based on analyzing bins of IR and beta-distributions), why does BBQ not
    produce worse metrics on the highest scoring samples? Can we isoalte these samples
    and perform tests on them? (Instead of just testing "top 10%").

    ########
    The highest likelihood model in the BBQ-ensemble produces AUC-ROC and MSE comparable
    to the entire ensemble (AUC-ROC 0.69953551015217341 and MSE 0.043266913296352877
    for the single model vs. AUC-ROC 0.70775821180884224 and MSE 0.043393969841217683 for
    the entire ensemble). While the AUC-ROC is slightly worse, the MSE is slightly better.
    The ECE 0.010919855957577638 for one model (vs. BBQ's 0.0076235109580911394) isn't bad
    either (BIR's 0.009823043745131722, IR's 0.0087272850603866657, k=100).
    AUC-ROC: IR 0.70987550204211003, BIR 0.71003757391577715
    MSE: IR 0.043217400882335832, BIR 0.043216607313252962

    Notes on BBQ
    The most important model made up for 93% of the ensemble (in one experiment). The second
    most important model (based on likelihood!) made up 6% of the ensemble, and the remaining
    460 models made up the last percentage (which of like 5 had any significance at all...).
    This in contrast to the likelihood (on testing set!) for weighted BBQ that puts >99%
    of the weight on one model.
"""
