from fizgrid.grid import Grid
from fizgrid.entities import Entity
from fizgrid.utils import Shape
import random, math


class Sniffer(Entity):
    def __init__(self, *args, **kwargs):
        # Override the init method of the Entity class to extend the functionality and add custom attributes.
        super().__init__(*args, **kwargs)
        self.goal_x = None
        self.goal_y = None
        self.tolerance = None
        self.speed = 1

    def get_dist_from_goal(self):
        # A function to calculate the distance from the entity to the goal.
        return (
            (self.goal_x - self.x_coord) ** 2
            + (self.goal_y - self.y_coord) ** 2
        ) ** 0.5

    def set_goal(self, x_coord, y_coord, tolerance=1):
        # Set the goal for the entity to reach and start the routing process.
        self.goal_x = x_coord
        self.goal_y = y_coord
        self.tolerance = tolerance
        self.add_next_route()

    def add_next_route(self):
        # Determine the next route for the entity to take.

        # This route is randomly calculated based on the angle and distance to the goal.
        #     - If the entity is within the tolerance of the goal, it will stop.
        #     - If the entity is not within the tolerance
        #         - It will calculate the target angle towards the goal.
        #         - It will calculate the target distance to the goal.
        #         - A random angle is generated by using a normal distribution centered on the goal angle.
        #         - A random distance is generated between 0 and the minimum of the distance to the goal and 5.
        distance_from_goal = self.get_dist_from_goal()
        if distance_from_goal < self.tolerance:
            return
        goal_angle_rad = math.atan2(
            self.goal_y - self.y_coord, self.goal_x - self.x_coord
        )
        random_angle = random.normalvariate(goal_angle_rad, math.pi / 2)
        distance = random.uniform(0, min(distance_from_goal, 5))
        x_shift = distance * math.cos(random_angle)
        y_shift = distance * math.sin(random_angle)

        self.add_route(
            waypoints=[
                (
                    self.x_coord + x_shift,
                    self.y_coord + y_shift,
                    distance * self.speed,
                ),
            ]
        )

    def on_realize(self, **kwargs):
        # This method is a blank method that should be overridden by the user.
        # It is called when an event for the entity is realized.

        # In this case, it is used to implment the goal checking logic and continuously add routes until the goal is reached.

        # The method checks if the entity has a goal set and if it is within the tolerance of the goal.
        # If the entity is not within the tolerance, it will add the next route.
        # If the entity is within the tolerance, it will set the goal to None.
        if self.goal_x is not None and self.goal_y is not None:
            if self.get_dist_from_goal() > self.tolerance:
                self.add_next_route()
            else:
                self.goal_x = None
                self.goal_y = None


forest = Grid(
    name="forest",
    x_size=1000,
    y_size=1000,
    add_exterior_walls=True,
)

# Add some pigs to the forest
truffle_pig_1 = forest.add_entity(
    Sniffer(
        name="Truffle_Pig_1",
        shape=Shape.rectangle(x_len=1, y_len=1, round_to=2),
        x_coord=450,
        y_coord=500,
    )
)
truffle_pig_2 = forest.add_entity(
    Sniffer(
        name="Truffle_Pig_2",
        shape=Shape.rectangle(x_len=1, y_len=1, round_to=2),
        x_coord=500,
        y_coord=450,
    )
)

truffle_pig_1.set_goal(
    x_coord=550,
    y_coord=500,
)

truffle_pig_2.set_goal(
    x_coord=500,
    y_coord=550,
)

# Run the sim
forest.simulate()

# print({
#     'Name': truffle_pig_1.name,
#     'x_coord': truffle_pig_1.x_coord,
#     'y_coord': truffle_pig_1.y_coord,
# })
# print({
#     'Name': truffle_pig_2.name,
#     'x_coord': truffle_pig_2.x_coord,
#     'y_coord': truffle_pig_2.y_coord,
# })

# Example Output
# {'Name': 'Truffle_Pig_1', 'x_coord': 550.6344767270211, 'y_coord': 499.90150976359365}
# {'Name': 'Truffle_Pig_2', 'x_coord': 499.2754299204255, 'y_coord': 550.6818320798665}


# Checks for success
success = True
truffle_pig_1.goal_x = 550
truffle_pig_1.goal_y = 500

truffle_pig_2.goal_x = 500
truffle_pig_2.goal_y = 550

if truffle_pig_1.get_dist_from_goal() > 1:
    success = False
if truffle_pig_2.get_dist_from_goal() > 1:
    success = False

if success:
    print("test_05.py: passed")
else:
    print("test_05.py: failed")
