import matplotlib.pyplot as plt


class PlotTrajectories:
    def __init__(self, time_range, list_trajectories=None):
        self.fig = None
        self.ax = None
        self.list_trajectories = list_trajectories if list_trajectories is not None else []
        self.time_range = time_range
        self._xlabel = "t"
        self._ylabel = "x"
        self._title = "Trayectorias"

    def _lazy_init_fig_ax(self):
        """
        To use lazy initialization pattern of the fig and ax.
        :return:
        """
        if self.fig is not None and self.ax is not None:
            return None
        self.fig, self.ax = plt.subplots()
        self.ax.set_xlabel(self._xlabel)
        self.ax.set_ylabel(self._ylabel)
        self.ax.set_title(self._title)

    def plot(self, color="b", alpha=0.2):
        """
        Plot the current trajectories
        :param alpha: the transparency.
        :param color: the color
        :return: itself
        """
        self._lazy_init_fig_ax()  # Init the fig and ax
        for traj in self.list_trajectories:
            self.ax.plot(self.time_range, traj, color, alpha=alpha)
        return self

    def add_trajectory(self, trajectory):
        """
        Add a new trajectory to plot.
        :param trajectory: a trajectory.
        :return: itself.
        """
        self.list_trajectories.append(trajectory)
        return self

    def add_trajectories(self, list_trajectories):
        """
        Add a list of trajectories to the current list of trajectories.
        :param list_trajectories: a list of trajectories.
        :return: itself.
        """
        self.list_trajectories += list_trajectories
        return self

    def clean_list_trajectories(self):
        """
        Clean the list of trajectories.
        :return: itself.
        """
        self.list_trajectories = []
        return self

    @property
    def xlabel(self):
        return self._xlabel

    @xlabel.setter
    def xlabel(self, xlabel):
        self._xlabel = xlabel
        self._lazy_init_fig_ax()
        self.ax.set_xlabel(xlabel)

    @property
    def ylabel(self):
        return self._ylabel

    @ylabel.setter
    def ylabel(self, ylabel):
        self._ylabel = ylabel
        self._lazy_init_fig_ax()
        self.ax.set_ylabel(ylabel)

    @property
    def title(self):
        return self._title

    @title.setter
    def title(self, title):
        self._title = title
        self._lazy_init_fig_ax()
        self.ax.set_title(title)
