"""
Useful class for Experiment tracking, and ensuring code is
saved alongside files.
"""
import datetime
import os
import shlex
import shutil
import subprocess
import typing
from pathlib import Path

import randomname


class Experiment:
    """This class contains utilities for managing experiments.
    It is a context manager, that when you enter it, changes
    your directory to a specified experiment folder (which
    optionally can have an automatically generated experiment
    name, or a specified one), and changes the CUDA device used
    to the specified device (or devices).

    Parameters
    ----------
    exp_directory : str
        Folder where all experiments are saved, by default "runs/".
    exp_name : str, optional
        Name of the experiment, by default uses the current time, date, and
        hostname to save.
    """

    def __init__(
        self,
        exp_directory: str = "runs/",
        exp_name: str = None,
    ):
        if exp_name is None:
            exp_name = self.generate_exp_name()
        exp_dir = Path(exp_directory) / exp_name
        exp_dir.mkdir(parents=True, exist_ok=True)

        self.exp_dir = exp_dir
        self.exp_name = exp_name
        self.git_tracked_files = (
            subprocess.check_output(
                shlex.split("git ls-tree --full-tree --name-only -r HEAD")
            )
            .decode("utf-8")
            .splitlines()
        )
        self.parent_directory = Path(".").absolute()

    def __enter__(self):
        self.prev_dir = os.getcwd()
        os.chdir(self.exp_dir)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        os.chdir(self.prev_dir)

    @staticmethod
    def generate_exp_name():
        """Generates a random experiment name based on the date
        and a randomly generated adjective-noun tuple.

        Returns
        -------
        str
            Randomly generated experiment name.
        """
        date = datetime.datetime.now().strftime("%y%m%d")
        name = f"{date}-{randomname.get_name()}"
        return name

    def snapshot(self, filter_fn: typing.Callable = lambda f: True):
        """Captures a full snapshot of all the files tracked by git at the time
        the experiment is run. It also captures the diff against the committed
        code as a separate file.

        Parameters
        ----------
        filter_fn : typing.Callable, optional
            Function that can be used to exclude some files
            from the snapshot, by default accepts all files
        """
        for f in self.git_tracked_files:
            if filter_fn(f):
                Path(f).parent.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(self.parent_directory / f, f)