import random

from mrtommy.scene import Scene
from mrtommy.solvers.prm import PRMSolver
from mrtommy.solvers.rrt_star import RRTStarSolver


class DRRTStarSolver(RRTStarSolver):
    def __init__(
        self,
        prm_landmarks: int,
        prm_k: int,
        max_iterations: int,
        max_distance: float,
        radius: float,
    ):
        self.prm_landmarks = prm_landmarks
        self.prm_k = prm_k
        super().__init__(max_iterations, max_distance, radius)

    def get_position(self, scene):
        return [random.choice(roadmap)[0] for roadmap in self.robot_roadmaps]

    def solve(self, scene):
        self.robot_roadmaps = []
        for robot in scene.robots:
            prm = PRMSolver(self.prm_landmarks, self.prm_k)
            prm.solve(Scene([robot], scene.obstacles))
            # TODO: currently the edges generated by PRM are unused
            self.robot_roadmaps.append(list(prm.roadmap.nodes))

        return super().solve(scene)

    def __repr__(self):
        return f"dRRT* with iterations={self.max_iterations}, η={self.max_distance}, radius={self.radius}, landmarks={self.prm_landmarks} and k={self.prm_k}"
