from __future__ import annotations
import numpy as np
import pandas as pd

from vivarium.framework.event import Event
from vivarium import Component
from vivarium.framework.engine import Builder
from vivarium.framework.population import SimulantData


class Movement(Component):
    ##############
    # Properties #
    ##############
    CONFIGURATION_DEFAULTS = {
        "field": {
            "width": 1000,
            "height": 1000,
        },
        "movement": {
            "max_speed": 2,
        },
    }

    #####################
    # Lifecycle methods #
    #####################

    def setup(self, builder: Builder) -> None:
        self.config = builder.configuration

        # docs-start: register_attribute_producer
        builder.value.register_attribute_producer(
            "acceleration", source=self.base_acceleration
        )
        # docs-end: register_attribute_producer
        self.randomness = builder.randomness.get_stream(self.name)
        builder.population.register_initializer(
            initializer=self.initialize_movement,
            columns=["x", "y", "vx", "vy"],
            required_resources=[self.randomness]
        )

    ##################################
    # Pipeline sources and modifiers #
    ##################################

    # docs-start: base_acceleration
    def base_acceleration(self, index: pd.Index[int]) -> pd.DataFrame:
        return pd.DataFrame(0.0, columns=["acc_x", "acc_y"], index=index)
    # docs-end: base_acceleration

    ########################
    # Event-driven methods #
    ########################

    def initialize_movement(self, pop_data: SimulantData) -> None:
        # Start randomly distributed, with random velocities
        new_population = pd.DataFrame(
            {
                "x": self.config.field.width * self.randomness.get_draw(pop_data.index, "x"),
                "y": self.config.field.height * self.randomness.get_draw(pop_data.index, "y"),
                "vx": ((2 * self.randomness.get_draw(pop_data.index, "vx")) - 1) * self.config.movement.max_speed,
                "vy": ((2 * self.randomness.get_draw(pop_data.index, "vy")) - 1) * self.config.movement.max_speed,
            },
            index=pop_data.index,
        )
        self.population_view.initialize(new_population)

    # docs-start: on_time_step
    def on_time_step(self, event: Event) -> None:
        def _apply_physics(pop: pd.DataFrame) -> pd.DataFrame:
            acceleration = self.population_view.get_frame(event.index, "acceleration")

            # Accelerate and limit velocity
            if not isinstance(acceleration, pd.DataFrame):
                raise ValueError("Acceleration must be a pd.DataFrame")
            pop[["vx", "vy"]] += acceleration.rename(columns=lambda c: c.replace("acc_", "v"))
            speed = np.sqrt(np.square(pop.vx) + np.square(pop.vy))
            velocity_scaling_factor = np.where(
                speed > self.config.movement.max_speed,
                self.config.movement.max_speed / speed,
                1.0,
            )
            pop["vx"] *= velocity_scaling_factor
            pop["vy"] *= velocity_scaling_factor

            # Move according to velocity
            pop["x"] += pop.vx
            pop["y"] += pop.vy

            # Loop around boundaries
            pop["x"] = pop.x % self.config.field.width
            pop["y"] = pop.y % self.config.field.height

            return pop

        self.population_view.update(["x", "y", "vx", "vy"], _apply_physics)
    # docs-end: on_time_step
