import copy
import sys
from typing import Any, Literal

import random
import math
import sb3_contrib

import gymnasium as gym
import numpy as np

from graph_jsp_env.disjunctive_graph_jsp_env import DisjunctiveGraphJspEnv
from jsp_instance_utils.instances import ft06, ft06_makespan
from sb3_contrib.common.maskable.distributions import MaskableCategoricalDistribution
from sb3_contrib.common.maskable.policies import MaskableActorCriticPolicy
from sb3_contrib.common.wrappers import ActionMasker

from gymcts.gymcts_agent import GymctsAgent
from gymcts.gymcts_env_abc import GymctsABC
from gymcts.gymcts_node import GymctsNode

from gymcts.logger import log


class GraphJspNeuralGYMCTSWrapper(GymctsABC, gym.Wrapper):

    def __init__(self, env: DisjunctiveGraphJspEnv):
        gym.Wrapper.__init__(self, env)

    def load_state(self, state: Any) -> None:
        self.env.reset()
        for action in state:
            self.env.step(action)

    def is_terminal(self) -> bool:
        return self.env.unwrapped.is_terminal()

    def get_valid_actions(self) -> list[int]:
        return list(self.env.unwrapped.valid_actions())

    def rollout(self) -> float:
        terminal = env.is_terminal()

        if terminal:
            lower_bound = env.unwrapped.reward_function_parameters['scaling_divisor']
            return - env.unwrapped.get_makespan() / lower_bound + 2

        reward = 0
        while not terminal:
            action = random.choice(self.get_valid_actions())
            obs, reward, terminal, truncated, _ = env.step(action)

        return reward + 2

    def get_state(self) -> Any:
        return env.unwrapped.get_action_history()

    def action_masks(self) -> np.ndarray | None:
        """Return the action mask for the current state."""
        return self.env.unwrapped.valid_action_mask()


class GymctsNeuralNode(GymctsNode):
    PUCT_v3_mu = 0.95

    MuZero_c1 = 1.25
    MuZero_c2 = 19652.0

    """
    PUCT (Predictor + UCT) exploration terms:

    PUCT_v0:
        c * P(s, a) * √( N(s) / (1 + N(s,a)) )

    PUCT_v1:
        c * P(s, a) * √( 2 * ln(N(s)) / N(s,a) )

    PUCT_v2:
        c * P(s, a) * √( N(s) ) / N(s,a)

    PUCT_v3:
        c * P(s, a)^μ * √( N(s) / (1 + N(s,a)) )

    PUCT_v4:
        c * ( P(s, a) / (1 + N(s,a)) )

    PUCT_v5:
        c * P(s, a) * ( √(N(s)) + 1 ) / (N(s,a) + 1)

    PUCT_v6:
        c * P(s, a) * N(s) / (1 + N(s,a))

    PUCT_v7:
        c * P(s, a) * ( √(N(s)) + ε ) / (N(s,a) + 1)

    PUCT_v8:
        c * P(s, a) * √( (ln(N(s)) + 1) / (1 + N(s,a)) )

    PUCT_v9:
        c * P(s, a) * √( N(s) / (1 + N(s,a)) )

    PUCT_v10:
        c * P(s, a) * √( ln(N(s)) / (1 + N(s,a)) )


    MuZero exploration terms:

    MuZero_v0:
        P(s, a) * √( N(s) / (1 + N(s,a)) ) * [ c₁ + ln( (N(s) + c₂ + 1) / c₂ ) ]

    MuZero_v1:
        P(s, a) * √( N(s) / (1 + N(s,a)) ) * [ c₁ + ln( (N(s) + c₂ + 1) / c₂ ) ]


    Where:
        - N(s):      number of times state s has been visited
        - N(s,a):    number of times action a was taken from state s
        - P(s,a):    prior probability of selecting action a from state s
        - c, c₁, c₂: exploration constants
        - μ:         exponent applied to P(s,a) in some variants
        - ε:         small constant to avoid division by zero (in PUCT 7)
    """
    score_variate: Literal[
        "PUCT_v0",
        "PUCT_v1",
        "PUCT_v2",
        "PUCT_v3",
        "PUCT_v4",
        "PUCT_v5",
        "PUCT_v6",
        "PUCT_v7",
        "PUCT_v8",
        "PUCT_v9",
        "PUCT_v10",
        "MuZero_v0",
        "MuZero_v1",
    ] = "PUCT_v0"

    def __init__(
            self,
            action: int,
            parent: 'GymctsNeuralNode',
            env_reference: GymctsABC,
            prior_selection_score: float,
            observation: np.ndarray | None = None,
    ):
        super().__init__(action, parent, env_reference)

        self._obs = observation
        self._selection_score_prior = prior_selection_score

    def tree_policy_score(self) -> float:
        # call the superclass (GymctsNode) for ucb_score
        c = GymctsNode.ubc_c
        # the way alpha zero does it
        # exploration_term = self._selection_score_prior * c * math.sqrt(math.log(self.parent.visit_count)) / (1 + self.visit_count)
        # the way the vanilla gymcts does it
        p_sa = self._selection_score_prior
        n_s = self.parent.visit_count
        n_sa = self.visit_count

        assert 0 <= GymctsNode.best_action_weight <= 1
        b = GymctsNode.best_action_weight
        exploitation_term = 0.0 if self.visit_count == 0 else (1 - b) * self.mean_value + b * self.max_value

        if GymctsNeuralNode.score_variate == "PUCT_v0":
            return exploitation_term + c * p_sa * math.sqrt(n_s) / (1 + n_sa)
        elif GymctsNeuralNode.score_variate == "PUCT_v1":
            return exploitation_term + c * p_sa * math.sqrt(2 * math.log(n_s) / (n_sa))
        elif GymctsNeuralNode.score_variate == "PUCT_v2":
            if n_sa == 0:
                return float("inf")  # Avoid division by zero
            return exploitation_term + c * p_sa * math.sqrt(n_s) / n_sa
        elif GymctsNeuralNode.score_variate == "PUCT_v3":
            return exploitation_term + c * (p_sa ** GymctsNeuralNode.PUCT_v3_mu) * math.sqrt(n_s / (1 + n_sa))
        elif GymctsNeuralNode.score_variate == "PUCT_v4":
            return exploitation_term + c * (p_sa / (1 + n_sa))
        elif GymctsNeuralNode.score_variate == "PUCT_v5":
            return exploitation_term + c * p_sa * (math.sqrt(n_s) + 1) / (n_sa + 1)
        elif GymctsNeuralNode.score_variate == "PUCT_v6":
            return exploitation_term + c * p_sa * n_s / (1 + n_sa)
        elif GymctsNeuralNode.score_variate == "PUCT_v7":
            epsilon = 1e-8
            return exploitation_term + c * p_sa * (math.sqrt(n_s) + epsilon) / (n_sa + 1)
        elif GymctsNeuralNode.score_variate == "PUCT_v8":
            return exploitation_term + c * p_sa * math.sqrt((math.log(n_s) + 1) / (1 + n_sa))
        elif GymctsNeuralNode.score_variate == "PUCT_v9":
            return exploitation_term + c * p_sa * math.sqrt(n_s / (1 + n_sa))
        elif GymctsNeuralNode.score_variate == "PUCT_v10":
            return exploitation_term + c * p_sa * math.sqrt(math.log(n_s) / (1 + n_sa))
        elif GymctsNeuralNode.score_variate == "MuZero_v0":
            c1 = GymctsNeuralNode.MuZero_c1
            c2 = GymctsNeuralNode.MuZero_c2
            return exploitation_term + c * p_sa * math.sqrt(n_s) / (1 + n_sa) * (c1 + math.log((n_s + c2 + 1) / c2))
        elif GymctsNeuralNode.score_variate == "MuZero_v1":
            c1 = GymctsNeuralNode.MuZero_c1
            c2 = GymctsNeuralNode.MuZero_c2
            return exploitation_term + c * p_sa * math.sqrt(n_s) / (1 + n_sa) * (c1 + math.log((n_s + c2 + 1) / c2))

        exploration_term = self._selection_score_prior * c * math.sqrt(
            math.log(self.parent.visit_count) / (self.visit_count)) if self.visit_count > 0 else float("inf")
        return self.mean_value + exploration_term

    def get_best_action(self) -> int:
        """
        Returns the best action of the node. The best action is the action with the highest score.
        The best action is the action that has the highest score.

        :return: the best action of the node.
        """
        return max(self.children.values(), key=lambda child: child.max_value).action

    def __str__(self, colored=False, action_space_n=None) -> str:
        """
        Returns a string representation of the node. The string representation is used for visualisation purposes.
        It is used for example in the mcts tree visualisation functionality.

        :param colored: true if the string representation should be colored, false otherwise. (ture is used by the mcts tree visualisation)
        :param action_space_n: the number of actions in the action space. This is used for coloring the action in the string representation.
        :return: a potentially colored string representation of the node.
        """
        if not colored:

            if not self.is_root():
                return f"(a={self.action}, N={self.visit_count}, Q_v={self.mean_value:.2f}, best={self.max_value:.2f}, {GymctsNeuralNode.score_variate}={self.tree_policy_score():.2f})"
            else:
                return f"(N={self.visit_count}, Q_v={self.mean_value:.2f}, best={self.max_value:.2f}) [root]"

        import gymcts.colorful_console_utils as ccu

        if self.is_root():
            return f"({ccu.CYELLOW}N{ccu.CEND}={self.visit_count}, {ccu.CYELLOW}Q_v{ccu.CEND}={self.mean_value:.2f}, {ccu.CYELLOW}best{ccu.CEND}={self.max_value:.2f})"

        if action_space_n is None:
            raise ValueError("action_space_n must be provided if colored is True")

        p = ccu.CYELLOW
        e = ccu.CEND
        v = ccu.CCYAN

        def colorful_value(value: float | int | None) -> str:
            if value == None:
                return f"{ccu.CGREY}None{e}"
            color = ccu.CCYAN
            if value == 0:
                color = ccu.CRED
            if value == float("inf"):
                color = ccu.CGREY
            if value == -float("inf"):
                color = ccu.CGREY

            if isinstance(value, float):
                return f"{color}{value:.2f}{e}"

            if isinstance(value, int):
                return f"{color}{value}{e}"

        root_node = self.get_root()
        mean_val = f"{self.mean_value:.2f}"

        return ((f"("
                 f"{p}a{e}={ccu.wrap_evenly_spaced_color(s=self.action, n_of_item=self.action, n_classes=action_space_n)}, "
                 f"{p}N{e}={colorful_value(self.visit_count)}, "
                 f"{p}Q_v{e}={ccu.wrap_with_color_scale(s=mean_val, value=self.mean_value, min_val=root_node.min_value, max_val=root_node.max_value)}, "
                 f"{p}best{e}={colorful_value(self.max_value)}") +
                (
                    f", {p}{GymctsNeuralNode.score_variate}{e}={colorful_value(self.tree_policy_score())})" if not self.is_root() else ")"))


class GymctsNeuralAgent(GymctsAgent):

    def __init__(self,
                 env: GymctsABC,
                 *args,
                 model_kwargs=None,
                 score_variate: Literal[
                     "PUCT_v0",
                     "PUCT_v1",
                     "PUCT_v2",
                     "PUCT_v3",
                     "PUCT_v4",
                     "PUCT_v5",
                     "PUCT_v6",
                     "PUCT_v7",
                     "PUCT_v8",
                     "PUCT_v9",
                     "PUCT_v10",
                     "MuZero_v0",
                     "MuZero_v1",
                 ] = "PUCT_v0",
                 **kwargs
                 ):

        # init super class
        super().__init__(
            env=env,
            *args,
            **kwargs
        )
        if score_variate not in [
            "PUCT_v0",
            "PUCT_v1",
            "PUCT_v2",
            "PUCT_v3",
            "PUCT_v4",
            "PUCT_v5",
            "PUCT_v6",
            "PUCT_v7",
            "PUCT_v8",
            "PUCT_v9",
            "PUCT_v10",
            "MuZero_v0",
            "MuZero_v1",
        ]:
            raise ValueError(f"Invalid score_variate: {score_variate}. Must be one of: "
                             f"['PUCT_v0', 'PUCT_v1', 'PUCT_v2', 'PUCT_v3', 'PUCT_v4', 'PUCT_v5', "
                             f"'PUCT_v6', 'PUCT_v7', 'PUCT_v8', 'PUCT_v9', 'PUCT_v10', 'MuZero_v0', 'MuZero_v1']")
        GymctsNeuralNode.score_variate = score_variate

        if model_kwargs is None:
            model_kwargs = {}
        obs, info = env.reset()

        self.search_root_node = GymctsNeuralNode(
            action=None,
            parent=None,
            env_reference=env,
            observation=obs,
            prior_selection_score=1.0,
        )

        def mask_fn(env: gym.Env) -> np.ndarray:
            mask = env.action_masks()
            if mask is None:
                mask = np.ones(env.action_space.n, dtype=np.float32)
            return mask

        env = ActionMasker(env, action_mask_fn=mask_fn)

        model_kwargs = {
                           "policy": MaskableActorCriticPolicy,
                           "env": env,
                           "verbose": 1,
                       } | model_kwargs

        self._model = sb3_contrib.MaskablePPO(**model_kwargs)

    def learn(self, total_timesteps: int, **kwargs) -> None:
        """Learn from the environment using the MaskablePPO model."""
        self._model.learn(total_timesteps=total_timesteps, **kwargs)

    def expand_node(self, node: GymctsNeuralNode) -> None:
        log.debug(f"expanding node: {node}")
        # EXPANSION STRATEGY
        # expand all children

        child_dict = {}

        self._load_state(node)

        obs_tensor, vectorized_env = self._model.policy.obs_to_tensor(np.array([node._obs]))
        action_masks = np.array([self.env.action_masks()])
        distribution = self._model.policy.get_distribution(obs=obs_tensor, action_masks=action_masks)
        unwrapped_distribution = distribution.distribution.probs[0]

        # print(f'valid actions: {node.valid_actions}')
        # print(f'env mask: {self.env.action_masks()}')
        # print(f'env valid actions: {self.env.get_valid_actions()}')
        """
                for action in node.valid_actions:
            # reconstruct state
            # load state of leaf node
            self._load_state(node)

            obs, reward, terminal, truncated, _ = self.env.step(action)
            child_dict[action] = GymctsNeuralNode(
                action=action,
                parent=node,
                env_reference=self.env,
                observation=obs,
                prior_selection_score=1.0,
            )
        node.children = child_dict
        return
        """

        for action, prob in enumerate(unwrapped_distribution):
            self._load_state(node)

            log.debug(f"Probabily for action {action}: {prob}")

            if prob == 0.0:
                continue

            assert action in node.valid_actions, f"Action {action} is not in valid actions: {node.valid_actions}"

            obs, reward, terminal, truncated, _ = self.env.step(action)
            child_dict[action] = GymctsNeuralNode(
                action=action,
                parent=node,
                observation=copy.deepcopy(obs),
                env_reference=self.env,
                prior_selection_score=float(prob)
            )

        node.children = child_dict
        # print(f"Expanded node {node} with {len(node.children)} children.")


if __name__ == '__main__':
    log.setLevel(20)

    env_kwargs = {
        "jps_instance": ft06,
        "default_visualisations": ["gantt_console", "graph_console"],
        "reward_function_parameters": {
            "scaling_divisor": ft06_makespan
        },
        "reward_function": "nasuta",
    }

    env = DisjunctiveGraphJspEnv(**env_kwargs)
    env.reset()

    env = GraphJspNeuralGYMCTSWrapper(env)

    import torch

    model_kwargs = {
        "gamma": 0.99013,
        "gae_lambda": 0.9,
        "normalize_advantage": True,
        "n_epochs": 28,
        "n_steps": 432,
        "max_grad_norm": 0.5,
        "learning_rate": 6e-4,
        "policy_kwargs": {
            "net_arch": {
                "pi": [90, 90],
                "vf": [90, 90],
            },
            "ortho_init": True,
            "activation_fn": torch.nn.ELU,
            "optimizer_kwargs": {
                "eps": 1e-7
            }
        }
    }

    agent = GymctsNeuralAgent(
        env=env,
        render_tree_after_step=True,
        render_tree_max_depth=3,
        exclude_unvisited_nodes_from_render=False,
        number_of_simulations_per_step=15,
        # clear_mcts_tree_after_step = False,
        model_kwargs=model_kwargs
    )

    agent.learn(total_timesteps=10_000)

    agent.solve()

    actions = agent.solve(render_tree_after_step=True)
    for a in actions:
        obs, rew, term, trun, info = env.step(a)

    env.render()
    makespan = env.unwrapped.get_makespan()
    print(f"makespan: {makespan}")
