from ensemble_support import *
from alad_simple import *
from weight_inference import *
from aatp_iterative_gradient import *
from query_model import *


class MetricsStructure(object):
    def __init__(self, train_aucs=None, test_aucs=None, train_precs=None, test_precs=None,
                 train_aprs=None, test_aprs=None, train_n_at_top=None, test_n_at_top=None,
                 all_weights=None, queried=None):
        self.train_aucs = train_aucs
        self.test_aucs = test_aucs
        self.train_precs = train_precs
        self.test_precs = test_precs
        self.train_aprs = train_aprs
        self.test_aprs = test_aprs
        self.train_n_at_top = train_n_at_top
        self.test_n_at_top = test_n_at_top
        self.all_weights = all_weights
        self.queried = queried
        self.test_indexes = []


def get_alad_metrics_structure(budget, opts):
    metrics = MetricsStructure(
        train_aucs=np.zeros(shape=(1, budget)),
        # for precision@k first two columns are fid,k
        train_precs=[],
        train_aprs=np.zeros(shape=(1, budget)),
        train_n_at_top=[],
        all_weights=[],
        queried=[]
    )
    for k in range(len(opts.precision_k)):
        metrics.train_precs.append(np.zeros(shape=(1, budget)))
        metrics.train_n_at_top.append(np.zeros(shape=(1, budget)))
    return metrics


def save_alad_metrics(metrics, opts):
    cansave = (opts.resultsdir != "" and os.path.isdir(opts.resultsdir))
    if cansave:
        save(metrics, filepath=opts.get_metrics_path())


def load_alad_metrics(opts):
    metrics = None
    fpath = opts.get_metrics_path()
    canload = (opts.resultsdir != "" and os.path.isfile(fpath))
    if canload:
        # print "Loading metrics" + fpath
        metrics = load(fpath)
    else:
        print "Cannot load " + fpath
    return metrics


class MetricsCollection(object):
    def __init__(self, fids, runidxs, metrics):
        self.fids = fids
        self.runidxs = runidxs
        self.metrics = metrics


def consolidate_alad_metrics(fids, runidxs, opts):
    """

    :param fids:
    :param runidxs:
    :param opts:
    :return: MetricsCollection
    """
    metrics_struct = MetricsCollection(fids=fids, runidxs=runidxs, metrics=[])
    for i in range(len(fids)):
        fmetrics = []
        for j in range(len(runidxs)):
            opts.set_multi_run_options(fids[i], runidxs[j])
            fmetrics.append(load_alad_metrics(opts))
        metrics_struct.metrics.append(fmetrics)
    return metrics_struct


class SequentialResults(object):
    def __init__(self, num_seen=None, num_seen_baseline=None,
                 true_queried_indexes=None, true_queried_indexes_baseline=None,
                 stream_window=None, stream_window_baseline=None,
                 aucs=None):
        self.num_seen = num_seen
        self.num_seen_baseline = num_seen_baseline
        self.true_queried_indexes = true_queried_indexes
        se
lf.true_queried_indexes_baseline = true_queried_indexes_baseline
        self.stream_window = stream_window
        self.stream_window_baseline = stream_window_baseline
        self.aucs = aucs


def summarize_ensemble_num_seen(ensemble, metrics, fid=0, runidx=0):
    """
    IMPORTANT: returned queried_indexes and queried_indexes_baseline a
re 1-indexed (NOT 0-indexed)
    """
    nqueried = len(metrics.queried)
    num_seen = np.zeros(shape=(1, nqueried + 2))
    num_seen_baseline = np.zeros(shape=(1, nqueried + 2))

    num_seen[0, 0:2] = [fid, runidx]
    num_seen[0, 2:(num_seen.shape[1])] = np.cumsum(ensemble.labels[metrics.queried])

    qlbls = ensemble.labels[ensemble.ordered_anom_idxs[0:nqueried]]
    num_seen_baseline[0, 0:2] = [fid, runidx]
    num_seen_baseline[0, 2:(num_seen_baseline.shape[1])] = np.cumsum(qlbls)

    # the ensembles store samples in sorted order of default anomaly
    # scores. The corresponding indexes are stored in ensemble.original_indexes
    true_queried_indexes = np.zeros(shape=(1, nqueried + 2))
    true_queried_indexes[0, 0:2] = [fid, runidx]
    # Note: make the queried indexes relative 1 (NOT zero)
    true_queried_indexes[0, 2:(true_queried_indexes.shape[1])] = ensemble.original_indexes[metrics.queried] + 1

    # the ensembles store samples in sorted order of default anomaly
    # scores. The corresponding indexes are stored in ensemble.original_indexes
    true_queried_indexes_baseline = np.zeros(shape=(1, nqueried + 2))
    true_queried_indexes_baseline[0, 0:2] = [fid, runidx]
    # Note: make the queried indexes relative 1 (NOT zero)
    true_queried_indexes_baseline[0, 2:(true_queried_indexes_baseline.shape[1])] = \
        ensemble.original_indexes[np.arange(nqueried)] + 1

    return num_seen, num_seen_baseline, true_queried_indexes, true_queried_indexes_baseline


def summarize_alad_metrics(ensembles, metrics_struct):
    nqueried = len(metrics_struct.metrics[0][0].queried)
    num_seen = np.zeros(shape=(0, nqueried+2))
    num_seen_baseline = np.zeros(shape=(0, nqueried+2))
    true_queried_indexes = np.zeros(shape=(0, nqueried+2))
    true_queried_indexes_baseline = np.zeros(shape=(0, nqueried + 2))
    for i in range(len(metrics_struct.metrics)):
        # file level
        submetrics = metrics_struct.metrics[i]
        subensemble = ensembles[i]
        for j in range(len(submetrics)):
            # rerun level
            queried = submetrics[j].queried
            lbls = subensemble[j].labels

            nseen = np.zeros(shape=(1, nqueried+2))
            nseen[0, 0:2] = [metrics_struct.fids[i], metrics_struct.runidxs[j]]
            nseen[0, 2:(nseen.shape[1])] = np.cumsum(lbls[queried])
            num_seen = rbind(num_seen, nseen)

            qlbls = subensemble[j].labels[subensemble[j].ordered_anom_idxs[0:nqueried]]
            nseen = np.zeros(shape=(1, nqueried+2))
            nseen[0, 0:2] = [metrics_struct.fids[i], metrics_struct.runidxs[j]]
            nseen[0, 2:(nseen.shape[1])] = np.cumsum(qlbls)
            num_seen_baseline = rbind(num_seen_baseline, nseen)

            # the ensembles store samples in sorted order of default anomaly
            # scores. The corresponding indexes are stored in ensemble.original_indexes
            t_idx = np.zeros(shape=(1, nqueried + 2))
            t_idx[0, 0:2] = [metrics_struct.fids[i], metrics_struct.runidxs[j]]
            t_idx[0, 2:(t_idx.shape[1])] = subensemble[j].original_indexes[queried]
            # Note: make the queried indexes realive 1 (NOT zero)
            true_queried_indexes = rbind(true_queried_indexes, t_idx + 1)

            # the ensembles store samples in sorted order of default anomaly
            # scores. The corresponding indexes are stored in ensemble.original_indexes
            b_idx = np.zeros(shape=(1, nqueried + 2))
            b_idx[0, 0:2] = [metrics_struct.fids[i], metrics_struct.runidxs[j]]
            b_idx[0, 2:(b_idx.shape[1])] = subensemble[j].original_indexes[np.arange(nqueried)]
            # Note: make the queried indexes realive 1 (NOT zero)
            true_queried_indexes_baseline = rbind(true_queried_indexes_baseline, b_idx + 1)

    return SequentialResults(num_seen=num_seen, num_seen_baseline=num_seen_baseline,
                             true_queried_indexes=true_queried_indexes,
                             true_queried_indexes_baseline=true_queried_indexes_baseline)


def save_alad_summary(alad_summary, opts):
    cansave = opts.resultsdir != "" and os.path.isdir(opts.resultsdir)
    if cansave:
        save(alad_summary, filepath=opts.get_metrics_summary_path())


def load_alad_summary(opts):
    alad_summary = None
    fpath = opts.get_metrics_summary_path()
    canload = opts.resultsdir != "" and os.path.isfile(fpath)
    if canload:
        alad_summary = load(fpath)
    else:
        print ("Cannot load %s" % fpath)
    return alad_summary


class Budget(object):
    def __init__(self, topK, budget):
        self.topK = topK
        self.budget = budget


def get_budget_topK(n, opts):
    # set topK as per tau or input topK
    topK = opts.topK
    if topK <= 0:
        topK = int(np.round(opts.tau * n))  # function of total number of instances
    budget = opts.budget
  