import pandas as pd
import numpy as np


def create_map_reverse(arr):
    m = {}
    rm = {}
    for i in range(len(arr)):
        m[arr[i]] = i
        rm[i] = arr[i]
    return m, rm


def get_data_prop(df, nspmap, sp, prop, nos):
    nos = int(nos * prop)
    sp = [nspmap[x] for x in sp]
    data = df[df["homology_species"].isin(sp)]
    if nos < len(data):
        random_indexes = np.random.permutation(len(data))
        data = data.loc[data.index.values[random_indexes[:nos]]]
        return data
    else:
        return data


def create_balanced_dataset_paralog(
        df,
        matrix,
        spnmap,
        nspmap,
        nos,
        hom_type,
        spname):
    df = df[df["homology_type"] == hom_type]
    dist = matrix[spnmap[spname]]
    dist_sort = np.argsort(dist)
    sp_far = dist_sort[-5:]
    sp_near = dist_sort[1:6]
    # get the records for the species which are far away
    df_far = get_data_prop(df, nspmap, sp_far, 0.2, nos)
    # get the records for the nearby species
    df_near = get_data_prop(df, nspmap, sp_near, 0.2, nos)
    df_dist = pd.concat([df_far, df_near])
    nos = nos - len(df_dist)
    df = df.drop(df_dist.index.values)
    random_ind = np.random.permutation(len(df))
    df_r = df.loc[df.index.values[random_ind[:nos]]]
    df = pd.concat([df_r, df_dist])
    return df


def select_data_goc(df, prop, nos):
    df = df[df["goc_score"] == 0.0]
    nos = int(prop * nos)
    rind = np.random.permutation(len(df))
    if len(df) < nos:
        return df
    else:
        return df.loc[df.index.values[rind[:nos]]]


def create_balanced_dataset_ortholog(
        df,
        matrix,
        spnmap,
        nspmap,
        nos,
        hom_type,
        spname):
    df = df[df["homology_type"] == hom_type]
    dist = matrix[spnmap[spname]]
    dist_sort = np.argsort(dist)
    sp_far = dist_sort[-5:]
    sp_near = dist_sort[1:6]
    # get the records for the species which are far away
    df_far = get_data_prop(df, nspmap, sp_far, 0.2, nos)
    # get the records for the nearby species
    df_near = get_data_prop(df, nspmap, sp_near, 0.2, nos)
    df_dist = pd.concat([df_far, df_near])
    len(df_dist)
    df = df.drop(df_dist.index.values)
    df_goc = select_data_goc(df, 0.1, nos)
    len(df_goc)
    df = df.drop(df_goc.index.values)
    nos = nos - len(df_dist) - len(df_goc)
    random_ind = np.random.permutation(len(df))
    df_r = df.loc[df.index.values[random_ind[:nos]]]
    df = pd.concat([df_r, df_goc, df_dist])
    return df


def select(df, nos, matrix, spnmap, nspmap, sp):
    nos_p = int((0.5 * nos) / 2)
    nos_o = int((0.5 * nos) / 3)

    # get the paralogy data
    df_p1 = create_balanced_dataset_paralog(
        df, matrix, spnmap, nspmap, nos_p, "within_species_paralog", sp)
    df_p2 = create_balanced_dataset_paralog(
        df, matrix, spnmap, nspmap, nos_p, "other_paralog", sp)
    # get the orthology data
    df_o1 = create_balanced_dataset_ortholog(
        df, matrix, spnmap, nspmap, nos_o, "ortholog_one2many", sp)
    df_o2 = create_balanced_dataset_ortholog(
        df, matrix, spnmap, nspmap, nos_o, "ortholog_many2many", sp)
    df_o3 = create_balanced_dataset_ortholog(
        df, matrix, spnmap, nspmap, nos_o, "ortholog_one2one", sp)

    # concatenate everything
    df = pd.concat([df_o1, df_o2, df_o3, df_p1, df_p2])
    return df