# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from typing import Iterator, Mapping, Type, Optional, Tuple, List
import math
import copy
import random
import os

import torch
import tensorwatch as tw
from torch.utils.data.dataloader import DataLoader
import yaml

from overrides import overrides

from archai.common.common import logger
from archai.common.checkpoint import CheckPoint
from archai.common.config import Config
from archai.nas.model_desc_builder import ModelDescBuilder
from archai.nas.arch_trainer import TArchTrainer
from archai.nas import nas_utils
from archai.nas.model_desc import CellType, ModelDesc
from archai.common.trainer import Trainer
from archai.datasets import data
from archai.nas.model import Model
from archai.common.metrics import EpochMetrics, Metrics
from archai.common import utils
from archai.nas.finalizers import Finalizers
from archai.nas.searcher import ModelMetrics, Searcher, SearchResult
from archai.nas import nas_utils


class SearchCombinations(Searcher):
    @overrides
    def search(self, conf_search:Config, model_desc_builder:ModelDescBuilder,
               trainer_class:TArchTrainer, finalizers:Finalizers)->SearchResult:

        # region config vars
        conf_model_desc = conf_search['model_desc']
        conf_post_train = conf_search['post_train']

        conf_checkpoint = conf_search['checkpoint']
        resume = conf_search['resume']
        # endregion

        self._checkpoint = nas_utils.create_checkpoint(conf_checkpoint, resume)

        macro_combinations = list(self.get_combinations(conf_search))
        start_macro_i, best_search_result = self.restore_checkpoint(conf_search,
                                                                     macro_combinations)
        best_macro_comb = -1,-1,-1 # reductions, cells, nodes

        for macro_comb_i in range(start_macro_i, len(macro_combinations)):
            reductions, cells, nodes = macro_combinations[macro_comb_i]
            logger.pushd(f'r{reductions}.c{cells}.n{nodes}')

            # build model description that we will search on
            model_desc = self.build_model_desc(model_desc_builder, conf_model_desc,
                                               reductions, cells, nodes)

            # perform search on model description
            model_desc, search_metrics = self.search_model_desc(conf_search,
                model_desc, trainer_class, finalizers)

            # train searched model for few epochs to get some perf metrics
            model_metrics = self.train_model_desc(model_desc,
                                                  conf_post_train)

            assert model_metrics is not None, "'post_train' section in yaml should have non-zero epochs if running combinations search"

            # save result
            self.save_trained(conf_search, reductions, cells, nodes, model_metrics)

            # update the best result so far
            if self.is_better_metrics(best_search_result.search_metrics,
                                      model_metrics.metrics):
                best_search_result = SearchResult(model_desc, search_metrics,
                                                  model_metrics.metrics)
                best_macro_comb = reductions, cells, nodes

            # checkpoint
            assert best_search_result is not None
            self.record_checkpoint(macro_comb_i, best_search_result)
            logger.popd() # reductions, cells, nodes

        assert best_search_result is not None
        self.clean_log_result(conf_search, best_search_result)
        logger.info({'best_macro_comb':best_macro_comb})

        return best_search_result

    def is_better_metrics(self, metrics1:Optional[Metrics],
                          metrics2:Optional[Metrics])->bool:
        if metrics1 is None or metrics2 is None:
            return True
        return metrics2.best_val_top1() >= metrics1.best_val_top1()

    def restore_checkpoint(self, conf_search:Config, macro_combinations)\
            ->Tuple[int, Optional[SearchResult]]:

        conf_pareto = conf_search['pareto']
        pareto_summary_filename = conf_pareto['summary_filename']

        summary_filepath = utils.full_path(pareto_summary_filename)

        # if checkpoint is available then restart from last combination we were running
        checkpoint_avail = self._checkpoint is not None
        resumed, state = False, None
        start_macro_i, best_result = 0, None
        if checkpoint_avail:
            state = self._checkpoint.get('search', None)
            if state is not None:
                start_macro_i = state['start_macro_i']
                assert start_macro_i >= 0 and start_macro_i < len(macro_combinations)

                best_result = yaml.load(state['best_result'], Loader=yaml.Loader)

                start_macro_i += 1 # resume after the last checkpoint
                resumed = True

        if not resumed:
            # erase previous file left over from run
            utils.zero_file(summary_filepath)

        logger.warn({'resumed': resumed, 'checkpoint_avail': checkpoint_avail,
                     'checkpoint_val': state is not None,
                     'start_macro_i': start_macro_i,
                     'total_macro_combinations': len(macro_combinations)})
        return start_macro_i, best_result

    def record_checkpoint(self, macro_comb_i:int, best_result:SearchResult)->None:
        if self._checkpoint is not None:
            state = {'start_macro_i': macro_comb_i,
                     'best_result': yaml.dump(best_result)}
            self._checkpoint.new()
            self._checkpoint['search'] = state
            self._checkpoint.commit()

    def get_combinations(self, conf_search:Config)->Iterator[Tuple[int, int, int]]:
        conf_pareto = conf_search['pareto']
        conf_model_desc = conf_search['model_desc']

        min_cells = conf_model_desc['n_cells']
        min_reductions = conf_model_desc['n_reductions']
        min_nodes = conf_model_desc['cell']['n_nodes']
        max_cells = conf_pareto['max_cells']
        max_reductions = conf_pareto['max_reductions']
        max_nodes = conf_pareto['max_nodes']

        logger.info({'min_reductions': min_reductions,
                     'min_cells': min_cells,
                     'min_nodes': min_nodes,
                     'max_reductions': max_reductions,
                     'max_cells': max_cells,
                     'max_nodes': max_nodes
                     })

        # TODO: what happens when reductions is 3 but cells is 2? have to step
        # through code and check
        for reductions in range(min_reductions, max_reductions+1):
            for cells in range(min_cells, max_cells+1):
                for nodes in range(min_nodes, max_nodes+1):
                    yield reductions, cells, nodes

    def save_trained(self, conf_search:Config, reductions:int, cells:int, nodes:int,
                      model_metrics:ModelMetrics)->None:
        """Save the model and metric info into a log file"""

        metrics_dir = conf_search['metrics_dir']

        # construct path where we will save
        subdir = utils.full_path(metrics_dir.format(**vars()), create=True)

        model_stats = nas_utils.get_model_stats(model_metrics.model)

        # save model_stats in its own file
        model_stats_filepath = os.path.join(subdir, 'model_stats.yaml')
        if model_stats_filepath:
            with open(model_stats_filepath, 'w') as f:
                yaml.dump(model_stats, f)

        # save just metrics separately for convinience
        metrics_filepath = os.path.join(subdir, 'metrics.yaml')
        if metrics_filepath:
            with open(metrics_filepath, 'w') as f:
                yaml.dump(model_stats.metrics, f)

        logger.info({'model_stats_filepath': model_stats_filepath,
                     'metrics_filepath': metrics_filepath})

        # append key info in root pareto data
        if self._summary_filepath:
            train_top1 = val_top1 = train_epoch = val_epoch = math.nan
            # extract metrics
            if model_metrics.metrics:
                best_metrics = model_metrics.metrics.run_metrics.best_epoch()
                train_top1 = best_metrics[0].top1.avg
                train_epoch = best_metrics[0].index
                if best_metrics[1]:
                    val_top1 = best_metrics[1].top1.avg if len(best_metrics)>1 else math.nan
                    val_epoch = best_metrics[1].index if len(best_metrics)>1 else math.nan

            # extract model stats
            flops = model_stats.Flops
            parameters = model_stats.parameters
            inference_memory = model_stats.inference_memory
            inference_duration = model_stats.duration

            utils.append_csv_file(self._summary_filepath, [
                ('reductions', reductions),
                ('cells', cells),
                ('nodes', nodes),
                ('train_top1', train_top1),
                ('train_epoch', train_epoch),
                ('val_top1', val_top1),
                ('val_epoch', val_epoch),
                ('flops', flops),
                ('params', parameters),
                ('inference_memory', inference_memory),
                ('inference_duration', inference_duration)
                ])



