"""
Trainer tests
"""
import pytest

import numpy as np
import tempfile
from os.path import isfile

import torch
from torch.nn import Linear

from nequip.data import AtomicDataDict
from nequip.train.trainer import Trainer
from nequip.utils.savenload import load_file
from nequip.nn import GraphModuleMixin

# set up two config to test
DEBUG = False
NATOMS = 3
NFRAMES = 10
minimal_config = dict(
    n_train=4,
    n_val=4,
    exclude_keys=["sth"],
    max_epochs=2,
    batch_size=2,
    learning_rate=1e-2,
    optimizer="Adam",
    seed=0,
    restart=False,
    T_0=50,
    T_mult=2,
    loss_coeffs={"forces": 2},
    early_stopping_patiences={"loss": 50},
    early_stopping_lower_bounds={"LR": 1e-10},
)
configs_to_test = [dict(), minimal_config]
loop_config = pytest.mark.parametrize("trainer", configs_to_test, indirect=True)
one_config_test = pytest.mark.parametrize("trainer", [minimal_config], indirect=True)


@pytest.fixture(scope="class")
def trainer(request):
    """
    Generate a class instance with minimal configurations
    """
    params = request.param

    model = DummyNet(3)
    with tempfile.TemporaryDirectory(prefix="output") as path:
        params["root"] = path
        c = Trainer(model=model, **params)
        yield c


class TestTrainerSetUp:
    """
    test initialization
    """

    @one_config_test
    def test_init(self, trainer):
        assert isinstance(trainer, Trainer)


class TestDuplicateError:
    def test_duplicate_id_2(self, temp_data):

        minimal_config["root"] = temp_data

        model = DummyNet(3)
        c1 = Trainer(model=model, **minimal_config)
        logfile1 = c1.logfile

        c2 = Trainer(model=model, **minimal_config)
        logfile2 = c2.logfile

        assert c1.root == c2.root
        assert c1.workdir != c2.workdir
        assert c1.logfile.endswith("log")
        assert c2.logfile.endswith("log")


class TestInit:
    @one_config_test
    def test_init_model(self, trainer):
        trainer.init_model()


class TestSaveLoad:
    @loop_config
    @pytest.mark.parametrize("state_dict", [True, False])
    @pytest.mark.parametrize("training_progress", [True, False])
    def test_as_dict(self, trainer, state_dict, training_progress):

        dictionary = trainer.as_dict(
            state_dict=state_dict, training_progress=training_progress
        )

        assert "optimizer_kwargs" in dictionary
        assert state_dict == ("state_dict" in dictionary)
        assert training_progress == ("progress" in dictionary)
        assert len(dictionary["optimizer_kwargs"]) > 1

    @loop_config
    @pytest.mark.parametrize("format, suffix", [("torch", "pth"), ("yaml", "yaml")])
    def test_save(self, trainer, format, suffix):

        with tempfile.NamedTemporaryFile(suffix="." + suffix) as tmp:
            file_name = tmp.name
            trainer.save(file_name, format=format)
            assert isfile(file_name), "fail to save to file"
            assert suffix in file_name

    @loop_config
    @pytest.mark.parametrize("append", [True, False])
    def test_from_dict(self, trainer, append):

        # torch.save(trainer.model, trainer.best_model_path)

        dictionary = trainer.as_dict(state_dict=True, training_progress=True)
        trainer1 = Trainer.from_dict(dictionary, append=append)

        for key in [
            "best_model_path",
            "last_model_path",
            "logfile",
            "epoch_log",
            "batch_log",
            "workdir",
        ]:
            v1 = getattr(trainer, key, None)
            v2 = getattr(trainer1, key, None)
            print(key, v1, v2)
            assert append == (v1 == v2)

    @loop_config
    @pytest.mark.parametrize("append", [True, False])
    def test_from_file(self, trainer, append):

        format = "torch"

        with tempfile.NamedTemporaryFile(suffix=".pth") as tmp:

            file_name = tmp.name
            file_name = trainer.save(file_name, format=format)

            assert isfile(trainer.last_model_path)

            trainer1 = Trainer.from_file(file_name, format=format, append=append)

            for key in [
                "best_model_path",
                "last_model_path",
                "logfile",
                "epoch_log",
                "batch_log",
                "workdir",
            ]:
                v1 = getattr(trainer, key, None)
                v2 = getattr(trainer1, key, None)
                print(key, v1, v2)
                assert append == (v1 == v2)

            for iparam, group1 in enumerate(trainer.optim.param_groups):
                group2 = trainer1.optim.param_groups[iparam]

                for key in group1:
                    v1 = group1[key]
                    v2 = group2[key]
                    if key != "params":
                        assert v1 == v2


class TestData:
    @one_config_test
    @pytest.mark.parametrize("mode", ["random", "sequential"])
    def test_split(self, trainer, nequip_dataset, mode):

        trainer.train_val_split = mode
        trainer.set_dataset(nequip_dataset)
        for i, batch in enumerate(trainer.dl_train):
            print(i, batch)


class TestTrain:
    @one_config_test
    def test_train(self, trainer, nequip_dataset):

        v0 = get_param(trainer.model)
        trainer.set_dataset(nequip_dataset)
        trainer.train()
        v1 = get_param(trainer.model)

        assert not np.allclose(v0, v1), "fail to train parameters"
        assert isfile(trainer.last_model_path), "fail to save best model"

    @one_config_test
    def test_load_w_revision(self, trainer):

        with tempfile.TemporaryDirectory() as folder:

            file_name = trainer.save(f"{folder}/a.pth")

            dictionary = load_file(
                supported_formats=dict(torch=["pt", "pth"]),
                filename=file_name,
                enforced_format="torch",
            )

            dictionary["root"] = folder
            dictionary["progress"]["stop_arg"] = None
            dictionary["max_epochs"] *= 2
            trainer1 = Trainer.from_dict(dictionary, append=False)
            assert trainer1.iepoch == trainer.iepoch
            assert trainer1.max_epochs == minimal_config["max_epochs"] * 2

    @one_config_test
    def test_restart_training(self, trainer, nequip_dataset):

        model = trainer.model
        device = trainer.device
        optimizer = trainer.optim
        trainer.set_dataset(nequip_dataset)
        trainer.train()

        v0 = get_param(trainer.model)
        with tempfile.TemporaryDirectory() as folder:
            file_name = trainer.save(f"{folder}/hello.pth")

            dictionary = load_file(
                filename=file_name,
                supported_formats=dict(torch=["pt", "pth"]),
                enforced_format="torch",
            )

            dictionary["progress"]["stop_arg"] = None
            dictionary["max_epochs"] *= 2
            dictionary["root"] = folder

            trainer1 = Trainer.from_dict(dictionary, append=False)
            trainer1.stop_arg = None
            trainer1.max_epochs *= 2
            trainer1.set_dataset(nequip_dataset)
            trainer1.train()
            v1 = get_param(trainer1.model)

            assert not np.allclose(v0, v1), "fail to train parameters"


class TestRescale:
    def test_scaling(self, scale_train):

        trainer = scale_train

        data = trainer.dataset_train[0]

        def get_loss_val(validation):
            trainer.ibatch = 0
            trainer.n_batches = 10
            trainer.reset_metrics()
            trainer.batch_step(
                data=data,
                validation=validation,
            )
            metrics, _ = trainer.metrics.flatten_metrics(
                trainer.batch_metrics,
            )
            return trainer.batch_losses, metrics

        ref_loss, ref_met = get_loss_val(True)
        loss, met = get_loss_val(False)

        # metrics should always be in real unit
        for k in ref_met:
            assert np.allclose(ref_met[k], met[k])
        # loss should always be in normalized unit
        for k in ref_loss:
            assert np.allclose(ref_loss[k], loss[k])

        assert np.allclose(np.sqrt(loss["loss_f"]) * trainer.scale, met["f_rmse"])


def get_param(model):
    v = []
    for p in model.parameters():
        v += [np.hstack(p.data.tolist())]
    v = np.hstack(v)
    return v


class DummyNet(GraphModuleMixin, torch.nn.Module):
    def __init__(self, ndim, nydim=1) -> None:
        super().__init__()
        self.ndim = ndim
        self.nydim = nydim
        self.linear1 = Linear(ndim, nydim)
        self.linear2 = Linear(ndim, nydim * 3)
        self._init_irreps(
            irreps_in={"pos": "1x1o"},
            irreps_out={
                AtomicDataDict.FORCE_KEY: "1x1o",
                AtomicDataDict.TOTAL_ENERGY_KEY: "1x0e",
            },
        )

    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        data = data.copy()
        x = data["pos"]
        data.update(
            {
                AtomicDataDict.FORCE_KEY: self.linear2(x),
                AtomicDataDict.TOTAL_ENERGY_KEY: self.linear1(x),
            }
        )
        return data


class DummyScale(torch.nn.Module):
    """ mimic the rescale model"""

    def __init__(self, key, scale, shift) -> None:
        super().__init__()
        self.key = key
        self.scale_by = torch.as_tensor(scale, dtype=torch.get_default_dtype())
        self.shift_by = torch.as_tensor(shift, dtype=torch.get_default_dtype())
        self.linear2 = Linear(3, 3)

    def forward(self, data):
        out = self.linear2(data["pos"])
        if not self.training:
            out = out * self.scale_by
            out = out + self.shift_by
        return {self.key: out}

    def scale(self, data, force_process=False):
        data = data.copy()
        if force_process or not self.training:
            data[self.key] = data[self.key] * self.scale_by
            data[self.key] = data[self.key] + self.shift_by
        return data

    def unscale(self, data, force_process=False):
        data = data.copy()
        if force_process or self.training:
            data[self.key] = data[self.key] - self.shift_by
            data[self.key] = data[self.key] / self.scale_by
        return data


@pytest.fixture(scope="class")
def scale_train(nequip_dataset):
    with tempfile.TemporaryDirectory(prefix="output") as path:
        trainer = Trainer(
            model=DummyScale(AtomicDataDict.FORCE_KEY, scale=1.3, shift=1),
            n_train=4,
            n_val=4,
            max_epochs=0,
            batch_size=2,
            loss_coeffs=AtomicDataDict.FORCE_KEY,
            root=path,
        )
        trainer.set_dataset(nequip_dataset)
        trainer.train()
        trainer.scale = 1.3
        yield trainer
