import logging
import os
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from kliff.descriptors.descriptor import Descriptor
from kliff.utils import create_directory, seed_all, to_path

logger = logging.getLogger(__name__)


class ModelTorch(nn.Module):
    """
    Base class for machine learning models.

    Typically, a user will not directly use this.

    Args:
        descriptor: atomic environment descriptor for computing configuration
            fingerprints. See :meth:`~kliff.descriptors.SymmetryFunction` and
            :meth:`~kliff.descriptors.Bispectrum`.
        seed: random seed.
    """

    def __init__(self, descriptor: Descriptor, seed: int = 35):
        super(ModelTorch, self).__init__()

        self.descriptor = descriptor
        self.seed = seed
        seed_all(seed)

        dtype = self.descriptor.get_dtype()

        if dtype == np.float32:
            self.dtype = torch.float32
        elif dtype == np.float64:
            self.dtype = torch.float64
        else:
            raise ModelTorchError(f"Not support dtype {dtype}.")

        self.save_prefix = None
        self.save_start = None
        self.save_frequency = None

    def forward(self, x):
        raise NotImplementedError("`forward` not implemented.")

    def fit(self, path: Path):
        """
        Fit the model using analytic solution.

        Args:
            path: path to the fingerprints generated by the descriptor.
        """
        raise ModelTorchError(
            "Analytic fitting not supported for this model. Minimize a loss function to "
            "train the model instead."
        )

    def set_save_metadata(self, prefix: Path, start: int, frequency: int = 1):
        """
        Set metadata that controls how the model are saved during training.

        If this function is called before minimization starts, the model will be
        saved to the directory specified by `prefix` every `frequency` epochs,
        beginning at the `start` epoch.

        Args:
            prefix: Path to the directory where the models are saved.
                Model will be named as `{prefix}/model_epoch{ep}.pt`, where `ep` is
                the epoch number.
            start: Epoch number at which begins to save the model.
            frequency: Save the model every `frequency` epochs.
        """
        self.save_prefix = str(prefix)
        self.save_start = int(start)
        self.save_frequency = int(frequency)

    def save(self, filename: Path):
        """
        Save a model to disk.

        Args:
            filename: Path to store the model.
        """
        filename = to_path(filename)
        create_directory(filename)
        torch.save(self.state_dict(), str(filename))

        # save descriptor mean and stdev
        fname = filename.parent.joinpath("mean_and_stdev.pkl")
        self.descriptor.dump_mean_stdev(fname)

    def load(self, filename: Path, mode: str = "train"):
        """
        Load a model on disk into memory.

        Args:
            filename: Path where the model is stored.
            mode: Purpose of the loaded model. Should be either `train` or `eval`.
        """
        filename = to_path(filename)
        self.load_state_dict(torch.load(str(filename)))

        if mode == "train":
            self.train()
        elif mode == "eval":
            self.eval()
        else:
            raise ModelTorchError(
                'Unrecognized mode "{}" in model.load().'.format(mode)
            )

        # load descriptor mean and stdev
        fname = filename.parent.joinpath("mean_and_stdev.pkl")
        self.descriptor.load_mean_stdev(fname)

    def write_kim_model(self, path: Path = None):
        raise NotImplementedError("`write_kim_model` not implemented.")


class ModelTorchError(Exception):
    def __init__(self, msg):
        super(ModelTorchError, self).__init__(msg)
        self.msg = msg
