from datetime import datetime
import math
from typing import List, Tuple
from matplotlib.axes import Axes
from matplotlib.figure import Figure
import matplotlib.pyplot as plt

from planner.cbFunctions import heuristicFunction, cbf
from planner.drone import DroneGoalState, DroneState
from planner.flyZone import FlyZone
from planner.helperFunctions import reverse_insort
from planner.movement import Movement

from planner.obstacle import Obstacle
from planner.plotting import plotPathToCurrentDroneStateLive, plotlive
from position import Position
import os

import glob
from PIL import Image

from trajectory import Trajectory



class Planner:
  CONTROLL_BARRIER_FUNCTIONS = {
    "heuristic": heuristicFunction,
    "cbf": cbf
  }

  HUMAN = Obstacle(0, 0, 0.25)

  def __init__(self, dt: float, obstacles: List["Obstacle"], flyZone: "FlyZone", verboseLevel: int, cbfKey: str) -> None:
    self.dt = dt
    Movement.POSSIBLE_ACCELERATIONS = [acc * dt for acc in Movement.POSSIBLE_ACCELERATIONS]
    self.obstacles = obstacles
    self.flyZone = flyZone
    self.verboseLevel = verboseLevel
    self.cbfKey = cbfKey
    self.cbf = self.CONTROLL_BARRIER_FUNCTIONS[cbfKey]

    self.iterations = 0
    self.currentGoalState = None
    self.currentStartState = None
    self.numIterationsForLastPath = None
    self.movements = []
    self.visitedMovements = set()

     # Animation
    self.saveAnimation = False
    self.lengthOfAnimation = 10
    self.animationFrameDuration = 0.25  # In seconds
    self.numAnimationFrames = self.lengthOfAnimation / self.animationFrameDuration
    self.currentAnimationFrame = 0
    self.animationKey = None


  @plotlive
  def plotStateLive(self, droneState: "DroneState", ax: Axes, fig: Figure, isFinalDroneState: bool) -> None:
    plotPathToCurrentDroneStateLive(
      droneState=droneState, 
      planner=self,
      ax=ax,
      fig=fig,
      isFinalDroneState=isFinalDroneState
    )


  def addMovements(self, associatedDroneState: "DroneState") -> Tuple[bool, "Movement"]:
    for angle in Movement.POSSIBLE_ANGLES:
      for acceleration in Movement.POSSIBLE_ACCELERATIONS:
        newMovement = Movement(
          angle=angle,
          acceleration=acceleration,
          planner=self,
          associatedDroneState=associatedDroneState
        )

        if newMovement.willHaveReachedGoal:
          return True, newMovement
        elif newMovement.isValid and newMovement.key not in self.visitedMovements:
          reverse_insort(self.movements, newMovement) 
          self.visitedMovements.add(newMovement.key)
            
    return False, None


  def printMovements(self) -> None:
    print("\n-Movements-")
    for movement in self.movements:
      print(movement)
    print("---")


  def animationPlot(self, droneState: "DroneState", ax: Axes, fig: Figure, isFinalDroneState:bool=False) -> None:
    if self.numIterationsForLastPath is not None:
      divider = self.numAnimationFrames - 3
      plotAt = self.numIterationsForLastPath // divider
    
      if plotAt == 0.0:
        plotAt = 1
      
      plotNow = self.iterations % plotAt == 0 or self.iterations == 0

      if plotNow or isFinalDroneState:
        self.currentAnimationFrame += 1
        self.plotStateLive(droneState, ax, fig, isFinalDroneState)


  def printStatus(self, droneState: "DroneState") -> None:
    print()
    print(f"#{self.iterations} | C={droneState.selectedMovement.cost}")


  def getYawForInitDroneState(self, initDronePos: "Position", goalState: "DroneGoalState"):
    initYaw = math.atan2(goalState.y - initDronePos.y, goalState.x - initDronePos.x)
    return initYaw


  def getFinalDroneState(self, movementThatReachedGoal: "Movement"):
    finalDroneState = movementThatReachedGoal.associatedDroneState
    finalDroneState.selectedMovement = movementThatReachedGoal

    # finalDroneState = DroneState(penultimateDroneState)
    return finalDroneState


  def findPathToGoal(self, initDroneState: "DroneState", goalState: "DroneGoalState"):
    self.currentGoalState = goalState
    self.currentStartState = initDroneState

    reachedGoal, movementThatReachedGoal = self.addMovements(initDroneState)

    if (self.saveAnimation and self.numIterationsForLastPath is not None) or self.verboseLevel == 1:
      fig, ax = plt.subplots()
    else:
      fig, ax = None, None


    if self.saveAnimation:
      self.animationPlot(initDroneState, ax, fig)

    if self.verboseLevel == 1:
      self.printStatus(associatedDroneState)
      self.plotStateLive(newDroneState, ax, fig, False)

    if reachedGoal:
      finalDroneState = self.getFinalDroneState(movementThatReachedGoal)
      return finalDroneState, True

    while len(self.movements) > 0:
      self.iterations += 1

      cheapestMovement = self.movements.pop()

      associatedDroneState = cheapestMovement.associatedDroneState

      associatedDroneState.selectedMovement = cheapestMovement

      newDroneState = DroneState(associatedDroneState)

      if self.verboseLevel == 1 and self.iterations % 1 == 0 and self.iterations > 20000:
        self.printStatus(associatedDroneState)
        self.plotStateLive(newDroneState, ax, fig, False)
      elif self.verboseLevel == 2 and self.iterations % 1000 == 0:
        self.printStatus(associatedDroneState)
      

      if self.saveAnimation:
        self.animationPlot(newDroneState, ax, fig)

      reachedGoal, movementThatReachedGoal = self.addMovements(newDroneState)

      if reachedGoal:
        finalDroneState = self.getFinalDroneState(movementThatReachedGoal)
          
        if self.saveAnimation:
          self.animationPlot(finalDroneState, ax, fig, True)

        return finalDroneState, True

    return None, False


  def setKey(self, finalDroneState: "DroneState", isFinalTrajectory:bool=False) -> bool:
    startX, startY = round(self.currentStartState.x, 2), round(self.currentStartState.y, 2)
    startPos = f"S({startX},{startY})"

    goalX, goalY = round(self.currentGoalState.x, 2), round(self.currentGoalState.y, 2)
    goalPos = f"G({goalX},{goalY})"
    dateString = str(datetime.now()).split(".")[0]
    self.animationKey = f"{dateString} | T={self.lengthOfAnimation} | fps={round(1/self.animationFrameDuration)} | cbf={self.cbfKey} | {startPos}-->{goalPos} | D={finalDroneState.depth} | i_tot={self.iterations}"
    
    if isFinalTrajectory:
      self.animationKey += " | Final trajectory"


  def resetParams(self):
    self.currentAnimationFrame = 0
    self.iterations = 0
    self.movements = []
    self.visitedMovements = set()


def makeAnimationFromTrajectory(planner: "Planner", dt:float=None, lastFrameDuration:int=5000):
  def getImageNum(imageName: str) -> int:
    return int(imageName.split("/")[-1][:-4])

  if dt is None:
    dt = planner.animationFrameDuration * 1000

  images = glob.glob(f"savedTrajectories/{planner.animationKey}/animationFrames/*.png")

  images.sort(key=getImageNum)

  durations = []
  numImages = len(images)

  for _ in range(numImages-1):
    durations.append(dt)

  durations.append(lastFrameDuration)

  imagePath = f"savedTrajectories/{planner.animationKey}/animation.gif"

  frames = [Image.open(image) for image in images]

  frames[0].save(
    imagePath,
    save_all=True, 
    append_images=frames[1:], 
    optimize=False, 
    duration=durations, 
    loop=0
  )


def recordAnimation(initDronePosition: "Position", planner: "Planner", goalStates: List["DroneGoalState"]):
  planner.saveAnimation = True
  numGoalStates = len(goalStates)
  currentDroneState = initDronePosition

  for goalStateIdx in range(numGoalStates):
    currentGoalState = goalStates[goalStateIdx]
    
    initDroneState = DroneState(
        parent=None, 
        x=currentDroneState.x, 
        y=currentDroneState.y,
        yaw=planner.getYawForInitDroneState(currentDroneState, currentGoalState)
      )

    planner.currentStartState = initDroneState

    for i in range(2):
      finalDroneState, foundGoalState = planner.findPathToGoal(initDroneState, currentGoalState)
      currentDroneState = finalDroneState

      if foundGoalState:
        print("Reached the goal! :)")
        
      else:
        print("No path found :(")
        break

      if i == 0:
        planner.setKey(currentDroneState)

        if planner.saveAnimation:
          os.mkdir(f"animationFrames/{planner.animationKey}")

        planner.numIterationsForLastPath = planner.iterations
        planner.iterations = 0
        planner.movements = []
        planner.visitedMovements = set()
      
    if planner.saveAnimation:
      planner.makeAnimationFromTrajectory()
    
    if goalStateIdx < numGoalStates - 1:
      planner.numIterationsForLastPath = None
      planner.currentAnimationFrame = 0
      planner.iterations = 0
      planner.movements = []
      planner.visitedMovements = set()
      
  
  planner.saveAnimation = False


def recordFinalTrajectory(finalDroneState: "DroneState", planner: "Planner"):
  planner.saveAnimation = True

  statesInFinalTrajectory = [finalDroneState]

  currentDroneState = finalDroneState.parent

  while currentDroneState is not None:
    statesInFinalTrajectory.append(currentDroneState)
    currentDroneState = currentDroneState.parent
  
  numDroneStatesInFinalTrajectory = len(statesInFinalTrajectory)
  fig, ax = plt.subplots()

  for droneStateIdx in range(numDroneStatesInFinalTrajectory-1, 0, -1):
    currentDroneState = statesInFinalTrajectory[droneStateIdx]
    planner.plotStateLive(currentDroneState, ax, fig, False)
    planner.currentAnimationFrame += 1

  planner.plotStateLive(statesInFinalTrajectory[0], ax, fig, True)

  makeAnimationFromTrajectory(planner, dt=planner.dt * 1000, lastFrameDuration=2000)


from datetime import datetime
import math
from typing import List, Tuple
from matplotlib.axes import Axes
from matplotlib.figure import Figure
import matplotlib.pyplot as plt

from planner.cbFunctions import heuristicFunction, cbf
from planner.drone import DroneGoalState, DroneState
from planner.flyZone import FlyZone
from planner.helperFunctions import reverse_insort
from planner.movement import Movement

from planner.obstacle import Obstacle
from planner.plotting import plotPathToCurrentDroneStateLive, plotlive
from position import Position
import os

import glob
from PIL import Image

from trajectory import Trajectory



class Planner:
  CONTROLL_BARRIER_FUNCTIONS = {
    "heuristic": heuristicFunction,
    "cbf": cbf
  }

  HUMAN = Obstacle(0, 0, 0.25)

  def __init__(self, dt: float, obstacles: List["Obstacle"], flyZone: "FlyZone", verboseLevel: int, cbfKey: str) -> None:
    self.dt = dt
    Movement.POSSIBLE_ACCELERATIONS = [acc * dt for acc in Movement.POSSIBLE_ACCELERATIONS]
    self.obstacles = obstacles
    self.flyZone = flyZone
    self.verboseLevel = verboseLevel
    self.cbfKey = cbfKey
    self.cbf = self.CONTROLL_BARRIER_FUNCTIONS[cbfKey]

    self.iterations = 0
    self.currentGoalState = None
    self.currentStartState = None
    self.numIterationsForLastPath = None
    self.movements = []
    self.visitedMovements = set()

     # Animation
    self.saveAnimation = False
    self.lengthOfAnimation = 10
    self.animationFrameDuration = 0.25  # In seconds
    self.numAnimationFrames = self.lengthOfAnimation / self.animationFrameDuration
    self.currentAnimationFrame = 0
    self.animationKey = None


  @plotlive
  def plotStateLive(self, droneState: "DroneState", ax: Axes, fig: Figure, isFinalDroneState: bool) -> None:
    plotPathToCurrentDroneStateLive(
      droneState=droneState, 
      planner=self,
      ax=ax,
      fig=fig,
      isFinalDroneState=isFinalDroneState
    )


  def addMovements(self, associatedDroneState: "DroneState") -> Tuple[bool, "Movement"]:
    for angle in Movement.POSSIBLE_ANGLES:
      for acceleration in Movement.POSSIBLE_ACCELERATIONS:
        newMovement = Movement(
          angle=angle,
          acceleration=acceleration,
          planner=self,
          associatedDroneState=associatedDroneState
        )

        if newMovement.willHaveReachedGoal:
          return True, newMovement
        elif newMovement.isValid and newMovement.key not in self.visitedMovements:
          reverse_insort(self.movements, newMovement) 
          self.visitedMovements.add(newMovement.key)
            
    return False, None


  def printMovements(self) -> None:
    print("\n-Movements-")
    for movement in self.movements:
      print(movement)
    print("---")


  def animationPlot(self, droneState: "DroneState", ax: Axes, fig: Figure, isFinalDroneState:bool=False) -> None:
    if self.numIterationsForLastPath is not None:
      divider = self.numAnimationFrames - 3
      plotAt = self.numIterationsForLastPath // divider
    
      if plotAt == 0.0:
        plotAt = 1
      
      plotNow = self.iterations % plotAt == 0 or self.iterations == 0

      if plotNow or isFinalDroneState:
        self.currentAnimationFrame += 1
        self.plotStateLive(droneState, ax, fig, isFinalDroneState)


  def printStatus(self, droneState: "DroneState") -> None:
    print()
    print(f"#{self.iterations} | C={droneState.selectedMovement.cost}")


  def getYawForInitDroneState(self, initDronePos: "Position", goalState: "DroneGoalState"):
    initYaw = math.atan2(goalState.y - initDronePos.y, goalState.x - initDronePos.x)
    return initYaw


  def getFinalDroneState(self, movementThatReachedGoal: "Movement"):
    finalDroneState = movementThatReachedGoal.associatedDroneState
    finalDroneState.selectedMovement = movementThatReachedGoal

    # finalDroneState = DroneState(penultimateDroneState)
    return finalDroneState


  def findPathToGoal(self, initDroneState: "DroneState", goalState: "DroneGoalState"):
    self.currentGoalState = goalState
    self.currentStartState = initDroneState

    reachedGoal, movementThatReachedGoal = self.addMovements(initDroneState)

    if (self.saveAnimation and self.numIterationsForLastPath is not None) or self.verboseLevel == 1:
      fig, ax = plt.subplots()
    else:
      fig, ax = None, None


    if self.saveAnimation:
      self.animationPlot(initDroneState, ax, fig)

    if self.verboseLevel == 1:
      self.printStatus(associatedDroneState)
      self.plotStateLive(newDroneState, ax, fig, False)

    if reachedGoal:
      finalDroneState = self.getFinalDroneState(movementThatReachedGoal)
      return finalDroneState, True

    while len(self.movements) > 0:
      self.iterations += 1

      cheapestMovement = self.movements.pop()

      associatedDroneState = cheapestMovement.associatedDroneState

      associatedDroneState.selectedMovement = cheapestMovement

      newDroneState = DroneState(associatedDroneState)

      if self.verboseLevel == 1 and self.iterations % 1 == 0 and self.iterations > 20000:
        self.printStatus(associatedDroneState)
        self.plotStateLive(newDroneState, ax, fig, False)
      elif self.verboseLevel == 2 and self.iterations % 1000 == 0:
        self.printStatus(associatedDroneState)
      

      if self.saveAnimation:
        self.animationPlot(newDroneState, ax, fig)

      reachedGoal, movementThatReachedGoal = self.addMovements(newDroneState)

      if reachedGoal:
        finalDroneState = self.getFinalDroneState(movementThatReachedGoal)
          
        if self.saveAnimation:
          self.animationPlot(finalDroneState, ax, fig, True)

        return finalDroneState, True

    return None, False


  def setKey(self, finalDroneState: "DroneState", isFinalTrajectory:bool=False) -> bool:
    startX, startY = round(self.currentStartState.x, 2), round(self.currentStartState.y, 2)
    startPos = f"S({startX},{startY})"

    goalX, goalY = round(self.currentGoalState.x, 2), round(self.currentGoalState.y, 2)
    goalPos = f"G({goalX},{goalY})"
    dateString = str(datetime.now()).split(".")[0]
    self.animationKey = f"{dateString} | T={self.lengthOfAnimation} | fps={round(1/self.animationFrameDuration)} | cbf={self.cbfKey} | {startPos}-->{goalPos} | D={finalDroneState.depth} | i_tot={self.iterations}"
    
    if isFinalTrajectory:
      self.animationKey += " | Final trajectory"


  def resetParams(self):
    self.currentAnimationFrame = 0
    self.iterations = 0
    self.movements = []
    self.visitedMovements = set()


def makeAnimationFromTrajectory(planner: "Planner", dt:float=None, lastFrameDuration:int=5000):
  def getImageNum(imageName: str) -> int:
    return int(imageName.split("/")[-1][:-4])

  if dt is None:
    dt = planner.animationFrameDuration * 1000

  images = glob.glob(f"savedTrajectories/{planner.animationKey}/animationFrames/*.png")

  images.sort(key=getImageNum)

  durations = []
  numImages = len(images)

  for _ in range(numImages-1):
    durations.append(dt)

  durations.append(lastFrameDuration)

  imagePath = f"savedTrajectories/{planner.animationKey}/animation.gif"

  frames = [Image.open(image) for image in images]

  frames[0].save(
    imagePath,
    save_all=True, 
    append_images=frames[1:], 
    optimize=False, 
    duration=durations, 
    loop=0
  )


def recordAnimation(initDronePosition: "Position", planner: "Planner", goalStates: List["DroneGoalState"]):
  planner.saveAnimation = True
  numGoalStates = len(goalStates)
  currentDroneState = initDronePosition

  for goalStateIdx in range(numGoalStates):
    currentGoalState = goalStates[goalStateIdx]
    
    initDroneState = DroneState(
        parent=None, 
        x=currentDroneState.x, 
        y=currentDroneState.y,
        yaw=planner.getYawForInitDroneState(currentDroneState, currentGoalState)
      )

    planner.currentStartState = initDroneState

    for i in range(2):
      finalDroneState, foundGoalState = planner.findPathToGoal(initDroneState, currentGoalState)
      currentDroneState = finalDroneState

      if foundGoalState:
        print("Reached the goal! :)")
        
      else:
        print("No path found :(")
        break

      if i == 0:
        planner.setKey(currentDroneState)

        if planner.saveAnimation:
          os.mkdir(f"animationFrames/{planner.animationKey}")

        planner.numIterationsForLastPath = planner.iterations
        planner.iterations = 0
        planner.movements = []
        planner.visitedMovements = set()
      
    if planner.saveAnimation:
      planner.makeAnimationFromTrajectory()
    
    if goalStateIdx < numGoalStates - 1:
      planner.numIterationsForLastPath = None
      planner.currentAnimationFrame = 0
      planner.iterations = 0
      planner.movements = []
      planner.visitedMovements = set()
      
  
  planner.saveAnimation = False


def recordFinalTrajectory(finalDroneState: "DroneState", planner: "Planner"):
  planner.saveAnimation = True

  statesInFinalTrajectory = [finalDroneState]

  currentDroneState = finalDroneState.parent

  while currentDroneState is not None:
    statesInFinalTrajectory.append(currentDroneState)
    currentDroneState = currentDroneState.parent
  
  numDroneStatesInFinalTrajectory = len(statesInFinalTrajectory)
  fig, ax = plt.subplots()

  for droneStateIdx in range(numDroneStatesInFinalTrajectory-1, 0, -1):
    currentDroneState = statesInFinalTrajectory[droneStateIdx]
    planner.plotStateLive(currentDroneState, ax, fig, False)
    planner.currentAnimationFrame += 1

  planner.plotStateLive(statesInFinalTrajectory[0], ax, fig, True)

  makeAnimationFromTrajectory(planner, dt=planner.dt * 1000, lastFrameDuration=2000)


# Animation # In seconds # finalDroneState = DroneState(penultimateDroneState)