from typing import Union

import numpy as np
from sortedcontainers import SortedDict

from quadrotor_environment.quadrotor_model import QuadrotorModel, SysState


class ModelDelayWrapper:
    """A wrapper around QuadrotorModel that delays observations, actions and introduces jitter to the controller
    execution.
    """
    __slots__ = "delay_scale", "model", "controller_jitter", "observation_delay", "observation_jitter", \
                "action_delay", "action_jitter", "controller_period", "past_states", "future_actions", "time",

    def __init__(self, model: QuadrotorModel,
                 delay_scale: float = 1,
                 controller_jitter: Union[float, str] = 0,
                 controller_period: float = 0.01,
                 observation_delay: float = 0,
                 observation_jitter: Union[float, str] = 0,
                 action_delay: float = 0,
                 action_jitter: Union[float, str] = 0):
        """
        Creates a new ModelDelayWrapper wrapping the specified model.

        All times specified in this class are in seconds.

        Note:
            All jitters can either be specified as a float, in which case they are interpreted as a standard
            deviation of a mean centered gauss distribution from which to draw them, or as a path to a text file that
            can be loaded using np.loadtxt. The jitters are then generated by drawing random elements from the resulting
            numpy array. This feature is useful if you measured the jitter in your real hardware system and want to
            reproduce the jitter distribution in simulation as closely as possible.

        :param model: The model to wrap.
        :param delay_scale: Global scaling factor by which all other delays are multiplied. Default is 1.
        :param controller_jitter: The amount of jitter to apply to the control period.
        :param controller_period: The duration between calls to the controller. The simulation will advance this much on
        each call to the step method.
        :param observation_delay: The age of the state that is to be observed.
        :param observation_jitter: The amount of jitter in the age of the state that is observed.
        :param action_delay: The amount of time that has to pass before an action is applied to the system.
        :param action_jitter: The jitter to be added to the action delay.
        """

        if isinstance(observation_jitter, str):
            observation_jitter = np.loadtxt(observation_jitter)
        if isinstance(action_jitter, str):
            action_jitter = np.loadtxt(action_jitter)
        if isinstance(controller_jitter, str):
            controller_jitter = np.loadtxt(controller_jitter)

        observation_jitter = np.asarray(observation_jitter)
        action_jitter = np.asarray(action_jitter)
        controller_jitter = np.asarray(controller_jitter)

        assert observation_delay >= 0
        assert action_delay >= 0
        assert np.all(observation_jitter >= 0)
        assert np.all(action_jitter >= 0)
        assert delay_scale >= 0

        self.delay_scale = delay_scale
        self.model = model
        self.controller_jitter = controller_jitter
        self.observation_delay = observation_delay
        self.observation_jitter = observation_jitter
        self.action_delay = action_delay
        self.action_jitter = action_jitter
        self.controller_period = controller_period

        self.past_states = SortedDict()
        self.future_actions = SortedDict()
        self.time = None

    def reset(self, initial_action: np.ndarray, initial_state: SysState) -> (SysState, float, float):
        """
        Resets the state of the model delay wrapper by clearing all past states, future actions and initializing with
        the specified initial action and state.
        
        :param initial_action: The initial action to assume.
        :param initial_state: The initial state to assume
        :return: A triple consisting of
            1. the initial observed state (which is jus the initial state that as passed),
            2. the initial controller period and
            3. the initial observation age.
        """
        self.past_states = SortedDict()
        self.future_actions = SortedDict()

        observation_delay = self._sample_observation_delay()
        self.time = observation_delay
        self.past_states[0] = (initial_action, initial_state)
        return initial_state, self._sample_controller_period(), observation_delay

    def compute_past_state(self, time: float) -> SysState:
        """
        Computes the state at some arbitrary point in the past.

        :param time: The time at which to compute the state.
        :return: The state at the time.
        """
        # Find last computed state just before the observed state in the state history
        # (sbt = state before t = most recent state before t)
        sbt_index = self.past_states.bisect(time) - 1
        assert sbt_index >= 0
        t_sbt, (sbt_action, sbt) = self.past_states.peekitem(sbt_index)
        assert time >= t_sbt

        return self.model.next_state(sbt, sbt_action, time - t_sbt)

    def compute_current_state(self) -> SysState:
        """
        Computes the current state of the system.

        :return: The current state of the system.
        """
        return self.compute_past_state(self.time)

    def step(self, action: np.ndarray) -> (SysState, float, float):
        """
        Applies an action to the delayed system and advances the time by the controller period.

        :param action: The action to execute (probably in the future because it is delayed).
        :return: A triple consisting of
            1. the observed state (probably from the past),
            2. the actual controller period (with jitter)
            3. The age of the observed state.
        """
        # Insert future action to the action schedule
        self.future_actions[self.time + self._sample_action_delay()] = action

        current_controller_period = self._sample_controller_period()
        self.time += current_controller_period
        self._materialize_past_states()

        # Sample time of observation
        t_obs = np.clip(self.time - self._sample_observation_delay(), 0, self.time)

        # Compute observed state
        observed_state = self.compute_past_state(t_obs)

        return observed_state, current_controller_period, self.time - t_obs

    def _materialize_past_states(self):
        # Compute states that now lie in the past
        while len(self.future_actions) > 0 and self.future_actions.peekitem(0)[0] <= self.time:
            t_a, new_action = self.future_actions.popitem(0)
            t_sba, (old_action, sba) = self.past_states.peekitem()
            assert t_a - t_sba >= 0
            new_state = self.model.next_state(sba, old_action, t_a - t_sba)
            self.past_states[t_a] = (new_action, new_state)
    
    def _sample_action_delay(self):
        if self.action_jitter.ndim == 0:
            return np.clip(self.action_delay + np.random.randn() * self.action_jitter, 0, np.inf) * self.delay_scale
        else:
            return np.clip(self.action_delay + np.random.choice(self.action_jitter), 0, np.inf) * self.delay_scale
        
    def _sample_observation_delay(self):
        if self.observation_jitter.ndim == 0:
            return np.clip(self.observation_delay + np.random.randn() * self.observation_jitter, 0, np.inf) \
                   * self.delay_scale
        else:
            return np.clip(self.observation_delay + np.random.choice(self.observation_jitter), 0, np.inf) \
                   * self.delay_scale

    def _sample_controller_period(self):
        if self.controller_jitter.ndim == 0:
            return np.clip(self.controller_period + np.random.randn() * self.controller_jitter, 0, np.inf)
        else:
            return np.clip(self.controller_period + np.random.choice(self.controller_jitter), 0, np.inf)
