import numpy as np

from .base_plots import TimePlot


class StatePlot(TimePlot):
    """Plot to display the environments states and their references."""

    _default_limit_line_cfg = {"color": "red", "linestyle": "--", "linewidth": 0.75}

    # Labels for each state variable.
    state_labels = {
        "omega": r"$\omega$/(1/s)",
        "torque": "$T$/Nm",
        "i": "$i$/A",
        "i_a": r"$i_{\mathrm{a}}$/A",
        "i_e": r"$i_{\mathrm{e}}$/A",
        "i_b": r"$i_{\mathrm{b}}$/A",
        "i_c": r"$i_{\mathrm{c}}$/A",
        "i_sq": r"$i_{\mathrm{sq}}$/A",
        "i_sd": r"$i_{\mathrm{sd}}$/A",
        "u": "$u$/V",
        "u_a": r"$u_{\mathrm{a}}$/V",
        "u_b": r"$u_{\mathrm{b}}$/V",
        "u_c": r"$u_{\mathrm{c}}$/V",
        "u_sq": r"$u_{\mathrm{sq}}$/V",
        "u_sd": r"$u_{\mathrm{sd}}$/V",
        "u_e": r"$u_{\mathrm{e}}$/V",
        "u_sup": r"$u_{\mathrm{sup}}$/V",
        "epsilon": r"$\epsilon$/rad",
    }

    @property
    def state(self):
        return self._state

    def __init__(self, state):
        """
        Args:
            state(str): Name of the state to plot
        """
        super().__init__()

        self._state_line_config = self._default_time_line_cfg.copy()
        self._ref_line_config = self._default_time_line_cfg.copy()
        self._limit_line_config = self._default_limit_line_cfg.copy()

        #: State space of the plotted variable
        self._state_space = None
        #: State name of the plotted variable
        self._state = state
        #: Index in the state array of the plotted variable
        self._state_idx = None
        #: Maximal value of the plotted variable
        self._limits = None
        # Bool: Flag if the plotted variable is referenced.
        self._referenced = None

        # matplotlib-Lines for the state and reference
        self._state_line = None
        self._reference_line = None

        # Data containers
        self._state_data = []
        self._ref_data = []

        # Flag, if the passed data is normalized
        self._normalized = True

        self._scale_plots_to_data = False

    def set_env(self, env):
        # Docstring of superclass
        super().set_env(env)
        ps = env.physical_system
        rg = env.reference_generator
        # Save the index of the state.
        self._state_idx = ps.state_positions[self._state]
        # The maximal values of the state.
        self._limits = ps.limits[self._state_idx]
        self._state_space = (
            ps.state_space.low[self._state_idx],
            ps.state_space.high[self._state_idx],
        )
        # Bool: if the state is referenced.
        self._referenced = rg.referenced_states[self._state_idx]
        # Bool: if the data is already normalized to an interval of [-1, 1]
        self._normalized = self._limits != self._state_space[1]
        self.reset_data()

        min_limit = self._limits * self._state_space[0] if self._normalized else self._state_space[0]
        max_limit = self._limits * self._state_space[1] if self._normalized else self._state_space[1]
        spacing = 0.1 * (max_limit - min_limit)

        # Set the y-axis limits to fixed initital values
        self._y_lim = (min_limit - spacing, max_limit + spacing)

        # Set the y-axis label
        self._label = self.state_labels.get(self._state, self._state)

        self._scale_plots_to_data = env.scale_plots

    def reset_data(self):
        super().reset_data()
        # Initialize the data containers
        self._state_data = np.full(shape=self._x_data.shape, fill_value=np.nan)
        self._ref_data = np.full(shape=self._x_data.shape, fill_value=np.nan)

    def initialize(self, axis):
        # Docstring of superclass
        super().initialize(axis)

        # Line to plot the state data
        (self._state_line,) = self._axis.plot(self._x_data, self._state_data, **self._state_line_config, zorder=2)
        self._lines = [self._state_line]

        # If the state is referenced plot also the reference line
        if self._referenced:
            (self._reference_line,) = self._axis.plot(self._x_data, self._ref_data, **self._ref_line_config, zorder=1)
            self._lines.append(self._reference_line)
        min_limit = self._limits * self._state_space[0] if self._normalized else self._state_space[0]
        max_limit = self._limits * self._state_space[1] if self._normalized else self._state_space[1]
        if self._state_space[0] < 0:
            self._axis.axhline(min_limit, **self._limit_line_config)
        lim = self._axis.axhline(max_limit, **self._limit_line_config)

        y_label = self._label
        unit_split = y_label.find("/")
        if unit_split == -1:
            unit_split = len(y_label)
        limit_label = y_label[:unit_split] + r"$_{\mathrm{max}}$" + y_label[unit_split:]

        if self._referenced:
            ref_label = y_label[:unit_split] + r"$^*$" + y_label[unit_split:]
            self._axis.legend(
                (self._state_line, self._reference_line, lim),
                (y_label, ref_label, limit_label),
                loc="upper left",
                numpoints=20,
            )
        else:
            self._axis.legend(
                (self._state_line, lim),
                (y_label, limit_label),
                loc="upper left",
                numpoints=20,
            )

        self._y_data = [self._state_data, self._ref_data]

    def on_step_end(self, k, state, reference, reward, terminated):
        super().on_step_end(k, state, reference, reward, terminated)
        # Write the data to the data containers
        state_ = state[self._state_idx]
        ref = reference[self._state_idx]
        idx = self.data_idx
        self._x_data[idx] = self._t
        self._state_data[idx] = state_ * self._limits if self._normalized else state_
        if self._referenced:
            self._ref_data[idx] = ref * self._limits if self._normalized else ref
