"""
base.py
"""

from dataclasses import dataclass
from typing import Tuple, Callable

import networkx as nx
import numpy as np
from pydrake.geometry import SceneGraph
from pydrake.multibody.plant import MultibodyPlant
from pydrake.multibody.tree import ModelInstanceIndex

# Internal Imports
from brom_drake.motion_planning.algorithms.motion_planner import MotionPlanner


@dataclass
class BaseRRTPlannerConfig:
    random_seed: int = 23
    steering_step_size: float = 0.05
    prob_sample_goal: float = 0.25
    max_iterations: int = int(1e4)
    convergence_threshold: float = 1e-3


class BaseRRTPlanner(MotionPlanner):
    def __init__(
        self,
        robot_model_idx: ModelInstanceIndex,
        plant: MultibodyPlant,
        scene_graph: SceneGraph,
        config: BaseRRTPlannerConfig = None,
    ):
        # Input Processing

        # Setup
        # self.dim_q = plant.num_actuated_dofs(robot_model_idx)

        self.config = config
        if self.config is None:
            self.config = BaseRRTPlannerConfig()

        # Use the parent class constructor
        super().__init__(
            robot_model_idx,
            plant,
            scene_graph,
            random_seed=self.config.random_seed,
        )

    def find_nearest_node(
        self, rrt: nx.DiGraph, q_random: np.ndarray
    ) -> Tuple[nx.classes.reportviews.NodeView, int]:
        """
        Description:
            This function finds the nearest node in the RRT to the random configuration.
        """
        nearest_node = None
        min_distance = float("inf")

        for node in rrt.nodes:
            distance = np.linalg.norm(rrt.nodes[node]["q"] - q_random)
            if distance < min_distance:
                min_distance = distance
                nearest_node_idx = node
                nearest_node = rrt.nodes[node]

        return nearest_node, nearest_node_idx

    def plan(
        self,
        q_start: np.ndarray,
        q_goal: np.ndarray,
        collision_check_fcn: Callable[[np.ndarray], bool] = None,
    ) -> Tuple[nx.DiGraph, bool]:
        """
        Description:
            This function executes the RRT planning algorithm.
        """
        # Input Processing
        q_start = q_start.flatten()
        q_goal = q_goal.flatten()
        if q_start.shape[0] != self.dim_q or q_goal.shape[0] != self.dim_q:
            raise ValueError(
                f"Start configuration shape ({q_start.shape}) and goal configuration "
                + f"shape ({q_goal.shape}) must match the dimension of the robot ({self.dim_q})."
            )

        if collision_check_fcn is None:
            collision_check_fcn = self.check_collision_in_config

        # Setup
        rrt = nx.DiGraph()
        rrt.add_node(0, q=q_start)  # Make this graph contain all configurations
        prob_sample_goal = self.config.prob_sample_goal

        # RRT Planner
        for iteration in range(self.config.max_iterations):
            # Sample random configuration OR goal
            if np.random.rand() < prob_sample_goal:
                q_random = q_goal.copy()
            else:
                q_random = self.sample_random_configuration()

            # Find nearest node in the tree
            nearest_node, nearest_node_idx = self.find_nearest_node(rrt, q_random)

            # Steer from nearest node to random configuration
            q_new = self.steer(nearest_node["q"], q_random)

            # Check for collisions
            if collision_check_fcn(q_new):
                continue

            # Add new node to the tree
            rrt.add_node(rrt.number_of_nodes(), q=q_new)
            rrt.add_edge(nearest_node_idx, rrt.number_of_nodes() - 1)

            # Check if we have reached the goal
            if np.linalg.norm(q_new - q_goal) < self.config.convergence_threshold:
                print(f"Goal reached after {iteration} iterations!")
                print(f"Check all points in path: {rrt.nodes}")
                for node in rrt.nodes:
                    print(f"Node {node}: {rrt.nodes[node]}")
                    collision_check_value = collision_check_fcn(rrt.nodes[node]["q"])
                    if collision_check_value:
                        print(
                            f"Collision check node: {collision_check_fcn(rrt.nodes[node]['q'])}"
                        )
                return rrt, rrt.number_of_nodes() - 1

        # If we exit the loop without finding a path to the goal,
        # return the RRT and indicate failure
        print("Max iterations reached without finding a path to the goal.")
        return rrt, -1

    def steer(self, nearest_node, q_random: np.ndarray) -> np.ndarray:
        """
        Description
        -----------
        This function steers from the nearest node towards the random configuration.
        """
        # Setup
        step_size = self.config.steering_step_size
        q_nearest = nearest_node

        # Calculate the direction vector
        direction = q_random - q_nearest
        distance = np.linalg.norm(direction)

        # Either:
        # 1. Move a full step towards the random configuration
        # 2. Or, if the distance is less than the step size, return the random configuration
        if distance < step_size:
            return q_random

        return q_nearest + step_size * (direction / distance)
