# coding: utf-8
import sys
import unittest

import numpy as np
import torch
from torch.autograd import Variable

sys.path.append("../")
from ..loss import MySymkl2d, Symkl2d


def kld(p, q):
    """Calculates Kullback–Leibler divergence"""
    p = np.array(p)
    q = np.array(q)
    return np.sum(p * np.log(p / q), axis=(p.ndim - 1))


def jsd(p, q):
    """Calculates Jensen-Shannon Divergence"""
    p = np.array(p)
    q = np.array(q)
    m = 0.5 * (p + q)
    return 0.5 * kld(p, m) + 0.5 * kld(q, m)


def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()


class TestKLD(unittest.TestCase):
    """test class for loss.py
    """

    def test_one_d(self):
        np1_one_dim = np.array([0.7, 0.2, 0.1])
        np2_one_dim = np.array([0.1, 0.8, 0.1])

        t1_one_dim = Variable(torch.FloatTensor(np1_one_dim))
        t2_one_dim = Variable(torch.FloatTensor(np2_one_dim))

        in1 = softmax(np1_one_dim)
        in2 = softmax(np2_one_dim)
        actual = 0.5 * (kld(in2, in1) + kld(in1, in2))

        mysymkl = MySymkl2d()
        nmlsymkl = Symkl2d(size_average=False)
        averaged_symkl = Symkl2d(size_average=True)

        pred_mysymkl = mysymkl(t1_one_dim, t2_one_dim).data[0]
        pred_nmlsymkl = nmlsymkl(t1_one_dim, t2_one_dim).data[0]
        pred_averaged_symkl = averaged_symkl(t1_one_dim, t2_one_dim).data[0]

        self.assertAlmostEqual(pred_mysymkl * len(np1_one_dim), actual)
        self.assertAlmostEqual(pred_nmlsymkl, actual)
        self.assertAlmostEqual(pred_averaged_symkl, pred_mysymkl)

    def test_four_d(self):
        batch_size, n_ch, w, h = 16, 3, 4, 5

        np1_four_dim = np.random.random([batch_size, n_ch, w, h])
        np2_four_dim = np.random.random([batch_size, n_ch, w, h])

        t1_one_dim = Variable(torch.FloatTensor(np1_four_dim))
        t2_one_dim = Variable(torch.FloatTensor(np2_four_dim))

        in1 = softmax(np1_four_dim)  # TODO: Need to be fixed
        in2 = softmax(np2_four_dim)  # TODO: Need to be fixed
        actual = 0.5 * (kld(in2, in1) + kld(in1, in2))  # TODO: Need to be fixed

        mysymkl = MySymkl2d()
        nmlsymkl = Symkl2d(size_average=False)
        averaged_symkl = Symkl2d(size_average=True)

        pred_mysymkl = mysymkl(t1_one_dim, t2_one_dim).data[0]
        pred_nmlsymkl = nmlsymkl(t1_one_dim, t2_one_dim).data[0]
        pred_symkl = averaged_symkl(t1_one_dim, t2_one_dim).data[0]

        # self.assertAlmostEqual(pred_mysymkl * 3, actual) # TODO: Need to be fixed
        # self.assertAlmostEqual(pred_nmlsymkl, actual) # TODO: Need to be fixed
        self.assertAlmostEqual(pred_symkl, pred_mysymkl)
        self.assertAlmostEqual(pred_symkl, pred_nmlsymkl / (batch_size * n_ch * w * h))


if __name__ == "__main__":
    unittest.main()
