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 mpl_toolkits.mplot3d import Axes3D

WIDTH = int(70)
HEIGHT = int(35)
OBSTACLE_PATH = "Assets/Data_image/obstacle.pkl"
POSITION_PATH1 = ue.Application.dataPath + '/Data_position/80/Walk1.csv'
POSITION_PATH2 = ue.Application.dataPath + '/Data_position/80/Walk2.csv'
POSITION_PATH3 = ue.Application.dataPath + '/Data_position/80/Walk4.csv'
HEATMAP_PATH = "Assets/Data_image/80/heatmap3DMultiple1.png"
SAVEFILE = "Assets/Resources/Json/outputMultiple.json"

def get_data_hight(path):
    dataset = pd.read_csv(path, sep=';', usecols=["Delta Time", "Position x", "Position z"], decimal=',', dtype={'Delta Time': float, 'Position x': float, 'Position z': float})
    dataset["Position x"] = dataset["Position x"].round(0)
    dataset["Position z"] = dataset["Position z"].round(0)
    dataset["Delta Time"] = dataset["Delta Time"].round(2) 

    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
data1 = pd.read_csv(POSITION_PATH1, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
data1 = data1.round(0)
data2 = pd.read_csv(POSITION_PATH2, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
data2 = data2.round(0)
data3 = pd.read_csv(POSITION_PATH3, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
data3 = data3.round(0)

# 2. Group by positions and count appearance
data_count1 = data1.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
data_count2 = data2.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
data_count3 = data3.groupby(['Position x', 'Position z']).size().reset_index(name='counts')

# 3.1 Assign x, y, z, width, depth, height
x1 = data_count1["Position x"].tolist()
y1 = data_count1["Position z"].tolist()
z1 = np.zeros_like(len(x1))
# dz1 = data_count1["counts"].tolist() # Change

x2 = data_count2["Position x"].tolist()
y2 = data_count2["Position z"].tolist()
z2 = np.zeros_like(len(x2))
# dz2 = data_count2["counts"].tolist() # Change

x3 = data_count3["Position x"].tolist()
y3 = data_count3["Position z"].tolist()
z3 = np.zeros_like(len(x3))
# dz3 = data_count3["counts"].tolist() # Change

# 3.2 Assign height of the bars
dz_dynamic1 = get_data_hight(POSITION_PATH1)
dz1 = []
for i in range(len(x1)):
    dz1.append(dz_dynamic1[int(y1[i])][int(x1[i])])

dz_dynamic2 = get_data_hight(POSITION_PATH2)
dz2 = []
for i in range(len(x2)):
    dz2.append(dz_dynamic2[int(y2[i])][int(x2[i])])

dz_dynamic3 = get_data_hight(POSITION_PATH3)
dz3 = []
for i in range(len(x3)):
    dz3.append(dz_dynamic3[int(y3[i])][int(x3[i])])

# 3.3 Add offset to day2 and day3
x2[:] = [a+0.5 for a in x2[:]]
x3[:] = [a+0.5 for a in x3[:]]
y3[:] = [a+0.5 for a in y3[:]]

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

# 5.1 Create custom colormap Day1, Day2, Day3
cmap1 = LinearSegmentedColormap.from_list(name='day1', colors=[(0.40,0.76,0.65), (0.11,0.62,0.47)])
cmap2 = LinearSegmentedColormap.from_list(name='day2', colors=[(0.99,0.55,0.38), (0.85,0.37,0.01)])
cmap3 = 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_array1 = np.array(data_count1['counts'])
fracs1 = dz_array1.astype(float) / dz_array1.max()
color_values1 = cmap1(fracs1.tolist())

dz_array2 = np.array(data_count2['counts'])
fracs2 = dz_array2.astype(float) / dz_array2.max()
color_values2 = cmap2(fracs2.tolist())

dz_array3 = np.array(data_count3['counts'])
fracs3 = dz_array3.astype(float) / dz_array3.max()
color_values3 = cmap3(fracs3.tolist())

# 6. Create the bars
img = ax.bar3d(x1, y1, z1, 0.5, 0.5, dz1, color=color_values1, shade=False)
img = ax.bar3d(x2, y2, z2, 0.5, 0.5, dz2, color=color_values2, shade=False)
img = ax.bar3d(x3, y3, z3, 0.5, 0.5, dz3, color=color_values3, shade=False)

# 7. Create Colorbar
color_intense1 = data_count1['counts'].tolist()
color_map1 = cm.ScalarMappable(cmap=cmap1)
color_map1.set_array(color_intense1)
fig.colorbar(color_map1)

color_intense2 = data_count2['counts'].tolist()
color_map2 = cm.ScalarMappable(cmap=cmap2)
color_map2.set_array(color_intense2)
fig.colorbar(color_map2)

color_intense3 = data_count3['counts'].tolist()
color_map3 = cm.ScalarMappable(cmap=cmap3)
color_map3.set_array(color_intense3)
fig.colorbar(color_map3)

# 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
dz1 = (np.array(dz1)/np.array(dz1).max()).tolist()
dz2 = (np.array(dz2)/np.array(dz2).max()).tolist()
dz3 = (np.array(dz3)/np.array(dz3).max()).tolist()

jsonData = {"x1" : x1,
            "y1" : y1,
            "z1" : z1.tolist(),
            "dz1" : dz1,
            "color_values1" : color_values1.flatten().tolist(),
            "color_intense1" : color_intense1,
            "x2" : x2,
            "y2" : y2,
            "z2" : z2.tolist(),
            "dz2" : dz2,
            "color_values2" : color_values2.flatten().tolist(),
            "color_intense2" : color_intense2,
            "x3" : x3,
            "y3" : y3,
            "z3" : z3.tolist(),
            "dz3" : dz3,
            "color_values3" : color_values3.flatten().tolist(),
            "color_intense3" : color_intense3}

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