#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import time
import sympy as sp
import numpy as np
import pandas as pd
from tqdm import tqdm
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.collections import EllipseCollection
from matplotlib import patches
from fractions import Fraction

# ------------------------------------------------------------------------------
# logger
# ------------------------------------------------------------------------------

# logging.basicConfig(format='%(levelname)s %(asctime)s @%(name)s # %(message)s',
#                     datefmt='%m/%d/%Y %I:%M:%S %p',
#                     # filename='example.log',
#                     level=logging.DEBUG)

# ------------------------------------------------------------------------------
# utilities
# ------------------------------------------------------------------------------


def mat_show(mat):
    """Graphical visualization of a 2D matrix.
    """
    fig = plt.figure()
    ax = fig.add_subplot(111)
    cax = ax.matshow(to_np_mat(mat), interpolation='nearest')
    fig.colorbar(cax)


def is_symmetric(a, tol=1e-8):
    """Check if matrix is symmetric.

    """
    return np.allclose(a, a.T, atol=tol)


def mat(array):
    """For a given 2D array return a numpy matrix.
    """
    return np.matrix(array)


def vec(vector):
    """Construct a column vector of type numpy matrix.
    """
    return np.matrix(vector).reshape(-1, 1)


def to_np_mat(sympy_mat):
    """Cast sympy Matrix to numpy matrix of float type.

    Parameters
    ----------
    m: sympy 2D matrix

    Returns
    -------
    a numpy asmatrix

    """
    return np.asmatrix(sympy_mat.tolist(), dtype=np.float)


def to_np_array(sympy_mat):
    """Cast sympy Matrix to numpy matrix of float type. Works for N-D matrices as
    compared to to_np_mat().

    Parameters
    ----------
    m: sympy 2D matrix

    Returns
    -------
    a numpy asmatrix

    """
    return np.asarray(sympy_mat.tolist(), dtype=np.float)


def to_np_vec(sympy_vec):
    """Transforms a 1D sympy vector (e.g. 5 x 1) to numpy array (e.g. (5,)).

    Parameters
    ----------
    v: 1D sympy vector

    Returns
    -------
    a 1D numpy array

    """
    return np.asarray(sp.flatten(sympy_vec), dtype=np.float)


def is_pd(A):
    """Check if matrix is positive definite
    """
    return np.all(np.linalg.eigvals(A) > 0)


def nullspace(A, atol=1e-13, rtol=0):
    """Compute an approximate basis for the nullspace of A.

    The algorithm used by this function is based on the singular value
    decomposition of `A`.

    Parameters
    ----------
    A : ndarray
        A should be at most 2-D.  A 1-D array with length k will be treated
        as a 2-D with shape (1, k)
    atol : float
        The absolute tolerance for a zero singular value.  Singular values
        smaller than `atol` are considered to be zero.
    rtol : float
        The relative tolerance.  Singular values less than rtol*smax are
        considered to be zero, where smax is the largest singular value.

    If both `atol` and `rtol` are positive, the combined tolerance is the
    maximum of the two; that is::
        tol = max(atol, rtol * smax)
    Singular values smaller than `tol` are considered to be zero.

    Return value
    ------------
    ns : ndarray
        If `A` is an array with shape (m, k), then `ns` will be an array
        with shape (k, n), where n is the estimated dimension of the
        nullspace of `A`.  The columns of `ns` are a basis for the
        nullspace; each element in numpy.dot(A, ns) will be approximately
        zero.
    """

    A = np.atleast_2d(A)
    u, s, vh = np.linalg.svd(A)
    tol = max(atol, rtol * s[0])
    nnz = (s >= tol).sum()
    ns = vh[nnz:].conj().T
    return ns


def lrs_inequality_vertex_enumeration(A, b):
    """Find the vertices given an inequality system A * x <= b. This function
    depends on lrs library.

    Parameters
    ----------

    A: numpy array [m x n]

    b: numpy array [m]

    Returns
    -------

    v: numpy array [k x n]
        the vertices of the polytope

    """
    # export H-representation
    with open('temp.ine', 'w') as file_handle:
        file_handle.write('Feasible_Set\n')
        file_handle.write('H-representation\n')
        file_handle.write('begin\n')
        file_handle.write(str(A.shape[0]) + ' ' +
                          str(A.shape[1] + 1) + ' rational\n')
        for i in range(0, A.shape[0]):
            file_handle.write(str(Fraction(b[i])))
            for j in range(0, A.shape[1]):
                file_handle.write(' ' + str(Fraction(-A[i, j])))

            file_handle.write('\n')

        file_handle.write('end\n')

    # call lrs
    try:
        os.system('lrs temp.ine > temp.ext')
    except OSError as e:
        raise RuntimeError(e)

    # read the V-representation
    vertices = []
    with open('temp.ext', 'r') as file_handle:
        begin = False
        for line in file_handle:
            if begin:
                if 'end' in line:
                    break

                comp = line.split()
                try:
                    v_type = comp.pop(0)
                except:
                    print('No feasible solution')

                if v_type is '1':
                    v = [float(Fraction(i)) for i in comp]
                    vertices.append(v)

            else:
                if 'begin' in line:
                    begin = True

    # delete temporary files
    try:
        os.system('rm temp.ine temp.ext')
    except OSError as e:
        pass

    return vertices


def ccd_inequality_vertex_enumeration(A, b):
    """Find the vertices given an inequality system A * x <= b. This function
    depends on pycddlib (cdd).

    Parameters
    ----------

    A: numpy array [m x n]

    b: numpy array [m]

    Returns
    -------

    v: numpy array [k x n]
        the vertices of the polytope

    """
    import cdd
    # try floating point, if problem fails try exact arithmetics (slow)
    try:
        M = cdd.Matrix(np.hstack((b.reshape(-1, 1), -A)),
                       number_type='float')
        M.rep_type = cdd.RepType.INEQUALITY
        p = cdd.Polyhedron(M)
    except:
        print('Warning: switch to exact arithmetics')
        M = cdd.Matrix(np.hstack((b.reshape(-1, 1), -A)),
                       number_type='fraction')
        M.rep_type = cdd.RepType.INEQUALITY
        p = cdd.Polyhedron(M)

    G = np.array(p.get_generators())

    if not G.shape[0] == 0:
        return G[np.where(G[:, 0] == 1.0)[0], 1:].tolist()
    else:
        raise ValueError('Infeasible Inequality')


def optimization_based_sampling(A, b, optimziation_samples,
                                closiness_tolerance, max_opt_iterations):
    """Efficient method for sampling the feasible set for a large system of
    inequalities (A x <= b). When the dimension of the set (x) is large,
    deterministic and pure randomized techniques fail to solve this problem
    efficiently.

    This method uses constrained optimization in order to find n solutions that
    satisfy the inequality. Each iteration new randomized objective function is
    assigned so that the optimization will find a different solution.

    Parameters
    ----------
    A: numpy array [m x n]

    b: numpy array [m]

    optimziation_samples: integer
        number of samples to generate

    closiness_tolerance: float
        accept solution if distance is larger than the provided tolerance

    max_opt_iterations: integer
        maximum iteration of the optimization algorithm

    Returns
    -------

    solutions: list

    """
    from scipy.optimize import minimize

    nullity = A.shape[1]
    solutions = []
    j = 0
    while j < optimziation_samples:
        # change objective function randomly (Dirichlet distribution ensures
        # that w sums to 1)
        w = np.random.dirichlet(np.ones(nullity), size=1)

        def objective(x):
            return np.sum((w * x) ** 2)

        # solve the optimization_based_sampling
        def inequality_constraint(x):
            return A.dot(x) - b

        x0 = np.random.uniform(-1, 1, nullity)
        bounds = tuple([(None, None) for i in range(0, nullity)])
        constraints = ({'type': 'ineq', 'fun': inequality_constraint})
        options = {'maxiter': max_opt_iterations}
        sol = minimize(objective, x0, method='SLSQP',
                       bounds=bounds,
                       constraints=constraints,
                       options=options)

        x = sol.x
        # check if solution satisfies the system
        if np.all(A.dot(x) <= b):
            # if solution is not close to the rest then accept
            close = False
            for xs in solutions:
                if np.linalg.norm(xs - x, 2) < closiness_tolerance:
                    close = True
                    break

            if not close:
                solutions.append(x)
                j = j + 1

    return solutions


def convex_bounded_vertex_enumeration(A,
                                      b,
                                      convex_combination_passes=1,
                                      method='lrs'):
    """Sample a convex, bounded inequality system A * x <= b. The vertices of the
    convex polytope are first determined. Then the convexity property is used to
    generate additional solutions within the polytope.

    Parameters
    ----------

    A: numpy array [m x n]

    b: numpy array [m]

    convex_combination_passes: int (default 1)
        recombine vertices to generate additional solutions using the convex
        property

    method: str (lrs or cdd or rnd)

    Returns
    -------

    v: numpy array [k x n]
        solutions within the convex polytope

    """
    # find polytope vertices
    if method == 'lrs':
        solutions = lrs_inequality_vertex_enumeration(A, b)
    elif method == 'cdd':
        solutions = ccd_inequality_vertex_enumeration(A, b)
    elif method == 'rnd':
        solutions = optimization_based_sampling(A, b,
                                                A.shape[0] ** 2,
                                                1e-3, 1000)
    else:
        raise RuntimeError('Unsupported method: choose "lrs" or "cdd" or "rnd"')

    # since the feasible space is a convex set we can find additional solution
    # in the form z = a * x_i + (1-a) x_j
    for g in range(0, convex_combination_passes):
        n = len(solutions)
        for i in range(0, n):
            for j in range(0, n):
                if i == j:
                    continue

                a = 0.5
                x1 = np.array(solutions[i])
                x2 = np.array(solutions[j])
                z = a * x1 + (1 - a) * x2
                solutions.append(z.tolist())

    # remove duplicates from 2D list
    solutions = [list(t) for t in set(tuple(element) for element in solutions)]

    return np.array(solutions, np.float)


def test_inequality_sampling(d):
    A = np.array([[0, 0, -1],
                  [0, -1, 0],
                  [1, 0, 0],
                  [-1, 0, 0],
                  [0, 1, 0],
                  [0, 0, 1]])
    b = 0.5 * np.ones(6)
    print(A, b)
    t1 = time.time()
    # solutions = optimization_based_sampling(A, b, 20, 0.3, 1000)
    solutions = convex_bounded_vertex_enumeration(A, b, d)
    t2 = time.time()
    print('Execution time: ' + str(t2 - t1) + 's')

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(solutions[:, 0], solutions[:, 1], solutions[:, 2])
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    plt.show()
    fig.tight_layout()
    fig.savefig('results/inequality_sampling/inequality_sampling_d' + str(d) + '.png',
                format='png', dpi=300)
    fig.savefig('results/inequality_sampling/inequality_sampling_d' + str(d) + '.pdf',
                format='pdf', dpi=300)
    fig.savefig('results/inequality_sampling/inequality_sampling_d' + str(d) + '.eps',
                format='eps', dpi=300)


# test_inequality_sampling(0)
# test_inequality_sampling(1)
# test_inequality_sampling(2)

def tensor3_vector_product(T, v):
    """Implements a product of a rank-3 tensor (3D array) with a vector using
    tensor product and tensor contraction.

    Parameters
    ----------

    T: sp.Array of dimensions n x m x k

    v: sp.Array of dimensions k x 1

    Returns
    -------

    A: sp.Array of dimensions n x m

    Example
    -------

    >>>T = sp.Array([[[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]],
                     [[13, 16, 19, 22], [14, 17, 20, 23], [15, 18, 21, 24]]])
    ⎡⎡1  4  7  10⎤  ⎡13  16  19  22⎤⎤
    ⎢⎢           ⎥  ⎢              ⎥⎥
    ⎢⎢2  5  8  11⎥  ⎢14  17  20  23⎥⎥
    ⎢⎢           ⎥  ⎢              ⎥⎥
    ⎣⎣3  6  9  12⎦  ⎣15  18  21  24⎦⎦
    >>>v = sp.Array([1, 2, 3, 4]).reshape(4, 1)
    ⎡1⎤
    ⎢ ⎥
    ⎢2⎥
    ⎢ ⎥
    ⎢3⎥
    ⎢ ⎥
    ⎣4⎦
    >>>tensor3_vector_product(T, v)
    ⎡⎡70⎤  ⎡190⎤⎤
    ⎢⎢  ⎥  ⎢   ⎥⎥
    ⎢⎢80⎥  ⎢200⎥⎥
    ⎢⎢  ⎥  ⎢   ⎥⎥
    ⎣⎣90⎦  ⎣210⎦⎦

    """
    import sympy as sp
    assert(T.rank() == 3)
    # reshape v to ensure 1D vector so that contraction do not contain x 1
    # dimension
    v.reshape(v.shape[0], )
    p = sp.tensorproduct(T, v)
    return sp.tensorcontraction(p, (2, 3))


def test_tensor_product():
    T = sp.Array([[[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]],
                  [[13, 16, 19, 22], [14, 17, 20, 23], [15, 18, 21, 24]]])
    v = sp.Array([1, 2, 3, 4]).reshape(4, 1)
    display(T, v)
    display(tensor3_vector_product(T, v))


# test_tensor_product()


def draw_ellipse(ax, xc, A, scale=1.0, show_axis=False):
    """Construct an ellipse representation of a 2x2 matrix.

    Parameters
    ----------
    ax: plot axis

    xc: np.array 2 x 1
        center of the ellipse
    mat: np.array 2 x 2

    scale: float (default=1.0)
        scale factor of the principle axes

    """
    eigen_values, eigen_vectors = np.linalg.eig(A)
    idx = np.abs(eigen_values).argsort()[::-1]
    eigen_values = eigen_values[idx]
    eigen_vectors = eigen_vectors[:, idx]
    phi = np.rad2deg(np.arctan2(eigen_vectors[1, 0], eigen_vectors[0, 0]))

    ellipse = patches.Ellipse(xy=(xc[0, 0], xc[1, 0]),
                              width=2 * scale * eigen_values[0],
                              height=2 * scale * eigen_values[1],
                              angle=phi,
                              linewidth=2, fill=False)
    ax.add_patch(ellipse)

    # axis
    if show_axis:
        x_axis = np.array([[xc[0, 0], xc[1, 0]],
                           [xc[0, 0] + scale * np.abs(eigen_values[0]) * eigen_vectors[0, 0],
                            xc[1, 0] + scale * np.abs(eigen_values[0]) * eigen_vectors[1, 0]]])
        y_axis = np.array([[xc[0, 0], xc[1, 0]],
                           [xc[0, 0] + scale * eigen_values[1] * eigen_vectors[0, 1],
                            xc[1, 0] + scale * eigen_values[1] * eigen_vectors[1, 1]]])
        ax.plot(x_axis[:, 0], x_axis[:, 1], '-r', label='x-axis')
        ax.plot(y_axis[:, 0], y_axis[:, 1], '-g', label='y-axis')

    return phi, eigen_values, eigen_vectors


def test_ellipse():
    fig, ax = plt.subplots()
    xc = vec([0, 0])
    M = mat([[2, 1], [1, 2]])
    # M = mat([[-2.75032375, -11.82938331], [-11.82938331, -53.5627191]])
    print np.linalg.matrix_rank(M)
    phi, l, v = draw_ellipse(ax, xc, M, 1, True)
    print(phi, l, v)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.axis('equal')
    ax.legend()
    fig.show()


# test_ellipse()


def calculate_feasible_muscle_set(feasible_muscle_set_analysis, base_name,
                                  t_start, t_end, dt, speed):
    """ Calculates the feasible muscle space of a simulation.

    Parameters
    ----------

    feasible_muscle_set_analysis: FeasibleMuscleSetAnalysis

    base_name: base name of simulation files

    t_start: t start

    t_end: t end

    dt: time interval for reporting

    speed: speed of animation

    """
    print('Calculating feasible muscle set ...')
    time = np.linspace(t_start, t_end, t_end / dt + 1, endpoint=True)
    for i, t in enumerate(tqdm(time)):
        visualize_feasible_muscle_set(feasible_muscle_set_analysis, t,
                                      base_name + str(i).zfill(6), 'png')

    command = 'convert -delay ' + \
              str(speed * dt) + ' -loop 0 ' + base_name + \
              '*.png ' + base_name + 'anim.gif'

    print(command)
    try:
        os.system(command)
    except:
        print('unable to execute command')


def visualize_feasible_muscle_set(feasible_muscle_set_analysis, t,
                                  fig_name='fig/feasible_muscle_set', format='png'):
    """ Visualization of the feasible muscle space.

    Parameters
    ----------

    feasible_muscle_set_analysis: FeasibleMuscleSetAnalysis
    t: time instance to evaluate the feasible
    fig_name: figure name for saving
    format: format (e.g. .png, .pdf, .eps)

    """
    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
    feasible_muscle_set_analysis.visualize_simple_muscle(t, ax)
    fig.suptitle('t = ' + str(np.around(t, decimals=2)),
                 y=1.00, fontsize=12, fontweight='bold')
    fig.tight_layout()
    fig.savefig(fig_name + '.' + format, format=format, dpi=300)
    fig.savefig(fig_name + '.pdf', format='pdf', dpi=300)
    fig.savefig(fig_name + '.eps', format='eps', dpi=300)


def apply_generalized_force(f):
    """Applies a generalized force (f) in a manner that is consistent with Newton's
    3rd law.

    Parameters
    ----------
    f: generalized force

    """
    n = len(f)
    tau = []
    for i in range(0, n):
        if i == n - 1:
            tau.append(f[i])
        else:
            tau.append(f[i] - f[i + 1])

    return tau


def custom_exponent(q, A, k, q_lim):
    """ Sympy representation of custom exponent function.

    f(q) = A e^(k (q - q_lim)) / (150) ** k

    """
    return A * sp.exp(k * (q - q_lim)) / (148.42) ** k


def coordinate_limiting_force(q, q_low, q_up, a, b):
    """A continuous coordinate limiting force for a rotational joint.

    It applies an exponential force when approximating a limit. The convention
    is that positive force is generated when approaching the lower limit and
    negative when approaching the upper. For a = 1, F ~= 1N at the limits.

    Parameters
    ----------
    q: generalized coordinate
    q_low: lower limit
    q_up: upper limit
    a: force at limits
    b: rate of the slop

    Note: q, q_low, q_up must have the same units (e.g. rad)

    """
    return custom_exponent(q_low + 5, a, b, q) - custom_exponent(q, a, b, q_up - 5)


def test_limiting_force():
    """
    """
    q = np.linspace(0, np.pi / 4, 100, endpoint=True)
    f = [coordinate_limiting_force(qq, 0, np.pi / 4, 1, 50) for qq in q]

    plt.plot(q, np.array(f))
    plt.show()


def gaussian(x, a, m, s):
    """Gaussian function.

    f(x) = a e^(-(x - m)^2 / (2 s ^2))

    Parameters
    ----------
    x: x
    a: peak
    m: mean
    s: standard deviation

    For a good approximation of an impulse at t = 0.3 [x, 1, 0.3, 0.01].

    """
    return a * np.exp(- (x - m) ** 2 / (2 * s ** 2))


def test_gaussian():
    """
    """
    t = np.linspace(0, 2, 200)
    y = [gaussian(tt, 0.4, 0.4, 0.01) for tt in t]
    plt.plot(t, y)
    plt.show()


def rotate(origin, point, angle):
    """Rotate a point counterclockwise by a given angle around a given origin.

    The angle should be given in radians.

    """
    R = np.asmatrix([[np.cos(angle), - np.sin(angle)],
                     [np.sin(angle), np.cos(angle)]])

    q = origin + R * (point - origin)

    return q


def sigmoid(t, t0, A, B):
    """Implementation of smooth sigmoid function.

    Parameters
    ----------
    t: time to be evalutaed
    t0: delay
    A: magnitude
    B: slope

    Returns
    -------
    (y, y', y'')

    """
    return (A * (np.tanh(B * (t - t0)) + 1) / 2,
            A * B * (- np.tanh(B * (t - t0)) ** 2 + 1) / 2,
            - A * B ** 2 * (- np.tanh(B * (t - t0)) ** 2 + 1) * np.tanh(B * (t - t0)))


def test_sigmoid():
    """
    """
    t, A, B, t0 = sp.symbols('t A B t0')
    y = A / 2 * (sp.tanh(B * (t - t0 - 1)) + 1)
    yd = sp.diff(y, t)
    ydd = sp.diff(yd, t)
    print('\n', y, '\n', yd, '\n', ydd)

    tt = np.linspace(-2, 2, 100)
    yy = np.array([sigmoid(x, 0.5, 2, 5) for x in tt])
    plt.plot(tt, yy)
    plt.show()


def plot_corr_ellipses(data, ax=None, **kwargs):
    """For a given correlation matrix "data", plot the correlation matrix in terms
    of ellipses.

    parameters
    ----------
    data: Pandas dataframe containing the correlation of the data (df.corr())
    ax: axis (e.g. fig, ax = plt.subplots(1, 1))
    kwards: keywords arguments (cmap="Greens")

    https://stackoverflow.com/questions/34556180/
    how-can-i-plot-a-correlation-matrix-as-a-set-of-ellipses-similar-to-the-r-open

    """
    M = np.array(data)
    if not M.ndim == 2:
        raise ValueError('data must be a 2D array')
    if ax is None:
        fig, ax = plt.subplots(1, 1, subplot_kw={'aspect': 'equal'})
        ax.set_xlim(-0.5, M.shape[1] - 0.5)
        ax.set_ylim(-0.5, M.shape[0] - 0.5)

    # xy locations of each ellipse center
    xy = np.indices(M.shape)[::-1].reshape(2, -1).T

    # set the relative sizes of the major/minor axes according to the strength of
    # the positive/negative correlation
    w = np.ones_like(M).ravel()
    h = 1 - np.abs(M).ravel()
    a = 45 * np.sign(M).ravel()

    ec = EllipseCollection(widths=w, heights=h, angles=a, units='x', offsets=xy,
                           transOffset=ax.transData, array=M.ravel(), **kwargs)
    ax.add_collection(ec)

    # if data is a DataFrame, use the row/column names as tick labels
    if isinstance(data, pd.DataFrame):
        ax.set_xticks(np.arange(M.shape[1]))
        ax.set_xticklabels(data.columns, rotation=90)
        ax.set_yticks(np.arange(M.shape[0]))
        ax.set_yticklabels(data.index)

    return ec


def get_cmap(n, name='hsv'):
    """Returns a function that maps each index in 0, 1, ..., n-1 to a distinct RGB
    color; the keyword argument name must be a standard mpl colormap name.

    """
    return plt.cm.get_cmap(name, n)


def assert_if_same(A, B):
    """Assert whether two quantities (value, vector, matrix) are the same."""
    assert np.isclose(
        np.array(A).astype(np.float64),
        np.array(B).astype(np.float64)).all() == True, 'quantities not equal'


def christoffel_symbols(M, q, i, j, k):
    """
    M [n x n]: inertia mass matrix
    q [n x 1]: generalized coordinates
    i, j, k  : the indexies to be computed
    """
    return sp.Rational(1, 2) * (sp.diff(M[i, j], q[k]) + sp.diff(M[i, k], q[j]) -
                                sp.diff(M[k, j], q[i]))


def coriolis_matrix(M, q, dq):
    """
    Coriolis matrix C(q, qdot) [n x n]
    Coriolis forces are computed as  C(q, qdot) *  qdot [n x 1]
    """
    n = M.shape[0]
    C = sp.zeros(n, n)
    for i in range(0, n):
        for j in range(0, n):
            for k in range(0, n):
                C[i, j] = C[i, j] + christoffel_symbols(M, q, i, j, k) * dq[k]
    return C


def substitute(symbols, constants):
    """For a given symbolic sequence substitute symbols."""
    return np.array([substitute(exp, constants)
                     if hasattr(exp, '__iter__')
                     else exp.subs(constants) for exp in symbols])
