#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import itertools
import logging
import random
import unittest
from collections import defaultdict
from itertools import permutations

import numpy as np
import pytest
import reagent.core.types as rlt
import torch
import torch.nn.functional as F
from parameterized import parameterized
from reagent.model_utils.seq2slate_utils import (
    DECODER_START_SYMBOL,
    Seq2SlateMode,
    Seq2SlateOutputArch,
    mask_logits_by_idx,
    per_symbol_to_per_seq_log_probs,
    per_symbol_to_per_seq_probs,
    subsequent_mask,
    pytorch_decoder_mask,
)
from reagent.test.ranking.seq2slate_utils import (
    MODEL_TRANSFORMER,
    ON_POLICY,
    create_batch,
    create_seq2slate_net,
    rank_on_policy,
    run_seq2slate_tsp,
)


logger = logging.getLogger(__name__)


output_arch_list = [
    Seq2SlateOutputArch.FRECHET_SORT,
    Seq2SlateOutputArch.AUTOREGRESSIVE,
]
temperature_list = [1.0, 2.0]


class TestSeq2SlateOnPolicy(unittest.TestCase):
    def setUp(self):
        np.random.seed(0)
        random.seed(0)
        torch.manual_seed(0)

    def test_pytorch_decoder_mask(self):
        batch_size = 3
        src_seq_len = 4
        num_heads = 2

        memory = torch.randn(batch_size, src_seq_len, num_heads)
        tgt_in_idx = torch.tensor([[1, 2, 3], [1, 4, 2], [1, 5, 4]]).long()
        tgt_tgt_mask, tgt_src_mask = pytorch_decoder_mask(memory, tgt_in_idx, num_heads)

        expected_tgt_tgt_mask = (
            torch.tensor(
                [
                    [False, True, True],
                    [False, False, True],
                    [False, False, False],
                ],
            )
            .unsqueeze(0)
            .repeat(batch_size * num_heads, 1, 1)
        )
        expected_tgt_src_mask = torch.tensor(
            [
                [
                    [False, False, False, False],
                    [True, False, False, False],
                    [True, True, False, False],
                ],
                [
                    [False, False, False, False],
                    [False, False, True, False],
                    [True, False, True, False],
                ],
                [
                    [False, False, False, False],
                    [False, False, False, True],
                    [False, False, True, True],
                ],
            ]
        ).repeat_interleave(num_heads, dim=0)
        assert torch.all(tgt_tgt_mask == expected_tgt_tgt_mask)
        assert torch.all(tgt_src_mask == expected_tgt_src_mask)

    def test_per_symbol_to_per_seq_log_probs(self):
        """
        Test per_symbol_to_per_seq_log_probs method
        """
        batch_size = 1
        seq_len = 3
        candidate_size = seq_len + 2

        tgt_out_idx = torch.tensor([[0, 2, 1]]) + 2
        per_symbol_log_probs = torch.randn(batch_size, seq_len, candidate_size)
        per_symbol_log_probs[0, :, :2] = float("-inf")
        per_symbol_log_probs[0, 1, 2] = float("-inf")
        per_symbol_log_probs[0, 2, 2] = float("-inf")
        per_symbol_log_probs[0, 2, 4] = float("-inf")
        per_symbol_log_probs = F.log_softmax(per_symbol_log_probs, dim=2)

        expect_per_seq_log_probs = (
            per_symbol_log_probs[0, 0, 2]
            + per_symbol_log_probs[0, 1, 4]
            + per_symbol_log_probs[0, 2, 3]
        )
        computed_per_seq_log_probs = per_symbol_to_per_seq_log_probs(
            per_symbol_log_probs, tgt_out_idx
        )
        np.testing.assert_allclose(
            expect_per_seq_log_probs, computed_per_seq_log_probs, atol=0.001, rtol=0.0
        )

    def test_per_symbol_to_per_seq_probs(self):
        batch_size = 1
        seq_len = 3
        candidate_size = seq_len + 2

        tgt_out_idx = torch.tensor([[0, 2, 1]]) + 2
        per_symbol_log_probs = torch.randn(batch_size, seq_len, candidate_size)
        per_symbol_log_probs[0, :, :2] = float("-inf")
        per_symbol_log_probs[0, 1, 2] = float("-inf")
        per_symbol_log_probs[0, 2, 2] = float("-inf")
        per_symbol_log_probs[0, 2, 4] = float("-inf")
        per_symbol_log_probs = F.log_softmax(per_symbol_log_probs, dim=2)
        per_symbol_probs = torch.exp(per_symbol_log_probs)

        expect_per_seq_probs = (
            per_symbol_probs[0, 0, 2]
            * per_symbol_probs[0, 1, 4]
            * per_symbol_probs[0, 2, 3]
        )
        computed_per_seq_probs = per_symbol_to_per_seq_probs(
            per_symbol_probs, tgt_out_idx
        )
        np.testing.assert_allclose(
            expect_per_seq_probs, computed_per_seq_probs, atol=0.001, rtol=0.0
        )

    def test_subsequent_mask(self):
        expect_mask = torch.tensor([[1, 0, 0], [1, 1, 0], [1, 1, 1]])
        mask = subsequent_mask(3, torch.device("cpu"))
        assert torch.all(torch.eq(mask, expect_mask))

    def test_mask_logits_by_idx(self):
        logits = torch.tensor(
            [
                [
                    [1.0, 2.0, 3.0, 4.0, 5.0],
                    [2.0, 3.0, 4.0, 5.0, 6.0],
                    [3.0, 4.0, 5.0, 6.0, 7.0],
                ],
                [
                    [5.0, 4.0, 3.0, 2.0, 1.0],
                    [6.0, 5.0, 4.0, 3.0, 2.0],
                    [7.0, 6.0, 5.0, 4.0, 3.0],
                ],
            ]
        )
        tgt_in_idx = torch.tensor(
            [[DECODER_START_SYMBOL, 2, 3], [DECODER_START_SYMBOL, 4, 3]]
        )
        masked_logits = mask_logits_by_idx(logits, tgt_in_idx)
        expected_logits = torch.tensor(
            [
                [
                    [float("-inf"), float("-inf"), 3.0, 4.0, 5.0],
                    [float("-inf"), float("-inf"), float("-inf"), 5.0, 6.0],
                    [float("-inf"), float("-inf"), float("-inf"), float("-inf"), 7.0],
                ],
                [
                    [float("-inf"), float("-inf"), 3.0, 2.0, 1.0],
                    [float("-inf"), float("-inf"), 4.0, 3.0, float("-inf")],
                    [float("-inf"), float("-inf"), 5.0, float("-inf"), float("-inf")],
                ],
            ]
        )
        assert torch.all(torch.eq(masked_logits, expected_logits))

    @parameterized.expand(itertools.product(output_arch_list, temperature_list))
    @torch.no_grad()
    def test_seq2slate_transformer_propensity_computation(
        self, output_arch, temperature
    ):
        """
        Test propensity computation of seq2slate net
        """
        candidate_num = 4
        candidate_dim = 2
        hidden_size = 32
        all_perm = torch.tensor(
            list(permutations(torch.arange(candidate_num), candidate_num))
        )
        batch_size = len(all_perm)
        device = torch.device("cpu")

        seq2slate_net = create_seq2slate_net(
            MODEL_TRANSFORMER,
            candidate_num,
            candidate_dim,
            hidden_size,
            output_arch,
            temperature,
            device,
        )
        batch = create_batch(
            batch_size,
            candidate_num,
            candidate_dim,
            device,
            ON_POLICY,
            diverse_input=False,
        )
        batch = rlt.PreprocessedRankingInput.from_input(
            state=batch.state.float_features,
            candidates=batch.src_seq.float_features,
            device=device,
            action=all_perm,
        )
        per_symbol_log_prob = seq2slate_net(
            batch, mode=Seq2SlateMode.PER_SYMBOL_LOG_PROB_DIST_MODE
        ).log_probs
        per_seq_log_prob = seq2slate_net(
            batch, mode=Seq2SlateMode.PER_SEQ_LOG_PROB_MODE
        ).log_probs
        per_seq_log_prob_computed = per_symbol_to_per_seq_log_probs(
            per_symbol_log_prob, all_perm + 2
        )
        # probabilities of two modes should match
        np.testing.assert_allclose(
            per_seq_log_prob, per_seq_log_prob_computed, atol=0.00001
        )
        # probabilities of all possible permutations should sum up to 1
        np.testing.assert_allclose(
            torch.sum(torch.exp(per_seq_log_prob)), 1.0, atol=0.00001
        )

    @parameterized.expand(itertools.product(output_arch_list, temperature_list))
    def test_seq2slate_transformer_onpolicy_basic_logic(self, output_arch, temperature):
        """
        Test basic logic of seq2slate on policy sampling
        """
        device = torch.device("cpu")
        candidate_num = 4
        candidate_dim = 2
        batch_size = 4096
        hidden_size = 32
        seq2slate_net = create_seq2slate_net(
            MODEL_TRANSFORMER,
            candidate_num,
            candidate_dim,
            hidden_size,
            output_arch,
            temperature,
            device,
        )
        batch = create_batch(
            batch_size,
            candidate_num,
            candidate_dim,
            device,
            ON_POLICY,
            diverse_input=False,
        )

        action_to_propensity_map = {}
        action_count = defaultdict(int)
        total_count = 0
        for i in range(50):
            model_propensity, model_action = rank_on_policy(
                seq2slate_net, batch, candidate_num, greedy=False
            )
            for propensity, action in zip(model_propensity, model_action):
                action_str = ",".join(map(str, action.numpy().tolist()))

                # Same action always leads to same propensity
                if action_to_propensity_map.get(action_str) is None:
                    action_to_propensity_map[action_str] = float(propensity)
                else:
                    np.testing.assert_allclose(
                        action_to_propensity_map[action_str],
                        float(propensity),
                        atol=0.001,
                        rtol=0.0,
                    )

                action_count[action_str] += 1
                total_count += 1

            logger.info(f"Finish {i} round, {total_count} data counts")

        # Check action distribution
        for action_str, count in action_count.items():
            empirical_propensity = count / total_count
            computed_propensity = action_to_propensity_map[action_str]
            logger.info(
                f"action={action_str}, empirical propensity={empirical_propensity}, "
                f"computed propensity={computed_propensity}"
            )
            np.testing.assert_allclose(
                computed_propensity, empirical_propensity, atol=0.01, rtol=0.0
            )

    def test_seq2slate_transformer_on_policy_simple_tsp(self):
        """
        Solve Traveling Salesman Problem. Cities comes from a fixed set of nodes (cities).
        Easily hit reward threshold after one batch training
        """
        device = torch.device("cpu")
        batch_size = 4096
        epochs = 1
        num_batches = 50
        expect_reward_threshold = 1.12
        hidden_size = 32
        num_candidates = 6
        diverse_input = False
        learning_rate = 0.001
        learning_method = ON_POLICY
        policy_gradient_interval = 1
        run_seq2slate_tsp(
            MODEL_TRANSFORMER,
            batch_size,
            epochs,
            num_candidates,
            num_batches,
            hidden_size,
            diverse_input,
            learning_rate,
            expect_reward_threshold,
            learning_method,
            policy_gradient_interval,
            device,
        )

    @pytest.mark.seq2slate_long
    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
    def test_seq2slate_transformer_on_policy_hard_tsp(self):
        """
        Solve Traveling Salesman Problem. Data comes from different sets of cities.
        """
        device = torch.device("cuda")
        batch_size = 4096
        epochs = 3
        num_batches = 300
        expect_reward_threshold = 1.05
        hidden_size = 32
        num_candidates = 6
        diverse_input = True
        learning_rate = 0.001
        learning_method = ON_POLICY
        policy_gradient_interval = 1
        run_seq2slate_tsp(
            MODEL_TRANSFORMER,
            batch_size,
            epochs,
            num_candidates,
            num_batches,
            hidden_size,
            diverse_input,
            learning_rate,
            expect_reward_threshold,
            learning_method,
            policy_gradient_interval,
            device,
        )
