import gc
import time
import ray

import numpy as np

from openea.modules.finding.similarity import sim
from openea.modules.utils.util import task_divide


def calculate_rank_(idx, sim_mat, top_k, accurate, total_num):
    assert 1 in top_k
    mr = 0
    mrr = 0
    hits = [0] * len(top_k)
    hits1_rest = set()
    for i in range(len(idx)):
        gold = idx[i]
        if accurate:
            rank = (-sim_mat[i, :]).argsort()
        else:
            rank = np.argpartition(-sim_mat[i, :], np.array(top_k) - 1)
        hits1_rest.add((gold, rank[0]))
        assert gold in rank
        rank_index = np.where(rank == gold)[0][0]
        mr += (rank_index + 1)
        mrr += 1 / (rank_index + 1)
        for j in range(len(top_k)):
            if rank_index < top_k[j]:
                hits[j] += 1
    mr /= total_num
    mrr /= total_num
    return mr, mrr, hits, hits1_rest


@ray.remote(num_cpus=1)
def calculate_rank(idx, sim_mat, top_k, accurate, total_num):
    assert 1 in top_k
    mr = 0
    mrr = 0
    hits = [0] * len(top_k)
    hits1_rest = set()
    for i in range(len(idx)):
        gold = idx[i]
        if accurate:
            rank = (-sim_mat[i, :]).argsort()
        else:
            rank = np.argpartition(-sim_mat[i, :], np.array(top_k) - 1)
        hits1_rest.add((gold, rank[0]))
        assert gold in rank
        rank_index = np.where(rank == gold)[0][0]
        mr += (rank_index + 1)
        mrr += 1 / (rank_index + 1)
        for j in range(len(top_k)):
            if rank_index < top_k[j]:
                hits[j] += 1
    mr /= total_num
    mrr /= total_num
    return mr, mrr, hits, hits1_rest


def greedy_alignment(embed1, embed2, top_k, nums_threads, metric, normalize, csls_k, accurate, is_print):
    t = time.time()
    sim_mat = sim(embed1, embed2, metric=metric, normalize=normalize, csls_k=csls_k)
    num = sim_mat.shape[0]
    if nums_threads > 1:
        hits = [0] * len(top_k)
        mr, mrr = 0, 0
        alignment_rest = set()
        rests = list()
        search_tasks = task_divide(np.array(range(num)), nums_threads)
        for task in search_tasks:
            mat = sim_mat[task, :]
            res = calculate_rank.remote(task, mat, top_k, accurate, num)
            rests.append(res)
        for rest in ray.get(rests):
            sub_mr, sub_mrr, sub_hits, sub_hits1_rest = rest
            mr += sub_mr
            mrr += sub_mrr
            hits += np.array(sub_hits)
            alignment_rest |= sub_hits1_rest
    else:
        mr, mrr, hits, alignment_rest = calculate_rank_(list(range(num)), sim_mat, top_k, accurate, num)
    assert len(alignment_rest) == num
    hits = np.array(hits) / num
    for i in range(len(hits)):
        hits[i] = round(hits[i], 3)
    cost = time.time() - t
    if is_print:
        if accurate:
            if csls_k > 0:
                print("accurate results with csls: csls={}, hits@{} = {}, mr = {:.3f}, mrr = {:.6f}, time = {:.3f} s ".
                      format(csls_k, top_k, hits, mr, mrr, cost))
            else:
                print("accurate results: hits@{} = {}, mr = {:.3f}, mrr = {:.6f}, time = {:.3f} s ".
                      format(top_k, hits, mr, mrr, cost))
        else:
            if csls_k > 0:
                print("quick results with csls: csls={}, hits@{} = {}, time = {:.3f} s ".format(csls_k, top_k, hits,
                                                                                                cost))
            else:
                print("quick results: hits@{} = {}, time = {:.3f} s ".format(top_k, hits, cost))

    sim_list = []
    for i, j in alignment_rest:
        sim_list.append(sim_mat[i, j])
    del embed1
    del embed2
    del sim_mat
    gc.collect()
    return alignment_rest, hits, mr, mrr, sim_list


def valid(embeds1, embeds2, mapping, top_k, threads_num, metric='inner', normalize=False, csls_k=0, accurate=False,
          is_print=True):
    if mapping is None:
        _, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(embeds1, embeds2, top_k, threads_num,
                                                                metric, normalize, csls_k, accurate, is_print)
    else:
        test_embeds1_mapped = np.matmul(embeds1, mapping)
        _, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(test_embeds1_mapped, embeds2, top_k, threads_num,
                                                                metric, normalize, csls_k, accurate, is_print)
    return hits1_12, mrr_12, sim_list


def test(embeds1, embeds2, mapping, top_k, threads_num, metric='inner', normalize=False, csls_k=0, accurate=True,
         is_print=True):
    if mapping is None:
        alignment_rest_12, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(embeds1, embeds2, top_k, threads_num,
                                                                                metric, normalize, csls_k, accurate,
                                                                                is_print)
    else:
        test_embeds1_mapped = np.matmul(embeds1, mapping)
        alignment_rest_12, hits1_12, mr_12, mrr_12, sim_list = greedy_alignment(test_embeds1_mapped, embeds2, top_k,
                                                                                threads_num,
                                                                                metric, normalize, csls_k, accurate,
                                                                                is_print)
    return alignment_rest_12, hits1_12, mrr_12, sim_list