import contextlib
import os
import tempfile
import unittest
from copy import deepcopy
from io import StringIO
from unittest.mock import Mock, patch

import torch
from omegaconf import OmegaConf

from mmf.common.registry import registry
from mmf.models.base_model import BaseModel
from mmf.utils.checkpoint import Checkpoint
from mmf.utils.configuration import load_yaml
from mmf.utils.early_stopping import EarlyStopping
from mmf.utils.file_io import PathManager
from mmf.utils.logger import Logger
from tests.test_utils import compare_state_dicts


@contextlib.contextmanager
def mock_env_with_temp():
    d = tempfile.TemporaryDirectory()
    patched = patch("mmf.utils.checkpoint.get_mmf_env", return_value=d.name)
    patched.start()
    yield d.name
    d.cleanup()
    patched.stop()


class SimpleModule(BaseModel):
    def __init__(self, config={}):
        super().__init__(config)
        self.base = torch.nn.Sequential(
            torch.nn.Linear(5, 4), torch.nn.Tanh(), torch.nn.Linear(4, 5)
        )

        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(5, 4), torch.nn.Tanh(), torch.nn.Linear(4, 5)
        )

        self.loss = torch.nn.CrossEntropyLoss()

    def forward(self, x, target):
        x = self.classifier(self.base(x))
        return {"losses": {"total_loss": self.loss(x, target)}}


class OnlyBase(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.base_test = torch.nn.Sequential(
            torch.nn.Linear(5, 4), torch.nn.Tanh(), torch.nn.Linear(4, 5)
        )

    def format_state_key(self, key):
        return key


class TestUtilsCheckpoint(unittest.TestCase):
    def setUp(self):
        import argparse

        torch.manual_seed(1234)
        # An easy way to get a AttributeDict object
        self.trainer = argparse.Namespace()
        self.config = OmegaConf.create(
            {
                "model": "simple",
                "model_config": {},
                "checkpoint": {
                    "save_git_details": False,
                    "reset": {"optimizer": False, "counts": False, "all": False},
                    "pretrained_state_mapping": {"base_test": "base"},
                },
                "config_override": "test",
            }
        )
        # Keep original copy for testing purposes
        self.trainer.config = deepcopy(self.config)
        self.trainer.writer = Mock(spec=Logger)

        self.trainer.model = SimpleModule()

        self.trainer.optimizer = torch.optim.Adam(
            self.trainer.model.parameters(), lr=1e-01
        )

    def test_save_config(self):
        with mock_env_with_temp() as d:
            Checkpoint(self.trainer)
            config = load_yaml(os.path.join(d, "config.yaml"))
            self.assertFalse(config == self.config)
            self.assertTrue(config == self.trainer.config)

    def test_save_and_load_state_dict(self):
        with mock_env_with_temp() as d:
            checkpoint = Checkpoint(self.trainer)
            self._init_early_stopping(checkpoint)
            self._do_a_pass()
            # Test normal case
            checkpoint.save(1500)

            self.assertTrue(
                PathManager.exists(os.path.join(d, "models", "model_1500.ckpt"))
            )
            self.assertTrue(PathManager.exists(os.path.join(d, "current.ckpt")))
            self.assertFalse(PathManager.exists(os.path.join(d, "best.ckpt")))
            os.remove(os.path.join(d, "models", "model_1500.ckpt"))
            os.remove(os.path.join(d, "current.ckpt"))

            best_model = deepcopy(self.trainer.model)
            best_optimizer = deepcopy(self.trainer.optimizer)
            # Test with update_best
            checkpoint.save(2000, update_best=True)

            self.assertTrue(
                PathManager.exists(os.path.join(d, "models", "model_2000.ckpt"))
            )
            self.assertTrue(PathManager.exists(os.path.join(d, "best.ckpt")))
            self.assertTrue(PathManager.exists(os.path.join(d, "current.ckpt")))

            self._do_a_pass()
            checkpoint.save(2500)

            # Test resume
            self.trainer.config.checkpoint.resume = True

            current_model = deepcopy(self.trainer.model)
            current_optimizer = deepcopy(self.trainer.optimizer)
            checkpoint.load_state_dict()

            self.assertFalse(
                compare_state_dicts(
                    self.trainer.model.state_dict(), best_model.state_dict()
                )
            )
            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), current_model.state_dict()
                )
            )
            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, best_optimizer)
            )
            self.assertFalse(
                self._compare_optimizers(
                    self.trainer.optimizer, best_optimizer, skip_keys=True
                )
            )
            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, current_optimizer)
            )
            self.assertTrue(
                self._compare_optimizers(
                    self.trainer.optimizer, current_optimizer, skip_keys=True
                )
            )

            base_0_weight_current = self.trainer.model.base[0].weight.data.clone()

            # Test resume_best
            self.trainer.config.checkpoint.resume = True
            self.trainer.config.checkpoint.resume_best = True

            checkpoint.load_state_dict()

            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), best_model.state_dict()
                )
            )
            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, best_optimizer)
            )
            self.assertTrue(
                self._compare_optimizers(
                    self.trainer.optimizer, best_optimizer, skip_keys=True
                )
            )
            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, current_optimizer)
            )
            self.assertFalse(
                self._compare_optimizers(
                    self.trainer.optimizer, current_optimizer, skip_keys=True
                )
            )
            base_0_weight_best = self.trainer.model.base[0].weight.data.clone()

            self.trainer.config.checkpoint.resume_best = False
            # Test distributed settings
            self.trainer.model = torch.nn.DataParallel(self.trainer.model)
            checkpoint.load_state_dict()

            self.assertTrue(
                torch.equal(
                    self.trainer.model.module.base[0].weight, base_0_weight_current
                )
            )
            self.assertFalse(
                torch.equal(
                    self.trainer.model.module.base[0].weight, base_0_weight_best
                )
            )

    def test_finalize_and_restore_from_it(self):
        with mock_env_with_temp():
            checkpoint = Checkpoint(self.trainer)
            self._init_early_stopping(checkpoint)
            original_model = deepcopy(self.trainer.model)
            self._do_a_pass()
            model_1500 = deepcopy(self.trainer.model)
            checkpoint.save(1500)

            swap = self.trainer.model
            self.trainer.model = original_model
            checkpoint.restore()
            # First test without best.ckpt
            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), original_model.state_dict()
                )
            )
            self.assertFalse(
                compare_state_dicts(
                    self.trainer.model.state_dict(), model_1500.state_dict()
                )
            )

            self.trainer.model = swap

            self._do_a_pass()
            model_2000 = deepcopy(self.trainer.model)
            checkpoint.save(2000, update_best=True)

            self._do_a_pass()
            model_2500 = deepcopy(self.trainer.model)
            checkpoint.save(2500)

            checkpoint.restore()

            self.assertFalse(
                compare_state_dicts(
                    self.trainer.model.state_dict(), original_model.state_dict()
                )
            )
            self.assertFalse(
                compare_state_dicts(
                    self.trainer.model.state_dict(), model_1500.state_dict()
                )
            )
            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), model_2000.state_dict()
                )
            )
            self.assertFalse(
                compare_state_dicts(
                    self.trainer.model.state_dict(), model_2500.state_dict()
                )
            )

    def test_finalize_and_resume_file(self):
        with mock_env_with_temp() as d:
            checkpoint = Checkpoint(self.trainer)
            self._init_early_stopping(checkpoint)
            self._do_a_pass()
            checkpoint.finalize()
            original = deepcopy(self.trainer.model)
            pth_path = os.path.join(d, "simple_final.pth")
            self.assertTrue(PathManager.exists(pth_path))

            self._do_a_pass()

            after_a_pass = deepcopy(self.trainer.model)
            original_optimizer = deepcopy(self.trainer.optimizer)
            self.trainer.config.checkpoint.resume_file = pth_path

            with contextlib.redirect_stdout(StringIO()):
                checkpoint.load_state_dict()
            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), original.state_dict()
                )
            )
            self.assertFalse(
                compare_state_dicts(
                    self.trainer.model.state_dict(), after_a_pass.state_dict()
                )
            )
            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, original_optimizer)
            )
            # Keys will not be same as we just updated the model
            self.assertTrue(
                self._compare_optimizers(
                    self.trainer.optimizer, original_optimizer, skip_keys=True
                )
            )

    def test_resets(self):
        with mock_env_with_temp():
            checkpoint = Checkpoint(self.trainer)
            self._init_early_stopping(checkpoint)
            self._do_a_pass()

            original_optimizer = deepcopy(self.trainer.optimizer)
            original_model = deepcopy(self.trainer.model)

            self.trainer.current_epoch = 3
            checkpoint.save(2000, update_best=True)
            self.trainer.current_epoch = 4
            # Test reset all
            self.trainer.config.checkpoint.resume = True
            self.trainer.config.checkpoint.reset.all = True
            checkpoint.load_state_dict()

            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), original_model.state_dict()
                )
            )

            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, original_optimizer)
            )
            self.assertTrue(
                self._compare_optimizers(
                    self.trainer.optimizer, original_optimizer, skip_keys=True
                )
            )
            self.assertEqual(self.trainer.num_updates, 0)
            self.assertEqual(self.trainer.current_iteration, 0)
            self.assertEqual(self.trainer.current_epoch, 4)

            # Test reset_optimizer
            self._init_early_stopping(checkpoint)
            self.trainer.config.checkpoint.reset.all = False
            self.trainer.config.checkpoint.reset.optimizer = True
            checkpoint.load_state_dict()

            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), original_model.state_dict()
                )
            )

            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, original_optimizer)
            )
            self.assertTrue(
                self._compare_optimizers(
                    self.trainer.optimizer, original_optimizer, skip_keys=True
                )
            )

            self.assertEqual(self.trainer.num_updates, 2000)
            self.assertEqual(self.trainer.current_iteration, 2000)
            self.assertEqual(self.trainer.current_epoch, 3)

            self._init_early_stopping(checkpoint)
            # Test reset_counts
            self.trainer.config.checkpoint.reset.all = False
            self.trainer.config.checkpoint.reset.optimizer = False
            self.trainer.config.checkpoint.reset.counts = True
            checkpoint.load_state_dict()

            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), original_model.state_dict()
                )
            )

            self.assertTrue(
                self._compare_optimizers(
                    self.trainer.optimizer, original_optimizer, skip_keys=True
                )
            )
            self.assertEqual(self.trainer.num_updates, 0)
            self.assertEqual(self.trainer.current_iteration, 0)
            self.assertEqual(self.trainer.current_epoch, 2)

            # Test with resume_best
            self._do_a_pass()
            checkpoint.save(3000)
            self._init_early_stopping(checkpoint)
            self.trainer.config.checkpoint.reset.all = False
            self.trainer.config.checkpoint.resume_best = True
            self.trainer.config.checkpoint.reset.optimizer = True
            self.trainer.config.checkpoint.reset.counts = False
            checkpoint.load_state_dict()

            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.state_dict(), original_model.state_dict()
                )
            )

            self.assertFalse(
                self._compare_optimizers(self.trainer.optimizer, original_optimizer)
            )
            self.assertFalse(
                self._compare_optimizers(
                    self.trainer.optimizer, original_optimizer, skip_keys=True
                )
            )
            self.assertEqual(self.trainer.num_updates, 1000)
            self.assertEqual(self.trainer.current_iteration, 1000)
            self.assertEqual(self.trainer.current_epoch, 3)

    def test_zoo_load(self):
        with mock_env_with_temp():
            checkpoint = Checkpoint(self.trainer)
            self._init_early_stopping(checkpoint)
            self._do_a_pass()

            original_model = deepcopy(self.trainer.model)
            ret_load_pretrained_zoo = {
                "config": self.config.model_config,
                "checkpoint": deepcopy(self.trainer.model.state_dict()),
                "full_config": self.config,
            }

            self._do_a_pass()

            with patch(
                "mmf.utils.checkpoint.load_pretrained_model",
                return_value=ret_load_pretrained_zoo,
            ):
                self.trainer.config.checkpoint.resume_zoo = "random"
                with contextlib.redirect_stdout(StringIO()):
                    checkpoint.load_state_dict()
                self.assertTrue(
                    compare_state_dicts(
                        self.trainer.model.state_dict(), original_model.state_dict()
                    )
                )

                # Now, test zoo override
                self.trainer.config.checkpoint.zoo_config_override = True
                SimpleModule.from_pretrained = Mock(
                    return_value=deepcopy(original_model)
                )
                registry.register_model("simple")(SimpleModule)
                with contextlib.redirect_stdout(StringIO()):
                    checkpoint.load_state_dict()
                self.assertTrue(
                    compare_state_dicts(
                        self.trainer.model.state_dict(), original_model.state_dict()
                    )
                )

    def test_pretrained_load(self):
        with mock_env_with_temp() as d:
            checkpoint = Checkpoint(self.trainer)
            self._init_early_stopping(checkpoint)
            self._do_a_pass()
            original_model = deepcopy(self.trainer.model)
            # Test with zoo now
            ret_load_pretrained_zoo = {
                "config": self.config.model_config,
                "checkpoint": deepcopy(self.trainer.model.state_dict()),
                "full_config": self.config,
            }

            checkpoint.save(2000)
            self.trainer.config.checkpoint.resume_file = os.path.join(d, "current.ckpt")
            self.trainer.config.checkpoint.resume_pretrained = True
            self.trainer.model = OnlyBase()
            checkpoint.load_state_dict()

            self.assertTrue(
                compare_state_dicts(
                    self.trainer.model.base_test.state_dict(),
                    original_model.base.state_dict(),
                )
            )

            with patch(
                "mmf.utils.checkpoint.load_pretrained_model",
                return_value=ret_load_pretrained_zoo,
            ):
                self.trainer.config.checkpoint.resume_zoo = "random"
                self.trainer.config.checkpoint.resume_file = None
                self.trainer.model = OnlyBase()
                checkpoint.load_state_dict()

                self.assertTrue(
                    compare_state_dicts(
                        self.trainer.model.base_test.state_dict(),
                        original_model.base.state_dict(),
                    )
                )

    def _init_early_stopping(self, checkpoint):
        self.trainer.num_updates = 0
        self.trainer.current_iteration = 0
        self.trainer.current_epoch = 0
        self.trainer.early_stopping = EarlyStopping(self.trainer.model, checkpoint)
        self.trainer.early_stopping.best_monitored_iteration = 1000
        self.trainer.early_stopping.best_monitored_update = 1000
        self.trainer.early_stopping.best_monitored_value = 0.1
        self.trainer.current_epoch = 2

    def _do_a_pass(self):
        self.trainer.optimizer.zero_grad()
        self.trainer.model.train()
        with contextlib.redirect_stdout(StringIO()):
            loss = self.trainer.model(
                torch.rand(5, 5, requires_grad=True),
                torch.empty(5, dtype=torch.long).random_(5),
            )

        loss["losses"]["total_loss"].sum().backward()
        self.trainer.optimizer.step()

    def _compare_optimizers(self, a, b, skip_keys=False):
        state_dict_a = a.state_dict()
        state_dict_b = b.state_dict()
        state_a = state_dict_a["state"]
        state_b = state_dict_b["state"]

        same = True
        if not skip_keys:
            same = same and list(state_a.keys()) == list(state_b.keys())
            same = same and state_dict_a["param_groups"] == state_dict_b["param_groups"]

        for item1, item2 in zip(state_a.values(), state_b.values()):
            same = same and compare_state_dicts(item1, item2)

        return same
