ps_solution.py

import numpy as np
from worlds import Tree, Gridmap, generate_some_worlds, generate_worlds
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
import tikzplotlib as tplt
import cloudpickle
import pickle
from drone_params import Drone
import os
import sys
from scipy import integrate
from scipy.interpolate import lagrange
from scipy.interpolate import CubicSpline


class PsSolution:
    """
    The PsSolution class is used to evaluate a solution generated by PsControl class.
    """

    def __init__(self):

        self.n_simpson = 10
        self.state = []
        self.control = []
        self.time = []
        self.obstacles = []
        self.n_col = []
        self.n_seg = 0
        self.sampled = False
        self.relative_error_max = None
        self.relative_deflection = None
        self.max_error = None
        self.time_error = False
        self.time_eq_error = False
        self.objective = -1

    def new_solution(self, m, drone=None, obs=None, fit='polyfit', n_simpson=10, fix_solution = False):
        """
        Generate solution based on data from solved Pyomo model, Drone class with parametes in world obs.
        
        Parameters
        ----------
        m (object): The Pyomo model.
        drone (object): The drone class with parameters and dynamics.
        obs (object): The world class with parameters in world obs.
        fit (str): The type of fitting to use. Default is 'polyfit'.
            Supported values are 'polyfit', 'lagrange' and 'spline'.
        n_simpson (int): The number of points to use for the Simpson integration. Default is 10.
        fix_solution (bool): Whether to fix the solution. !not fully implemented! Default is False.
        """

        # set parameters
        self.fit = fit
        self.n_simpson = int(n_simpson)
        self.n_col = m.n_col
        self.n_seg = len(self.n_col)

        # load drone
        if drone is None:
            from drone_params import Drone
            self.drone = Drone
        else:
            self.drone = drone
            
        # import drone dynamics for evaluation
        self.__eval_dynamics = drone.eval_dynamics
        
        if hasattr(m, 'psm_approx'):    
            self.psm_approx = m.psm_approx
        else:
            self.psm_approx = 'chebyshev'

        # extract solving time
        if hasattr(m, 'presolve_time'):
            self.presolve_time = m.presolve_time
            self.solve_time = m.solve_time
            self.psm_solve_time = self.presolve_time + self.solve_time
        else:
            self.presolve_time = 0
            self.solve_time = 0
            self.psm_solve_time = 0

        self.objective = m.obj.expr()  # extract objective value

        # load obstacles
        self.obs = obs
        if obs is not None and self.obs.n_obstacles is None:
            if self.obs.dim == 3:
                self.obs.n_obstacles = self.obs.map3d.shape[0]
            elif self.obs.dim == 2:
                self.obs.n_obstacles = self.obs.map2d.shape[0]

        # dimension of optimal control variables
        self.nx = m.nx.at(-1)+1
        self.nu = m.nu.at(-1)+1

        # initialize optimal control variables in collocation points
        self.time = np.zeros(m.n_col_sum)
        self.state = np.zeros((self.nx, m.n_col_sum))
        self.control = np.zeros((self.nu, m.n_col_sum))

        # calculate gradual number of collocation points
        self.n_col_add = np.zeros(self.n_col.shape[0]+1, dtype=int)
        for i, cols in enumerate(self.n_col):
            self.n_col_add[i+1] = self.n_col_add[i] + cols

        # load time in collocation points from solution and check the time continuity
        for i in m.t:
            self.time[i] = m.t[i].value
            if i > 0:
                if self.time[i] < self.time[i-1]:
                    print(f"[ERROR] Time continuum violation at {i}")
                    self.time_error = True
                elif self.time[i] == self.time[i-1] and i not in self.n_col_add:
                    print(f"[WARNING] Time continuum equility at {i}")
                    self.time_eq_error = True

        # load state in collocation points
        for i in m.x:
            self.state[i] = m.x[i].value

        # load control in collocation points
        for i in m.u:
            self.control[i] = m.u[i].value

        # build list with start and end time of each segment
        self.multiseg_time = []
        for i, seg_start in enumerate(self.n_col_add[:-1]):
            self.multiseg_time.append(self.time[seg_start:self.n_col_add[i+1]])
            
        if (self.time_error or self.time_eq_error) and fix_solution:
            self.fix_solution()
        
    def fix_solution(self):
        """
        Fix the solution by checking for errors in collocation points and
        fitting the solution to the original number of collocation points.
        """
        
        # copy original solution
        time = self.time.copy()
        state = self.state.copy()
        control = self.control.copy()
        n_col = self.n_col.copy()
        n_col_old = self.n_col.copy()
        n_col_add = self.n_col_add.copy()
        n_col_add_old = self.n_col_add.copy()
        deleted_segment = []
        
        # search and delete errors
        i = 1
        len_time = len(time)
        while i < len_time:
            if (time[i] < time[i-1] or time[i] == time[i-1]) and i not in n_col_add:
                # delete wrong point
                print(f"Fixing point {i}")
                time = np.delete(time, i)
                state = np.delete(state, i, axis=1)
                control = np.delete(control, i, axis=1)
                # find segment number
                if n_col_add[0] > i:
                    seg = 0
                else:
                    seg = np.where(n_col_add <= i)[0][-1]
                n_col[seg] = n_col[seg]-1
                if n_col[seg] <= 0:
                    # delete segment without good collocation points
                    # TODO: test segment deletion: 
                    raise Exception("Whole segment {seg} is wrong.")
                    #  new multiseg_time won't fit to old collocation scheme
                    deleted_segment.append([seg, n_col[seg], self.multiseg_time[seg][0], self.multiseg_time[seg][-1]])
                    n_col_old = np.delete(n_col_old, seg)
                    
                    if seg == 0:
                        n_col_old[0] = n_col_old[0] + n_col[seg]
                    else:
                        n_col_old[seg-1] = n_col_old[seg-1] + n_col[seg]
                    n_col = np.delete(n_col, seg)
                    n_col_add = np.delete(n_col_add, seg)
                i = i-1
                len_time = len_time-1
            else:
                i = i+1
        
        self.time = time
        self.state = state
        self.control = control
        self.n_col = n_col
        # calculate gradual number of collocation points
        self.n_col_add = np.zeros(self.n_col.shape[0]+1, dtype=int)
        for i, cols in enumerate(self.n_col):
            self.n_col_add[i+1] = self.n_col_add[i] + cols
        # get t0 and tf for each segment
        self.multiseg_time = []
        for i, seg_start in enumerate(self.n_col_add[:-1]):
            self.multiseg_time.append(self.time[seg_start:self.n_col_add[i+1]])
        
        if self.psm_approx == 'chebyshev':
            from chebyshev import cheb_scaled
        elif self.psm_approx == 'legendre':
            from legendre import legendre_scaled as cheb_scaled
            
        # fit solution to original number of collocation points
        new_multiseg_time = []
        for i, cols in enumerate(n_col_old):
            tt, __, __ = cheb_scaled(cols-1, [self.multiseg_time[i][0],self.multiseg_time[i][-1]])
            new_multiseg_time.append(tt)
        self.time, self.state, self.control = self.resample_solution(new_multiseg_time)
        self.n_col = n_col_old
        self.n_col_add = n_col_add_old
        self.multiseg_time = new_multiseg_time
        
        self.time_error = False
        self.time_eq_error = False
        

    def plot(self, save=False, img_folder="img/", dpi = 600, figsize = (6, 4)):
        """
        Plots the collocation points, state and control of the drone over time. 

        Args:
            save (bool): Whether to save the plots as image files or not.
            img_folder (str): The folder path where the image files should be saved.
        """

        # time vs collocation points
        plt.figure(figsize=figsize)
        plt.plot(self.time, '.')
        plt.xlabel('collocation points [k]')
        plt.ylabel('$t$ [s]')
        if save:
            plt.savefig(img_folder+'time.eps')
            plt.savefig(img_folder+'time.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'time.tex')

        # position
        plt.figure(figsize=figsize)
        for i in range(0, 3):
            plt.plot(self.time, self.state[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$r$ [m]')
        plt.legend(['$r_x$', '$r_y$', '$r_z$'])
        if save:
            plt.savefig(img_folder+'x_position.eps')
            plt.savefig(img_folder+'x_position.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_position.tex')

        # speed
        plt.figure(figsize=figsize)
        for i in range(3, 6):
            plt.plot(self.time, self.state[i, :])
        plt.xlabel('t[s]')
        plt.ylabel('$v$ [m/s]')
        plt.legend(['$v_x$', '$v_y$', '$v_z$'])
        if save:
            plt.savefig(img_folder+'x_velocity.eps')
            plt.savefig(img_folder+'x_velocity.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_velocity.tex')

        # quaternion
        plt.figure(figsize=figsize)
        for i in range(6, 10):
            plt.plot(self.time, self.state[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$q$')
        plt.legend(['$q_w$', '$q_x$', '$q_y$', '$q_z$'])
        if save:
            plt.savefig(img_folder+'x_quaternion.eps')
            plt.savefig(img_folder+'x_quaternion.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_quaternion.tex')

        # angular rate
        plt.figure(figsize=figsize)
        for i in range(10, 13):
            plt.plot(self.time, self.state[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$\omega$ [rad/s]')
        plt.legend(['$\omega_x$', '$\omega_y$', '$\omega_z$'])
        if save:
            plt.savefig(img_folder+'x_angular_rate.eps')
            plt.savefig(img_folder+'x_angular_rate.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_angular_rate.tex')

        # thrust
        plt.figure(figsize=figsize)
        plt.plot(self.time, self.control[0, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$T$ [N]')
        if save:
            plt.savefig(img_folder+'u_thrust.eps')
            plt.savefig(img_folder+'u_thrust.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'u_thrust.tex')

        # control
        plt.figure(figsize=figsize)
        for i in range(1, 4):
            plt.plot(self.time, self.control[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel(r'$\tau$ [N$\cdot$m]')
        plt.legend([r'$\tau_x$', r'$\tau_y$', r'$\tau_z$'])
        if save:
            plt.savefig(img_folder+'u_torque.eps')
            plt.savefig(img_folder+'u_torque.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'u_torque.tex')

        fig = plt.figure(figsize=figsize)
        fig.set_figwidth(8)
        fig.set_figheight(8)
        if self.obs.dim == 3:
            ax = plt.axes(projection="3d")
            #plt.title('3D grid map')
            ax.scatter3D(self.obs.map3d[:, 0],
                         self.obs.map3d[:, 1],
                         self.obs.map3d[:, 2],
                         marker='s',
                         linewidths=self.obs.space**2)
            ax.set_zlabel("$z$ [m]")
            ax.scatter3D(self.drone.x0[0],
                         self.drone.x0[1],
                         self.drone.x0[2],
                         marker='o',
                         linewidths=self.obs.space,
                         color='green')
            ax.scatter3D(self.drone.xf[0],
                         self.drone.xf[1],
                         self.drone.xf[2],
                         marker='o',
                         linewidths=self.obs.space,
                         color='red')
            ax.plot3D(self.state[0, :],
                      self.state[1, :],
                      self.state[2, :],
                      color='orange')
        else:
            ax = plt.axes()
            ax.scatter(self.obs.map2d[:, 0],
                       self.obs.map2d[:, 1],
                       marker='s',
                       linewidths=self.obs.space**2, color='blue')
            ax.scatter(self.drone.x0[0],
                       self.drone.x0[1],
                       marker='o',
                       linewidths=self.obs.space,
                       color='green')
            ax.scatter(self.drone.xf[0],
                       self.drone.xf[1],
                       marker='o',
                       linewidths=self.obs.space,
                       color='red')
            ax.plot(self.state[0, :],
                    self.state[1, :],
                    color='orange')
        ax.set_xlabel("$x$ [m]")
        ax.set_ylabel("$y$ [m]")

        fig.subplots_adjust(left=0, bottom=0, right=1,
                            top=1, wspace=0, hspace=0)
        if save:
            fig.savefig(img_folder+'3d_position.eps')
            plt.savefig(img_folder+'3d_position.png', dpi=dpi, bbox_inches="tight")
            # 3D plots DO NOT WORK in tikzplotlib!!!

        plt.show()

    def plot_sampled(self, save=False, img_folder="img/", dt=0.01, dpi = 600, figsize = (6,4)):
        """
        Plots trajectories resampled according to sampling parameter dt[s].

        Args:
            save (bool): Whether to save the plots as image files or not.
            img_folder (str): The folder path where the image files should be saved.
            dt (float): Sampling parameter

        Returns:
            None
        """
        time_array = [np.arange(self.time[0], self.time[-1], dt)]
        self.__get_polynomial_solution()
        time_sampled, state_sampled, control_sampled = self.resample_solution(
            time_array)
        time_sampled, state_sampled, control_sampled = time_sampled[
            0], state_sampled[0], control_sampled[0]
        # time vs collocation points
        plt.figure(figsize=figsize)
        plt.plot(self.time, '.')
        plt.xlabel('collocation points [k]')
        plt.ylabel('$t$ [s]')
        if save:
            plt.savefig(img_folder+'time.eps')
            plt.savefig(img_folder+'time.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'time.tex')

        # position
        plt.figure(figsize=figsize)
        for i in range(0, 3):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$r$ [m]')
        plt.legend(['$r_x$', '$r_y$', '$r_z$'])
        if save:
            plt.savefig(img_folder+'x_position.eps')
            plt.savefig(img_folder+'x_position.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_position.tex')

        # speed
        plt.figure(figsize=figsize)
        for i in range(3, 6):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$v[m/s]$')
        plt.legend(['$v_x$', '$v_y$', '$v_z$'])
        if save:
            plt.savefig(img_folder+'x_velocity.eps')
            plt.savefig(img_folder+'x_velocity.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_velocity.tex')

        # quaternion
        plt.figure(figsize=figsize)
        for i in range(6, 10):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$q$')
        plt.legend(['$q_w$', '$q_x$', '$q_y$', '$q_z$'])
        if save:
            plt.savefig(img_folder+'x_quaternion.eps')
            plt.savefig(img_folder+'x_quaternion.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_quaternion.tex')

        # angular rate
        plt.figure(figsize=figsize)
        for i in range(10, 13):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$\omega$ [rad/s]')
        plt.legend(['$\omega_x$', '$\omega_y$', '$\omega_z$'])
        if save:
            plt.savefig(img_folder+'x_angular_rate.eps')
            plt.savefig(img_folder+'x_angular_rate.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_angular_rate.tex')

        # thrust
        plt.figure(figsize=figsize)
        plt.plot(time_sampled, control_sampled[0, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$T$ [N]')
        if save:
            plt.savefig(img_folder+'u_thrust.eps')
            plt.savefig(img_folder+'u_thrust.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'u_thrust.tex')

        # control
        plt.figure(figsize=figsize)
        for i in range(1, 4):
            plt.plot(time_sampled, control_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel(r'$\tau$ [N$\cdot$m]')
        plt.legend([r'$\tau_x$', r'$\tau_y$', r'$\tau_z$'])
        if save:
            plt.savefig(img_folder+'u_torque.eps')
            plt.savefig(img_folder+'u_torque.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'u_torque.tex')

        fig = plt.figure(figsize=figsize)
        fig.set_figwidth(8)
        fig.set_figheight(8)
        if self.obs.dim == 3:
            ax = plt.axes(projection="3d")
            #plt.title('3D grid map')
            ax.scatter3D(self.obs.map3d[:, 0],
                         self.obs.map3d[:, 1],
                         self.obs.map3d[:, 2],
                         marker='s',
                         linewidths=self.obs.space**2)
            ax.set_zlabel("$z$ [m]")
            ax.scatter3D(self.drone.x0[0],
                         self.drone.x0[1],
                         self.drone.x0[2],
                         marker='o',
                         linewidths=self.obs.space,
                         color='green')
            ax.scatter3D(self.drone.xf[0],
                         self.drone.xf[1],
                         self.drone.xf[2],
                         marker='o',
                         linewidths=self.obs.space,
                         color='red')
            ax.plot3D(state_sampled[0, :],
                      state_sampled[1, :],
                      state_sampled[2, :],
                      color='orange')
        else:
            ax = plt.axes()
            ax.scatter(self.obs.map2d[:, 0],
                       self.obs.map2d[:, 1],
                       marker='s',
                       linewidths=self.obs.space**2)
            ax.scatter(self.drone.x0[0],
                       self.drone.x0[1],
                       marker='o',
                       linewidths=self.obs.space,
                       color='green')
            ax.scatter(self.drone.xf[0],
                       self.drone.xf[1],
                       marker='o',
                       linewidths=self.obs.space,
                       color='red')
            ax.plot(self.state[0, :],
                    self.state[1, :],
                    color='orange')
        ax.set_xlabel("$x$ [m]")
        ax.set_ylabel("$y$ [m]")

        fig.subplots_adjust(left=0, bottom=0, right=1,
                            top=1, wspace=0, hspace=0)
        if save:
            fig.savefig(img_folder+'3d_position.eps')
            plt.savefig(img_folder+'3d_position.png', dpi=dpi, bbox_inches="tight")
            # 3D plots DO NOT WORK in tikzplotlib!!!

        plt.show()
        
    def plot_sampled_with_col(self, save=False, img_folder="img/", dt=0.01, dpi = 600, figsize = (6,4)):
        """
        Plots trajectories resampled according to sampling parameter dt[s] with collocation points.

        Args:
            save (bool): Whether to save the plots as image files or not.
            img_folder (str): The folder path where the image files should be saved.
            dt (float): Sampling parameter

        Returns:
            None
        """
        # Get the default property cycle
        prop_cycle = plt.rcParams['axes.prop_cycle']
        default_colors = prop_cycle.by_key()['color']
        
        time_array = [np.arange(self.time[0], self.time[-1], dt)]
        self.__get_polynomial_solution()
        time_sampled, state_sampled, control_sampled = self.resample_solution(
            time_array)
        time_sampled, state_sampled, control_sampled = time_sampled[
            0], state_sampled[0], control_sampled[0]
        # time vs collocation points
        plt.figure(figsize=figsize)
        plt.plot(self.time,np.zeros(len(self.time)), '.')
        plt.ylabel('placement of collocation points [k]')
        plt.xlabel('$t$ [s]')
        if save:
            plt.savefig(img_folder+'time.eps')
            plt.savefig(img_folder+'time.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'time.tex')

        # position
        plt.figure(figsize=figsize)
        for i in range(0, 3):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$r$ [m]')
        plt.legend(['$r_x$', '$r_y$', '$r_z$'])
        for i in range(0, 3):
            plt.plot(self.time, self.state[i, :],'o', color = default_colors[i])
        if save:
            plt.savefig(img_folder+'x_position.eps')
            plt.savefig(img_folder+'x_position.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_position.tex')

        # speed
        plt.figure(figsize=figsize)
        for i in range(3, 6):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$v$ [m/s]')
        plt.legend(['$v_x$', '$v_y$', '$v_z$'])
        for i in range(3, 6):
            plt.plot(self.time, self.state[i, :],'o', color = default_colors[i-3])
        if save:
            plt.savefig(img_folder+'x_velocity.eps')
            plt.savefig(img_folder+'x_velocity.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_velocity.tex')

        # quaternion
        plt.figure(figsize=figsize)
        for i in range(7, 10):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$q$')
        plt.legend(['$q_x$', '$q_y$', '$q_z$'])
        for i in range(7, 10):
            plt.plot(self.time, self.state[i, :],'o', color = default_colors[i-7])
        if save:
            plt.savefig(img_folder+'x_quaternion.eps')
            plt.savefig(img_folder+'x_quaternion.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_quaternion.tex')

        # angular rate
        plt.figure(figsize=figsize)
        for i in range(10, 13):
            plt.plot(time_sampled, state_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$\omega$ [rad/s]')
        plt.legend(['$\omega_x$', '$\omega_y$', '$\omega_z$'])
        for i in range(10, 13):
            plt.plot(self.time, self.state[i, :],'o', color = default_colors[i-10])
        if save:
            plt.savefig(img_folder+'x_angular_rate.eps')
            plt.savefig(img_folder+'x_angular_rate.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'x_angular_rate.tex')

        # thrust
        plt.figure(figsize=figsize)
        plt.plot(time_sampled, control_sampled[0, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel('$T$ [N]')
        plt.plot(self.time, self.control[0, :],'o', color = default_colors[0])
        if save:
            plt.savefig(img_folder+'u_thrust.eps')
            plt.savefig(img_folder+'u_thrust.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'u_thrust.tex')

        # control
        plt.figure(figsize=figsize)
        for i in range(1, 4):
            plt.plot(time_sampled, control_sampled[i, :])
        plt.xlabel('$t$ [s]')
        plt.ylabel(r'$\tau$ [N$\cdot$m]')
        plt.legend([r'$\tau_x$', r'$\tau_y$', r'$\tau_z$'])
        for i in range(1, 4):
            plt.plot(self.time, self.control[i, :],'o', color = default_colors[i-1])
        if save:
            plt.savefig(img_folder+'u_torque.eps')
            plt.savefig(img_folder+'u_torque.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'u_torque.tex')

        fig = plt.figure(figsize=(figsize[0],figsize[0]))
        if self.obs.dim == 3:
            ax = plt.axes(projection="3d")
            self.plot_3d_obstacles(ax, self.obs.map3d, self.obs.space*np.sqrt(3)/2+self.drone.safe_radius)
            #plt.title('3D grid map')
            ax.scatter3D(self.obs.map3d[:, 0],
                         self.obs.map3d[:, 1],
                         self.obs.map3d[:, 2],
                         marker='s',
                         linewidths=self.obs.space**2)
            ax.set_zlabel("$z$ [m]")
            ax.scatter3D(self.drone.x0[0],
                         self.drone.x0[1],
                         self.drone.x0[2],
                         marker='o',
                         linewidths=self.obs.space,
                         color='green')
            ax.scatter3D(self.drone.xf[0],
                         self.drone.xf[1],
                         self.drone.xf[2],
                         marker='o',
                         linewidths=self.obs.space,
                         color='red')
            ax.plot3D(state_sampled[0, :],
                      state_sampled[1, :],
                      state_sampled[2, :],
                      color='orange')
            ax.plot3D(self.state[0, :],
                      self.state[1, :],
                      self.state[2, :],
                      marker='o', linestyle='None', color='orange')
            ax.set_box_aspect([1,1,1])  # Equal aspect ratio
            # Set the limits of the axes to fit the range of your data
            xmin, xmax = np.min(self.obs.map3d[:, 0]), np.max(self.obs.map3d[:, 0])
            ymin, ymax = np.min(self.obs.map3d[:, 1]), np.max(self.obs.map3d[:, 1])
            zmin, zmax = np.min(self.obs.map3d[:, 2]), np.max(self.obs.map3d[:, 2])
            ax.set_xlim([xmin, xmax])
            ax.set_ylim([ymin, ymax])
            ax.set_zlim([zmin, zmax])
        else:
            ax = plt.axes()
            # Plot obstacles as circles
            for obs_pos in self.obs.map2d:
                circle = Circle(obs_pos, self.obs.space*np.sqrt(2)/2+self.drone.safe_radius,
                                color='blue', fill=False)  # Create a circle around the obstacle
                ax.add_patch(circle)  # Add the circle to the axes
            ax.scatter(self.obs.map2d[:, 0],
                       self.obs.map2d[:, 1],
                       marker='s',
                       linewidths=self.obs.space**2)
            ax.scatter(self.drone.x0[0],
                       self.drone.x0[1],
                       marker='o',
                       linewidths=self.obs.space,
                       color='green')
            ax.scatter(self.drone.xf[0],
                       self.drone.xf[1],
                       marker='o',
                       linewidths=self.obs.space,
                       color='red')
            ax.plot(state_sampled[0, :],
                    state_sampled[1, :],
                    color='orange')
            ax.plot(self.state[0, :],
                    self.state[1, :],
                    marker='o', linestyle='None', color='orange')
        ax.set_xlabel("$x$ [m]")
        ax.set_ylabel("$y$ [m]")

        # fig.subplots_adjust(left=0, bottom=0, right=1,
        #                     top=1, wspace=0, hspace=0)
        if save:
            fig.savefig(img_folder+'3d_position.eps')
            plt.savefig(img_folder+'3d_position.png', dpi=dpi, bbox_inches="tight")
            # 3D plots DO NOT WORK in tikzplotlib!!!

        plt.show()
        
    def plot_3d_obstacles(self, ax, centers, radius):
        """
        Plots 3D spherical obstacles given centers and radius.

        Args:
        - ax: Axes3D object to plot on.
        - centers: numpy array of shape (N, 3) containing obstacle centers.
        - radius: Radius of the spherical obstacles.
        """
        u = np.linspace(0, 2 * np.pi, 30)
        v = np.linspace(0, np.pi, 30)
        x = radius * np.outer(np.cos(u), np.sin(v))
        y = radius * np.outer(np.sin(u), np.sin(v))
        z = radius * np.outer(np.ones(np.size(u)), np.cos(v))
        
        for center in centers:
            ax.plot_surface(x + center[0], y + center[1], z + center[2], color='b', alpha=0.3)
        
    def get_sampled_trajectory(self, dt = 0.01):
        """Get sampled trajectory"""
        
        time_array = [np.arange(self.time[0], self.time[-1], dt)]
        self.__get_polynomial_solution()
        time_sampled, state_sampled, control_sampled = self.resample_solution(
            time_array)
        time_sampled, state_sampled, control_sampled = time_sampled[
            0], state_sampled[0], control_sampled[0]
        
        return time_sampled, state_sampled, control_sampled

    def __get_polynomial(self, y, t, n):
        """Return polynomial of n-th order"""
        return np.polynomial.Polynomial.fit(t, y, n-1)
    
    def __get_polynomial_lagrange(self, y, t, n):
        """Return polynomial of n-th order"""
        return lagrange(t, y)
    
    def __get_spline(self, y, t, n):
        """Return spline"""
        return CubicSpline(t, y, n)

    def __get_polynomial_cheby(self, y, t, n):
        """Return chebyshev polynomial of n-th order"""
        # TODO: rescaling to <-1,1> needed due to numerical stability
        return np.polynomial.chebyshev.Chebyshev.fit(t, y, n-1)

    def __get_polynomials(self, y, t, n_col):
        """Return array of polynomials for each segment"""
        p = []
        n_col_add = np.zeros(n_col.shape[0]+1, dtype=int)
        for i, cols in enumerate(n_col):
            n_col_add[i+1] = n_col_add[i] + cols
            if self.fit == 'cheby':
                p.append(self.__get_polynomial_cheby(
                    y[n_col_add[i]:n_col_add[i+1]], t[n_col_add[i]:n_col_add[i+1]], cols))
            elif self.fit == 'lagrange':
                p.append(self.__get_polynomial_lagrange(
                    y[n_col_add[i]:n_col_add[i+1]], t[n_col_add[i]:n_col_add[i+1]], cols))
            elif self.fit == 'spline':
                p.append(self.__get_spline(
                    y[n_col_add[i]:n_col_add[i+1]], t[n_col_add[i]:n_col_add[i+1]], cols))
            else:
                p.append(self.__get_polynomial(
                    y[n_col_add[i]:n_col_add[i+1]], t[n_col_add[i]:n_col_add[i+1]], cols))
        return p

    def __get_n_polynomials(self, Y, t, n_col):
        """Returns array of N functions Y[N,t] for each segment """
        pY = []
        for y in Y:
            pY.append(self.__get_polynomials(y, t, n_col))
        return pY

    def __get_polynomial_solution(self):
        """Calculate polynomial representation of solution"""

        if not hasattr(self, 'state_poly'):
            self.state_poly = self.__get_n_polynomials(
                self.state, self.time, self.n_col)
        if not hasattr(self, 'control_poly'):
            self.control_poly = self.__get_n_polynomials(
                self.control, self.time, self.n_col)

        return

    def __sample_solution(self):
        """ 
        This function generates a solution sampled on the selected n_simpson points.
        It generates time, state, state derivative, and control samples for each segment.
        The generated samples are stored as attributes of the class.
        """
        # Generate time samples for each segment
        t_sample = []
        for t in self.multiseg_time:
            t_sample_seg = []
            for i in range(t.shape[0]-1):
                t_sample_seg.append(np.linspace(
                    t[i], t[i+1], self.n_simpson))
            t_sample.append(t_sample_seg)

        # Generate state and state derivative samples for each segment
        state_sample = []
        state_derivative_sample = []
        for p_array in self.state_poly:
            # go though state elements
            state_sample_seg = []
            state_derivative_sample_seg = []
            for i_seg, p in enumerate(p_array):
                # go through segments
                for t in t_sample[i_seg]:
                    # go through collocation points neghborhood
                    state_sample_seg.append(p(t))
                    if self.fit == 'spline':
                        state_derivative_sample_seg.append(p.derivative(1)(t))
                    else:
                        state_derivative_sample_seg.append(p.deriv(1)(t))
            state_sample.append(state_sample_seg)
            state_derivative_sample.append(state_derivative_sample_seg)

        # Generate control samples for each segment
        control_sample = []
        for p_array in self.control_poly:
            control_sample_seg = []
            for i_seg, p in enumerate(p_array):
                for t in t_sample[i_seg]:
                    control_sample_seg.append(p(t))
            control_sample.append(control_sample_seg)

        # Save generated samples as attributes of the class
        self.__state_sample = state_sample
        self.__control_sample = control_sample
        self.__t_sample = t_sample
        self.__state_derivative_sample = state_derivative_sample

        self.sampled = True

    def resample_solution(self, new_multiseg_time):
        """
        Resamples the solution at the given times.

        Args:
            new_multiseg_time (list): A list of arrays, where each array contains 
                                    the new times for each segment.

        Returns:
            tuple: A tuple containing:
            - time_sample (array): The new time array.
            - state_sample (list): A list of arrays, where each array contains the 
                                state samples for each segment.
            - control_sample (list): A list of arrays, where each array contains 
                                    the control samples for each segment.
        """
        if not hasattr('self','state_poly'):
            self.__get_polynomial_solution()
        
        time_sample = new_multiseg_time

        # Adherence to time continuity
        for i in range(len(time_sample[:-1])):
            time_sample[i+1][0] = time_sample[i][-1]

        # Generate state and state derivative samples for each segment
        state_sample = []
        control_sample = []
        i = 0
        for new_seg_time in new_multiseg_time:
            n = len(new_seg_time)
            state_sample_seg = np.zeros([self.nx, int(n)])
            control_sample_seg = np.zeros([self.nu, int(n)])

            for j, t_new in enumerate(new_seg_time):
                while t_new > self.multiseg_time[i][-1] and i < self.n_seg-1:
                    # check if time instant is within current segment
                    i += 1
                for k, p in enumerate(self.state_poly):
                    state_sample_seg[k, j] = p[i](t_new)
                for k, p in enumerate(self.control_poly):
                    control_sample_seg[k, j] = p[i](t_new)

            state_sample.append(state_sample_seg)
            control_sample.append(control_sample_seg)

        return time_sample, state_sample, control_sample

    def evaluate_solution(self, n_simpson = 10):
        """ 
        This function evaluates the dynamics error of the solution.
        It generates the sample solution using the polynomial solutions
        obtained from __get_polynomial_solution() function.
        """
        # Set the number of poinst for Simpson's rule
        self.n_simpson = n_simpson
        self.sampled = False
        # Sample the solution
        self.__sample_solution()
        # Fit solution for state and control as polynomial
        self.__get_polynomial_solution()
        # Calculate the dynamics error
        self.__dynamics_error()
        # Calculate the adherence to constraints
        self.__constraints_error()

    def get_deflection(self):
        """
        Calculates the deflection of the local relative error based on the relative error values
        stored in `self.relative_error_max`.
        Returns a list of deflection values for each segment in `self.relative_error_max`.
        """
        # Ensure that the relative error is calculated
        if self.relative_error_max is None:
            self.get_relative_error()

        deflection = []
        # Calculate deflection from self.relative_error_max
        for relative_error_max_seg in self.relative_error_max:
            # Convert the list to a NumPy array
            relative_error_max_seg = np.array(relative_error_max_seg)

            # Calculate the difference between consecutive local errors
            differences = np.diff(relative_error_max_seg)

            # Compute the absolute value of the differences to obtain the deflection values
            deflection_seg = np.abs(differences)

            deflection.append(deflection_seg.tolist())

        self.relative_deflection = deflection

        return deflection

    def plot_deflection(self, figsize=(6,4), save = False, img_folder = 'img/', dpi = 600):
        """
        Plots the deflection of the local relative error based on the deflection values
        stored in `self.deflection` and the original solution time in `self.time`.
        """
        # Ensure that the deflection is calculated
        if self.relative_deflection is None:
            self.get_deflection()

        # Flatten the deflection list and add None between segments to separate them
        flat_deflection = []
        for i, deflection_seg in enumerate(self.relative_deflection):
            flat_deflection.extend(deflection_seg)
            if i < len(self.relative_deflection) - 1:
                flat_deflection.append(None)

        # Flatten the time list and add None between segments to separate them
        flat_time = np.delete(self.time, self.n_col_add[1:]-1)

        # Plot the deflection values
        plt.figure(figsize=figsize)
        plt.step(flat_time[:-1], flat_deflection, linestyle='-', where='pre')
        plt.xlabel('$t$ [s]')
        plt.ylabel('$\Delta\epsilon_{r\max,i}$')
        if save:
            plt.savefig(img_folder+'error_deflection.eps')
            plt.savefig(img_folder+'error_deflection.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'error_deflection.tex')
        plt.show()

    def plot_relative_error_max(self, figsize=(6,4), save = False, img_folder = 'img/', dpi = 600):
        """
        Plots the maximum relative error based on the values stored in `self.relative_error_max`
        and the original solution time in `self.time`.
        """
        # Ensure that the relative error is calculated
        if self.relative_error_max is None:
            self.get_relative_error()

        # Flatten the relative_error_max list and add None between segments to separate them
        flat_relative_error_max = []
        for i, relative_error_max_seg in enumerate(self.relative_error_max):
            flat_relative_error_max.extend(relative_error_max_seg)
            if i < len(self.relative_error_max) - 1:
                flat_relative_error_max.append(self.relative_error_max[i+1][0])
            #     flat_relative_error_max.append(None)
            #     pass

        # Flatten the time list and add None between segments to separate them
        flat_time = self.time.copy()
        plt.figure(figsize=figsize)
        # Plot the maximum relative error values
        plt.step(flat_time[:-1], flat_relative_error_max, linestyle='-', where='pre')
        plt.xlabel('$t$ [s]')
        plt.ylabel('$\epsilon_{r\max,i}$')
        if save:
            plt.savefig(img_folder+'error_relative.eps')
            plt.savefig(img_folder+'error_relative.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'error_relative.tex')
        plt.show()

    def plot_absolute_error(self, figsize=(6,4), save = False, img_folder = 'img/', dpi = 600):
        """
        Plots the maximum error based on the values stored in `self.max_error`
        and the original solution time in `self.time`.
        """
        # Ensure that the maximum error is calculated
        if self.max_error is None:
            self.get_absolute_error()

        # Flatten the max_error list and add None between segments to separate them
        flat_max_error = []
        for i, max_error_seg in enumerate(self.max_error):
            flat_max_error.extend(max_error_seg)
            if i < len(self.max_error) - 1:
                # flat_max_error.append(None)
                flat_max_error.append(self.max_error[i+1][0])

        # Flatten the time list and add None between segments to separate them
        flat_time = self.time.copy()

        # Plot the maximum error values
        plt.figure(figsize=figsize)
        plt.step(flat_time[:-1], flat_max_error, linestyle='-', where='pre')
        plt.yscale('log')
        plt.xlabel('$t$ [s]')
        plt.ylabel('$\epsilon_{a\max,i}$')
        if save:
            plt.savefig(img_folder+'error_absolute.eps')
            plt.savefig(img_folder+'error_absolute.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'error_absolute.tex')
        plt.show()
        
    def plot_obstacles_error(self, figsize=(6,4), save = False, img_folder = 'img/', dpi = 600):
        """
        Plots the collisions with obstacles measure based on the values stored in `self.max_error`
        and the original solution time in `self.time`.
        """
        # Ensure that the maximum error is calculated
        if self.integral_obstacle_error is None:
            self.get_constraints_error()

        # Flatten the max_error list and add None between segments to separate them
        flat_obstacle_error = []
        for i, integral_obstacle_error_seg in enumerate(self.integral_obstacle_error):
            flat_obstacle_error.extend(integral_obstacle_error_seg)
            if i < len(self.integral_obstacle_error) - 1:
                # flat_max_error.append(None)
                flat_obstacle_error.append(self.integral_obstacle_error[i+1][0])

        # Flatten the time list and add None between segments to separate them
        flat_time = self.time.copy()

        # Plot the maximum error values
        plt.figure(figsize=figsize)
        plt.step(flat_time[:-1], flat_obstacle_error, linestyle='-', where='pre')
        # plt.yscale('log')
        plt.xlabel('$t$ [s]')
        plt.ylabel('Collisions with Obstacles')
        if save:
            plt.savefig(img_folder+'error_obstacle.eps')
            plt.savefig(img_folder+'error_obstacle.png', dpi=dpi, bbox_inches="tight")
            tplt.save(img_folder+'error_obstacle.tex')
        plt.show()
        
    def get_sampled_state_np(self):
        
         # Ensure that the solution is sampled
        if not self.sampled:
            self.__sample_solution()

        # Flatten the max_error list and add None between segments to separate them
        flat_state = []
        for j, state_element in enumerate(self.__state_sample):
            flat_state_element = []
            for i, state_seg in enumerate(state_element):
                flat_state_element.extend(state_seg)
                # if i < len(self.integral_obstacle_error) - 1:
                    # flat_max_error.append(None)
                    # flat_state_element.append(self.__state_sample[i+1][0])
            flat_state.append(flat_state_element)
            
        # Flatten the time list and add None between segments to separate them
        state_sample_np = np.array(flat_state)
        return state_sample_np
        
    def plot_sampled_path(self, figsize=(6,4)):
        """
        Plots the collisions with obstacles measure based on the values stored in `self.max_error`
        and the original solution time in `self.time`.
        """
        
        state_sample_np = self.get_sampled_state_np()
        # Plot the maximum error values
        plt.figure(figsize=figsize)
        plt.plot(state_sample_np[0,:], state_sample_np[1,:], linestyle='-')
        # plt.yscale('log')
        plt.xlabel('x')
        plt.ylabel('y')
        self.plot_sampled_with_col()
        plt.show()

    def get_absolute_error(self):
        """
        Calculates the maximum error from the integral error stored in `self.integral_error`.
        Returns a list of maximum errors for each segment in `self.integral_error`.
        """
        # Evaluate the dynamics error
        if not hasattr(self, 'integral_error'):
            self.__dynamics_error()

        max_error = []
        # Calculate maximum error from self.integral_error
        for error_seg in self.integral_error:
            max_error_seg = []
            for error in error_seg:
                max_error_seg.append(np.max(error))
            max_error.append(max_error_seg)

        self.max_error = max_error

        return max_error

    def get_relative_error(self):
        """
        Calculates the relative error from the integral error stored in `self.integral_error`.
        Returns a list of maximum relative errors for each segment in `self.integral_error`.
        """
        # Evaluate the dynamics error
        if not hasattr(self, 'integral_error'):
            self.__dynamics_error()
        if not hasattr(self, 'max_error'):
            self.get_absolute_error()
        if self.max_error is None:
            self.get_absolute_error()

        relative_error_max = []
        # Calculate relative error from self.integral_error
        for error_seg in self.max_error:
            relative_error_seg_max = []
            mean_error_seg_max = np.mean(error_seg)
            for error in error_seg:
                relative_error_seg_max.append(error/mean_error_seg_max)
            relative_error_max.append(relative_error_seg_max)

        self.relative_error_max = relative_error_max

        return relative_error_max

    def get_constraints_error(self):
        """
        Calculates the adherence to constraints.
        Returns a list of integral errors for each segment.
        """
        self.__constraints_error()

        constraints_error_state = np.sum(np.concatenate(self.integral_state_error))
        constraints_error_control = np.sum(np.concatenate(self.integral_control_error))
        constraints_error_obstacles = np.sum(
            np.concatenate(self.integral_obstacle_error))
        constraints_error_total = constraints_error_state + \
            constraints_error_control+constraints_error_obstacles

        return constraints_error_total, constraints_error_state, constraints_error_control, constraints_error_obstacles

    def __dynamics_error(self):
        # Fit solution for state and control as polynomial
        if not hasattr(self, 'state_poly'):
            self.__get_polynomial_solution()
        # Sample the solution
        if not self.sampled:
            self.__sample_solution()
        # Evaluate error in dynamics
        self.__sample_dynamics_error()
        # Evaluate integral error in dynamics
        self.__integral_dynamics_error()
        
    def access_element(self, i, j):
        if not hasattr(self, 'n_col_add_eval'):
            # calculate gradual number of collocation points
            self.n_col_add_eval = np.zeros(self.n_col.shape[0]+1, dtype=int)
            for i, cols in enumerate(self.n_col):
                self.n_col_add_eval[i+1] = self.n_col_add_eval[i] + cols-1
                
        # Check if i is valid
        if i < 0 or i >= len(self.n_col):
            raise IndexError("Segment index out of range")
        # Check if j is valid
        if j < 0 or (i < len(self.n_col) - 1 and j >= self.n_col[i]-1):
            raise IndexError("Collocation point index out of range")
        # Calculate the index for the 1D array
        index = self.n_col_add_eval[i] + j
        return index
    
    def __sample_dynamics_error(self):
        """ 
        This function calculates and stores the dynamics error for a generated sample solution 
        by looping through each segment and time sample of the trajectory, 
        and calculating the absolute difference between the state derivatives 
        obtained from the polynomial solution and the dynamics function.
        """
                
        # Initialize an empty list to store the dynamics error
        self.dynamics_error = []
        # Loop through each segment and time sample of the trajectory
        for i in range(self.n_seg):
            dynamics_error_seg = []
            for j in range(self.n_col[i]-1):
                dynamics_error_subseg = []
                for k in range(self.n_simpson):
                    cur_point = self.access_element(i, j)
                    state_derivative_l = [self.__state_derivative_sample[l][cur_point][k] for l in range(self.nx)]
                    state_l = [self.__state_sample[l][cur_point][k] for l in range(self.nx)]
                    control_l = [self.__control_sample[l][cur_point][k] for l in range(self.nu)]
                    # Calculate dynamics error
                    state_derivative_np = np.array(state_derivative_l)  
                    dynamics_np = np.array(self.__eval_dynamics(state_l, control_l))
                    dynamics_error_subseg.append(np.abs(state_derivative_np - dynamics_np))
                    np.abs(state_derivative_np - dynamics_np)

                # Store the calculated dynamics error for the current time sample as a list
                dynamics_error_seg.append(dynamics_error_subseg)

            # Store the calculated dynamics error for the current segment as a nested list
            self.dynamics_error.append(dynamics_error_seg)

    def __integral_dynamics_error(self):
        """ 
        This function compute the integral of the dynamics error over time using composite Simpson's rule.
        """
        self.integral_error = []
        # self.integral_error_np = []

        # Loop through each segment and collocation point of the trajectory
        for i in range(self.n_seg):
            integral_seg = []
            self.integral_seg = np.array((self.nx, self.n_seg))
            for j in range(self.n_col[i]-1):
                integral_state = []
                for k in range(self.nx):
                    # Extract n_simpson samples of state k in segment i from collocation point j to col. point j+j
                    dynamics_error_temp = [
                        self.dynamics_error[i][j][l][k] for l in range(self.n_simpson)]
                    # Calculate the integral of the current state k using composite Simpson's rule.
                    subinterval_integral = integrate.simpson(
                        dynamics_error_temp, self.__t_sample[i][j])
                    if subinterval_integral < 0:
                        print("WARNING: Integral underflow")
                        subinterval_integral = np.inf
                    integral_state.append(subinterval_integral)
                integral_seg.append(integral_state)
            # Store the calculated dynamics error for the current segment as a nested list
            self.integral_error.append(integral_seg)

    def __constraints_error(self):
        # Fit solution for state and control as polynomial
        if not hasattr(self, 'state_poly'):
            self.__get_polynomial_solution()
        # Sample the solution
        if not self.sampled:
            self.__sample_solution()
        # Evaluate adherence to constraints
        self.__sample_constraints_error()
        # Evaluate integral adherence to constraints
        self.__integral_constraints_error()
    
    def recalc_constraints_error(self, new_n_simpson = None):
        if new_n_simpson is not None:
            self.n_simpson = new_n_simpson
        self.__get_polynomial_solution()
        self.__sample_solution()
        self.__sample_constraints_error()
        self.__integral_constraints_error()

    def __sample_constraints_error(self):
        """ 
        This function calculates and stores the constraints error for a generated sample solution 
        by looping through each segment and time sample of the trajectory. 
        The constraints error is calculated as the amount by which each element 
        in the state and control vectors exceed the corresponding bounds. 
        Moreover trajectory collisions with spherical or cylindrical obstacles are evaluated.
        """
        # Initialize empty lists to store the constraints error
        self.constraints_state_error = []
        self.constraints_control_error = []
        self.constraints_obstacle_error = []
        idx = 0
        # Loop through each segment and time sample of the trajectory
        for i in range(self.n_seg):
            constraints_state_error_seg = []
            constraints_control_error_seg = []
            constraints_obstacle_error_seg = []
            for j in range(self.n_col[i]-1):
                constraints_state_error_subseg = []
                constraints_control_error_subseg = []
                constraints_obstacle_error_subseg = []
                for k in range(self.n_simpson):
                    cur_point = self.access_element(i, j)
                    state_l = np.array([self.__state_sample[l][cur_point][k] for l in range(self.nx)])
                    control_l = np.array([self.__control_sample[l][cur_point][k] for l in range(self.nu)])

                    # Calculate the amount by which each element in the state and control vectors exceed the corresponding bounds
                    state_constraint_error = self.__get_box_excess(
                        state_l, self.drone.x_min, self.drone.x_max)
                    control_constraint_error = self.__get_box_excess(
                        control_l, self.drone.u_min, self.drone.u_max)

                    # Calculate the amount by which each element in the state vectors exceed the boundaries of obstacles
                    obstacle_constraint_error = []
                    if self.obs.dim == 3:
                        # spherical obstacles
                        radius = self.obs.space*np.sqrt(3)/2+self.drone.safe_radius
                        for obstacle in self.obs.map2d:
                            obstacle_constraint_error.append(self.__get_sphere_collision(
                                np.array(state_l[:3]), obstacle, radius))
                    else:
                        # cylindrical obstacles
                        radius = self.obs.space*np.sqrt(2)/2+self.drone.safe_radius
                        for obstacle in self.obs.map2d:
                            obstacle_constraint_error.append(self.__get_sphere_collision(
                                np.array(state_l[:2]), obstacle, radius))

                    # Append the state and control constraints errors for the current time sample to the constraints error subsegment
                    constraints_state_error_subseg.append(
                        state_constraint_error)
                    constraints_control_error_subseg.append(
                        control_constraint_error)
                    constraints_obstacle_error_subseg.append(
                        obstacle_constraint_error)

                # Store the calculated constraints error for the current time sample as a nested list
                constraints_state_error_seg.append(
                    constraints_state_error_subseg)
                constraints_control_error_seg.append(
                    constraints_control_error_subseg)
                constraints_obstacle_error_seg.append(
                    constraints_obstacle_error_subseg)

            # Store the calculated constraints error for the current segment as a nested list
            self.constraints_state_error.append(constraints_state_error_seg)
            self.constraints_control_error.append(
                constraints_control_error_seg)
            self.constraints_obstacle_error.append(
                constraints_obstacle_error_seg)

    def __integral_constraints_error(self):
        """ 
        This function compute the integral of the constraints error over time using composite Simpson's rule.
        """
        self.integral_state_error = []
        self.integral_control_error = []
        self.integral_obstacle_error = []
        # Loop through each segment and collocation point of the trajectory
        for i in range(self.n_seg):
            integral_state_seg = []
            self.integral_state_seg = np.array((self.nx, self.n_seg))
            integral_control_seg = []
            self.integral_control_seg = np.array((self.nu, self.n_seg))
            integral_obstacle_seg = []
            self.integral_obstacle_seg = np.array(
                (self.obs.n_obstacles, self.n_seg))
            for j in range(self.n_col[i]-1):
                # error in adherence of the state trajectory
                integral_state = []
                integral_control = []
                integral_obstacle = []
                for k in range(self.nx):
                    # Extract n_simpson samples of state k in segment i from collocation point j to col. point j+j
                    constraint_error_temp = [
                        self.constraints_state_error[i][j][l][k] for l in range(self.n_simpson)]
                    # Calculate the integral of the current state k using composite Simpson's rule.
                    subinterval_state_integral = integrate.simpson(
                        constraint_error_temp, self.__t_sample[i][j])
                    integral_state.append(subinterval_state_integral)
                integral_state_seg.append(integral_state)

                # error in adherence of the control trajectory
                for k in range(self.nu):
                    # Extract n_simpson samples of control k in segment i from collocation point j to col. point j+j
                    constraint_error_temp = [
                        self.constraints_control_error[i][j][l][k] for l in range(self.n_simpson)]
                    # Calculate the integral of the current control k using composite Simpson's rule.
                    subinterval_control_integral = integrate.simpson(
                        constraint_error_temp, self.__t_sample[i][j])
                    integral_control.append(subinterval_control_integral)
                integral_control_seg.append(integral_control)

                # collisions with obstacles
                for k in range(self.obs.n_obstacles):
                    # Extract n_simpson samples of obstacle collisions in segment i from collocation point j to col. point j+j
                    constraint_error_temp = [
                        self.constraints_obstacle_error[i][j][l][k] for l in range(self.n_simpson)]
                    # Calculate the integral of the current obstacle k using composite Simpson's rule.
                    subinterval_obstacle_integral = integrate.simpson(
                        constraint_error_temp, self.__t_sample[i][j])
                    integral_obstacle.append(subinterval_obstacle_integral)
                integral_obstacle_seg.append(integral_obstacle)

            # Store the integrals for the current segment as a nested list
            self.integral_state_error.append(integral_state_seg)
            self.integral_control_error.append(integral_control_seg)
            self.integral_obstacle_error.append(integral_obstacle_seg)

    def __get_box_excess(self, val, val_min, val_max):
        """
        Calculates how much each element in val exceeds the box constraints.

        Parameters:
            val (list): List of values to be checked.
            val_min (list): List of lower bounds for corresponding element in val.
            val_max (list): List of upper bounds for corresponding element in val.

        Returns:
            list: List of the excesses for each element in val. If an element is within the bounds,
                its value will be 0. If it exceeds the bounds, its value will be the amount it exceeds
                the bound by (positive for exceeding the upper bound, negative for exceeding the lower bound).
        """
        excesses = []
        for i in range(len(val)):
            if val[i] < val_min[i]:
                excesses.append(np.abs(val[i] - val_min[i]))
            elif val[i] > val_max[i]:
                excesses.append(np.abs(val[i] - val_max[i]))
            else:
                excesses.append(0)
        return excesses

    def __get_sphere_collision(self, pos, center, r):
        """
        Calculates how much each element in pos exceeds the spherical/cylindrical constraint represented by (cx, cy, cz) and r.

        Parameters:
            pos (list): List of position (x, y, z) to be checked.
            center (list): the center of the sphere.
            r (float): Radius of the sphere.

        Returns:
            list: Excess of obstacles border. If an element is within the constraint,
                its value will be 0. If it exceeds the constraint, its value will be the amount it exceeds
                the bound by (positive for collision with the bound).    
        """
        # distance to center
        dist_squared = np.sum((pos - center)**2)
        dist = np.sqrt(dist_squared)

        # comparison of distance with radius
        if dist < r:
            excess = r - dist
        else:
            excess = 0.0

        return excess
    
    def get_sphere_collision(self, pos, center, r):
        return self.__get_sphere_collision(pos, center, r)

def example():
    current_folder = os.path.dirname(sys.argv[0])

    folder = 'data_traj'
    i = -1
    file_names = ['orchard', 'columns', 'random_spheres',
                  'forest', 'walls', 'example', 'simple']

    # with open(current_folder+'/'+folder+'/ps_testing_with_warmstart.pkl', mode='rb') as file:
    #     m = cloudpickle.load(file)

    with open(current_folder+'/'+folder+'/ps_testing_with_warmstart_multiseg.pkl', mode='rb') as file:
        m = cloudpickle.load(file)

    # set drone
    drone = Drone()
    # Open the file in binary mode
    with open(current_folder+'/'+'worlds'+'/'+file_names[-1]+'.pkl', 'rb') as file:
        # Call load method to deserialze
        world = pickle.load(file)

    # world limits, start and goal
    drone.x_min[0] = world.x_min
    drone.x_max[0] = world.x_max
    drone.x_min[1] = world.y_min
    drone.x_max[1] = world.y_max
    if world.dim == 2:
        drone.x0[0:2] = world.start
        drone.xf[0:2] = world.goal
        drone.x0[2] = 1.0
        drone.x0[2] = 1.0
    else:
        drone.x0[0:3] = world.start
        drone.xf[0:3] = world.goal
        drone.x_min[2] = world.z_min
        drone.x_max[2] = world.z_max

    sln = PsSolution()

    sln.new_solution(m, drone, world)
    rel_err = sln.get_relative_error()
    sln.plot_relative_error_max()
    print(np.shape(rel_err))
    sln.plot_absolute_error()
    deflection_err = sln.get_deflection()
    print(np.shape(deflection_err))
    sln.plot_deflection()
    constraint_error = sln.get_constraints_error()
    print(constraint_error)
    # sln.plot()
    print(sln.multiseg_time)
    sln.resample_solution(sln.multiseg_time)

    print("The End")


def example_plot():
    current_folder = os.path.dirname(sys.argv[0])

    folder = 'data_traj'
    i = -1
    file_names = ['orchard', 'columns', 'random_spheres',
                  'forest', 'walls', 'example', 'simple']

    # with open(current_folder+'/'+folder+'/ps_testing_with_warmstart.pkl', mode='rb') as file:
    #     m = cloudpickle.load(file)
    
    with open(current_folder+'/'+folder+'/ps_testing_with_warmstart_multiseg.pkl', mode='rb') as file:
        m = cloudpickle.load(file)

    # set drone
    drone = Drone()
    # Open the file in binary mode
    with open(current_folder+'/'+'worlds'+'/'+file_names[-1]+'.pkl', 'rb') as file:
        # Call load method to deserialze
        world = pickle.load(file)

    # world limits, start and goal
    drone.x_min[0] = world.x_min
    drone.x_max[0] = world.x_max
    drone.x_min[1] = world.y_min
    drone.x_max[1] = world.y_max
    if world.dim == 2:
        drone.x0[0:2] = world.start
        drone.xf[0:2] = world.goal
        drone.x0[2] = 1.0
        drone.x0[2] = 1.0
    else:
        drone.x0[0:3] = world.start
        drone.xf[0:3] = world.goal
        drone.x_min[2] = world.z_min
        drone.x_max[2] = world.z_max

    sln = PsSolution()

    sln.new_solution(m, 10, drone, world)
    sln.plot_sampled_with_col(dt=0.01)

    print("The End")
    
def show_solution(path = None):
    if path is None:
        return
    
    with open(path, mode='rb') as file:
        sln = cloudpickle.load(file)
    sln.plot_sampled_with_col(dt=0.01)
    
def example_errors(path = None, new_n_simpson = 10):
    if path is None:
        return
    with open(path, mode='rb') as file:
        sol = cloudpickle.load(file)
    # sol.plot_obstacles_error()
    del(sol.state_poly)
    del(sol.integral_error)
    sol.sampled = False
    sol.recalc_constraints_error(new_n_simpson)
    constraints_error = sol.get_constraints_error() #  "Total Violation", "State Violation", "Control Violation", "Obstacles Violation"
    sol.get_absolute_error()
    abs_error = np.max(np.concatenate(sol.max_error))
    print(abs_error, constraints_error)
    
    # state_sample_np = sol.get_sampled_state_np().T
    # excess = 0
    # for obs in sol.obs.map2d:
    #     for traj_point in state_sample_np:
    #         excess += sol.get_sphere_collision(traj_point[:2], obs, sol.obs.space*np.sqrt(2.0)/2+sol.drone.safe_radius)
    # if excess > 0:
    #     print(excess)

if __name__ == '__main__':
    # example()
    # example_plot()
    
    # show_solution('/home/zboucek/mygit/default/disertace-prog/traj_plan/results/data/trajs/ps_sol_simple_boundon_conston_none_ctrloff_isingle_multi_chebyshev_with_warmstart.pkl')
    example_errors('/home/zboucek/mygit/default/disertace-prog/traj_plan/results/data/trajs/ps_sol_simple_boundon_conston_none_ctrloff_isingle_multi_chebyshev_with_warmstart.pkl',100)