
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any

import numpy as np
import pandas as pd

from vivarium import Component
from vivarium.framework.engine import Builder


class Force(Component, ABC):
    ##############
    # Properties #
    ##############
    @property
    def configuration_defaults(self) -> dict[str, Any]:
        return {
            self.__class__.__name__.lower(): {
                "max_force": 0.03,
            },
        }

    columns_required = []

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

    def setup(self, builder: Builder) -> None:
        self.config = builder.configuration[self.__class__.__name__.lower()]
        self.max_speed = builder.configuration.movement.max_speed

        self.neighbors = builder.value.get_value("neighbors")

        builder.value.register_value_modifier(
            "acceleration",
            modifier=self.apply_force,
            required_resources=self.columns_required + [self.neighbors],
        )

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

    def apply_force(self, index: pd.Index[int], acceleration: pd.DataFrame) -> pd.DataFrame:
        neighbors = self.neighbors(index)
        pop = self.population_view.get(index)
        pairs = self._get_pairs(neighbors, pop)

        raw_force = self.calculate_force(pairs)
        force = self._normalize_and_limit_force(
            force=raw_force,
            velocity=pop[["vx", "vy"]],
            max_force=self.config.max_force,
            max_speed=self.max_speed,
        )

        acceleration.loc[force.index, ["x", "y"]] += force[["x", "y"]]
        return acceleration

    ##################
    # Helper methods #
    ##################

    @abstractmethod
    def calculate_force(self, neighbors: pd.DataFrame) -> pd.DataFrame:
        pass

    def _get_pairs(self, neighbors: pd.Series[int], pop: pd.DataFrame) -> pd.DataFrame:
        pairs = (
            pop.join(neighbors.rename("neighbors"))
            .reset_index()
            .explode("neighbors")
            .merge(
                pop.reset_index(),
                left_on="neighbors",
                right_index=True,
                suffixes=("_self", "_other"),
            )
        )
        pairs["distance_x"] = pairs.x_other - pairs.x_self
        pairs["distance_y"] = pairs.y_other - pairs.y_self
        pairs["distance"] = np.sqrt(pairs.distance_x**2 + pairs.distance_y**2)

        return pairs

    def _normalize_and_limit_force(
        self,
        force: pd.DataFrame,
        velocity: pd.DataFrame,
        max_force: float,
        max_speed: float,
    ) -> pd.DataFrame:
        normalization_factor = np.where(
            (force.x != 0) | (force.y != 0),
            max_speed / self._magnitude(force),
            1.0,
        )
        force["x"] *= normalization_factor
        force["y"] *= normalization_factor
        force["x"] -= velocity.loc[force.index, "vx"]
        force["y"] -= velocity.loc[force.index, "vy"]
        magnitude = self._magnitude(force)
        limit_scaling_factor = np.where(
            magnitude > max_force,
            max_force / magnitude,
            1.0,
        )
        force["x"] *= limit_scaling_factor
        force["y"] *= limit_scaling_factor
        return force[["x", "y"]]

    def _magnitude(self, df: pd.DataFrame) -> pd.Series[float]:
        return pd.Series(np.sqrt(np.square(df.x) + np.square(df.y)), dtype=float)


class Separation(Force):
    """Push boids apart when they get too close."""

    configuration_defaults = {
        "separation": {
            "distance": 30,
            "max_force": 0.03,
        },
    }

    def calculate_force(self, neighbors: pd.DataFrame) -> pd.DataFrame:
        # Push boids apart when they get too close
        separation_neighbors = neighbors[neighbors.distance < self.config.distance].copy()
        force_scaling_factor = np.where(
            separation_neighbors.distance > 0,
            ((-1 / separation_neighbors.distance) / separation_neighbors.distance),
            1.0,
        )
        separation_neighbors["force_x"] = (
            separation_neighbors["distance_x"] * force_scaling_factor
        )
        separation_neighbors["force_y"] = (
            separation_neighbors["distance_y"] * force_scaling_factor
        )

        force: pd.DataFrame = (
            separation_neighbors.groupby("index_self")[["force_x", "force_y"]]
            .sum()
            .rename(columns=lambda c: c.replace("force_", ""))
        )

        return force


class Cohesion(Force):
    """Push boids together."""

    def calculate_force(self, pairs: pd.DataFrame) -> pd.DataFrame:
        return (
            pairs.groupby("index_self")[["distance_x", "distance_y"]]
            .sum()
            .rename(columns=lambda c: c.replace("distance_", ""))
        )


class Alignment(Force):
    """Push boids toward where others are going."""

    def calculate_force(self, pairs: pd.DataFrame) -> pd.DataFrame:
        return (
            pairs.groupby("index_self")[["vx_other", "vy_other"]]
            .sum()
            .rename(columns=lambda c: c.replace("v", "").replace("_other", ""))
        )
