import tarfile
import re
from PIL import Image
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import scipy.cluster.vq as vq
import numpy as np
from scipy.io import loadmat
from scipy.special import ellipeinc, ellipe
from os import path


def plot_toy_slice(x, y, drop_dim):
    """
    Plot a slice of the 4D input data in our toy example.

    A 3D scatter plot is generated by droping the dimension `drop_dim`. The
    color of each plotted point is determined using the labels `y`.

    Parameters
    ----------

    x: (n, 4) array-like
    y: (n,) array-like
        The labels.
    drop_dim: int
        The dimension to cut, 1 indexed.
    """
    x = np.array(x)
    x_3d = np.delete(x, drop_dim - 1, 1)
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    y = y.astype(int)
    ax.scatter(x_3d[:,0], x_3d[:,1], x_3d[:,2], c=_COLORS[y])
    plt.show()


def plot_toy_3d(x, y):
    """
    Show a 3D scatter plot of our toy data.

    Parameters
    ----------

    x: (n, 3) array-like
    y: (n,) array-like
        The labels
    """
    y = y.astype(int)
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(x[:,0], x[:,1], x[:,2], c=_COLORS[y])
    ax.set_aspect('equal')
    plt.show()


def plot_toy_2d(x, y):
    """
    Plot the low-dimensional representation of the toy data.

    The function uses a colored scatter plot where each point is
    from `x` and the color is determined by the labels `y`.

    Parameters
    ----------

    x: (n, 2) array-like
        The two dimensional representation.
    y: (n,) array-like
       The labels.
    """
    x = np.array(x)
    y = np.array(y).reshape(-1).astype(int)
    if len(x.shape) != 2 or len(y) != x.shape[0] or x.shape[1] != 2:
        raise ValueError("Invalid input shape")
    theta = _find_orientation(x, y)
    x = _rot_mat(theta).dot(x.T).T
    plt.scatter(x[:,0], x[:,1], c=_COLORS[y], s=30)
    plt.gca().set_aspect('equal')
    plt.show()
    return x


def hog_test_data():
    """
    Load test data for the HOG features.

    Can be used to check the implementation.

    Returns
    -------

    image: array with shape (25,50)
    steps: dictionary with values of intermediate steps
    """
    data = loadmat(path.join(path.dirname(__file__), 'hog_ref.mat'))
    print("Number of bins: {}".format(data.pop('n_bins')[0,0]))
    print("Cell size: {}".format(data.pop('cell_size')[0,0]))
    print("Block size: {}".format(data.pop('blk_size')[0,0]))
    print("Unsigned directions: {}".format(data.pop('unsigned')[0,0] == 1))
    print("Clip value: {}".format(data.pop('clip_val')[0,0]))
    for key in list(data.keys()):
        if key.startswith('__'):
            del data[key]

    image = data.pop('image')
    return image, data


def map_on_ellipse(xq, a=4, b=1, gap_angle=90):
    """Map 2D points on bend ellipse in 3D.

    Parameters
    ----------

    xq: (n, 2) array-like
        The 2D points
    a, b: non-negative, scalar
        Ellipsis axis
    gap_angle: scalar in [0,360)
        Degree of the gap the bend ellipse should have.
        E.g. for 0 the data is mapped onto a closed ellipse.

    Returns
    -------

    xyz: (n, 3) array-like
        Mapped 3D points
    """
    gap_angle = np.deg2rad(gap_angle) / 2
    # Range of the data
    x_min = np.min(xq[:,0])
    x_max = np.max(xq[:,0])
    width = x_max - x_min

    # Compute perimeter of the open ellipse
    a = abs(a)
    b = abs(b)
    major, minor = max(a, b), min(a, b)
    ecc = (1 - minor**2 / major**2) ** 2
    offset = ellipeinc(gap_angle, ecc)
    diameter = major * (4 * ellipe(ecc) - offset)
    # Scale ellipse to match the span
    scale = width / diameter
    a *= scale
    b *= scale

    # Find the angles. We use the canonical parametrization.
    x_zeroed = xq[:,0] - x_min
    alpha = (2 * np.pi - 2 * gap_angle) * x_zeroed / (x_max - x_min) + gap_angle  # initial guess
    for _ in range(6):
        # 6 Newton steps to solve inverse elliptic integral.
        # 6 Newton steps ought to be enough for anybody.
        f = major * ellipeinc(alpha, ecc) - (x_zeroed + major * offset)
        f_prime = major * np.sqrt(1 - ecc * np.sin(alpha)**2)
        alpha = alpha - f / f_prime
    xy = np.vstack((a * np.cos(alpha), b * np.sin(alpha))).T
    z = xq[:,1]  # z-axis is the previous y-axis, we only bend in the xy-plane
    xyz = np.hstack((xy, z.reshape(-1, 1)))
    return xyz


def load_pedestrian_images(split, label):
    """
    Load test/train images from the pedestrian dataset.


    Parameters
    ----------

    split: str
        Either 'train' or 'test', to specify if training or testing data shall
        be returned.
    label: bool
        Specifies whether the images with or without pedestrians shall
        be returned.

    Returns
    -------

    images: ndarray
        A NumPy array with shape (N, 100, 50, 3) representing pixel values of
        the images. Each image is 100x50 and has three color channels (RGB).
        The pixel values are in the range [0, 255].
    """
    if split not in ('train', 'test'):
        raise ValueError("The split must be 'train' or 'test'")
    folder = 'ped/{}/{}'.format(split, int(bool(label)))
    images = []
    for member in _PED_TARFILE.getmembers():
        if not member.isfile() or not path.dirname(member.name) == folder:
            continue
        if not re.match(_IMAGE_FILE_RE, path.basename(member.name)):
            continue
        image = Image.open(_PED_TARFILE.extractfile(member))
        images.append(np.asarray(image))
    return np.array(images, dtype="float64")


_COLORS = np.load(path.join(path.dirname(__file__), 'pca_toy_colors.npy'))
_PED_TARFILE = tarfile.open(path.join(path.dirname(__file__), "ped.tar.gz"), "r:gz")
_IMAGE_FILE_RE = re.compile(r'^\d{3}\.jpg$')


def _rot_mat(theta):
    """Compute 2D rotatation matrix."""
    c, s = np.cos(theta), np.sin(theta)
    R = np.array([[c, -s], [s, c]])
    return R


def _find_orientation(x, y, eye_color_index=4):
    """Find the orientation of the face."""
    old = np.seterr(all='raise')
    try:
        eyes, _ = vq.kmeans2(x[y == eye_color_index], 2)
    except:
        return 0
    finally:
        np.seterr(**old)
    eye_line = eyes[0] - eyes[1]
    rad = np.arctan2(*eye_line) + np.pi / 2
    eyes_rot = _rot_mat(rad).dot(eyes.T).T
    if eyes_rot[0,1] < 0:
        rad += np.pi
    return rad
