# -*- coding: utf-8 -*-
"""Test for KL functions"""

import unittest

import torch

from trinity.algorithm.kl_fn import KL_FN


class KLFnTest(unittest.TestCase):
    def setUp(self):
        seed = 42
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

        shape = (4, 10)
        self.logprob = 2 * torch.rand(shape) - 1
        self.ref_logprob = 2 * torch.rand(shape) - 1
        self.old_logprob = 2 * torch.rand(shape) - 1
        self.response_mask = torch.rand(shape) > 0.5

    def test_k1_kl_fn(self):
        kl_fn_cls = KL_FN.get("k1")
        kl_fn = kl_fn_cls(kl_coef=0.01)
        kl = kl_fn.calculate_kl(self.logprob, self.ref_logprob)
        expected_kl = self.logprob - self.ref_logprob
        self.assertTrue(torch.allclose(kl, expected_kl))

    def test_k2_kl_fn(self):
        kl_fn_cls = KL_FN.get("k2")
        kl_fn = kl_fn_cls(kl_coef=0.01)
        kl = kl_fn.calculate_kl(self.logprob, self.ref_logprob)
        expected_kl = (self.logprob - self.ref_logprob).square() * 0.5
        self.assertTrue(torch.allclose(kl, expected_kl))

    def test_k3_kl_fn(self):
        kl_fn_cls = KL_FN.get("k3")
        kl_fn = kl_fn_cls(kl_coef=0.01)
        kl = kl_fn.calculate_kl(self.logprob, self.ref_logprob)
        logr = self.ref_logprob - self.logprob
        expected_kl = logr.exp() - 1 - logr
        self.assertTrue(torch.allclose(kl, expected_kl))

    def test_abs_kl_fn(self):
        kl_fn_cls = KL_FN.get("abs")
        kl_fn = kl_fn_cls(kl_coef=0.01)
        kl = kl_fn.calculate_kl(self.logprob, self.ref_logprob)
        expected_kl = torch.abs(self.logprob - self.ref_logprob)
        self.assertTrue(torch.allclose(kl, expected_kl))

    def test_low_var_kl_fn(self):
        kl_fn_cls = KL_FN.get("low_var_kl")
        kl_fn = kl_fn_cls(kl_coef=0.01)
        kl = kl_fn.calculate_kl(self.logprob, self.ref_logprob)
        kl_intermediate = self.ref_logprob - self.logprob
        kl_intermediate = torch.clamp(kl_intermediate, min=-20, max=20)
        ratio = torch.exp(kl_intermediate)
        expected_kl = torch.clamp((ratio - kl_intermediate - 1).contiguous(), min=-10, max=10)
        self.assertTrue(torch.allclose(kl, expected_kl))

    def test_dummy_kl_fn(self):
        kl_fn_cls = KL_FN.get("none")
        kl_fn = kl_fn_cls(kl_coef=0.01)
        kl = kl_fn.calculate_kl(self.logprob, self.ref_logprob)
        expected_kl = torch.zeros_like(self.logprob)
        self.assertTrue(torch.allclose(kl, expected_kl))

    def test_corrected_k3_fallback(self):
        k3_fn = KL_FN.get("k3")(kl_coef=0.01)
        corrected_k3_fn = KL_FN.get("corrected_k3")(kl_coef=0.01)
        kl_standard = k3_fn.calculate_kl(self.logprob, self.ref_logprob)
        kl_corrected_no_old = corrected_k3_fn.calculate_kl(
            self.logprob, self.ref_logprob, old_logprob=None
        )
        self.assertTrue(torch.allclose(kl_standard, kl_corrected_no_old))

    def test_corrected_k3_with_old_logprob(self):
        corrected_k3_fn = KL_FN.get("corrected_k3")(kl_coef=0.01)
        kl_corrected = corrected_k3_fn.calculate_kl(
            self.logprob, self.ref_logprob, self.old_logprob
        )
        logr = self.ref_logprob - self.logprob
        kl_standard = logr.exp() - 1 - logr
        log_ratio_is = self.logprob - self.old_logprob
        ratio_is = log_ratio_is.exp()
        ratio_is = torch.clamp(ratio_is, min=0.0, max=2.0)
        expected_kl = ratio_is * kl_standard
        self.assertTrue(torch.allclose(kl_corrected, expected_kl))

    def test_corrected_k3_same_policy(self):
        k3_fn = KL_FN.get("k3")(kl_coef=0.01)
        corrected_k3_fn = KL_FN.get("corrected_k3")(kl_coef=0.01)
        kl_standard = k3_fn.calculate_kl(self.logprob, self.ref_logprob)
        kl_corrected = corrected_k3_fn.calculate_kl(self.logprob, self.ref_logprob, self.logprob)
        self.assertTrue(torch.allclose(kl_standard, kl_corrected, rtol=1e-4, atol=1e-6))

    def test_corrected_k3_loss(self):
        corrected_k3_fn = KL_FN.get("corrected_k3")(kl_coef=0.01)
        kl_loss, metrics = corrected_k3_fn.calculate_kl_loss(
            logprob=self.logprob,
            ref_logprob=self.ref_logprob,
            response_mask=self.response_mask,
            loss_agg_mode="token-mean",
            old_logprob=self.old_logprob,
        )
        self.assertEqual(kl_loss.dim(), 0)
        self.assertIn("kl_loss", metrics)
        self.assertIn("kl_coef", metrics)
        self.assertEqual(metrics["kl_coef"], 0.01)

    def test_kl_loss_aggregation_modes(self):
        corrected_k3_fn = KL_FN.get("corrected_k3")(kl_coef=0.01)
        kl_loss_mean, _ = corrected_k3_fn.calculate_kl_loss(
            logprob=self.logprob,
            ref_logprob=self.ref_logprob,
            response_mask=self.response_mask,
            loss_agg_mode="token-mean",
            old_logprob=self.old_logprob,
        )
        kl_loss_sum, _ = corrected_k3_fn.calculate_kl_loss(
            logprob=self.logprob,
            ref_logprob=self.ref_logprob,
            response_mask=self.response_mask,
            loss_agg_mode="seq-mean-token-sum",
            old_logprob=self.old_logprob,
        )
        self.assertGreater(kl_loss_sum.item(), kl_loss_mean.item())
