from argparse import Namespace
from pathlib import Path

import torch
from torch import nn
from torch.nn import ModuleList

from ml_lib.modules.util import LightningBaseModule
from ml_lib.utils.config import Config
from ml_lib.utils.model_io import SavedLightningModels
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, DatasetMixin,
                                BaseDataloadersMixin)


class Ensemble(DatasetMixin,
               BaseDataloadersMixin,
               BaseTrainMixin,
               BaseValMixin,
               BaseOptimizerMixin,
               LightningBaseModule
               ):

    def __init__(self, hparams):
        super(Ensemble, self).__init__(hparams)

        # Dataset
        # =============================================================================
        self.dataset = self.build_dataset()

        # Model Paramters
        # =============================================================================
        # Additional parameters
        self.in_shape = self.dataset.train_dataset.sample_shape
        self.conv_filters = self.params.filters
        self.criterion = nn.BCELoss()

        # Pre_trained_models
        out_path = Path('output') / self.params.secondary_type
        # exp_paths = list(out_path.rglob(f'*{self.params.exp_fingerprint}'))
        exp_paths = list(out_path.rglob('*e87b8f455ba134504b1ae17114ac2a2a'))
        config_ini_files = sum([list(exp_path.rglob('config.ini')) for exp_path in exp_paths], [])

        self.model_list = ModuleList()

        configs = [Config() for _ in range(len(config_ini_files))]
        for config, ini_file in zip(configs, config_ini_files):
            config.read_file(ini_file.open('r'))
            model = SavedLightningModels.load_checkpoint(models_root_path=config.exp_path / config.version).restore()
            self.model_list.append(model)

    def forward(self, batch, **kwargs):
        ys = [model(batch).main_out for model in self.model_list]
        tensor = torch.stack(ys).mean(dim=0)

        return Namespace(main_out=tensor)