import warnings
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Type, Union, Dict, List, Tuple

import numpy as np

from .model import Model
from .utils import get_config_params


class Agent(ABC):
    def __init__(self, model_cls: Type[Model], observation_space: Any, action_space: Any, config: dict = None,
                 *args, **kwargs) -> None:
        """
        This method MUST be called between (0.) and (4.) in subclasses for doing initialization works

        0. [IN '__init__' of SUBCLASSES] Define parameters, model configurations and other related variables
        1. If 'config' is not 'None', set specified configuration parameters (which appear after 'config') for agent or
        specified configurations for model
        2. Initialize model instances
        3. Build training part of computational graph
        4. [IN '__init__' of SUBCLASSES] Do other operations if necessary

        :param model_cls: Model class that agent adopts
        :param observation_space: Env observation space
        :param action_space: Env action space
        :param config: Configurations for agent and models
        :param args: Positional configurations for agent only (ignored if specified in 'config')
        :param kwargs: Keyword configurations for agent only (ignored if specified in 'config')
        """
        self.model_cls = model_cls
        self.observation_space = observation_space
        self.action_space = action_space

        # 1. Set configurations
        if config is not None:
            self.load_config(config)

        # 2. Initialize instances of 'model_cls'
        self.model_instances = None
        self._init_model_instances(config)

        # 3. Build training part of computational graph
        self.build()

    @abstractmethod
    def build(self) -> None:
        """Build computational graph for training"""
        pass

    @abstractmethod
    def preprocess(self, state: Any, *args, **kwargs) -> Any:
        """Preprocess the game state"""
        pass

    @abstractmethod
    def set_weights(self, *args, **kwargs) -> None:
        pass

    @abstractmethod
    def get_weights(self, *args, **kwargs) -> Any:
        pass

    @abstractmethod
    def save(self, path: Path, *args, **kwargs) -> None:
        """Save the checkpoint file"""
        pass

    @abstractmethod
    def load(self, path: Path, *args, **kwargs) -> None:
        """Load the checkpoint file"""
        pass

    @abstractmethod
    def learn(self, training_data: Dict[str, np.ndarray], *args, **kwargs) -> Union[Dict[str, float], None]:
        """
        Train the agent with data generated by 'prepare_update'
        :param training_data: A dictionary of lists of training_data, such as:
                {'state': [[1, 2], [3, 4]], 'action': [1, 0], 'value': [0.1, 0.3]}
        :param args: Optional positional arguments
        :param kwargs: Optional keyword arguments
        :return: Training statistics
        """
        pass

    @abstractmethod
    def prepare_training_data(self, trajectory: List[Tuple[Any, Any, Any, Any, Any, dict]]) -> Dict[str, np.ndarray]:
        """
        Prepare the training data from a trajectory
        :param trajectory: A list of transitions that consists of (state, action, reward, done, next_state, extra_data)
        :return: Prepared training data
        """
        pass

    @abstractmethod
    def post_process_training_data(self, training_data: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        """
        Process all training data concatenated from trajectories before sending them to learner(s)
        :param training_data: Concatenated training data
        :return: Post-processed data
        """
        pass

    @abstractmethod
    def update_sampling(self, current_step: int, total_steps: int, *args, **kwargs) -> None:
        """Do some updates according to the process of sampling"""
        pass

    @abstractmethod
    def update_training(self, current_step: int, total_steps: int, *args, **kwargs) -> None:
        """Do some updates according to the process of training"""
        pass

    def export_config(self) -> dict:
        """Export dictionary as configurations"""
        param_dict = {p: getattr(self, p) for p in get_config_params(self)}

        if len(self.model_instances) == 1:
            model_config = self.model_instances[0].export_config()
        else:
            model_config = [x.export_config() for x in self.model_instances]
        param_dict.update({'model': model_config})

        return param_dict

    def load_config(self, config: dict) -> None:
        """Load dictionary as configurations and initialize model instances"""
        for key, val in config.items():
            if key in get_config_params(self):
                self.__dict__[key] = val
            elif key != 'model':
                warnings.warn(f"Invalid config item '{key}' ignored", RuntimeWarning)

    def predict(self, state: Any, *args, **kwargs) -> Any:
        """Get the action distribution at specific state"""
        return self.model_instances[0].forward(state, *args, **kwargs)

    def policy(self, state: Any, *args, **kwargs) -> Any:
        """Choose action during exploitation"""
        return np.argmax(self.predict(state, *args, **kwargs)[0])

    def sample(self, state: Any, *args, **kwargs) -> Tuple[Any, Dict]:
        """Return action and other information (value, distribution et al) during exploration/sampling"""
        p = self.predict(state, *args, **kwargs)[0]
        return np.random.choice(len(p), p=p), {}

    def _init_model_instances(self, config: Union[dict, None]) -> None:
        """Initialize model instances"""
        self.model_instances = []

        def create_model_instance(_c: dict):
            ret = {}
            for k, v in _c.items():
                if k in valid_config:
                    ret[k] = v
                else:
                    warnings.warn(f"Invalid config item '{k}' ignored", RuntimeWarning)
            self.model_instances.append(self.model_cls(self.observation_space, self.action_space, **ret))

        if config is not None and 'model' in config:
            model_config = config['model']
            valid_config = get_config_params(self.model_cls)

            if isinstance(model_config, list):
                for _, c in enumerate(model_config):
                    create_model_instance(c)
            elif isinstance(model_config, dict):
                create_model_instance(model_config)
        else:
            self.model_instances.append(self.model_cls(self.observation_space, self.action_space))
