# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from scipy.integrate import odeint


class Planets(object):
    """
    Implements a 2D environments where there are N bodies (planets) that attract each other according to a 1/r law.

    We assume the mass of each body is 1.
    """

    # For each dimension of the hypercube
    MIN_POS = 0.  # if box exists
    MAX_POS = 1.  # if box exists
    INIT_MAX_VEL = 1.
    GRAVITATIONAL_CONSTANT = 1.

    def __init__(self, num_bodies, num_dimensions=2, dt=0.01, contained_in_a_box=True):
        self.num_bodies = num_bodies
        self.num_dimensions = num_dimensions
        self.dt = dt
        self.contained_in_a_box = contained_in_a_box

        # state variables
        self.body_positions = None
        self.body_velocities = None
        self.reset()

    def reset(self):
        self.body_positions = np.random.uniform(self.MIN_POS, self.MAX_POS, size=(self.num_bodies, self.num_dimensions))
        self.body_velocities = self.INIT_MAX_VEL * np.random.uniform(-1, 1, size=(self.num_bodies, self.num_dimensions))

    @property
    def state(self):
        return np.concatenate((self.body_positions, self.body_velocities), axis=1)  # (N, 2D)

    def step(self):

        # Helper functions since ode solver requires flattened inputs
        def flatten(positions, velocities):  # positions shape (N, D); velocities shape (N, D)
            system_state = np.concatenate((positions, velocities), axis=1)  # (N, 2D)
            system_state_flat = system_state.flatten()  # ode solver requires flat, (N*2D,)
            return system_state_flat

        def unflatten(system_state_flat):  # system_state_flat shape (N*2*D,)
            system_state = system_state_flat.reshape(self.num_bodies, 2 * self.num_dimensions)  # (N, 2*D)
            positions = system_state[:, :self.num_dimensions]  # (N, D)
            velocities = system_state[:, self.num_dimensions:]  # (N, D)
            return positions, velocities

        # ODE function
        def system_first_order_ode(system_state_flat, _):

            positions, velocities = unflatten(system_state_flat)
            accelerations = np.zeros_like(velocities)  # init (N, D)

            for i in range(self.num_bodies):
                relative_positions = positions - positions[i]  # (N, D)
                distances = np.linalg.norm(relative_positions, axis=1, keepdims=True)  # (N, 1)
                distances[i] = 1.  # bodies don't affect themselves, and we don't want to divide by zero next

                # forces (see https://en.wikipedia.org/wiki/Numerical_model_of_the_Solar_System)
                force_vectors = self.GRAVITATIONAL_CONSTANT * relative_positions / (distances**self.num_dimensions)  # (N,D)
                force_vector = np.sum(force_vectors, axis=0)  # (D,)
                accelerations[i] = force_vector  # assuming mass 1.

            d_system_state_flat = flatten(velocities, accelerations)
            return d_system_state_flat

        # integrate + update
        current_system_state_flat = flatten(self.body_positions, self.body_velocities)  # (N*2*D,)
        _, next_system_state_flat = odeint(system_first_order_ode, current_system_state_flat, [0., self.dt])  # (N*2*D,)
        self.body_positions, self.body_velocities = unflatten(next_system_state_flat)  # (N, D), (N, D)

        # bounce off boundaries of box
        if self.contained_in_a_box:
            ind_below_min = self.body_positions < self.MIN_POS
            ind_above_max = self.body_positions > self.MAX_POS
            self.body_positions[ind_below_min] += 2. * (self.MIN_POS - self.body_positions[ind_below_min])
            self.body_positions[ind_above_max] += 2. * (self.MAX_POS - self.body_positions[ind_above_max])
            self.body_velocities[ind_below_min] *= -1.
            self.body_velocities[ind_above_max] *= -1.
            self.assert_bodies_in_box()  # check for bugs

    def animate(self, file_name=None, frames=1000, pixel_length=None, tight_format=True):
        """
        Animation function for visual debugging.
        """
        if self.num_dimensions is not 2:
            raise NotImplementedError

        if pixel_length is None:
            fig = plt.figure()
        else:
            # matplotlib can't render if pixel_length is too small, so just run in the background id pixels specified
            import matplotlib
            matplotlib.use('Agg')
            my_dpi = 96  # find your screen's dpi here: https://www.infobyip.com/detectmonitordpi.php
            fig = plt.figure(facecolor='lightslategray', figsize=(pixel_length/my_dpi, pixel_length/my_dpi), dpi=my_dpi)

        ax = fig.add_subplot(1, 1, 1)
        plt.axis('off')
        if tight_format:
            plt.subplots_adjust(left=0, right=1, top=1, bottom=0, wspace=None, hspace=None)
        body_colors = np.random.uniform(size=self.num_bodies)

        def render(_):
            self.step()
            x = self.body_positions[:, 0]
            y = self.body_positions[:, 1]

            ax.clear()
            # if tight_format:
            #     plt.subplots_adjust(left=0., right=1., top=1., bottom=0.)
            ax.scatter(x, y, marker='o', c=body_colors, cmap='viridis')
            # ax.set_title(self.__class__.__name__ + "\n(temperature inside box: {:.1f})".format(self.temperature))
            ax.set_xlim(self.MIN_POS, self.MAX_POS)
            ax.set_ylim(self.MIN_POS, self.MAX_POS)
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_aspect('equal')
            # ax.axis('off')
            ax.set_facecolor('black')
            if tight_format:
                ax.margins(x=0., y=0.)

        interval_milliseconds = 1000 * self.dt
        anim = animation.FuncAnimation(fig, render, frames=frames, interval=interval_milliseconds)

        plt.pause(1)
        if file_name is None:
            file_name = self.__class__.__name__.lower() + '.gif'
        file_name = 'images/' + file_name
        print('Saving file {} ...'.format(file_name))
        anim.save(file_name, writer='imagemagick')
        plt.close(fig)

    def assert_bodies_in_box(self):
        """
        if the sim goes really fast, they can bounce one-step out of box. Let's just check for this for now, fix later
        """
        assert np.all(self.body_positions >= self.MIN_POS) and np.all(self.body_positions <= self.MAX_POS)

    @property
    def temperature(self):
        """
        Temperature is the average kinetic energy of system
        :return: float
        """
        average_kinetic_energy = 0.5 * np.mean(np.linalg.norm(self.body_velocities, axis=1))  # (N, D) --> (1,)
        return average_kinetic_energy


class Electrons(Planets):
    """
    Implements a 2D environments where there are N bodies (electrons) that repel each other according to a 1/r law.
    """

    # override
    GRAVITATIONAL_CONSTANT = -1.  # negative means they repel


class IdealGas(Planets):
    """
    Implements a 2D environments where there are N bodies (gas molecules) that do not interact with each other.
    """

    # override
    GRAVITATIONAL_CONSTANT = 0.  # zero means they don't interact