#!/usr/bin/env python3
import os
import cv2
import numpy as np
import pandas as pd
from tqdm import tqdm
from skimage import transform
from sklearn.preprocessing import normalize
from sklearn.metrics import roc_curve, auc

class Mxnet_model_interf:
    def __init__(self, model_file, layer="fc1", image_size=(112, 112)):
        import mxnet as mx

        self.mx = mx
        cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
        if len(cvd) > 0 and int(cvd) != -1:
            ctx = [self.mx.gpu(ii) for ii in range(len(cvd.split(",")))]
        else:
            ctx = [self.mx.cpu()]

        prefix, epoch = model_file.split(",")
        print(">>>> loading mxnet model:", prefix, epoch, ctx)
        sym, arg_params, aux_params = self.mx.model.load_checkpoint(prefix, int(epoch))
        all_layers = sym.get_internals()
        sym = all_layers[layer + "_output"]
        model = self.mx.mod.Module(symbol=sym, context=ctx, label_names=None)
        model.bind(data_shapes=[("data", (1, 3, image_size[0], image_size[1]))])
        model.set_params(arg_params, aux_params)
        self.model = model

    def __call__(self, imgs):
        # print(imgs.shape, imgs[0])
        imgs = imgs.transpose(0, 3, 1, 2)
        data = self.mx.nd.array(imgs)
        db = self.mx.io.DataBatch(data=(data,))
        self.model.forward(db, is_train=False)
        emb = self.model.get_outputs()[0].asnumpy()
        return emb

class Torch_model_interf:
    def __init__(self, model_file, image_size=(112, 112)):
        import torch

        self.torch = torch
        cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
        device_name = "cuda:0" if len(cvd) > 0 and int(cvd) != -1 else "cpu"
        self.device = self.torch.device(device_name)
        try:
            self.model = self.torch.jit.load(model_file, map_location=device_name)
        except:
            print("Error: %s is weights only, please load and save the entire model by `torch.jit.save`" % model_file)
            self.model = None

    def __call__(self, imgs):
        # print(imgs.shape, imgs[0])
        imgs = imgs.transpose(0, 3, 1, 2).copy().astype("float32")
        imgs = (imgs - 127.5) * 0.0078125
        output = self.model(self.torch.from_numpy(imgs).to(self.device).float())
        return output.cpu().detach().numpy()

class ONNX_model_interf:
    def __init__(self, model_file, image_size=(112, 112)):
        import onnxruntime as ort
        ort.set_default_logger_severity(3)
        self.ort_session = ort.InferenceSession(model_file)
        self.output_names = [self.ort_session.get_outputs()[0].name]
        self.input_name = self.ort_session.get_inputs()[0].name

    def __call__(self, imgs):
        imgs = imgs.transpose(0, 3, 1, 2).astype("float32")
        imgs = (imgs - 127.5) * 0.0078125
        outputs = self.ort_session.run(self.output_names, {self.input_name: imgs})
        return outputs[0]

def keras_model_interf(model_file):
    import tensorflow as tf
    from tensorflow_addons.layers import StochasticDepth

    for gpu in tf.config.experimental.list_physical_devices("GPU"):
        tf.config.experimental.set_memory_growth(gpu, True)

    mm = tf.keras.models.load_model(model_file, compile=False)
    return lambda imgs: mm((tf.cast(imgs, "float32") - 127.5) * 0.0078125).numpy()

def face_align_landmark(img, landmark, image_size=(112, 112), method="similar"):
    tform = transform.AffineTransform() if method == "affine" else transform.SimilarityTransform()
    src = np.array(
        [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.729904, 92.2041]], dtype=np.float32
    )
    tform.estimate(landmark, src)
    # ndimage = transform.warp(img, tform.inverse, output_shape=image_size)
    # ndimage = (ndimage * 255).astype(np.uint8)
    M = tform.params[0:2, :]
    ndimage = cv2.warpAffine(img, M, image_size, borderValue=0.0)
    if len(ndimage.shape) == 2:
        ndimage = np.stack([ndimage, ndimage, ndimage], -1)
    else:
        ndimage = cv2.cvtColor(ndimage, cv2.COLOR_BGR2RGB)
    return ndimage

def read_IJB_meta_columns_to_int(file_path, columns, sep=" ", skiprows=0, header=None):
    # meta = np.loadtxt(file_path, skiprows=skiprows, delimiter=sep)
    meta = pd.read_csv(file_path, sep=sep, skiprows=skiprows, header=header).values
    return (meta[:, ii].astype("int") for ii in columns)

def extract_IJB_data_11(data_path, subset, save_path=None, force_reload=False):
    if save_path == None:
        save_path = os.path.join(data_path, subset + "_backup.npz")
    if not force_reload and os.path.exists(save_path):
        print(">>>> Reload from backup: %s ..." % save_path)
        aa = np.load(save_path)
        return (
            aa["templates"],
            aa["medias"],
            aa["p1"],
            aa["p2"],
            aa["label"],
            aa["img_names"],
            aa["landmarks"],
            aa["face_scores"],
        )

    if subset == "IJBB":
        media_list_path = os.path.join(data_path, "IJBB/meta/ijbb_face_tid_mid.txt")
        pair_list_path = os.path.join(data_path, "IJBB/meta/ijbb_template_pair_label.txt")
        img_path = os.path.join(data_path, "IJBB/loose_crop")
        img_list_path = os.path.join(data_path, "IJBB/meta/ijbb_name_5pts_score.txt")
    else:
        media_list_path = os.path.join(data_path, "IJBC/meta/ijbc_face_tid_mid.txt")
        pair_list_path = os.path.join(data_path, "IJBC/meta/ijbc_template_pair_label.txt")
        img_path = os.path.join(data_path, "IJBC/loose_crop")
        img_list_path = os.path.join(data_path, "IJBC/meta/ijbc_name_5pts_score.txt")

    print(">>>> Loading templates and medias...")
    templates, medias = read_IJB_meta_columns_to_int(media_list_path, columns=[1, 2])  # ['1.jpg', '1', '69544']
    print("templates: %s, medias: %s, unique templates: %s" % (templates.shape, medias.shape, np.unique(templates).shape))
    # templates: (227630,), medias: (227630,), unique templates: (12115,)

    print(">>>> Loading pairs...")
    p1, p2, label = read_IJB_meta_columns_to_int(pair_list_path, columns=[0, 1, 2])  # ['1', '11065', '1']
    print("p1: %s, unique p1: %s" % (p1.shape, np.unique(p1).shape))
    print("p2: %s, unique p2: %s" % (p2.shape, np.unique(p2).shape))
    print("label: %s, label value counts: %s" % (label.shape, dict(zip(*np.unique(label, return_counts=True)))))
    # p1: (8010270,), unique p1: (1845,)
    # p2: (8010270,), unique p2: (10270,) # 10270 + 1845 = 12115 --> np.unique(templates).shape
    # label: (8010270,), label value counts: {0: 8000000, 1: 10270}

    print(">>>> Loading images...")
    with open(img_list_path, "r") as ff:
        # 1.jpg 46.060 62.026 87.785 60.323 68.851 77.656 52.162 99.875 86.450 98.648 0.999
        img_records = np.array([ii.strip().split(" ") for ii in ff.readlines()])

    img_names = np.array([os.path.join(img_path, ii) for ii in img_records[:, 0]])
    landmarks = img_records[:, 1:-1].astype("float32").reshape(-1, 5, 2)
    face_scores = img_records[:, -1].astype("float32")
    print("img_names: %s, landmarks: %s, face_scores: %s" % (img_names.shape, landmarks.shape, face_scores.shape))
    # img_names: (227630,), landmarks: (227630, 5, 2), face_scores: (227630,)
    print("face_scores value counts:", dict(zip(*np.histogram(face_scores, bins=9)[::-1])))
    # {0.1: 2515, 0.2: 0, 0.3: 62, 0.4: 94, 0.5: 136, 0.6: 197, 0.7: 291, 0.8: 538, 0.9: 223797}

    print(">>>> Saving backup to: %s ..." % save_path)
    np.savez(
        save_path,
        templates=templates,
        medias=medias,
        p1=p1,
        p2=p2,
        label=label,
        img_names=img_names,
        landmarks=landmarks,
        face_scores=face_scores,
    )
    print()
    return templates, medias, p1, p2, label, img_names, landmarks, face_scores

def extract_gallery_prob_data(data_path, subset, save_path=None, force_reload=False):
    if save_path == None:
        save_path = os.path.join(data_path, subset + "_gallery_prob_backup.npz")
    if not force_reload and os.path.exists(save_path):
        print(">>>> Reload from backup: %s ..." % save_path)
        aa = np.load(save_path)
        return (
            aa["s1_templates"],
            aa["s1_subject_ids"],
            aa["s2_templates"],
            aa["s2_subject_ids"],
            aa["probe_mixed_templates"],
            aa["probe_mixed_subject_ids"],
        )

    if subset == "IJBC":
        meta_dir = os.path.join(data_path, "IJBC/meta")
        gallery_s1_record = os.path.join(meta_dir, "ijbc_1N_gallery_G1.csv")
        gallery_s2_record = os.path.join(meta_dir, "ijbc_1N_gallery_G2.csv")
        probe_mixed_record = os.path.join(meta_dir, "ijbc_1N_probe_mixed.csv")
    else:
        meta_dir = os.path.join(data_path, "IJBB/meta")
        gallery_s1_record = os.path.join(meta_dir, "ijbb_1N_gallery_S1.csv")
        gallery_s2_record = os.path.join(meta_dir, "ijbb_1N_gallery_S2.csv")
        probe_mixed_record = os.path.join(meta_dir, "ijbb_1N_probe_mixed.csv")

    print(">>>> Loading gallery feature...")
    s1_templates, s1_subject_ids = read_IJB_meta_columns_to_int(gallery_s1_record, columns=[0, 1], skiprows=1, sep=",")
    s2_templates, s2_subject_ids = read_IJB_meta_columns_to_int(gallery_s2_record, columns=[0, 1], skiprows=1, sep=",")
    print("s1 gallery: %s, ids: %s, unique: %s" % (s1_templates.shape, s1_subject_ids.shape, np.unique(s1_templates).shape))
    print("s2 gallery: %s, ids: %s, unique: %s" % (s2_templates.shape, s2_subject_ids.shape, np.unique(s2_templates).shape))

    print(">>>> Loading prope feature...")
    probe_mixed_templates, probe_mixed_subject_ids = read_IJB_meta_columns_to_int(
        probe_mixed_record, columns=[0, 1], skiprows=1, sep=","
    )
    print("probe_mixed_templates: %s, unique: %s" % (probe_mixed_templates.shape, np.unique(probe_mixed_templates).shape))
    print("probe_mixed_subject_ids: %s, unique: %s" % (probe_mixed_subject_ids.shape, np.unique(probe_mixed_subject_ids).shape))

    print(">>>> Saving backup to: %s ..." % save_path)
    np.savez(
        save_path,
        s1_templates=s1_templates,
        s1_subject_ids=s1_subject_ids,
        s2_templates=s2_templates,
        s2_subject_ids=s2_subject_ids,
        probe_mixed_templates=probe_mixed_templates,
        probe_mixed_subject_ids=probe_mixed_subject_ids,
    )
    print()
    return s1_templates, s1_subject_ids, s2_templates, s2_subject_ids, probe_mixed_templates, probe_mixed_subject_ids

def get_embeddings(model_interf, img_names, landmarks, batch_size=64, flip=True):
    steps = int(np.ceil(len(img_names) / batch_size))
    embs, embs_f = [], []
    for batch_id in tqdm(range(0, len(img_names), batch_size), "Embedding", total=steps):
        batch_imgs, batch_landmarks = img_names[batch_id : batch_id + batch_size], landmarks[batch_id : batch_id + batch_size]
        ndimages = [face_align_landmark(cv2.imread(img), landmark) for img, landmark in zip(batch_imgs, batch_landmarks)]
        ndimages = np.stack(ndimages)
        embs.extend(model_interf(ndimages))
        if flip:
            embs_f.extend(model_interf(ndimages[:, :, ::-1, :]))
    return np.array(embs), np.array(embs_f)

def process_embeddings(embs, embs_f=[], use_flip_test=True, use_norm_score=False, use_detector_score=True, face_scores=None):
    print(">>>> process_embeddings: Norm {}, Detect_score {}, Flip {}".format(use_norm_score, use_detector_score, use_flip_test))
    if use_flip_test and len(embs_f) != 0:
        embs = embs + embs_f
    if use_norm_score:
        embs = normalize(embs)
    if use_detector_score and face_scores is not None:
        embs = embs * np.expand_dims(face_scores, -1)
    return embs

def image2template_feature(img_feats=None, templates=None, medias=None, choose_templates=None, choose_ids=None):
    if choose_templates is not None:  # 1:N
        unique_templates, indices = np.unique(choose_templates, return_index=True)
        unique_subjectids = choose_ids[indices]
    else:  # 1:1
        unique_templates = np.unique(templates)
        unique_subjectids = None

    # template_feats = np.zeros((len(unique_templates), img_feats.shape[1]), dtype=img_feats.dtype)
    template_feats = np.zeros((len(unique_templates), img_feats.shape[1]))
    for count_template, uqt in tqdm(enumerate(unique_templates), "Extract template feature", total=len(unique_templates)):
        (ind_t,) = np.where(templates == uqt)
        face_norm_feats = img_feats[ind_t]
        face_medias = medias[ind_t]
        unique_medias, unique_media_counts = np.unique(face_medias, return_counts=True)
        media_norm_feats = []
        for u, ct in zip(unique_medias, unique_media_counts):
            (ind_m,) = np.where(face_medias == u)
            if ct == 1:
                media_norm_feats += [face_norm_feats[ind_m]]
            else:  # image features from the same video will be aggregated into one feature
                media_norm_feats += [np.mean(face_norm_feats[ind_m], 0, keepdims=True)]
        media_norm_feats = np.array(media_norm_feats)
        # media_norm_feats = media_norm_feats / np.sqrt(np.sum(media_norm_feats ** 2, -1, keepdims=True))
        template_feats[count_template] = np.sum(media_norm_feats, 0)
    template_norm_feats = normalize(template_feats)
    return template_norm_feats, unique_templates, unique_subjectids

def verification_11(template_norm_feats=None, unique_templates=None, p1=None, p2=None, batch_size=10000):
    try:
        print(">>>> Trying cupy.")
        import cupy as cp

        template_norm_feats = cp.array(template_norm_feats)
        score_func = lambda feat1, feat2: cp.sum(feat1 * feat2, axis=-1).get()
        test = score_func(template_norm_feats[:batch_size], template_norm_feats[:batch_size])
    except:
        score_func = lambda feat1, feat2: np.sum(feat1 * feat2, -1)

    template2id = np.zeros(max(unique_templates) + 1, dtype=int)
    template2id[unique_templates] = np.arange(len(unique_templates))

    steps = int(np.ceil(len(p1) / batch_size))
    score = []
    for id in tqdm(range(steps), "Verification"):
        feat1 = template_norm_feats[template2id[p1[id * batch_size : (id + 1) * batch_size]]]
        feat2 = template_norm_feats[template2id[p2[id * batch_size : (id + 1) * batch_size]]]
        score.extend(score_func(feat1, feat2))
    return np.array(score)

def evaluation_1N(query_feats, gallery_feats, query_ids, reg_ids, fars=[0.01, 0.1]):
    print("query_feats: %s, gallery_feats: %s" % (query_feats.shape, gallery_feats.shape))
    similarity = np.dot(query_feats, gallery_feats.T)  # (19593, 3531)

    top_1_count, top_5_count, top_10_count = 0, 0, 0
    pos_sims, neg_sims, non_gallery_sims = [], [], []
    for index, query_id in enumerate(query_ids):
        if query_id in reg_ids:
            gallery_label = np.argwhere(reg_ids == query_id)[0, 0]
            index_sorted = np.argsort(similarity[index])[::-1]

            top_1_count += gallery_label in index_sorted[:1]
            top_5_count += gallery_label in index_sorted[:5]
            top_10_count += gallery_label in index_sorted[:10]

            pos_sims.append(similarity[index][reg_ids == query_id][0])
            neg_sims.append(similarity[index][reg_ids != query_id])
        else:
            non_gallery_sims.append(similarity[index])
    total_pos = len(pos_sims)
    pos_sims, neg_sims, non_gallery_sims = np.array(pos_sims), np.array(neg_sims), np.array(non_gallery_sims)
    print("pos_sims: %s, neg_sims: %s, non_gallery_sims: %s" % (pos_sims.shape, neg_sims.shape, non_gallery_sims.shape))
    print("top1: %f, top5: %f, top10: %f" % (top_1_count / total_pos, top_5_count / total_pos, top_10_count / total_pos))

    correct_pos_cond = pos_sims > neg_sims.max(1)
    non_gallery_sims_sorted = np.sort(non_gallery_sims.max(1))[::-1]
    threshes, recalls = [], []
    for far in fars:
        # thresh = non_gallery_sims_sorted[int(np.ceil(non_gallery_sims_sorted.shape[0] * far)) - 1]
        thresh = non_gallery_sims_sorted[max(int((non_gallery_sims_sorted.shape[0]) * far) - 1, 0)]
        recall = np.logical_and(correct_pos_cond, pos_sims > thresh).sum() / pos_sims.shape[0]
        threshes.append(thresh)
        recalls.append(recall)
        # print("FAR = {:.10f} TPIR = {:.10f} th = {:.10f}".format(far, recall, thresh))
    cmc_scores = list(zip(neg_sims, pos_sims.reshape(-1, 1))) + list(zip(non_gallery_sims, [None] * non_gallery_sims.shape[0]))
    return top_1_count, top_5_count, top_10_count, threshes, recalls, cmc_scores

class IJB_test:
    def __init__(self, model_file, data_path, subset, batch_size=64, force_reload=False, restore_embs=None):
        templates, medias, p1, p2, label, img_names, landmarks, face_scores = extract_IJB_data_11(
            data_path, subset, force_reload=force_reload
        )
        if model_file != None:
            if model_file.endswith(".h5"):
                interf_func = keras_model_interf(model_file)
            elif model_file.endswith(".pth") or model_file.endswith(".pt"):
                interf_func = Torch_model_interf(model_file)
            elif model_file.endswith(".onnx") or model_file.endswith(".ONNX"):
                interf_func = ONNX_model_interf(model_file)
            else:
                interf_func = Mxnet_model_interf(model_file)
            self.embs, self.embs_f = get_embeddings(interf_func, img_names, landmarks, batch_size=batch_size)
        elif restore_embs != None:
            print(">>>> Reload embeddings from:", restore_embs)
            aa = np.load(restore_embs)
            if "embs" in aa and "embs_f" in aa:
                self.embs, self.embs_f = aa["embs"], aa["embs_f"]
            else:
                print("ERROR: %s NOT containing embs / embs_f" % restore_embs)
                exit(1)
            print(">>>> Done.")
        self.data_path, self.subset, self.force_reload = data_path, subset, force_reload
        self.templates, self.medias, self.p1, self.p2, self.label = templates, medias, p1, p2, label
        self.face_scores = face_scores.astype(self.embs.dtype)

    def run_model_test_single(self, use_flip_test=True, use_norm_score=False, use_detector_score=True):
        img_input_feats = process_embeddings(
            self.embs,
            self.embs_f,
            use_flip_test=use_flip_test,
            use_norm_score=use_norm_score,
            use_detector_score=use_detector_score,
            face_scores=self.face_scores,
        )
        template_norm_feats, unique_templates, _ = image2template_feature(img_input_feats, self.templates, self.medias)
        score = verification_11(template_norm_feats, unique_templates, self.p1, self.p2)
        return score

    def run_model_test_bunch(self):
        from itertools import product

        scores, names = [], []
        for use_norm_score, use_detector_score, use_flip_test in product([True, False], [True, False], [True, False]):
            name = "N{:d}D{:d}F{:d}".format(use_norm_score, use_detector_score, use_flip_test)
            print(">>>>", name, use_norm_score, use_detector_score, use_flip_test)
            names.append(name)
            scores.append(self.run_model_test_single(use_flip_test, use_norm_score, use_detector_score))
        return scores, names

    def run_model_test_1N(self, npoints=100):
        fars_cal = [10 ** ii for ii in np.arange(-4, 0, 4 / npoints)] + [1]  # plot in range [10-4, 1]
        fars_show_idx = np.arange(len(fars_cal))[:: npoints // 4]  # npoints=100, fars_show=[0.0001, 0.001, 0.01, 0.1, 1.0]

        g1_templates, g1_ids, g2_templates, g2_ids, probe_mixed_templates, probe_mixed_ids = extract_gallery_prob_data(
            self.data_path, self.subset, force_reload=self.force_reload
        )
        img_input_feats = process_embeddings(
            self.embs,
            self.embs_f,
            use_flip_test=True,
            use_norm_score=False,
            use_detector_score=True,
            face_scores=self.face_scores,
        )
        g1_templates_feature, g1_unique_templates, g1_unique_ids = image2template_feature(
            img_input_feats, self.templates, self.medias, g1_templates, g1_ids
        )
        g2_templates_feature, g2_unique_templates, g2_unique_ids = image2template_feature(
            img_input_feats, self.templates, self.medias, g2_templates, g2_ids
        )
        probe_mixed_templates_feature, probe_mixed_unique_templates, probe_mixed_unique_subject_ids = image2template_feature(
            img_input_feats, self.templates, self.medias, probe_mixed_templates, probe_mixed_ids
        )
        print("g1_templates_feature:", g1_templates_feature.shape)  # (1772, 512)
        print("g2_templates_feature:", g2_templates_feature.shape)  # (1759, 512)

        print("probe_mixed_templates_feature:", probe_mixed_templates_feature.shape)  # (19593, 512)
        print("probe_mixed_unique_subject_ids:", probe_mixed_unique_subject_ids.shape)  # (19593,)

        print(">>>> Gallery 1")
        g1_top_1_count, g1_top_5_count, g1_top_10_count, g1_threshes, g1_recalls, g1_cmc_scores = evaluation_1N(
            probe_mixed_templates_feature, g1_templates_feature, probe_mixed_unique_subject_ids, g1_unique_ids, fars_cal
        )
        print(">>>> Gallery 2")
        g2_top_1_count, g2_top_5_count, g2_top_10_count, g2_threshes, g2_recalls, g2_cmc_scores = evaluation_1N(
            probe_mixed_templates_feature, g2_templates_feature, probe_mixed_unique_subject_ids, g2_unique_ids, fars_cal
        )
        print(">>>> Mean")
        query_num = probe_mixed_templates_feature.shape[0]
        top_1 = (g1_top_1_count + g2_top_1_count) / query_num
        top_5 = (g1_top_5_count + g2_top_5_count) / query_num
        top_10 = (g1_top_10_count + g2_top_10_count) / query_num
        print("[Mean] top1: %f, top5: %f, top10: %f" % (top_1, top_5, top_10))

        mean_tpirs = (np.array(g1_recalls) + np.array(g2_recalls)) / 2
        show_result = {}
        for id, far in enumerate(fars_cal):
            if id in fars_show_idx:
                show_result.setdefault("far", []).append(far)
                show_result.setdefault("g1_tpir", []).append(g1_recalls[id])
                show_result.setdefault("g1_thresh", []).append(g1_threshes[id])
                show_result.setdefault("g2_tpir", []).append(g2_recalls[id])
                show_result.setdefault("g2_thresh", []).append(g2_threshes[id])
                show_result.setdefault("mean_tpir", []).append(mean_tpirs[id])
        print(pd.DataFrame(show_result).set_index("far").to_markdown())
        return fars_cal, mean_tpirs, g1_cmc_scores, g2_cmc_scores

def plot_roc_and_calculate_tpr(scores, names=None, label=None):
    print(">>>> plot roc and calculate tpr...")
    score_dict = {}
    for id, score in enumerate(scores):
        name = None if names is None else names[id]
        if isinstance(score, str) and score.endswith(".npz"):
            aa = np.load(score)
            score = aa.get("scores", [])
            label = aa["label"] if label is None and "label" in aa else label
            score_name = aa.get("names", [])
            for ss, nn in zip(score, score_name):
                score_dict[nn] = ss
        elif isinstance(score, str) and score.endswith(".npy"):
            name = name if name is not None else os.path.splitext(os.path.basename(score))[0]
            score_dict[name] = np.load(score)
        elif isinstance(score, str) and score.endswith(".txt"):
            # IJB meta data like ijbb_template_pair_label.txt
            label = pd.read_csv(score, sep=" ", header=None).values[:, 2]
        else:
            name = name if name is not None else str(id)
            score_dict[name] = score
    if label is None:
        print("Error: Label data is not provided")
        return None, None

    x_labels = [10 ** (-ii) for ii in range(1, 7)[::-1]]
    fpr_dict, tpr_dict, roc_auc_dict, tpr_result = {}, {}, {}, {}
    for name, score in score_dict.items():
        fpr, tpr, _ = roc_curve(label, score)
        roc_auc = auc(fpr, tpr)
        fpr, tpr = np.flipud(fpr), np.flipud(tpr)  # select largest tpr at same fpr
        tpr_result[name] = [tpr[np.argmin(abs(fpr - ii))] for ii in x_labels]
        fpr_dict[name], tpr_dict[name], roc_auc_dict[name] = fpr, tpr, roc_auc
    tpr_result_df = pd.DataFrame(tpr_result, index=x_labels).T
    tpr_result_df['AUC'] = pd.Series(roc_auc_dict)
    tpr_result_df.columns.name = "Methods"
    print(tpr_result_df.to_markdown())
    # print(tpr_result_df)

    try:
        import matplotlib.pyplot as plt

        fig = plt.figure()
        for name in score_dict:
            plt.plot(fpr_dict[name], tpr_dict[name], lw=1, label="[%s (AUC = %0.4f%%)]" % (name, roc_auc_dict[name] * 100))
        title = "ROC on IJB" + name.split("IJB")[-1][0] if "IJB" in name else "ROC on IJB"

        plt.xlim([10 ** -6, 0.1])
        plt.xscale("log")
        plt.xticks(x_labels)
        plt.xlabel("False Positive Rate")
        plt.ylim([0.3, 1.0])
        plt.yticks(np.linspace(0.3, 1.0, 8, endpoint=True))
        plt.ylabel("True Positive Rate")

        plt.grid(linestyle="--", linewidth=1)
        plt.title(title)
        plt.legend(loc="lower right", fontsize='x-small')
        plt.tight_layout()
        plt.show()
    except:
        print("matplotlib plot failed")
        fig = None

    return tpr_result_df, fig

def plot_dir_far_cmc_scores(scores, names=None):
    try:
        import matplotlib.pyplot as plt

        fig = plt.figure()
        for id, score in enumerate(scores):
            name = None if names is None else names[id]
            if isinstance(score, str) and score.endswith(".npz"):
                aa = np.load(score)
                score, name = aa.get("scores")[0], aa.get("names")[0]
            fars, tpirs = score[0], score[1]
            name = name if name is not None else str(id)

            auc_value = auc(fars, tpirs)
            label = "[%s (AUC = %0.4f%%)]" % (name, auc_value * 100)
            plt.plot(fars, tpirs, lw=1, label=label)

        plt.xlabel("False Alarm Rate")
        plt.xlim([0.0001, 1])
        plt.xscale("log")
        plt.ylabel("Detection & Identification Rate (%)")
        plt.ylim([0, 1])

        plt.grid(linestyle="--", linewidth=1)
        plt.legend(fontsize='x-small')
        plt.tight_layout()
        plt.show()
    except:
        print("matplotlib plot failed")
        fig = None

    return fig

def parse_arguments(argv):
    import argparse

    default_save_result_name = "IJB_result/{model_name}_{subset}_{type}.npz"
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-m", "--model_file", type=str, default=None, help="Saved model, keras h5 / pytorch jit pth / onnx / mxnet")
    parser.add_argument("-d", "--data_path", type=str, default="./", help="Dataset path containing IJBB and IJBC sub folder")
    parser.add_argument("-s", "--subset", type=str, default="IJBC", help="Subset test target, could be IJBB / IJBC")
    parser.add_argument("-b", "--batch_size", type=int, default=128, help="Batch size for get_embeddings")
    parser.add_argument(
        "-R", "--save_result", type=str, default=default_save_result_name, help="Filename for saving / restore result"
    )
    parser.add_argument("-L", "--save_label", action="store_true", help="Save label data, useful for plot only")
    parser.add_argument("-E", "--save_embeddings", action="store_true", help="Save embeddings data")
    parser.add_argument("-B", "--is_bunch", action="store_true", help="Run all 8 tests N{0,1}D{0,1}F{0,1}")
    parser.add_argument("-N", "--is_one_2_N", action="store_true", help="Run 1:N test instead of 1:1")
    parser.add_argument("-F", "--force_reload", action="store_true", help="Force reload, instead of using cache")
    parser.add_argument("-P", "--plot_only", nargs="*", type=str, help="Plot saved results, Format 1 2 3 or 1, 2, 3 or *.npy")
    args = parser.parse_known_args(argv)[0]

    if args.plot_only != None and len(args.plot_only) != 0:
        # Plot only
        from glob2 import glob

        score_files = []
        for ss in args.plot_only:
            score_files.extend(glob(ss.replace(",", "").strip()))
        args.plot_only = score_files
    elif args.model_file == None and args.save_result == default_save_result_name:
        print("Please provide -m MODEL_FILE, see `--help` for usage.")
        exit(1)
    elif args.model_file != None:
        if args.model_file.endswith(".h5") or args.model_file.endswith(".pth") or args.model_file.endswith(".pt") or args.model_file.endswith(".onnx"):
            # Keras model file "model.h5", pytorch model ends with `.pth` or `.pt`, onnx model ends with `.onnx`
            model_name = os.path.splitext(os.path.basename(args.model_file))[0]
        else:
            # MXNet model file "models/r50-arcface-emore/model,1"
            model_name = os.path.basename(os.path.dirname(args.model_file))

        if args.save_result == default_save_result_name:
            type = "1N" if args.is_one_2_N else "11"
            args.save_result = default_save_result_name.format(model_name=model_name, subset=args.subset, type=type)
    return args

if __name__ == "__main__":
    import sys

    args = parse_arguments(sys.argv[1:])
    if args.plot_only != None and len(args.plot_only) != 0:
        if args.is_one_2_N:
            plot_dir_far_cmc_scores(args.plot_only)
        else:
            plot_roc_and_calculate_tpr(args.plot_only)
    else:
        save_name = os.path.splitext(os.path.basename(args.save_result))[0]
        save_items = {}
        save_path = os.path.dirname(args.save_result)
        if len(save_path) != 0 and not os.path.exists(save_path):
            os.makedirs(save_path)

        tt = IJB_test(args.model_file, args.data_path, args.subset, args.batch_size, args.force_reload, args.save_result)
        if args.save_embeddings:  # Save embeddings first, in case of any error happens later...
            np.savez(args.save_result, embs=tt.embs, embs_f=tt.embs_f)

        if args.is_one_2_N:  # 1:N test
            fars, tpirs, _, _ = tt.run_model_test_1N()
            scores = [(fars, tpirs)]
            names = [save_name]
            save_items.update({"scores": scores, "names": names})
        elif args.is_bunch:  # All 8 tests N{0,1}D{0,1}F{0,1}
            scores, names = tt.run_model_test_bunch()
            names = [save_name + "_" + ii for ii in names]
            label = tt.label
            save_items.update({"scores": scores, "names": names})
        else:  # Basic 1:1 N0D1F1 test
            score = tt.run_model_test_single()
            scores, names, label = [score], [save_name], tt.label
            save_items.update({"scores": scores, "names": names})

        if args.save_embeddings:
            save_items.update({"embs": tt.embs, "embs_f": tt.embs_f})
        if args.save_label:
            save_items.update({"label": label})

        if args.model_file != None or args.save_embeddings:  # embeddings not restored from file or should save_embeddings again
            np.savez(args.save_result, **save_items)

        if args.is_one_2_N:
            plot_dir_far_cmc_scores(scores=scores, names=names)
        else:
            plot_roc_and_calculate_tpr(scores, names=names, label=label)


import collections
import numpy as np

ObsAct = collections.namedtuple('ObsAct', ('obs', 'action'))
ObsActExp = collections.namedtuple('ObsActExp', ('obs', 'action', 'exp'))

class FailurePredictor:
    def __init__(self,
                 q_model,
                 q_type='td3',
                 k_prev=3,
                 gap_between_k_prev=1,
                 req_consec_under=3,
                 backend='tf',
                 q_model_step=None,
                 use_delta_q=True,
                 discrim=None):

        self.q_model = q_model
        self.q_model_step = q_model_step
        self._use_delta_q = use_delta_q
        self._last_q = -1e100
        self.req_consec_under = req_consec_under
        self.q_type = q_type
        self.k_prev = k_prev  # how many precorrective sa pairs to add per each transition to expert data
        self.gap_between_k_prev = gap_between_k_prev
        self.backend = backend
        self.pre_corrective_pairs = []  # list containing ObsAct(o, a)
        self.episode_pairs = []  # list containg ObsActExp(o, a , e)
        self.q_values = np.array([])
        self.max_q = -1e100
        self.max_q_index = None
        self.tolerance = 0.01  # adjust the sensitivity, higher means less likely to indicate failure
        self._discrim = discrim
        self._num_consec_under = 0

    def _get_q_output(self, obs, action):
        if self.q_type == 'td3':
            if len(obs.shape) == 1:
                out, _ = self.q_model([obs], [action])
            else:
                out, _ = self.q_model(obs, action)
        return out

    def _get_batch(self):
        batch = ObsAct(*zip(*self.pre_corrective_pairs))
        return self._prep_batch_for_model(batch)

    def _get_batch_delta(self):
        batch_cur = ObsAct(*zip(*list(zip(*self.pre_corrective_pairs))[0]))
        batch_prev = ObsAct(*zip(*list(zip(*self.pre_corrective_pairs))[1]))
        return self._prep_batch_for_model(batch_cur), self._prep_batch_for_model(batch_prev)

    def _prep_obs_act_for_model(self, obs, action):
        if self.backend == 'tf':
            from tensorflow.contrib.eager.python import tfe as contrib_eager_python_tfe
            obs = contrib_eager_python_tfe.Variable(obs.astype('float32'))
            action = contrib_eager_python_tfe.Variable(action.astype('float32'))
        else:
            return NotImplementedError("Failure Predictor not implemented for backend %s" % self.backend)

        return obs, action

    def _prep_batch_for_model(self, batch):
        """
        Prepares batch of observations and actions for running through q model.
        :param batch: Should be an ObsAct containing all observations in obs and all actions in action
        :return:
        """
        return self._prep_obs_act_for_model(np.stack(batch.obs).astype('float32'),
                                            np.stack(batch.action).astype('float32'))

    def update_q_values(self):
        # user calls this after updating q model and before starting next episode
        if len(self.pre_corrective_pairs) > 0:
            if self._use_delta_q:
                self._update_q_values_delta()
            else:
                obs, action = self._get_batch()
                self.q_values = np.array(self._get_q_output(obs, action))
            self._update_max()

    def _update_q_values_delta(self):
        [obs, action], [prev_obs, prev_action] = self._get_batch_delta()
        q_values_cur = np.array(self._get_q_output(obs, action))
        q_values_prev = np.array(self._get_q_output(prev_obs, prev_action))
        self.q_values = q_values_cur - q_values_prev

    def _update_max(self):
        if len(self.pre_corrective_pairs) > 0:
            self.max_q = self.q_values.max()
            if self.max_q > 0:
                pass
                # import ipdb; ipdb.set_trace()
            if self._use_delta_q:
                self.max_q = min(self.max_q, 0)  # no threshold should be allowed to be above 0
            self.max_q_index = self.q_values.argmax()
        else:
            self.max_q = -1e100
            self.max_q_index = None

    def _push_back_pc(self, obs, action):
        self.pre_corrective_pairs.append(ObsAct(obs, action))

    def _push_back_pc_delta(self, obs, action, prev_obs, prev_action):
        self.pre_corrective_pairs.append([ObsAct(obs, action), ObsAct(prev_obs, prev_action)])

    def push_back_ep(self, obs, action, exp):
        # user should call this on every step of an episode
        # exp is a mask indicating whether action was expert or not, 1 for exp, 0 for not
        self.episode_pairs.append(ObsActExp(obs, action, exp))

    def update_pc_list_from_ep(self):
        # user calls this at the end of an episode
        test_REMOVE = False
        if len(self.episode_pairs) > 0:
            batch = ObsActExp(*zip(*self.episode_pairs))
            obs = np.stack(batch.obs)
            action = np.stack(batch.action)
            exp = np.stack(batch.exp)
            for i in range(obs.shape[0] - 1):
                if exp[i + 1] and not exp[i]:  # means we're at a transition to expert data
                    test_REMOVE = True
                    for k in range(self.k_prev):
                        ind = i - k * self.gap_between_k_prev
                        if exp[ind] or ind < 0: break
                        if self._use_delta_q and (exp[ind - 1] or ind - 1 < 0): break
                        if self._use_delta_q:
                            self._push_back_pc_delta(obs[ind], action[ind], obs[ind - 1], action[ind - 1])
                        else:
                            self._push_back_pc(obs[ind], action[ind])

            if test_REMOVE:
                pass
                # import ipdb; ipdb.set_trace()
            self.episode_pairs = []
            self._last_q = -1e100

    def worse_than_best(self, obs, action):
        # called by user to check if current o-a pair predicts failure
        # if self.q_model_step is not None and self.q_model_step < 300:
        #   return False
        prepped = self._prep_batch_for_model(ObsAct(obs, action))
        q_pred = np.array(self._get_q_output(prepped[0], prepped[1]))
        if self._use_delta_q:
            pred_delta_q = q_pred - self._last_q
            self._last_q = q_pred
            under_thresh = pred_delta_q < self.max_q - self.tolerance
        else:
            under_thresh = q_pred < self.max_q - self.tolerance

        self._num_consec_under = self._num_consec_under + 1 if under_thresh else 0
        if under_thresh:
            print('delq: %2.3f, max q: %2.3f, num_consec: %d' % (
            pred_delta_q, self.max_q - self.tolerance, self._num_consec_under))
        if self._use_delta_q:
            return under_thresh and self._num_consec_under >= self.req_consec_under, pred_delta_q, self.max_q
        else:
            return under_thresh and self._num_consec_under >= self.req_consec_under, q_pred, self.max_q

    def remove_best(self):
        if len(self.pre_corrective_pairs) > 0:
            # called by user if worse_than_best gave incorrect output
            self.pre_corrective_pairs.pop(self.max_q_index)
            self.q_values = np.delete(self.q_values, self.max_q_index)
            self._update_max()


# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
""" test transformer"""
import numpy as np
from mindspore import Tensor
from mindspore.common import dtype
from mindspore.nn.parallel import MultiHeadAttention, FeedForward, TransformerEncoderLayer, TransformerEncoder, \
    TransformerDecoder, TransformerDecoderLayer, Transformer
from mindspore.common.api import _executor

def test_transformer_encoder_only():
    model = Transformer(encoder_layers=2,
                        decoder_layers=0,
                        hidden_size=64,
                        ffn_hidden_size=64,
                        src_seq_length=16,
                        tgt_seq_length=32)

    encoder_input_value = Tensor(np.ones((2, 20, 64)), dtype.float32)
    encoder_input_mask = Tensor(np.ones((2, 1, 20, 20)), dtype.float16)

    _executor.compile(model, encoder_input_value, encoder_input_mask)

def test_encoder_and_decoder():
    model = Transformer(encoder_layers=1,
                        decoder_layers=2,
                        hidden_size=64,
                        ffn_hidden_size=64,
                        src_seq_length=20,
                        tgt_seq_length=20)

    encoder_input_value = Tensor(np.ones((2, 20, 64)), dtype.float32)
    encoder_input_mask = Tensor(np.ones((2, 1, 20, 20)), dtype.float16)

    decoder_input_value = Tensor(np.ones((2, 10, 64)), dtype.float32)
    decoder_input_mask = Tensor(np.ones((2, 1, 10, 10)), dtype.float16)
    memory_mask = Tensor(np.ones((2, 1, 10, 20)), dtype.float16)

    _executor.compile(model, encoder_input_value, encoder_input_mask,
                      decoder_input_value,
                      decoder_input_mask,
                      memory_mask)

def test_transformer_encoder():
    model = TransformerEncoder(num_layers=2,
                               hidden_size=8,
                               ffn_hidden_size=64,
                               seq_length=16,
                               num_heads=2)

    encoder_input_value = Tensor(np.ones((2, 16, 8)), dtype.float32)
    encoder_input_mask = Tensor(np.ones((2, 1, 16, 16)), dtype.float16)

    _executor.compile(model,
                      encoder_input_value,
                      encoder_input_mask)

def test_transformer_encoder_layer():
    model = TransformerEncoderLayer(hidden_size=8, ffn_hidden_size=64, seq_length=16,
                                    num_heads=2)

    encoder_input_value = Tensor(np.ones((2, 16, 8)), dtype.float32)
    encoder_input_mask = Tensor(np.ones((2, 1, 16, 16)), dtype.float16)

    _executor.compile(model,
                      encoder_input_value,
                      encoder_input_mask)

def test_transformer_encoder_layer_post_ture():
    model = TransformerEncoderLayer(hidden_size=8, ffn_hidden_size=64, seq_length=16,
                                    num_heads=2, post_layernorm_residual=True)

    encoder_input_value = Tensor(np.ones((2, 16, 8)), dtype.float32)
    encoder_input_mask = Tensor(np.ones((2, 1, 16, 16)), dtype.float16)

    _executor.compile(model,
                      encoder_input_value,
                      encoder_input_mask)

def test_transformer_decoder():
    model = TransformerDecoder(num_layers=1,
                               hidden_size=64,
                               ffn_hidden_size=64,
                               num_heads=2,
                               seq_length=10)

    encoder_input_value = Tensor(np.ones((2, 20, 64)), dtype.float32)

    decoder_input_value = Tensor(np.ones((2, 10, 64)), dtype.float32)
    decoder_input_mask = Tensor(np.ones((2, 1, 10, 10)), dtype.float16)
    memory_mask = Tensor(np.ones((2, 1, 10, 20)), dtype.float16)

    _executor.compile(model, decoder_input_value, decoder_input_mask,
                      encoder_input_value,
                      memory_mask)

def test_transformer_decoder_layer():
    model = TransformerDecoderLayer(
        hidden_size=64,
        ffn_hidden_size=64,
        num_heads=2,
        seq_length=10)

    encoder_input_value = Tensor(np.ones((2, 20, 64)), dtype.float32)

    decoder_input_value = Tensor(np.ones((2, 10, 64)), dtype.float32)
    decoder_input_mask = Tensor(np.ones((2, 1, 10, 10)), dtype.float16)
    memory_mask = Tensor(np.ones((2, 1, 10, 20)), dtype.float16)

    _executor.compile(model, decoder_input_value, decoder_input_mask,
                      encoder_input_value,
                      memory_mask)

def test_multihead_attention():
    model = MultiHeadAttention(hidden_size=15,
                               num_heads=3)
    from_tensor = Tensor(np.ones((2, 20, 15)), dtype.float32)
    to_tensor = Tensor(np.ones((2, 20, 15)), dtype.float16)
    attention_mask = Tensor(np.ones((2, 1, 20, 20)), dtype.float16)

    _executor.compile(model, from_tensor, to_tensor, attention_mask)

def test_feedforward_layer():
    model = FeedForward(hidden_size=15,
                        ffn_hidden_size=30,
                        dropout_rate=0.1,
                        hidden_act='relu')
    tensor = Tensor(np.ones((2, 20, 15)), dtype.float32)

    _executor.compile(model, tensor)


# -*- coding: utf-8 -*-
# SPDX-License-Identifier: BSD-3-Clause
# Author: Roman Feldbauer (version for scikit-hubness)
#         Daniel Obraczka (adaptation for kiez)
# PEP 563: Postponed Evaluation of Annotations

from __future__ import annotations

import numpy as np
from kiez.neighbors.neighbor_algorithm_base import NNAlgorithm

try:
    import nmslib
except ImportError:  # pragma: no cover
    nmslib = None

class HNSW(NNAlgorithm):
    """
    Wrapper for hierarchical navigable small world graphs based approximate nearest neighbor search

    Parameters
    ----------
    n_candidates: int
        number of nearest neighbors used in search
    metric: str, default = euclidean'
        distance measure used in search
        possible measures are found in :obj:`HNSW.valid_metrics`
    method: str, default = 'hnsw',
        ANN method to use. Currently, only 'hnsw' is supported.
    M: int, default = 16
        maximum number of neighbors in zero or above layers of hierarchical graph
    post_processing: int, default = 2
        More post processing means longer index creation,
        and higher retrieval accuracy.
    ef_construction: int, default = 200
        higher value improves quality of constructed graph but leads to longer indexing times
    n_jobs: int, default = 1
        Number of parallel jobs
    verbose: int, default = 0
        Verbosity level. If verbose > 0, show tqdm progress bar on indexing and querying.

    Notes
    -----
    See the nmslib documentation for more details: https://github.com/nmslib/nmslib/blob/master/manual/methods.md
    """

    valid_metrics = [
        "euclidean",
        "l2",
        "minkowski",
        "squared_euclidean",
        "sqeuclidean",
        "cosine",
        "cosinesimil",
    ]

    def __init__(
        self,
        n_candidates: int = 5,
        metric: str = "euclidean",
        method: str = "hnsw",
        M: int = 16,  # noqa: N803
        post_processing: int = 2,
        ef_construction: int = 200,
        n_jobs: int = 1,
        verbose: int = 0,
    ):

        self.space = None
        if nmslib is None:  # pragma: no cover
            raise ImportError(
                "Please install the `nmslib` package, before using this class.\n"
                "$ pip install nmslib"
            )

        if metric in [
            "euclidean",
            "l2",
            "minkowski",
            "squared_euclidean",
            "sqeuclidean",
        ]:
            if metric in ["squared_euclidean", "sqeuclidean"]:
                metric = "sqeuclidean"
            else:
                metric = "euclidean"
            self.space = "l2"
        elif metric in ["cosine", "cosinesimil"]:
            self.space = "cosinesimil"
        elif metric not in self.__class__.valid_metrics:
            raise ValueError(
                f"Unknown metric please try one of {self.__class__.valid_metrics}"
            )
        assert self.space in [
            "l2",
            "cosinesimil",
        ], f"Internal: self.space={self.space} not allowed"
        super().__init__(n_candidates=n_candidates, metric=metric, n_jobs=n_jobs)
        self.verbose = verbose
        self.method = method
        self.M = M  # noqa: N803
        self.post_processing = post_processing
        self.ef_construction = ef_construction

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(n_candidates={self.n_candidates}, "
            + f"verbose = {self.verbose}, "
            + f"method = {self.method}, "
            + f"M = {self.M}, "
            + f"space = {self.space}, "
            + f"post_processing = {self.post_processing}, "
            + f"ef_construction = {self.ef_construction}, "
            + f"n_jobs = {self.n_jobs}) "
        )

    def _fit(self, data, is_source: bool):
        method = self.method
        post_processing = self.post_processing
        big_m = self.M
        ef_construction = self.ef_construction

        hnsw_index = nmslib.init(method=method, space=self.space)
        hnsw_index.addDataPointBatch(data)
        hnsw_index.createIndex(
            {
                "M": big_m,
                "efConstruction": ef_construction,
                "post": post_processing,
                "indexThreadQty": self.n_jobs,
            },
            print_progress=(self.verbose >= 2),
        )
        return hnsw_index

    def _kneighbors(self, k, query, index, return_distance, is_self_querying):
        # Fetch the neighbor candidates
        neigh_ind_dist = index.knnQueryBatch(query, k=k, num_threads=self.n_jobs)

        # If fewer candidates than required are found for a query,
        # we save index=-1 and distance=NaN
        n_query = query.shape[0]
        neigh_ind = -np.ones((n_query, k), dtype=np.int32)
        neigh_dist = np.empty_like(neigh_ind, dtype=query.dtype) * np.nan

        for i, (ind, dist) in enumerate(neigh_ind_dist):
            neigh_ind[i, : ind.size] = ind
            neigh_dist[i, : dist.size] = dist

        # Convert cosine similarities to cosine distances
        if self.space == "cosinesimil":
            neigh_dist *= -1
            neigh_dist += 1
        elif self.space == "l2" and self.metric == "sqeuclidean":
            neigh_dist **= 2

        if return_distance:
            return neigh_dist, neigh_ind
        else:
            return neigh_ind


#!/usr/bin/env python
#
# Copyright (c) 2018 10X Genomics, Inc. All rights reserved.
#

"""Functions for summarizing a CountMatrix."""

from collections import OrderedDict
import h5py as h5
import numpy as np
import cellranger.constants as cr_constants
import cellranger.library_constants as lib_constants
import cellranger.rna.library as rna_library
import cellranger.matrix as cr_matrix
import cellranger.utils as cr_utils
import tenkit.stats as tk_stats

def _get_total_reads(summary, library_type):
    prefix = rna_library.get_library_type_metric_prefix(library_type)
    return int(summary.get(prefix + 'total_reads', 0))

def _get_conf_mapped_reads(summary, genomes, library_type):
    prefix = rna_library.get_library_type_metric_prefix(library_type)
    conf_mapped_metrics = [prefix + '_'.join([
        ref,
        cr_constants.TRANSCRIPTOME_REGION,
        cr_constants.CONF_MAPPED_READ_TYPE,
        'reads_frac']) for ref in genomes]
    total_reads = _get_total_reads(summary, library_type)
    return sum(float(summary.get(metric, 0)) * float(total_reads) for metric in conf_mapped_metrics)

def _get_barcode_summary_h5_indices(summary_h5, barcodes):
    """Get the indices into a barcode summary HDF5 file for the given barcodes.

    Does NOT preserve order of the barcodes argument.
    Args:
      summary_h5 (h5py.File): HDF5 file with datasets
      barcodes (iterable of str): List of barcode strings (e.g., "ACGT-1") to get.
    Returns:
      np.array(int): List of indices into the HDF5 arrays
    """
    if len(barcodes) == 0:
        return np.zeros(0, dtype=int)
    max_len = max(len(bc) for bc in barcodes)
    bc_arr = np.fromiter(barcodes, count=len(barcodes), dtype='S%d' % max_len)
    return np.flatnonzero(np.isin(summary_h5['bc_sequence'][:], bc_arr))

def _report_genome_agnostic_metrics(matrix, barcode_summary_h5, recovered_cells,
                                    cell_bc_seqs,
                                    total_reads,
                                    total_conf_mapped_reads,
                                    library_prefix):
    """Report metrics that are computed across all barcodes and all genomes"""
    d = {}

    matrix = matrix.view()

    genomes = matrix.get_genomes()
    if len(genomes) == 0:
        # For genomeless features, use "multi" in place of the genome prefix
        genomes = [lib_constants.MULTI_REFS_PREFIX]

    # Get number of cell bcs across all genomes
    cell_bcs_union = reduce(lambda a,x: a | set(x), cell_bc_seqs.itervalues(), set())
    n_cell_bcs_union = len(cell_bcs_union)

    d['filtered_bcs_transcriptome_union'] = n_cell_bcs_union
    d['%s_filtered_bcs' % lib_constants.MULTI_REFS_PREFIX] = n_cell_bcs_union

    # Report reads/cell across all genomes
    mean_reads_per_cell = tk_stats.robust_divide(total_reads, n_cell_bcs_union)
    d['%s_%s_total_raw_reads_per_filtered_bc' % (lib_constants.MULTI_REFS_PREFIX, cr_constants.TRANSCRIPTOME_REGION)] = mean_reads_per_cell
    # Create a feature-barcode dual whose name makes sense
    d['reads_per_cell'] = mean_reads_per_cell

    d['%s_%s_total_conf_mapped_reads_per_filtered_bc' % (lib_constants.MULTI_REFS_PREFIX, cr_constants.TRANSCRIPTOME_REGION)] = tk_stats.robust_divide(total_conf_mapped_reads, n_cell_bcs_union)

    # Split the matrix by genome
    if genomes[0] != lib_constants.MULTI_REFS_PREFIX:
        genome_matrices = OrderedDict(((g, matrix.select_features_by_genome(g)) for g in genomes))
    else:
        # Genomeless feature types
        genome_matrices = OrderedDict(((lib_constants.MULTI_REFS_PREFIX, matrix),))

    # Total UMI counts across all matrices and all filtered barcodes
    total_umi_counts = 0
    for mat in genome_matrices.itervalues():
        total_umi_counts += mat.select_barcodes_by_seq(cell_bcs_union).sum()

    # Deviation from cell load
    if recovered_cells is None:
        d['%s_filtered_bcs_difference_from_recovered_cells' % lib_constants.MULTI_REFS_PREFIX] = 0
        d['%s_filtered_bcs_relative_difference_from_recovered_cells' % lib_constants.MULTI_REFS_PREFIX] = 0
    else:
        d['%s_filtered_bcs_difference_from_recovered_cells' % lib_constants.MULTI_REFS_PREFIX] = int(n_cell_bcs_union) - int(recovered_cells)
        d['%s_filtered_bcs_relative_difference_from_recovered_cells' % lib_constants.MULTI_REFS_PREFIX] = tk_stats.robust_divide(n_cell_bcs_union - recovered_cells, recovered_cells)

    # Duplicate these metrics across genomes for backwards-compat
    for genome in genomes:
        d['%s_total_raw_reads_per_filtered_bc' % genome] = tk_stats.robust_divide(total_reads, n_cell_bcs_union)
        d['%s_total_conf_mapped_reads_per_filtered_bc' % genome] = tk_stats.robust_divide(total_conf_mapped_reads, n_cell_bcs_union)

        for read_type in cr_constants.MATRIX_REPORT_READ_TYPES:
            metric = '%s_total_%s_reads_per_filtered_bc' % (genome, read_type)
            if read_type in cr_constants.MATRIX_USE_MATRIX_FOR_READ_TYPE:
                n_reads = total_umi_counts
            else:
                # Note the extra underscore after the library_type prefix.
                # This is induced by the Reporter framework.
                h5_keys = [cr_utils.format_barcode_summary_h5_key(
                    library_prefix, g,
                    cr_constants.TRANSCRIPTOME_REGION,
                    read_type) for g in genomes]
                h5_keys = [x for x in h5_keys if x in barcode_summary_h5]
                n_reads = sum(np.array(barcode_summary_h5[h5_key]).sum() for h5_key in h5_keys)
            d[metric] = tk_stats.robust_divide(n_reads, n_cell_bcs_union)

    # Report frac reads in cells across all genomes
    total_conf_mapped_reads_in_cells = 0
    total_conf_mapped_barcoded_reads = 0

    for genome, g_mat in genome_matrices.iteritems():
        h5_key = '%s_%s_%s_%s_reads' % (library_prefix, genome, cr_constants.TRANSCRIPTOME_REGION,
                                        cr_constants.CONF_MAPPED_BC_READ_TYPE)
        cmb_reads = barcode_summary_h5[h5_key][:]
        cell_bc_indices = _get_barcode_summary_h5_indices(barcode_summary_h5,
                                                          cell_bcs_union)
        total_conf_mapped_reads_in_cells += cmb_reads[cell_bc_indices].sum()
        total_conf_mapped_barcoded_reads += cmb_reads.sum()
    frac_reads_in_cells = tk_stats.robust_divide(total_conf_mapped_reads_in_cells, total_conf_mapped_barcoded_reads)
    d['multi_filtered_bcs_conf_mapped_barcoded_reads_cum_frac'] = frac_reads_in_cells
    # Create a feature-barcode dual whose name makes sense
    d['feature_reads_in_cells'] = frac_reads_in_cells

    # Compute fraction of reads usable (conf mapped, barcoded, filtered barcode)
    unique_barcodes = set(cell_bcs_union)
    in_unique_barcodes_vectorized = np.vectorize(lambda x: x in unique_barcodes)
    filtered_bc_h5_row = in_unique_barcodes_vectorized(np.array(barcode_summary_h5['bc_sequence']))

    usable_reads = 0

    for genome in genomes:
        h5_key = cr_utils.format_barcode_summary_h5_key(
            library_prefix,
            genome,
            cr_constants.TRANSCRIPTOME_REGION,
            cr_constants.CONF_MAPPED_BC_READ_TYPE)

        if h5_key not in barcode_summary_h5:
            continue

        usable_reads += (filtered_bc_h5_row * np.array(barcode_summary_h5[h5_key])).sum()

    # Fraction reads usable
    d['%s_transcriptome_usable_reads_frac' % lib_constants.MULTI_REFS_PREFIX] = tk_stats.robust_divide(usable_reads, total_reads)
    # Create a feature barcoding dual whose name makes sense
    d['frac_feature_reads_usable'] = tk_stats.robust_divide(usable_reads, total_reads)

    # Usable reads
    d['%s_usable_reads' % lib_constants.MULTI_REFS_PREFIX] = usable_reads

    # Usable reads per cell
    reads_usable_per_cell = tk_stats.robust_divide(usable_reads, n_cell_bcs_union)
    d['%s_usable_reads_per_filtered_bc' % lib_constants.MULTI_REFS_PREFIX] = reads_usable_per_cell
    # Create a feature barcoding dual whose name makes sense
    d['feature_reads_usable_per_cell'] = reads_usable_per_cell

    # Compute matrix density
    filtered_mat = matrix.select_barcodes_by_seq(cell_bcs_union)
    total_nonzero_entries = filtered_mat.get_num_nonzero()
    filtered_shape = filtered_mat.get_shape()
    total_entries = filtered_shape[0] * filtered_shape[1]
    print total_entries, total_nonzero_entries, filtered_shape
    d['%s_filtered_gene_bc_matrix_density' % lib_constants.MULTI_REFS_PREFIX] = tk_stats.robust_divide(total_nonzero_entries, total_entries)

    return d

def _report(matrix, genome, barcode_summary_h5, recovered_cells, cell_bc_seqs,
            library_prefix):
    d = {}
    matrix = matrix.view()

    filtered_mat = matrix.select_barcodes_by_seq(cell_bc_seqs)
    filtered_mat_shape = filtered_mat.get_shape()
    cell_bc_indices = _get_barcode_summary_h5_indices(barcode_summary_h5, cell_bc_seqs)
    n_cell_bcs = len(cell_bc_seqs)

    # Don't compute metrics if no cells detected
    if n_cell_bcs == 0:
        return d

    # Compute matrix density
    d['filtered_gene_bc_matrix_density'] = tk_stats.robust_divide(filtered_mat.get_num_nonzero(),
                                                                  filtered_mat_shape[0]*filtered_mat_shape[1])

    counts_per_gene = filtered_mat.sum(axis=1)
    genes_top_n = min(cr_constants.TOP_N, len(counts_per_gene))
    top_genes_with_counts = {filtered_mat.int_to_feature_id(i): int(count) for i, count in cr_matrix.top_n(counts_per_gene, genes_top_n)}
    d['filtered_bcs_top_genes_with_reads'] = top_genes_with_counts

    unique_bcs_per_gene = filtered_mat.count_ge(axis=1, threshold=cr_constants.MIN_COUNTS_PER_BARCODE)
    top_genes_with_unique_bcs = {filtered_mat.int_to_feature_id(i): int(count) for i, count in cr_matrix.top_n(unique_bcs_per_gene, genes_top_n)}
    d['filtered_bcs_top_genes_with_unique_bcs'] = top_genes_with_unique_bcs

    # Total genes and counts
    total_genes_detected = np.count_nonzero(counts_per_gene)
    total_counts = int(counts_per_gene.sum())
    d['filtered_bcs_total_unique_genes_detected'] = total_genes_detected
    d['filtered_bcs_total_counts'] = total_counts

    def _summarize_per_barcode(a):
        mean = np.mean(a)
        stddev = np.std(a)
        return {
            'mean': mean,
            'median': np.median(a),
            'cv': tk_stats.robust_divide(float(stddev), float(mean)),
            'iqr': np.percentile(a, 75) - np.percentile(a, 25),
        }

    # Unique genes per bc
    unique_genes_per_bc = filtered_mat.count_ge(axis=0, threshold=cr_constants.MIN_COUNTS_PER_GENE)
    unique_genes_stats = _summarize_per_barcode(unique_genes_per_bc)
    for stat, value in unique_genes_stats.iteritems():
        d['filtered_bcs_%s_unique_genes_detected' % stat] = value

    # Counts per bc
    counts_per_bc = filtered_mat.sum(axis=0)
    counts_per_bc_stats = _summarize_per_barcode(counts_per_bc)
    for stat, value in counts_per_bc_stats.iteritems():
        d['filtered_bcs_%s_counts' % stat] = value

    # Cumulative fraction of counts going to top bcs
    filt_total_umis = filtered_mat.sum()
    raw_total_umis = matrix.sum()
    d['filtered_bcs_cum_frac'] = tk_stats.robust_divide(filt_total_umis, raw_total_umis)

    # cDNA PCR duplication in top bcs
    dupe_candidate_h5_key = cr_utils.format_barcode_summary_h5_key(library_prefix,
                                                                   genome,
                                                                   cr_constants.TRANSCRIPTOME_REGION,
                                                                   cr_constants.CONF_MAPPED_BC_READ_TYPE)
    if dupe_candidate_h5_key in barcode_summary_h5:
        n_reads = barcode_summary_h5[dupe_candidate_h5_key][:][cell_bc_indices].sum()
        n_deduped_reads = filt_total_umis
    else:
        n_reads = 0
        n_deduped_reads = 0
    d['filtered_bcs_%s_dupe_reads_frac' % cr_constants.CDNA_PCR_DUPE_TYPE] = 1 - tk_stats.robust_divide(n_deduped_reads, n_reads)

    # Reads per top bc for the various read types (computed over top bcs)
    for read_type in cr_constants.MATRIX_REPORT_READ_TYPES:
        # Compute (n_reads)/(n_bcs) over all bcs and over top bcs
        per_bc_metric = 'filtered_bcs_%s_reads_per_filtered_bc' % read_type

        # Cumulative fraction of reads going to top bcs
        frac_metric = 'filtered_bcs_%s_reads_cum_frac' % read_type

        if read_type in cr_constants.MATRIX_USE_MATRIX_FOR_READ_TYPE:
            n_reads = filt_total_umis
            n_all_reads = raw_total_umis
        else:
            h5_key = cr_utils.format_barcode_summary_h5_key(library_prefix,
                                                            genome,
                                                            cr_constants.TRANSCRIPTOME_REGION,
                                                            read_type)
            if h5_key in barcode_summary_h5:
                counts = barcode_summary_h5[h5_key][:]
                n_reads = counts[cell_bc_indices].sum()
                n_all_reads = counts.sum()
            else:
                n_reads = 0
                n_all_reads = 0
        d[per_bc_metric] = tk_stats.robust_divide(n_reads, n_cell_bcs)
        d[frac_metric] = tk_stats.robust_divide(n_reads, n_all_reads)
    return d

def report_genomes(matrix, reads_summary, barcode_summary_h5_path, recovered_cells, cell_bc_seqs):
    """Report on all genomes in this matrix"""
    barcode_summary_h5 = h5.File(barcode_summary_h5_path, 'r')

    metrics = {}

    genomes = matrix.get_genomes()
    assert len(cell_bc_seqs) == len(genomes)

    # Compute genome-agnostic metrics
    feature_types = sorted(list(set(f.feature_type for f in matrix.feature_ref.feature_defs)))
    for ftype in feature_types:
        total_reads = _get_total_reads(reads_summary, ftype)
        if total_reads == 0:
            continue

        genomes = matrix.get_genomes()
        conf_mapped_reads = _get_conf_mapped_reads(reads_summary, genomes, ftype)

        submatrix = matrix.view().select_features_by_type(ftype)
        prefix = rna_library.get_library_type_metric_prefix(ftype)

        m = _report_genome_agnostic_metrics(submatrix,
                                            barcode_summary_h5,
                                            recovered_cells,
                                            cell_bc_seqs,
                                            total_reads, conf_mapped_reads,
                                            prefix)

        if rna_library.has_genomes(ftype):
            for genome in genomes:
                # Compute genome-specific metrics
                genome_matrix = matrix.view().select_features_by_genome(genome)
                genome_summary = _report(genome_matrix,
                                         genome,
                                         barcode_summary_h5,
                                         recovered_cells,
                                         cell_bc_seqs[genome],
                                         prefix)

                for key, value in genome_summary.iteritems():
                    key = '_'.join([genome, key])
                    m[key] = value

        else:
            # This feature has no genomes
            cell_bcs_union = list(reduce(lambda a,x: a | set(x), cell_bc_seqs.itervalues(), set()))
            genome_summary = _report(submatrix,
                                     lib_constants.MULTI_REFS_PREFIX,
                                     barcode_summary_h5,
                                     recovered_cells,
                                     cell_bcs_union,
                                     prefix)
            for key, value in genome_summary.iteritems():
                key = '_'.join([lib_constants.MULTI_REFS_PREFIX, key])
                m[key] = value

        # Prepend feature type to metric keys
        m_prefixed = {(prefix+k): v for k, v in m.iteritems()}
        metrics.update(m_prefixed)

    return metrics


import sys
import os
import numpy as np
import pandas as pd
import fragility
from .utils import utils

import datetime
import time
from dateutil.relativedelta import relativedelta

class LoadPat(object):
    '''
    A class describing a dataset that is within our framework. This object will help set up the
    data for computation and any meta data necessary to link the computations together.

    This assumes that data is stored in a specific format!
    datadir/
        <patient1>
        <patient2>
            - <p>_rawnpy.npy
            - <p>_annotations.csv
            - <p>_chans.csv
            - <p>_headers.csv

    This mainly is the format that is the output of the dataconversion.py file from EDF files.

    Will need to reformat for other type of EEG files.
    '''
    def __init__(self, patient, datadir, clinoutcome=None):
        self.patient = patient
        self.datadir = datadir
        self.rawdatafile = patient + '_rawnpy.npy'
        self.annotationsfile = patient + '_annotations.csv'
        self.channelsfile = patient + '_chans.csv'
        self.headersfile = patient + '_headers.csv'
        self.clinoutcome = clinoutcome
        self.reference = None
        # get relevant channel data
        self.patid, self.seizid = utils.splitpatient(patient)
        self.included_chans, self.onsetchans, self.clinresult = utils.returnindices(
            self.patid, self.seizid)

        self.included_chans = self.included_chans.astype(int)
        if self.onsetchans is not None:
            self.onsetchans = np.asarray(
                [x.lower() for x in list(self.onsetchans)])

        # print("Loaded ", patient)
        self.loadchans_fromfile()
        self.loadheaders_fromfile()
        self.loadannotations_fromfile()

        if len(self.included_chans) == 0:
            self.included_chans = np.arange(
                0, len(self.chanlabels)).astype(int)

    def _metafilepath(self, filename):
        return os.path.join(self.datadir, self.patient, filename)

    def get_clindata(self):
        clindata = {
            'onset_time': self.onset_time,
            'offset_time': self.offset_time,
            'samplerate': self.samplerate,
            'reference': self.reference,
            'gender': self.gender,
            'age': self.age
        }
        return clindata

    def loadchans_fromfile(self):
        chansfile = self._metafilepath(self.channelsfile)
        # read in the csv file into pandas
        chanheaders = pd.read_csv(chansfile)
        # read in labels
        chanlabels = chanheaders['Labels']
        self.chanlabels = chanlabels.values

        # apply lower case transformation to keep things simple
        # return only included chans
        # if self.included_chans.size > 0:
        #     self.chanlabels[self.included_chans] = np.asarray([x.lower() for x in list(self.chanlabels[self.included_chans])])
        # print("Loaded channels file data")

    def loadrawdata_fromfile(self, reference=None):
        rawdatafile = self._metafilepath(self.rawdatafile)

        # adding in the ability to automatically read in clipped data
        # dictionaries...
        if os.path.exists(self._metafilepath(
                self.rawdatafile[0:-4] + '_clipped.npz')):
            sys.stderr.write(
                'READING IN RAW CLIPPED DATA WITH ONSET/OFFSET +/- 60 SECONDS')
            rawdatafile = self._metafilepath(
                self.rawdatafile[0:-4] + '_clipped.npz')
            data = np.load(rawdatafile)
            rawdata = data['rawdata']
        else:
            # load numpy array data
            rawdata = np.load(rawdatafile)

        self.reference = 'monopolar'
        # if signals by channels -> transpose
        if rawdata.shape[0] > rawdata.shape[1]:
            rawdata = rawdata.T

        # different montage/references
        if reference == 'avg':  # perform average referencing
            # average over each row
            avg = np.mean(rawdata, axis=1, keepdims=True)
            rawdata = rawdata - avg
            self.reference = 'average'
        elif reference == 'bipolar':  # perform bipolar montage
            self.reference = 'bipolar'
            samplefreq = self.samplerate

            # convert contacts into a list of tuples as data structure
            contacts = []
            try:
                chanlabels = self.chanlabels[self.included_chans]
            except:
                chanlabels = self.chanlabels

            for contact in chanlabels:
                thiscontactset = False
                for idx, s in enumerate(contact):
                    if s.isdigit() and not thiscontactset:
                        elec_label = contact[0:idx]
                        thiscontactset = True
                contacts.append((elec_label, int(contact[len(elec_label):])))
            # compute the bipolar scheme
            recording = fragility.preprocess.seegrecording.SeegRecording(
                contacts, rawdata, samplefreq)
            self.chanlabels = np.asarray(recording.get_channel_names_bipolar())
            rawdata = recording.get_data_bipolar()

        # return only included chans
        try:
            if self.included_chans.size > 0:
                rawdata = rawdata[self.included_chans,:]
        except IndexError as e:
            sys.stderr.write('Index error in this patient for loading raw data')
            # assert rawdata.shape[0] == len(self.included_chans)
            # print(e)

        # print("Loaded raw data")
        return rawdata

    def loadheaders_fromfile(self):
        headersfile = self._metafilepath(self.headersfile)
        # read in the file headers into pandas
        fileheaders = pd.read_csv(headersfile)

        # get important meta data (ADD ELEMENTS HERE TO ADD TO CLASS)
        birthdate = fileheaders['Birth Date']
        daterecording = fileheaders['Start Date (D-M-Y)']
        gender = fileheaders['Gender']
        equipment = fileheaders['Equipment']
        samplerate = fileheaders['Sample Frequency']
        recordduration = fileheaders['Data Record Duration (s)']

        # set metadata members of this class
        self.birthdate = birthdate.values[0]
        self.daterecording = daterecording.values[0]
        self.gender = gender.values[0]
        self.equipment = equipment.values[0]
        self.samplerate = samplerate.values[0]
        self.recordduration = recordduration.values[0]

        self.age = None
        # if not np.isnan(self.birthdate):
        try:
            now = datetime.datetime.now()
            birthdate = datetime.datetime.strptime(self.birthdate, '%d %b %Y')
            rdelta = relativedelta(now, birthdate)
            self.age = rdelta.years
            # print('Age in years - ', rdelta.years)
            # print('Age in months - ', rdelta.months)
            # print('Age in days - ', rdelta.days)
            # print("Loaded headers from file.")
        except TypeError:
            sys.stderr.write(
                "type error in birth date probably, so just setting to nan")
        except ValueError:
            sys.stderr.write(
                "value error in birth date probably, so just setting to nan")

    def loadannotations_fromfile(self):
        '''
        Here, we mainly are interested in the onset/offset times
        if they are present...
        '''
        annotationsfile = self._metafilepath(self.annotationsfile)
        # read in the clinical annotations into pandas
        annotations = pd.read_csv(annotationsfile)
        descriptions = annotations['Description'].values

        onset_time = []
        offset_time = []
        onsetset = False
        offsetset = False
        for idx, descrip in enumerate(descriptions):
            descrip = descrip.lower().split(' ')
            # print(descrip)
            if 'onset' in descrip \
                    or 'crise' in descrip \
                    or 'cgtc' in descrip \
                    or 'sz' in descrip or 'absence' in descrip:
                if not onsetset:
                    onset_time = annotations['Time (sec)'].values[idx]
                    onsetset = True
            if 'offset' in descrip \
                    or 'fin' in descrip \
                    or 'end' in descrip:
                if not offsetset:
                    offset_time = annotations['Time (sec)'].values[idx]
                    offsetset = True

        # set onset/offset times and markers
        # if np.array(onset_time).size > 0:
        try:
            self.onset_sec = onset_time
            self.onset_time = np.multiply(onset_time, 1000)
        except TypeError:
            print("no onset time!")
            self.onset_time = None
            self.onset_sec = None
        try:
            self.offset_sec = offset_time
            self.offset_time = np.multiply(offset_time, 1000)
        except TypeError:
            print("no offset time!")
            self.offset_time = None
            self.onset_sec
        # print('Ran setup of annotations data!')

    def readseiztimes(self, master_file):
        '''
        Helper function to read seizure times from a master excel file

        Mainly used to help in reading data from UMMC since it wasn't easily logged
        on the edf file.
        '''
        xl_file = pd.ExcelFile(master_file)
        center = 'UMMC'

        # load in the master excel sheet by a certain center
        dfs = {sheet_name: xl_file.parse(center)
               for sheet_name in xl_file.sheet_names}
        center_df = dfs[center]

        # get the patient row data
        center_df = center_df.apply(lambda x: x.astype(str).str.lower())
        pat_df = center_df.loc[center_df['Identifier'] == self.patient]

        def convert_totime(df_val): return datetime.datetime.strptime(
            df_val, '%H:%M:%S')

        # get the onset/offset electrographic time
        record_start = convert_totime(pat_df['Recording Start'].item())
        onset_time = (
            convert_totime(
                pat_df['E Onset'].item()) -
            record_start).total_seconds()
        offset_time = (
            convert_totime(
                pat_df['E Offset'].item()) -
            record_start).total_seconds()
        self.onset_time = np.multiply(onset_time, self.samplerate)
        self.offset_time = np.multiply(offset_time, self.samplerate)

    def clipdata(self, rawdata):
        '''
        Used for clipping data at +/- 60 seconds of their onset/offset times

        To reduce the amount of data that is needed to analyze
        '''
        if self.onset_time:
            firstind = int(
                self.onset_sec *
                self.samplerate -
                60 *
                self.samplerate)
        if self.offset_time:
            endind = int(
                self.offset_sec *
                self.samplerate +
                60 *
                self.samplerate)
        else:
            endind = firstind + 60*1000

        # firstind and endind needs to be within the rawdata bounds
        self.firstind = firstind
        self.endind = endind

        # clip the rawdata
        rawdata = rawdata[:, firstind:endind]
        clippedinds = firstind
        clippedsec = firstind / self.samplerate
        clippedms = firstind / self.samplerate * 1000

        # clip onsettime, offsettime
        self.onset_time = 60 * 1000
        self.offset_time = rawdata.shape[1] - 60 * 1000
        self.onset_sec = self.onset_sec - 60
        self.offset_sec = self.offset_sec - clippedsec

        return rawdata

    def clippeddatagen(self, rawdata, winsec=60):
        '''
        A data generator that goes through the raw data
        and returns clipped windows that have a predetermined
        number of seconds.

        To Do:
        - Test that beginind, endind always are continuous
        and there is no data cutoff.
        '''

        # clip data by the number of seconds in winsec
        winsize = winsec * self.samplerate
        numchans, numsamps = rawdata.shape

        ind = 0
        while True:
            beginind = ind * winsize
            endind = (ind + 1) * winsize

            # check if end of data will be reached
            if endind >= numsamps:
                yield rawdata[:, beginind:]
            else:
                yield rawdata[:, beginind:endind]


import sys, os
from sklearn.metrics import roc_auc_score, accuracy_score, precision_score, recall_score, matthews_corrcoef
import pandas as pd
import numpy as np
import json

EVALUATE_CAPITAL_ONLY=True

DO_SEQ_AVERAGING=False

#BINDING_CMAPS_PATH="/home/rsanchez/Tesis/rriPredMethod/data/bench5Data/newCodeData/computedFeatures/common/contactMapsBinding"
BINDING_CMAPS_PATH=None # if None, contact maps contained in the results file will be used, otherwise, new contacts will be added from the files in the path

def averageScores(scoresDf):
  labels_list=[]
  scores_list=[]
  for chain in scoresDf["chainId"].unique():
    df= scoresDf.loc[scoresDf["chainId"]==chain,:]
    result= [ ((int(elem), ""), (score, label)) if elem[-1].isdigit() else ((int(elem[:-1]), elem[-1]), (score, label))
                                for elem,score,label in zip(df["resId"], df["prediction"], df["categ"]) ]
    result.sort(key= lambda x: x[0])
    result= zip(* result)[1]
    scores, labels= zip(* result)
    scores= list(np.convolve(scores, np.array([1, 3, 1])/5.0, mode='same')+ np.array(scores))
    labels_list+= labels
    scores_list+= scores
  return scores_list, labels_list

def loadResults( resultsPath, fnameResults, cMapsPath=BINDING_CMAPS_PATH):
  prefix= fnameResults.split("_")[0].split(".")[0]
  if ".lig" in fnameResults:
    chainType="l"
  else:
    chainType="r"

  scoresDf= pd.read_table(os.path.join(resultsPath, fnameResults), comment="#", sep="\s+", dtype={"resId":str, "chainId":str})
  if not cMapsPath is None:
    newCmapSet=set([])
    for fname in os.listdir(cMapsPath):
      if ((chainType=="l" and "_l_" in fname) or (chainType=="r" and "_r_" in fname)) and fname.startswith(prefix):
        df= pd.read_table(os.path.join(cMapsPath,fname),sep='\s+', header='infer', comment="#",
                          dtype= {"chainIdL":str, "chainIdR":str, "resIdL":str, "resIdR":str,
                                        "chainId":str, "resId":str,  "resId":str})
        for i in range(df.shape[0]):
          chainId, resId, categ= df.iloc[i,:]
          if categ==1:
            newCmapSet.add((chainId, resId))
    for chainId, resId in newCmapSet:
      scoresDf.loc[(scoresDf["chainId"]==chainId) & (scoresDf["resId"]==resId),"categ"]=1
  return scoresDf

def get_single_chain_statistics(prefix, labels, scores):
  EVAL_PAIRS_AT= [ 2**3, 2**4]
  precisionAt=[]
  recallAt=[]
  scores= np.array(scores)
  labels= np.array(labels)
  try:
    roc_complex= roc_auc_score(labels, scores)
  except ValueError:
    roc_complex= np.nan
  probability_sorted_indexes = scores.argsort(axis=0)
  probability_sorted_indexes = probability_sorted_indexes[::-1]
  for evalPoint in EVAL_PAIRS_AT:
    try:
      label_predictions= np.ones(scores.shape[0], dtype=np.int32)* np.min( labels)

      label_predictions[probability_sorted_indexes[0 : evalPoint]]= np.repeat(1, evalPoint)
      precisionAt.append( precision_score(labels[probability_sorted_indexes[0 : evalPoint]],
                                        label_predictions[probability_sorted_indexes[0 : evalPoint]]))
      recallAt.append( recall_score(labels, label_predictions))
#      print(sum(labels==1), sum(label_predictions==1),precisionAt[-1], recallAt[-1])
    except IndexError:
      precisionAt.append( 0.0)
      recallAt.append( 0.0)
  summary= pd.DataFrame({"pdb":[prefix]})
  summary["auc_chains"]= [roc_complex]
  for evalPoint, precisionAt, recallAt in zip(EVAL_PAIRS_AT,precisionAt, recallAt):
    summary["prec_%d"%evalPoint]= [precisionAt]
    summary["reca_%d"%evalPoint]= [recallAt]
  return summary

def getRecallAtPrec(resultsPath, targetPrec=0.666, jsonFname=None, useSeqAvera= DO_SEQ_AVERAGING):
  allScores=[]
  allLabels=[]
  perComplexSummaries=[]
  for fname in sorted(os.listdir(resultsPath)):
    if EVALUATE_CAPITAL_ONLY and fname[1:4].islower(): continue
    print(fname)
    if ".rec" in fname or ".lig" in fname:
      results= loadResults( resultsPath, fname)
      if results is None: continue
      if useSeqAvera:
        scores, labels= averageScores(results)
      else:
        scores= list(results["prediction"].values)
        labels= list(results["categ"].values)

      summary= get_single_chain_statistics(fname, labels, scores)
      perComplexSummaries.append(summary)
#      print("%s %f"%(fname,roc_complex))
      allScores+= scores
      allLabels+= labels

  summary= pd.concat(perComplexSummaries, ignore_index=True)
  means= summary.mean(axis=0)
  summary= summary.append( summary.ix[summary.shape[0]-1,:],ignore_index=True )
  summary.ix[summary.shape[0]-1,0]=  "mean"
  summary.ix[summary.shape[0]-1, 1:]=  means

  roc_auc= roc_auc_score(allLabels, allScores)

  allScores= np.array(allScores)
  allScoresUnique= np.unique(np.round(allScores, decimals=4))
  indices= np.argsort(allScoresUnique)[1:-1]
  # import matplotlib.pyplot as plt
  #  plt.hist(allScoresUnique[:-2], bins=50)
  #  plt.show()
  bestThr=-1
  bestMcc= 0
  bestPrec= -1
  bestRec= -1
  thr_2_prec={"thr":[], "prec":[]}
  nonClassLabel= int(np.min(allLabels))
  for i in indices:
    tmpThr= allScoresUnique[i]
    binaryScores= np.array([1.0 if score>tmpThr else nonClassLabel for score in allScores])

    currentPrec= precision_score(allLabels, binaryScores)
    thr_2_prec["thr"].append(tmpThr)
    thr_2_prec["prec"].append(currentPrec)
    if abs(currentPrec- targetPrec) < abs(bestPrec- targetPrec):
      bestMcc= matthews_corrcoef(allLabels, binaryScores)
      bestThr= tmpThr
      bestPrec= precision_score(allLabels, binaryScores)
      bestRec= recall_score(allLabels, binaryScores)
  print(thr_2_prec)
  if jsonFname:
    with open(jsonFname,"w") as f:
      json.dump(thr_2_prec, f)
  print(summary.to_string(index=False))
  print("roc auc(mixed/mean): %f/%f Thr: %f mcc: %f prec: %f rec: %f"%(roc_auc,summary.iloc[-1,1], bestThr, bestMcc, bestPrec, bestRec))
  return bestThr

def countIfAboveThr(bestThr, resultsPath):
  totalNum=0
  numAbove=0
  for fname in sorted(os.listdir(resultsPath)):
    #print(fname)
    if ".rec" in fname or ".lig" in fname:
      results= loadResults( resultsPath, fname)
      totalNum+=1
      if results is None:
        continue
      else:
        scores= np.array(results["prediction"].values)
      if np.sum(scores>=bestThr):
        numAbove+=1
  print("Above thr %d / %d "%(numAbove,totalNum))

if __name__=="__main__":
  '''
   python evaluation/getPrecsAt_diffThr.py ~/Tesis/rriPredMethod/data/bench5Data/newCodeData/results/mixed_2/ ~/Tesis/rriPredMethod/data/bench5Data/newCodeData/results/thr_2_prec_b5_mixed_2.json

  '''
  if len(sys.argv)==3:
    resultsPath= os.path.expanduser(sys.argv[1])
    jsonFname= os.path.expanduser(sys.argv[2])
  else:
    raise ValueError("Bad arguments.")
  bestThr= getRecallAtPrec(resultsPath, jsonFname=jsonFname)
  countIfAboveThr(bestThr, resultsPath)