import UnityEngine as ue
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
# import seaborn as sns
# import matplotlib.patches as patches
# import matplotlib.colors as colors
# from sklearn import preprocessing
# from pylab import *
# from matplotlib import style
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"

# 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 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(figsize=(10,10))
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_map1 = cm.ScalarMappable(cmap=cmap1)
# color_map1.set_array(dz1)
# fig.colorbar(color_map1)

# color_map2 = cm.ScalarMappable(cmap=cmap2)
# color_map2.set_array(dz2)
# fig.colorbar(color_map2)

# color_map3 = cm.ScalarMappable(cmap=cmap3)
# color_map3.set_array(dz3)
# fig.colorbar(color_map3)

# 8. 
plt.show()

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