""":mod:`megastep.modules` are chunks of functionality that often turn up in megastep environments. 
""" 

import numpy as np
import torch
from rebar import arrdict
from . import spaces, geometry, cuda, plotting
import matplotlib.pyplot as plt

def to_local_frame(angles, p):
    """Converts a velocity vector in the global coordinate frame to one in the frame local to the agent"""
    a = np.pi/180*angles
    c, s = torch.cos(a), torch.sin(a)
    x, y = p[..., 0], p[..., 1]
    return torch.stack([c*x + s*y, -s*x + c*y], -1)

def to_global_frame(angles, p):
    """Converts a velocity vector in the local coordinate frame of the agent to one in the global frame."""
    a = np.pi/180*angles
    c, s = torch.cos(a), torch.sin(a)
    x, y = p[..., 0], p[..., 1]
    return torch.stack([c*x - s*y, s*x + c*y], -1)

class SimpleMovement:

    def __init__(self, core, speed=10, ang_speed=180, n_agents=None):
        """A simple movement system with no momentum. 

        There are seven actions in total: 
            * do nothing
            * forward/backward a fixed distance
            * strafe left/right a fixed distance
            * turn left/right a fixed angle

        :param core: The :class:`~megastep.core.Core` used by the environment.
        :param speed: The speed of the agent in its linear movements, in meters per second.
        :param ang_speed: The speed of the agent in its rotational movements, in degrees per second.
        :n_agents: The number of agents to output actions for. This is usually taken from the core; it can be usefully
            overridden in :ref:`multiagent environments <deathmatch-env>`.

        :var space: The :ref:`action space <spaces>` to present to the controlling network.
        """
        # noop, forward/backward, strafe left/right, turn left/right
        self.core = core
        velocity = torch.tensor([[0., 0.], [0., 1.], [0.,-1.], [1., 0.], [-1.,0.], [0., 0.], [0., 0.]])
        angvelocity = torch.tensor([0., 0., 0., 0., 0., +1., -1.])
        self._actionset = arrdict.arrdict(
            velocity=speed/core.fps*velocity,
            angvelocity=ang_speed/core.fps*angvelocity
        ).to(core.device)

        self.space = spaces.MultiDiscrete(n_agents or core.n_agents, 7)

    def __call__(self, decision):
        """Sets the velocity and angular velocity of the agents based on the actions given, then steps the physics.

        The ``decision['actions']`` are typically a (n_env, n_agent)-int tensor with values 0-7.

        :param decision: A :ref:`decision <decision-world>` dict with an ``action`` key conforming to ``.space``.
        :return: The result of the :func:`~megastep.cuda.physics` call.
        """
        core = self.core
        delta = self._actionset[decision.actions.long()]
        core.agents.angvelocity[:] = delta.angvelocity
        core.agents.velocity[:] = to_global_frame(core.agents.angles, delta.velocity)
        return cuda.physics(core.scenery, core.agents)

class MomentumMovement:

    def __init__(self, core, accel=5, ang_accel=180, decay=.125, n_agents=None):
        """A simple movement system *with* momentum. 

        There are seven actions in total: 
            * do nothing
            * accelerate forward/backward 
            * accelerate left/right 
            * torque left/right 

        TODO: Make the decay per second rather than per timestep.

        :param core: The :class:`~megastep.core.Core` used by the environment.
        :param accel: The acceleration of the agent in its linear movements, in meters per second squared.
        :param ang_accel: The acceleration of the agent in its rotational movements, in degrees per second squared.
        :param decay: The multiplicative decay of the agent's velocity per timestep. 1 means Newtonian motion; 0 means
            the only velocity is that generated by this timestep's acceleration.
        :n_agents: The number of agents to output actions for. This is usually taken from the core; it can be usefully
            overridden in :ref:`multiagent environments <deathmatch-env>`.

        :var space: The :ref:`action space <spaces>` to present to the controlling network.
        :var decay: The value of the decay parameter.
        """

        # noop, forward/backward, strafe left/right, turn left/right
        self.core = core
        velocity = torch.tensor([[0., 0.], [0., 1.], [0.,-1.], [1., 0.], [-1.,0.], [0., 0.], [0., 0.]])
        angvelocity = torch.tensor([0., 0., 0., 0., 0., +1., -1.])
        self._actionset = arrdict.arrdict(
            velocity=accel/core.fps*velocity,
            angvelocity=ang_accel/core.fps*angvelocity
        ).to(core.device)

        self.decay = decay

        self.space = spaces.MultiDiscrete(n_agents or core.n_agents, 7)

    def __call__(self, decision):
        """Sets the acceleration and angular acceleration of the agents based on the actions given, then steps the physics.

        The ``decision['actions']`` are typically a (n_env, n_agent)-int tensor with values 0-7.

        :param decision: A :ref:`decision <decision-world>` dict with an ``actions`` key conforming to ``.space``.
        :return: The result of the :func:`~megastep.cuda.physics` call.
        """
        core = self.core
        delta = self._actionset[decision.actions.long()]
        core.agents.angvelocity[:] = (1 - self.decay)*core.agents.angvelocity + delta.angvelocity
        core.agents.velocity[:] = (1 - self.decay)*core.agents.velocity + to_global_frame(core.agents.angles, delta.velocity)
        return cuda.physics(core.scenery, core.agents)

def unpack(d):
    """Unpacks :mod:`~megastep.cuda` datastructures into :ref:`arrdicts <dotdicts>` with the same attributes."""
    if isinstance(d, torch.Tensor):
        return d
    return arrdict.arrdict({k: unpack(getattr(d, k)) for k in dir(d) if not k.startswith('_')})

def render(core):
    """Calls :func:`~megastep.cuda.render`, turns the output into an :ref:`attrdict <dotdicts>`, 
    then converts the :attr:`~megastep.cuda.Render.screen` attribute into the kind of (batch, channel, height, width)
    tensor expected by PyTorch convolution modules.
    
    This is almost always what you want when rendering RGB observations.
    """
    render = unpack(cuda.render(core.scenery, core.agents))
    render = arrdict.arrdict({k: v.unsqueeze(2) for k, v in render.items()})
    render['screen'] = render.screen.permute(0, 1, 4, 2, 3)
    return render

def downsample(screen, subsample):
    """Factors a :func:`render`'d screen tensor along its final width dimension, returning something with shape
    (..., width/subsample, subsample). 
    
    Typically you chase this call by aggregating over the trailing dimension in some way; either mean or min or max
    or ``[..., 0]``.
    """
    return screen.view(*screen.shape[:-1], screen.shape[-1]//subsample, subsample)

class Depth:

    def __init__(self, core, n_agents=None, subsample=1, max_depth=10):
        """Generates depth observations.
        
        :param core: The :class:`~megastep.core.Core` used by the environment.
        :param n_agents: The number of agents to generate observations for. This is usually taken from the core; it can be usefully
            overridden in :ref:`multiagent environments <deathmatch-env>`.
        :param subsample: How many horizontal pixels to average when generating the observations. For example, if the 
            core is rendering at 256 pixels and ``subsample`` is 4, then 64-pixel observations will be returned. 
            A higher subsampling rate makes for slower rendering, but smoother observations. 
        :param max_depth: The maximum depth, corresponding to a zero in the observation. Given in meters.

        :var space: The :ref:`observation space <spaces>` to present to the controlling network.
        :var max_depth: The value of the ``max_depth`` parameter.
        :var subsample: The value of the ``subsample`` parameter.
        """
        n_agents = n_agents or core.n_agents
        self.core = core
        self.space = spaces.MultiImage(n_agents, 1, 1, core.res//subsample)
        self.max_depth = max_depth
        self.subsample = subsample

    def __call__(self, r=None):
        """Generates a depth observation. 

        This means calling :func:`~megastep.cuda.render` (if ``r`` isn't passed), then using it to create and return a
        depth tensor of values between 0 and 1. The values interpolate linearly between one at the near clipping plane
        (given by :attr:`megastep.core.AGENT_RADIUS`) and zero at ``max_depth``.

        :param r: The output of :func:`render`. :func:`render` will be called if this isn't passed.
            It's useful to pass ``r`` if you're doing other things with the render output.
        :return: A (n_env, n_agent, 1, res)-tensor of values between 0 and 1.
        """
        r = render(self.core) if r is None else r
        depth = 1 - ((r.distances - self.core.agent_radius)/self.max_depth).clamp(0, 1)
        self._last_obs = downsample(depth, self.subsample).mean(-1).unsqueeze(3)
        return self._last_obs
    
    def state(self, e=0):
        """The state of the module in sub-env ``e``, which is to say its last observation for ``e``. Useful in
        :ref:`plotting <plotting>`"""
        return self._last_obs[e].clone()

class RGB:

    def __init__(self, core, n_agents=None, subsample=1):
        """Generates RGB observations.
        
        :param core: The :class:`~megastep.core.Core` used by the environment.
        :param n_agents: The number of agents to generate observations for. This is usually taken from the core; it can be usefully
            overridden in :ref:`multiagent environments <deathmatch-env>`.
        :param subsample: How many horizontal pixels to average when generating the observations. For example, if the 
            core is rendering at 256 pixels and ``subsample`` is 4, then 64-pixel observations will be returned. 
            A higher subsampling rate makes for slower rendering, but smoother observations. 

        :var space: The :ref:`observation space <spaces>` to present to the controlling network.
        :var subsample: The value of the ``subsample`` parameter.
        """
        n_agents = n_agents or core.n_agents
        self.core = core
        self.space = spaces.MultiImage(n_agents, 3, 1, core.res//subsample)
        self.subsample = subsample

    def __call__(self, r=None):
        """Generates a depth observation. 

        This means calling :func:`~megastep.cuda.render` (if ``r`` isn't passed), then using it to create and return a
        RGB tensor. The RGB tensor needs :func:`~megastep.core.gamma_encode`'ing if you want to plot it.

        :param r: The output of :func:`render`. :func:`render` will be called if this isn't passed.
            It's useful to pass ``r`` if you're doing other things with the render output.
        :return: A (n_env, n_agent, 3, 1, res)-tensor of values between 0 and 1.
        """

        r = render(self.core) if r is None else r
        self._last_obs = downsample(r.screen, self.subsample).mean(-1)
        return self._last_obs
    
    def state(self, e=0):
        """The state of the module in sub-env ``e``, which is to say its last observation for ``e``. Useful in
        :ref:`plotting <plotting>`"""
        return self._last_obs[e].clone()

    @classmethod
    def plot_state(cls, state, axes=None):
        """Plots the state of this module using imshow. Make sure to :func:`~rebar.arrdict.numpyify` the state 
        before passing it here. Useful in :ref:`plotting <plotting>`."""
        n_agents = state.shape[0]
        axes = plt.subplots(n_agents, 1, squeeze=False) if axes is None else axes
        plotting.plot_images({'rgb': state}, axes)
        return axes

class IMU:

    def __init__(self, core, speed_scale=10., ang_scale=360., n_agents=None,):
        """Generate a linear-and-angular-velocity measurement. Kinda like a `inertial measurement unit
        <https://en.wikipedia.org/wiki/Inertial_measurement_unit>`_.

        :param core: The :class:`~megastep.core.Core` used by the environment.
        :param n_agents: The number of agents to generate observations for. This is usually taken from the core; it can be usefully
            overridden in :ref:`multiagent environments <deathmatch-env>`.
        :param speed_scale: The scale of speeds to use, with this value corresponding to an observation of 1. Given
            in meters per second.
        :param ang_scale: The scale of angular speeds to use, with this value corresponding to an observation of 1.
            Given in degrees per second.

        :var space: The :ref:`observation space <spaces>` to present to the controlling network.
        :var speed_scale: The value of the ``speed_scale`` parameter.
        :var ang_scale: The value of the ``ang_scale`` parameter.
        """
        self.core = core
        self.space = spaces.MultiVector(n_agents or core.n_agents, 3)
        self.speed_scale = speed_scale
        self.ang_scale = ang_scale

    def __call__(self):
        """Returns an IMU observation, which is to say a (n_env, n_agent, 3)-float tensor. The values along the final
        dimension are (angular_velocity, medial_velocity, lateral_velocity), with the ``angular_velocity`` scaled by 
        ``ang_scale`` and the linear velocities scaled by ``speed_scale``. 
        """ 
        return torch.cat([
            self.core.agents.angvelocity[..., None]/self.ang_scale,
            to_local_frame(self.core.agents.angles, self.core.agents.velocity)/self.speed_scale], -1)

def random_empty_positions(geometries, n_agents, n_points):
    """Returns a tensor of randomly-selected empty points in each :ref:`geometry <geometry>`.
    
    The returned tensor is a (n_geometries, n_agents, n_points, 2)-float tensor, with the coordinates given in meters.

    This is typcially used when you want to randomly move an agent to a new place, but *finding* an empty point at 
    each timestep is too expensive. So instead this is used to generate ``n_points`` empty points in advance, and then
    when you need one you can choose from the pre-generated options.
    """ 
    points = []
    for g in geometries:
        sample = np.stack((g.masks > 0).nonzero(), -1)

        # There might be fewer open points than we're asking for
        n_possible = min(len(sample)//n_agents, n_points)
        sample = sample[np.random.choice(np.arange(len(sample)), (n_possible, n_agents), replace=True)]

        # So repeat the sample until we've got enough
        sample = np.concatenate([sample]*int(n_points/len(sample)+1))[-n_points:]
        sample = np.random.permutation(sample)
        points.append(geometry.centers(sample, g.masks.shape, g.res).transpose(1, 0, 2))
    return arrdict.stack(points)
        
class RandomSpawns:

    def __init__(self, geometries, core, n_spawns=100):
        """Respawns agents to random empty locations in the geometry.

        :param geometries: The :ref:`geometry <geometry>` to use when calculating the spawn locations. Should be a list
            with one for each environment.
        :param core: The :class:`~megastep.core.Core` used by the environment.
        :param n_spawns: The number of spawns to choose between for each agent. This is precomputed when the class 
            is created, so that the respawns themselves are fast.
        """
        self.core = core

        positions = random_empty_positions(geometries, core.n_agents, n_spawns)
        angles = core.random.uniform(-180, +180, (len(geometries), core.n_agents, n_spawns))
        self._spawns = arrdict.torchify(arrdict.arrdict(positions=positions, angles=angles)).to(core.device)

    def __call__(self, reset):
        """Respawns the specified agents at random empty points.

        Respawned agents have their velocities zeroed.

        :param reset: A (n_env, n_agent) mask indicating which agents to reset. Trues will be reset; Falses will be
            left as is.
        """
        core = self.core
        required = reset.nonzero(as_tuple=True)
        choices = torch.randint_like(required[0], 0, self._spawns.angles.shape[1])
        core.agents.angles[required] = self._spawns.angles[(*required, choices)] 
        core.agents.positions[required] = self._spawns.positions[(*required, choices)] 
        core.agents.velocity[required] = 0.
        core.agents.angvelocity[required] = 0.

class RandomLifespans:

    def __init__(self, core, max_lifespan, min_lifespan=None):
        """Tracks how many steps each agent has been alive for, and indicates when they exceed a randomly-chosen
        lifespan.

        Lifespans are chosen randomly between ``min_lifespan`` and ``max_lifespan`` and re-rolled after each reset.
        This is useful when you want otherwise 'synchronous' environments to 'mix' so that you get a random
        distribution of behaviour in each batch, rather than one batch full of 'early life experience' and another of
        'late life experience'.
        
        :param core: The :class:`~megastep.core.Core` used by the environment.
        :param max_lifespan: The maximum lifespan.
        :type max_lifespan: int
        :param min_lifespan: The minimum lifespan; defaults to half of ``max_lifespan``.
        :type min_lifespan: int

        :var max_lifespan: Value of the ``max_lifespan`` parameter.
        :var min_lifespan: Value of the ``min_lifespan`` parameter.

        TODO: Test this now you've rewritten it.
        """
        min_lifespan = max_lifespan//2 if min_lifespan is None else min_lifespan
        self.min_lifespan = min_lifespan
        self.max_lifespan = max_lifespan
        self._max_lifespans = torch.zeros((core.n_envs, core.n_agents), dtype=torch.int, device=core.device)
        self._lifespans = torch.zeros_like(self._max_lifespans)
        self._reset(core.agent_full(True))

    def _reset(self, reset):
        self._lifespans[reset] = 0
        self._max_lifespans[reset] = torch.randint_like(self._max_lifespans, self.min_lifespan, self.max_lifespan)

    def __call__(self, reset=None):
        """Returns a tensor indicating which agents have outlived their randomly-generated lifespans.

        Every time this is called, an agent's time-lived is incremented by one. When it exceeds its max lifespan, 
        the returned mask will indicate this has happened, the time lived will be reset to zero, and the max 
        lifespan for that agent will be re-rolled.

        :param reset: An optional (n_env, n_agent)-bool tensor indicating which agents have been reset for other reasons.
            These agents will have their lifespans reset and max lifespans regenerated.
        :returns: An (n_env, n_agent)-bool tensor of which agents have exceeded their lifespan and should be reset.
        """
        self._lifespans += 1
        reset = torch.zeros_like(self._lifespans, dtype=torch.bool) if reset is None else reset
        reset = (self._lifespans >= self._max_lifespans) | reset
        self._reset(reset)
        return reset

    def state(self, e):
        """Returns the state of this module on sub-env ``e``. The state is a :ref:`arrdict <dotdicts>` of 
        the agents' lifespans and max lifespans as (n_agent,)-tensors."""
        return arrdict.arrdict(lifespan=self._lifespans[e], max_lifespans=self._max_lifespans[e]).clone()
