import pygame
import numpy as np
import os
import random
import matplotlib.pyplot as plt
from guided_mrmp.utils import Env, Roomba, Robot

from guided_mrmp.planners import RRT
from guided_mrmp.planners import RRTStar
from guided_mrmp.planners.multirobot.db_guided_mrmp import GuidedMRMP

class Simulator:
    def __init__(self, robots, dynamics_models, circle_obstacles, rectangle_obstacles, policy):
        """
        robots: list of Robot objects
        dynamics_models: list of DynamicsModel objects
        circle_obstacles: list of tuples (x,y,radius)
        rectangle_obstacles: list of tuples (x,y,width,height)
        policy: The policy that gives us the controls for each robot at a given time
        state: the current state of the world. This is a list of tuples, where each tuple is the state of a robot
        time: the current time
        """
        self.robots = robots
        self.circ_obstacles = circle_obstacles
        self.rect_obstacles = rectangle_obstacles
        self.policy = policy 

        self.state = [robot.current_position for robot in robots]

        self.num_robots = len(robots)
        self.dynamics_models = dynamics_models
        self.time = 0

    def advance(self, screen, dt):
        """
        Advance the simulation by dt seconds
        """
        controls = self.policy.advance(screen,self.state, self.time, dt)
        print(controls)
        for i in range(self.num_robots):
            new_state = self.dynamics_models[i].next_state(self.state[i], controls[i], dt)
            self.robots[i].current_position = new_state
            self.state[i] = new_state
        self.time += dt

    def draw_environment(self, screen):
        screen.fill((255,255,255))
        for obs in self.circ_obstacles:
            pygame.draw.circle(screen, (0,0,0), obs[0:2], obs[2])
        for obs in self.rect_obstacles:
            pygame.draw.rect(screen, (0,0,0), obs)

    def draw_robots(self, screen):
        for robot in self.robots:
            x,y,yaw = robot.current_position
            pygame.draw.circle(screen, robot.color, (x,y), robot.radius)
            # Plot direction marker
            dx = robot.radius * np.cos(robot.current_position[2])
            dy = robot.radius * np.sin(robot.current_position[2])
            pygame.draw.line(screen, (0,0,0), (x,y), (x+dx,y+dy), width=2)

    def run(self, screen):
        clock = pygame.time.Clock()
        pause = False
        while 1:
            clock.tick(30)
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    return
                if event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
                    return
                if event.type == pygame.KEYDOWN and event.key == pygame.K_p:
                    pause = not pause
                if event.type == pygame.KEYDOWN and event.key == pygame.K_RIGHT and pause:
                    self.draw_environment(screen)
                    self.policy.add_vis_guide_paths(screen)
                    self.draw_robots(screen)
                    
                    self.advance(screen,.3)
            if not pause:
                self.draw_environment(screen)
                self.policy.add_vis_guide_paths(screen)
                self.draw_robots(screen)
                
                self.advance(screen,.3)

            pygame.display.flip()