import UnityEngine as ue
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import json
from matplotlib.colors import LinearSegmentedColormap
from sklearn import preprocessing
from mpl_toolkits.mplot3d import Axes3D

WIDTH = int(70)
HEIGHT = int(35)
COND = "0"
INDEX = "1"
POSITION_PATH = ue.Application.dataPath + '/Data_position/'+ COND +'/Walk'+ INDEX +'.csv'
HEATMAP_PATH = 'Assets/Data_image/'+ COND +'/heatmap3D'+ INDEX +'.png'
SAVEFILE = 'Assets/Resources/Json/'+ COND +'/heatmap3D'+ INDEX +'.json'

def get_data_hight(dataset):
    dataset_human_on_same_spot = dataset.groupby(["Delta Time", "Position x", "Position z"]).size().reset_index(name='AmountOneSpot')
    dataset_human_amount_per_deltaTime = dataset_human_on_same_spot.groupby(['Delta Time']).size().reset_index(name='counts')
    
    solution = np.zeros((HEIGHT, WIDTH))
    old = np.zeros((HEIGHT, WIDTH))
    current = np.zeros((HEIGHT, WIDTH))

    minStart = 0
    for i in range(dataset_human_amount_per_deltaTime.shape[0]):
        for count in range(minStart, minStart + int(dataset_human_amount_per_deltaTime['counts'][i])):
            x = int(dataset_human_on_same_spot['Position x'][count])
            y = int(dataset_human_on_same_spot['Position z'][count])

            if i == 0:
                old[y][x] = dataset_human_on_same_spot['AmountOneSpot'][count]
            else :
                current[y][x] = dataset_human_on_same_spot['AmountOneSpot'][count]
        if i > 0:
            solution = solution + np.absolute(old - current)
            old = current
            current = np.zeros((HEIGHT, WIDTH))
        minStart += int(dataset_human_amount_per_deltaTime['counts'][i])

    return solution


# 1. Get position data from csv file
data = pd.read_csv(POSITION_PATH, sep=';', usecols=["Delta Time", "Position x", "Position z"], decimal=',', dtype={'Delta Time': float, 'Position x': float, 'Position z': float})
data["Position x"] = data["Position x"].round(0)
data["Position z"] = data["Position z"].round(0)
data["Delta Time"] = data["Delta Time"].round(2)

for i in range(len(data["Position x"].tolist())):
    if data["Position x"][i] % 2 != 0:
        data["Position x"][i] -= 1
    if data["Position z"][i] % 2 != 0:
        data["Position z"][i] -= 1

# 2. Group by positions and count appearance
data_count = data.groupby(['Position x', 'Position z']).size().reset_index(name='counts')

# 3.1 Assign x, y, z
x = data_count["Position x"].tolist()
y = data_count["Position z"].tolist()
z = np.zeros_like(len(x))

# 3.2 Assign height of the bars
dz_dynamic = get_data_hight(data)
dz = []
for i in range(len(x)):
    dz.append(dz_dynamic[int(y[i])][int(x[i])])

# 4. Create figure and axes
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 5.1 Create custom colormap Day1, Day2, Day3
cmap = LinearSegmentedColormap.from_list(name='day1', colors=[(0.40,0.76,0.65), (0.11,0.62,0.47)])
# cmap = LinearSegmentedColormap.from_list(name='day2', colors=[(0.99,0.55,0.38), (0.85,0.37,0.01)])
# cmap = LinearSegmentedColormap.from_list(name='day3', colors=[(0.55,0.63,0.80), (0.46,0.44,0.70)])

# 5.2 Initialize array for coloring the bars
dz_array = np.array(data_count['counts'])
fracs = dz_array.astype(float) / dz_array.max()
color_values = cmap(fracs.tolist())

# 5.3 calculate alpha value of the bars
alpha_values = np.array(data_count['counts'])
alpha_values = alpha_values.astype(np.float)
alpha_values = (alpha_values - alpha_values.min()) / (alpha_values.max() - alpha_values.min())

# 6. Create the bars
# for i in range(len(x)):
#     img = ax.bar3d(x[i], y[i], z, 1, 1, dz[i], color=color_values[i], shade=False)
img = ax.bar3d(x, y, z, 1, 1, dz, color=color_values, shade=False)

# 7. Create Colorbar
color_intense = data_count['counts'].tolist()
color_map = cm.ScalarMappable(cmap=cmap)
color_map.set_array(color_intense)
fig.colorbar(color_map)

# 8. 
ax.set_xlabel('Width')
ax.set_ylabel('Height')
ax.set_zlabel('Human dynamic')
plt.show()

# 10. Save 3D Heatmap
# heatmap.get_figure().savefig(HEATMAP_PATH, transparent=True)
# fig.savefig(HEATMAP_PATH, transparent=True)

# 11. Save the data from the Heatmap to create it in Unity
dz = (np.array(dz)/np.array(dz).max()).tolist()

jsonData = {"x1" : x,
            "y1" : y,
            "z1" : z.tolist(),
            "dz1" : dz,
            "alpha_values1" : alpha_values.flatten().tolist(),
            "x2" : 0,
            "y2" : 0,
            "z2" : 0,
            "dz2" : 0,
            "alpha_values2" : np.zeros_like(alpha_values).tolist(),
            "x3" : 0,
            "y3" : 0,
            "z3" : 0,
            "dz3" : 0,
            "alpha_values3" : np.zeros_like(alpha_values).tolist()}

with open(SAVEFILE, 'w') as filehandle:
    json.dump(jsonData, filehandle)