import time
import datetime
from multiprocessing import Process, Queue, cpu_count

import torch
import numpy as np
# from pytorch_transformers import BertModel
from transformers import BertModel

from . import utils
from . import stats


class ParallelSampler():
    def __init__(self, data, args, num_episodes=None):
        self.data = data
        self.args = args
        self.num_episodes = num_episodes

        self.all_classes = np.unique(self.data['label'])
        self.num_classes = len(self.all_classes)
        if self.num_classes < self.args.way:
            raise ValueError("Total number of classes is less than #way.")

        self.idx_list = []
        for y in self.all_classes:
            self.idx_list.append(
                np.squeeze(np.argwhere(self.data['label'] == y)))

        self.count = 0
        self.done_queue = Queue()

        self.num_cores = cpu_count() if args.n_workers is 0 else args.n_workers

        self.p_list = []
        for i in range(self.num_cores):
            self.p_list.append(
                Process(target=self.worker, args=(self.done_queue,)))

        for i in range(self.num_cores):
            self.p_list[i].start()

    def get_epoch(self):
        for _ in range(self.num_episodes):
            # wait until self.thread finishes
            support, query = self.done_queue.get()

            # convert to torch.tensor
            support = utils.to_tensor(support, self.args.cuda, ['raw'])
            query = utils.to_tensor(query, self.args.cuda, ['raw'])

            if self.args.meta_w_target:
                if self.args.meta_target_entropy:
                    w = stats.get_w_target(
                        support, self.data['vocab_size'],
                        self.data['avg_ebd'], self.args.meta_w_target_lam)
                else:  # use rr approxmation (this one is faster)
                    w = stats.get_w_target_rr(
                        support, self.data['vocab_size'],
                        self.data['avg_ebd'], self.args.meta_w_target_lam)
                support['w_target'] = w.detach()
                query['w_target'] = w.detach()

            support['is_support'] = True
            query['is_support'] = False

            yield support, query

    def worker(self, done_queue):
        '''
            Generate one task (support and query).
            Store into self.support[self.cur] and self.query[self.cur]
        '''
        while True:
            if done_queue.qsize() > 100:
                time.sleep(1)
                continue
            # sample ways
            sampled_classes = np.random.permutation(
                self.num_classes)[:self.args.way]

            source_classes = []
            for j in range(self.num_classes):
                if j not in sampled_classes:
                    source_classes.append(self.all_classes[j])
            source_classes = sorted(source_classes)

            # sample examples
            support_idx, query_idx = [], []
            for y in sampled_classes:
                tmp = np.random.permutation(len(self.idx_list[y]))
                support_idx.append(
                    self.idx_list[y][tmp[:self.args.shot]])
                query_idx.append(
                    self.idx_list[y][
                        tmp[self.args.shot:self.args.shot+self.args.query]])

            support_idx = np.concatenate(support_idx)
            query_idx = np.concatenate(query_idx)
            if self.args.mode == 'finetune' and len(query_idx) == 0:
                query_idx = support_idx

            # aggregate examples
            max_support_len = np.max(self.data['text_len'][support_idx])
            max_query_len = np.max(self.data['text_len'][query_idx])

            support = utils.select_subset(self.data, {}, ['text', 'text_len', 'label'],
                                          support_idx, max_support_len)
            query = utils.select_subset(self.data, {}, ['text', 'text_len', 'label'],
                                        query_idx, max_query_len)

            if self.args.embedding in ['idf', 'meta', 'meta_mlp']:
                # compute inverse document frequency over the meta-train set
                idf = stats.get_idf(self.data, source_classes)
                support['idf'] = idf
                query['idf'] = idf

            if self.args.embedding in ['iwf', 'meta', 'meta_mlp']:
                # compute SIF over the meta-train set
                iwf = stats.get_iwf(self.data, source_classes)
                support['iwf'] = iwf
                query['iwf'] = iwf

            if 'pos' in self.args.auxiliary:
                support = utils.select_subset(
                    self.data, support, ['head', 'tail'], support_idx)
                query = utils.select_subset(
                    self.data, query, ['head', 'tail'], query_idx)

            done_queue.put((support, query))

    def __del__(self):
        '''
            Need to terminate the processes when deleting the object
        '''
        for i in range(self.num_cores):
            self.p_list[i].terminate()

        del self.done_queue
