import numpy as np
import pylab as plt
import matplotlib as mpl
# from mpl_toolkits.mplot3d import Axes3D
import pandas as pd
from util import to_np_mat, plot_corr_ellipses,  \
    convex_bounded_vertex_enumeration, nullspace
from logger import Logger


# ------------------------------------------------------------------------
# FeasibleMuscleSetAnalysis
# ------------------------------------------------------------------------

def construct_muscle_space_inequality(NR, fm_par, Fmax):
    """Construct the feasible muscle space Z f_m0 <= B .

    Parameters
    ----------

    NR: moment arm null space matrix

    fm_par: particular muscle forces

    Fmax: maximum muscle force

    """
    Z0 = -NR
    Z1 = NR
    B0 = fm_par
    B1 = np.asmatrix(np.diag(Fmax)).reshape(Fmax.shape[0], 1) - fm_par
    Z = np.concatenate((Z0, Z1), axis=0)
    B = np.concatenate((B0, B1), axis=0)
    return Z, B


class FeasibleMuscleSetAnalysis:
    """Feasible muscle set analysis.

    The required command along with the state of the system are recorded. Then
    this information is used to compute the feasible muscle null space and
    visualize it.

    """

    def __init__(self, model, simulation_reporter):
        """
        """
        self.logger = Logger('FeasibleMuscleSetAnalsysis')
        self.model = model
        self.simulation_reporter = simulation_reporter

    def visualize_simple_muscle(self, t, ax=None):
        """Visualize the feasible force set at a particular time instance for a linear
        muscle.

        Parameters
        ----------

        t: time

        ax: 1 x 3 axis

        """
        m = self.model.md
        q, Z, B, NR, fm_par = self.calculate_simple_muscles(t)
        x_max = np.max(to_np_mat(self.model.Fmax))
        fm_set = self.generate_solutions(Z, B, NR, fm_par)
        dataframe = pd.DataFrame(fm_set, columns=['$m_' + str(i) + '$' for i in
                                                  range(1, m + 1)])

        # box plot
        if ax is None or ax.shape[0] < 3:
            fig, ax = plt.subplots(1, 3, figsize=(15, 5))

        # box plot
        dataframe.plot.box(ax=ax[0])
        ax[0].set_xlabel('muscle id')
        ax[0].set_ylabel('force $(N)$')
        ax[0].set_title('Muscle-Force Box Plot')
        ax[0].set_ylim([0, 1.1 * x_max])

        # correlation matrix
        cmap = mpl.cm.jet
        norm = mpl.colors.Normalize(vmin=-1, vmax=1)
        corr = dataframe.corr()
        m = plot_corr_ellipses(corr, ax=ax[1], norm=norm, cmap=cmap)
        cb = plt.colorbar(m, ax=ax[1], orientation='vertical', norm=norm,
                          cmap=cmap)
        cb.set_label('Correlation Coefficient')
        ax[1].margins(0.1)
        ax[1].set_xlabel('muscle id')
        ax[1].set_ylabel('muscle id')
        ax[1].set_title('Correlation Matrix')
        ax[1].axis('equal')

        # draw model
        self.model.draw_model(q, False, ax[2], scale=0.7, text=False)

    def calculate_simple_muscles(self, t):
        """Construct Z f_m0 <= B for the case of a linear muscle model for a particular
        time instance.

        Parameters
        ----------

        t: time

        """
        # find nearesrt index corresponding to t
        idx = np.abs(np.array(self.simulation_reporter.t) - t).argmin()
        t = self.simulation_reporter.t[idx]
        q = self.simulation_reporter.q[idx]
        u = self.simulation_reporter.u[idx]
        tau = self.simulation_reporter.tau[idx]
        pose = self.model.model_parameters(q=q, u=u)
        n = self.model.nd

        # calculate required variables
        R = to_np_mat(self.model.R.subs(pose))
        RBarT = np.asmatrix(np.linalg.pinv(R.T))
        # reduce to independent columns to avoid singularities (proposition 3)
        NR = nullspace(R.transpose())
        fm_par = np.asmatrix(-RBarT * tau.reshape((n, 1)))
        Fmax = to_np_mat(self.model.Fmax)

        Z, B = construct_muscle_space_inequality(NR, fm_par, Fmax)

        return q, Z, B, NR, fm_par

    def generate_solutions(self, A, b, NR, fm_par):
        """Sample the solution space that satisfy A x <= b.

        Parameters
        ----------

        A: matrix A

        b: column vector

        NR: moment arm nullspace

        fm_par: particular solution

        Returns
        -------

        muscle forces: a set of solutions that satisfy the problem

        """
        feasible_set = []
        fm0_set = convex_bounded_vertex_enumeration(np.array(A),
                                                    np.array(b).flatten(), 0)
        n = fm0_set.shape[0]
        for i in range(0, n):
            fm = fm_par + NR * np.matrix(fm0_set[i, :]).reshape(-1, 1)
            feasible_set.append(fm)

        return np.array(feasible_set).reshape(n, -1)


def test_feasible_set(model):
    feasible_set = FeasibleMuscleSetAnalysis(model)
    n = model.nd
    m = model.md
    feasible_set.record(1,
                        np.random.random((m, 1)),
                        np.random.random((m, m)),
                        np.random.random((n, 1)))
    fig, ax = plt.subplots(2, 3, figsize=(10, 10))
    feasible_set.visualize_simple_muscle(1, ax[0])
    feasible_set.visualize_simple_muscle(1, ax[1])
    plt.show()