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

# This file might cause import error for those who didn't install RL-related dependencies

import logging
import threading
from multiprocessing.pool import ThreadPool
from typing import Tuple

import gym
import numpy as np
import tianshou
import torch
import torch.nn as nn
import torch.nn.functional as F

from gym import spaces
from tianshou.data import to_torch
from tianshou.env.worker import EnvWorker

from typing_extensions import TypedDict

from nni.nas.execution import submit_models, wait_models
from nni.nas.execution.common import ModelStatus
from .utils import get_targeted_model


_logger = logging.getLogger(__name__)
_thread_lock = threading.Lock()


class MultiThreadEnvWorker(EnvWorker):
    def __init__(self, env_fn):
        self.env = env_fn()
        self.pool = ThreadPool(processes=1)
        super().__init__(env_fn)

    def get_env_attr(self, key):
        return getattr(self.env, key)

    def set_env_attr(self, key, value):
        return setattr(self.env, key, value)

    def __getattr__(self, key):
        if tianshou.__version__ >= '0.4.5':  # not a strict check here
            return super().__getattr__(key)  # https://github.com/thu-ml/tianshou/pull/478
        return getattr(self.env, key)

    def reset(self):
        return self.env.reset()

    def send(self, action):
        # for tianshou >= 0.4.6
        if action is None:
            self.result = self.pool.apply_async(self.env.reset)
        else:
            self.send_action(action)

    def recv(self):
        # for tianshou >= 0.4.6
        return self.result.get()

    @staticmethod
    def wait(*args, **kwargs):
        raise NotImplementedError('Async collect is not supported yet.')

    def send_action(self, action) -> None:
        # self.result is actually a handle
        self.result = self.pool.apply_async(self.env.step, (action,))

    def get_result(self):
        return self.result.get()

    def seed(self, seed):
        super().seed(seed)
        return self.env.seed(seed)

    def render(self, **kwargs):
        return self.env.render(**kwargs)

    def close_env(self) -> None:
        self.pool.terminate()
        return self.env.close()

class ObservationType(TypedDict):
    action_history: np.ndarray
    cur_step: int
    action_dim: int


class ModelEvaluationEnv(gym.Env[ObservationType, int]):
    def __init__(self, base_model, mutators, search_space):
        self.base_model = base_model
        self.mutators = mutators
        self.search_space = search_space
        self.ss_keys = list(self.search_space.keys())
        self.action_dim = max(map(lambda v: len(v), self.search_space.values()))
        self.num_steps = len(self.search_space)

    @property
    def observation_space(self):
        return spaces.Dict({
            'action_history': spaces.MultiDiscrete([self.action_dim] * self.num_steps),
            'cur_step': spaces.Discrete(self.num_steps + 1),
            'action_dim': spaces.Discrete(self.action_dim + 1)
        })

    @property
    def action_space(self):
        return spaces.Discrete(self.action_dim)

    def reset(self) -> ObservationType:
        self.action_history = np.zeros(self.num_steps, dtype=np.int32)
        self.cur_step = 0
        self.sample = {}
        return {
            'action_history': self.action_history,
            'cur_step': self.cur_step,
            'action_dim': len(self.search_space[self.ss_keys[self.cur_step]])
        }

    def step(self, action: int) -> Tuple[ObservationType, float, bool, dict]:
        cur_key = self.ss_keys[self.cur_step]
        assert action < len(self.search_space[cur_key]), \
            f'Current action {action} out of range {self.search_space[cur_key]}.'
        self.action_history[self.cur_step] = action
        self.sample[cur_key] = self.search_space[cur_key][action]
        self.cur_step += 1
        obs: ObservationType = {
            'action_history': self.action_history,
            'cur_step': self.cur_step,
            'action_dim': len(self.search_space[self.ss_keys[self.cur_step]]) \
                if self.cur_step < self.num_steps else self.action_dim
        }
        if self.cur_step == self.num_steps:
            with _thread_lock:
                model = get_targeted_model(self.base_model, self.mutators, self.sample)
            _logger.info(f'New model created: {self.sample}')
            submit_models(model)
            wait_models(model)
            if model.status == ModelStatus.Failed:
                return self.reset(), 0., False, {}
            rew = float(model.metric)  # type: ignore
            _logger.info(f'Model metric received as reward: {rew}')
            return obs, rew, True, {}
        else:

            return obs, 0., False, {}


class Preprocessor(nn.Module):
    def __init__(self, obs_space, hidden_dim=64, num_layers=1):
        super().__init__()
        self.action_dim = obs_space['action_history'].nvec[0]
        self.hidden_dim = hidden_dim
        # first token is [SOS]
        self.embedding = nn.Embedding(self.action_dim + 1, hidden_dim)
        self.rnn = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)

    def forward(self, obs):
        seq = F.pad(obs['action_history'] + 1, (1, 1))  # pad the start token and end token
        # end token is used to avoid out-of-range of v_s_. Will not actually affect BP.
        seq = self.embedding(seq.long())
        feature, _ = self.rnn(seq)
        return feature[torch.arange(len(feature), device=feature.device), obs['cur_step'].long() + 1]


class Actor(nn.Module):
    def __init__(self, action_space, preprocess):
        super().__init__()
        self.preprocess = preprocess
        self.action_dim = action_space.n
        self.linear = nn.Linear(self.preprocess.hidden_dim, self.action_dim)

    def forward(self, obs, **kwargs):
        obs = to_torch(obs, device=self.linear.weight.device)
        out = self.linear(self.preprocess(obs))
        # to take care of choices with different number of options
        mask = torch.arange(self.action_dim).expand(len(out), self.action_dim) >= obs['action_dim'].unsqueeze(1)
        out[mask.to(out.device)] = float('-inf')
        return F.softmax(out, dim=-1), kwargs.get('state', None)


class Critic(nn.Module):
    def __init__(self, preprocess):
        super().__init__()
        self.preprocess = preprocess
        self.linear = nn.Linear(self.preprocess.hidden_dim, 1)

    def forward(self, obs, **kwargs):
        obs = to_torch(obs, device=self.linear.weight.device)
        return self.linear(self.preprocess(obs)).squeeze(-1)
