import numpy as np
import pandas as pd
from scipy.stats import multivariate_normal
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import aaaaa

# checked; it is ok
# To generate distorsed covariance matrixes
def genS(sds, R, distor=1):
    # sds: standard deviations
    # R: correlation matrix
    # distor: level of distorsion
    sdsdis = sds * distor
    S = R * np.outer(sdsdis, sdsdis)
    np.fill_diagonal(S, sdsdis ** 2)
    return S

original_data_path = "/home/scpmaotj/Github/hcr/reference_data/"
synthetic_data_path = "/home/scpmaotj/Github/hcr/synthetic_data/"
plots_path = "/home/scpmaotj/Github/hcr/plots/"
nkey = 21

#Nivel distorsion
# 1 - Distorsionamos las desviaciones sin distorsionar la estructura de correlacion
# Funcion genS()
d = [1, 1.01, 1.05, 1.1, 1.2, 1.5, 1.7, 2]
# 2 - Distorsionamos las desviaciones y la estrucctura de correlacion
# Funcion???

# Intento sacar una mano al azar de una normal multivariante.
# Sacaremos m=100 observaciones
m = 100
#gesto <- 1

#for gesto in range(1, 43):
for gesto in range(1, 2):
    print(gesto)
    datXYZ = aaaaa.get_datXYZ(gesto, original_data_path)
    # Hacemos un grafico con ndib= 16 lecturas, no las 16 primeras sino desplazadas pos=4 posiciones
    ndib = 16
    pos = 4
    hands = datXYZ[pos:pos+ndib].reshape(ndib, nkey, 3)
    fig, axes = plt.subplots(4, 4, subplot_kw={'projection': '3d'})
    for i, ax in enumerate(axes.flat):
        ax.scatter(hands[i, :, 0], hands[i, :, 1], hands[i, :, 2], c='b', marker='o')
        ax.set_title(f"Original {i+pos}")
    plt.savefig(plots_path + f"{gesto}orig.pdf")
    plt.close()

    obsmean = np.mean(datXYZ, axis=0)
    obssd = np.std(datXYZ, axis=0)
    R = np.corrcoef(datXYZ, rowvar=False)

    id = len(d)
    np.random.seed(3663)

    for i in range(id):
        Sberria = genS(obssd, R, distor=d[i])
        z = multivariate_normal.rvs(mean=obsmean, cov=Sberria, size=m)
        np.savetxt(synthetic_data_path + f"gesto{gesto}d{i-1}.csv", z, delimiter=',', fmt='%.6f')
        simhands = z[pos:pos+ndib].reshape(ndib, nkey, 3).transpose(0, 2, 1)
        fig, axes = plt.subplots(4, 4, subplot_kw={'projection': '3d'})
        for j, ax in enumerate(axes.flat):
            ax.scatter(simhands[j, 0, :], simhands[j, 1, :], simhands[j, 2, :], c='b', marker='o')
            ax.set_title(f"Simulated, d={i-1}, It. {j+4}")
        plt.savefig(plots_path + f"gesto{gesto}_3coord_d{i-1}.pdf")
        plt.close()

# Selecciono un gesto, y miro algunas configuraciones simuladas segun grado de distorsion
gesto = 4
gd = 0
mirar = synthetic_data_path + f"gesto{gesto}d{gd}.csv"
z = np.loadtxt(mirar, delimiter=',')
m = z.shape[0]
simhands = np.zeros((nkey, 3, m))
for i in range(m):
    simhands[:, :, i] = z[i].reshape(nkey, 3)
simhands = simhands.transpose(2, 0, 1)
simhands = np.swapaxes(simhands, 1, 2)
fig, axes = plt.subplots(4, 4, subplot_kw={'projection': '3d'})
for i, ax in enumerate(axes.flat):
    ax.scatter(simhands[i, 0, :], simhands[i, 1, :], simhands[i, 2, :], c='b', marker='o')
    ax.set_title(str(i+1))
plt.show()