#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import division
from io import open
from sklearn.decomposition import PCA
from sklearn.preprocessing import scale
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation

import matplotlib.pyplot as plt
import matplotlib
import math
import numpy
import json
import argparse
# constantes PAR DEFAUT

# Affichage souhaité
FIGURE_2D = False
FIGURE_3D = False

# Fichier de sortie par defaut
OUTPUT = ''

# Nombre de dimensions de la PCA
# (toujours inférieur au nombre max de dimensions de la matrice)
NB_COMPONENTS = 3

# constantes redefinies par l'utilisateur
parser = argparse.ArgumentParser(description='This script make PCA of resource matrice.json file.')
parser.add_argument('-r','--resources', help='Path of resources JSON file [Required]', type=str, required=True)
parser.add_argument('-comp','--components', help='Number of components [Optionnal, default=3]', type=int, required=False)
parser.add_argument('-f2d','--figure2d', help='Show 2D PCA (GUI) [Optionnal, default=false]', action='store_true', required=False)
parser.add_argument('-f3d','--figure3d', help='Show 3D PCA (GUI) [Optionnal, default=false]', action='store_true', required=False)
parser.add_argument('-o','--output', help='Path of output file (record of 3D) [Optionnal, default=\'\']', type=str, required=False)
args = parser.parse_args()

if hasattr(args, 'components'):
    if args.components != None:
        NB_COMPONENTS = args.components
if hasattr(args, 'figure2d'):
    if args.figure2d != None:
        FIGURE_2D = args.figure2d
if hasattr(args, 'figure3d'):
    if args.figure3d != None:
        FIGURE_3D = args.figure3d
if hasattr(args, 'resources'):
    if args.resources != None:
        RESOURCES = json.loads(open(args.resources).read())
if hasattr(args, 'output'):
    if args.output != None:
        OUTPUT = args.output

# Liste des couleurs
# Exemple : 'rgb' => red blue green
COLORS = RESOURCES["colors"]

# Clés de chaques classes
# Exemple : [1, 2, 3]
KEYS = range(len(COLORS))

# Liste des noms des classes
# Exemple : numpy.array(['A', 'B', 'C'])
# len(CLASSES_NAMES) === LEN(COLORS)
CLASSES_NAMES = numpy.array(RESOURCES["classes_names"])

# Matrice numpy
# Exemple : numpy.matrix([[0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 2, 2], ...])
MATRIX = numpy.matrix(RESOURCES["matrix"])

# Tableau qui définit la classe de chaque vecteurs
# Exemple : numpy.array([1, 2, 3]) => Le premier vecteur [0, 0, 0, 0] appartient à la classe CLASSES_NAMES[1] -> 'A'
# len(TARGET) === len(MATRIX)
TARGETS = numpy.array(RESOURCES["targets"]) # Matrice numpy Tableau de classes de chaque individu

# PCA
data = scale(MATRIX)
pca = PCA(n_components=NB_COMPONENTS)
space = pca.fit(data)
matrix = space.transform(data)

# Dessine chaque dimension
if FIGURE_2D:
    plt.figure()
    for c, j, name in zip(COLORS, KEYS, CLASSES_NAMES):
        plt.scatter(matrix[TARGETS == j, 0], matrix[TARGETS == j, NB_COMPONENTS - 1], c=c, label=name)
    plt.legend()
    plt.title('PCA of my dataset axes [' + str(0) + '-' + str(NB_COMPONENTS - 1) + ']')
    for i in range(NB_COMPONENTS - 1):
        plt.figure()
        for c, j, name in zip(COLORS, KEYS, CLASSES_NAMES):
            plt.scatter(matrix[TARGETS == j, i], matrix[TARGETS == j, i + 1], c=c, label=name)
        plt.legend()
        plt.title('PCA of my dataset axes [' + str(i) + '-' + str(i + 1) + ']')
    plt.show()

# Creation de la vue 3D
if FIGURE_3D:
    figure3D = plt.figure()
    ax = figure3D.add_subplot(111, projection='3d')
    for c, j, name in zip(COLORS, KEYS, CLASSES_NAMES):
        ax.scatter(matrix[TARGETS == j, 0], matrix[TARGETS == j, 1], matrix[TARGETS == j, 2], c=c)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    scatters_proxy = []
    for color in COLORS:
        scatters_proxy.append(matplotlib.lines.Line2D([0],[0], linestyle="none", c=color, marker = 'o'))

    ax.legend(scatters_proxy, CLASSES_NAMES, numpoints = 1)

    plt.show()

# Creation de la video
if (len(OUTPUT) > 0):
    print "Recording..."
    # Create a figure and a 3D Axes
    scene = plt.figure()
    ax = Axes3D(scene)

    # Create an init function and the animate functions.
    # Both are explained in the tutorial. Since we are changing
    # the the elevation and azimuth and no objects are really
    # changed on the plot we don't have to return anything from
    # the init and animate function. (return value is explained
    # in the tutorial.
    def init():
        for c, j, name in zip(COLORS, KEYS, CLASSES_NAMES):
            ax.scatter(matrix[TARGETS == j, 0], matrix[TARGETS == j, 1], matrix[TARGETS == j, 2], c=c, s=20)
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_zlabel('Z')

    def animate(i):
        ax.view_init(elev=10., azim=i)

    # Animate
    anim = animation.FuncAnimation(scene, animate, init_func=init, frames=RESOURCES["video"]["frames"], interval=RESOURCES["video"]["interval"], blit=True)
    # Save
    anim.save(filename=OUTPUT, writer=RESOURCES["video"]["writer"], fps=RESOURCES["video"]["fps"], dpi=RESOURCES["video"]["dpi"])
    print "File saved at " + OUTPUT