"""Functions for managing DNNs"""

# pylint: disable=invalid-name

from pathlib import Path
from typing import Callable
import pandas as pd


def best_model(path: Path, metric: str = None):
    """
    Parses a logs.csv file generated by shrinkbench to get path to the best model from training.

    :param path: path to the root directory generated by shrinkbench, has a file called logs.csv
    :param metric: how to compare models to find the best one, it is {train_, val_} +
        {loss, acc1, acc5}.  If there is no metric provided, takes the model from the last epoch.
    :return: a path to the file that stores the model with the best metric
    """

    if not isinstance(path, Path):
        path = Path(path)

    csv_path = path / "logs.csv"
    df = pd.read_csv(csv_path)
    df = df[df["epoch"] >= 0]
    best_epoch = None

    if metric is not None:
        if "loss" in metric:
            best_epoch = int(df[df[metric] == df[metric].min()]["epoch"].iloc[0])
        if "acc" in metric:
            best_epoch = int(df[df[metric] == df[metric].max()]["epoch"].iloc[0])
    else:
        best_epoch = int(df['epoch'].max())

    return path / "checkpoints" / f"checkpoint-{best_epoch}.pt"


def get_hyperparameters(
    model_type: str, num_workers: int = 4, debug: int = None
) -> (dict, dict):
    """
    Returns the training and pruning hyperparameters for a model

    :param model_type: architecture of the model.
    :param num_workers: number of worker threads to use to load a dataset.
    :param debug: if specified, it is the number of epochs for training and
        finetuning

    :return: a tuple of (kwargs for training, kwargs for testing)
    """

    # for now all models use resnet parameters
    if "resnet" in model_type or 1:
        model_type = "resnet"

    model_kwargs = {
        "resnet": {
            "train": {
                "optim": "SGD",
                "epochs": 300 if not debug else debug,
                "lr": 1e-1,
                "weight_decay": 1e-4,
            },
            "prune": {
                "optim": "SGD",
                "epochs": 40 if not debug else debug,
                "lr": 1e-3,
                "weight_decay": 1e-4,
            },
        }
    }

    # kwargs used in both training and fine-tuning
    common_kwargs = {
        "dl_kwargs": {
            "batch_size": 128,
            "pin_memory": False,
            "num_workers": num_workers,
        },
        "save_freq": 1000,
        "early_stop_method": None,
    }

    train_kwargs = {
        **common_kwargs,
        "train_kwargs": {**model_kwargs[model_type]["train"]},
        "lr_schedule": get_lr_schedule(
            model_kwargs[model_type]["train"]["lr"],
            model_kwargs[model_type]["train"]["epochs"],
        ),
    }

    # use a fixed lr for finetuning
    prune_kwargs = {
        **common_kwargs,
        "train_kwargs": {**model_kwargs[model_type]["prune"]},
        "lr_schedule": get_lr_schedule(
            model_kwargs[model_type]["prune"]["lr"],
            model_kwargs[model_type]["prune"]["epochs"],
            "fixed",
        ),
    }

    return train_kwargs, prune_kwargs


def get_lr_schedule(
    intial_rate: float, num_epochs: int, method: str = None
) -> Callable:
    """
    Factory function to generate a learning rate schedule.

    :param intial_rate: the initial learning rate.
    :param num_epochs: the number of epochs used for training.
    :param method: the lr method.

    :return: a function with signature (epochs: int) that returns the multiplication factor
        for the initial learning rate given the epoch
    """

    def lr_schedule(epoch: int) -> float:
        """
        Given an epoch, get the multiplication factor for the learning rate.

        Keeps initial lr until halfway done with training, when lr decreases by a factor of 10.
        Then lr becomes 1/100 of initial lr at 75% of the way done with training.
        """
        if epoch < num_epochs * 0.5:
            return 1
        if num_epochs * 0.75 > epoch >= num_epochs * 0.5:
            return 0.1
        return 0.01

    if method == "fixed":
        return lambda x: 1
    return lr_schedule


if __name__ == '__main__':
    logs_path = r'C:\Users\Jonah\Desktop\Jonah\0Grad_1\Research\code\aicas\experiments\experiment_0' \
                r'\googlenet\CIFAR10\googlenet_GlobalMagGrad_2_compression_40_finetune_iterations' \
                r'\20211101-111936-I0Y6-8bd05137846f1a27442e8665fcd4d428'

    a = best_model(logs_path)
    print(a)