from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy

import numpy as np
import time
import utils.utils as utils
from collections import OrderedDict
import torch
from torch.autograd import Variable

import sys
sys.path.append("3rdparty/cider")
from pyciderevalcap.ciderD.ciderD import CiderD
sys.path.append("3rdparty/coco-caption")
from pycocoevalcap.bleu.bleu import Bleu

CiderD_scorer = None
Bleu_scorer = None

# CiderD_scorer = None
#CiderD_scorer = CiderD(df='corpus')

def init_cider_scorer(cached_tokens):
    global CiderD_scorer
    CiderD_scorer = CiderD_scorer or CiderD(df=cached_tokens)

def init_bleu_scorer():
    global Bleu_scorer
    Bleu_scorer = Bleu_scorer or Bleu(4)

def array_to_str(arr):
    out = ''
    for i in range(len(arr)):
        out += str(arr[i]) + ' '
        if arr[i] == 0:
            break
    return out.strip()

def get_self_critical_reward(model, att_feats, num_bbox, data, gen_result):
    batch_size = gen_result.size(0)# batch_size = sample_size * seq_per_img
    seq_per_img = batch_size // len(data['gts'])
    
    # get greedy decoding baseline
    greedy_res, _ = model.sample(att_feats, num_bbox)

    res = OrderedDict()
    
    gen_result = gen_result.cpu().numpy()
    greedy_res = greedy_res.cpu().numpy()
    for i in range(batch_size):
        res[i] = [array_to_str(gen_result[i])]
    for i in range(batch_size):
        res[batch_size + i] = [array_to_str(greedy_res[i])]

    gts = OrderedDict()
    for i in range(len(data['gts'])):
        gts[i] = [array_to_str(data['gts'][i][j]) for j in range(len(data['gts'][i]))]

    #_, scores = Bleu(4).compute_score(gts, res)
    #scores = np.array(scores[3])
    res = [{'image_id':i, 'caption': res[i]} for i in range(2 * batch_size)]
    gts = {i: gts[i % batch_size // seq_per_img] for i in range(2 * batch_size)}
    _, scores = CiderD_scorer.compute_score(gts, res)
    print('Cider scores:', _)

    scores = scores[:batch_size] - scores[batch_size:]

    rewards = np.repeat(scores[:, np.newaxis], gen_result.shape[1], 1)

    return rewards


def get_self_critical_reward_accu_steps(model, att_feats, num_bbox, data, gen_result, _grad_accu_step, grad_accu_steps):
    batch_size = gen_result.size(0)  # batch_size = sample_size * seq_per_img
    seq_per_img = batch_size* grad_accu_steps // len(data['gts'])
    sub_len_data_gts = len(data['gts'])//grad_accu_steps
    # print(batch_size, seq_per_img, sub_len_data_gts, len(data['gts']), grad_accu_steps)

    # get greedy decoding baseline
    greedy_res, _ = model.sample(Variable(att_feats.data, volatile=True), Variable(num_bbox.data, volatile=True))

    res = OrderedDict()

    gen_result = gen_result.cpu().numpy()
    greedy_res = greedy_res.cpu().numpy()
    for i in range(batch_size):
        res[i] = [array_to_str(gen_result[i])]
    for i in range(batch_size):
        res[batch_size + i] = [array_to_str(greedy_res[i])]

    gts = OrderedDict()
    for i in range(sub_len_data_gts):
        i_tmp = _grad_accu_step*sub_len_data_gts+i
        gts[i] = [array_to_str(data['gts'][i_tmp][j]) for j in range(len(data['gts'][i_tmp]))]

    # _, scores = Bleu(4).compute_score(gts, res)
    # scores = np.array(scores[3])
    res = [{'image_id': i, 'caption': res[i]} for i in range(2 * batch_size)]
    gts = {i: gts[i % batch_size // seq_per_img] for i in range(2 * batch_size)}
    _, scores = CiderD_scorer.compute_score(gts, res)
    print('Cider scores:', _)

    scores = scores[:batch_size] - scores[batch_size:]

    rewards = np.repeat(scores[:, np.newaxis], gen_result.shape[1], 1)

    return rewards

def get_self_critical_cider_bleu_reward(model, att_feats, num_bbox, data, gen_result):
    batch_size = gen_result.size(0)  # batch_size = sample_size * seq_per_img
    seq_per_img = batch_size // len(data['gts'])

    # get greedy decoding baseline
    greedy_res, _ = model.sample(Variable(att_feats.data, volatile=True), Variable(num_bbox.data, volatile=True))

    res = OrderedDict()

    gen_result = gen_result.cpu().numpy()
    greedy_res = greedy_res.cpu().numpy()
    for i in range(batch_size):
        res[i] = [array_to_str(gen_result[i])]
    for i in range(batch_size):
        res[batch_size + i] = [array_to_str(greedy_res[i])]

    gts = OrderedDict()
    for i in range(len(data['gts'])):
        gts[i] = [array_to_str(data['gts'][i][j]) for j in range(len(data['gts'][i]))]

    # _, scores = Bleu(4).compute_score(gts, res)
    # scores = np.array(scores[3])
    res = [{'image_id': i, 'caption': res[i]} for i in range(2 * batch_size)]
    gts = {i: gts[i % batch_size // seq_per_img] for i in range(2 * batch_size)}
    _, scores = CiderD_scorer.compute_score(gts, res)
    print('Cider scores:', _)

    scores = scores[:batch_size] - scores[batch_size:]

    rewards = np.repeat(scores[:, np.newaxis], gen_result.shape[1], 1)

    return rewards

def get_MC_reward(model, att_feats, num_bbox, data, gen_result, num_rollout):
    batch_size = gen_result.size(0)  # batch_size = sample_size * seq_per_img
    seq_per_img = batch_size // len(data['gts'])
    gts = OrderedDict()
    for i in range(len(data['gts'])):
        gts[i] = [array_to_str(data['gts'][i][j]) for j in range(len(data['gts'][i]))]

    gts = {i: gts[i % batch_size // seq_per_img] for i in range(batch_size)}
    rewards = []
    seq_len = gen_result.size(1)
    # sample init
    samples = 0
    for i in range(num_rollout):
        # greedy_res, _ = model.sample(Variable(att_feats.data, volatile=True), Variable(num_bbox.data, volatile=True),opt={'sample_max':0})
        # res = OrderedDict()
        # for j in range(batch_size):
        #     res[j] = [array_to_str(greedy_res[j])]
        # res = [{'image_id': j, 'caption': res[j]} for j in range(batch_size)]
        # _, scores = CiderD_scorer.compute_score(gts, res)
        # if i == 0:
        #     rewards.append(scores)
        # else:
        #     rewards[0] += scores
        for l in range(1, seq_len):
            data = gen_result[:, 0:l]
            samples, _ = model.sample_mc(Variable(data, volatile=True), Variable(att_feats.data, volatile=True), Variable(num_bbox.data, volatile=True),opt={'sample_max':0})
            samples_results = samples.cpu().numpy()
            res = OrderedDict()
            for j in range(batch_size):
                res[j] = [array_to_str(samples_results[j])]
            res = [{'image_id': j, 'caption': res[j]} for j in range(batch_size)]
            _, scores = CiderD_scorer.compute_score(gts, res)
            if i == 0:
                rewards.append(scores)
            else:
                rewards[l-1] += scores
        res = OrderedDict()
        for j in range(batch_size):
            res[j] = [array_to_str(gen_result[j])]
        res = [{'image_id': j, 'caption': res[j]} for j in range(batch_size)]
        _, scores = CiderD_scorer.compute_score(gts, res)
        if i == 0:
            rewards.append(scores)
        else:
            rewards[seq_len-1] += scores
    rewards = np.transpose(np.array(rewards)) / (1.0 * num_rollout)  # batch_size * (seq_len + 1)

    greedy_res, _ = model.sample(Variable(att_feats.data, volatile=True), Variable(num_bbox.data, volatile=True))

    res = OrderedDict()
    greedy_res = greedy_res.cpu().numpy()
    for i in range(batch_size):
        res[i] = [array_to_str(greedy_res[i])]

    res = [{'image_id': i, 'caption': res[i]} for i in range(batch_size)]
    _, scores = CiderD_scorer.compute_score(gts, res)
    # print('Cider scores:', _)
    # rewards = np.repeat(rewards[:, np.newaxis], gen_result.shape[1], 1)

    return rewards[:, :] - scores[:, np.newaxis]