import UnityEngine as ue
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
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_PATH = ue.Application.dataPath + "/Data_position/80/Walk1.csv"
HEATMAP_PATH = "Assets/Data_image/80/heatmap3D1.png"

def get_data_hight():
    dataset = pd.read_csv(POSITION_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
data = pd.read_csv(POSITION_PATH, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
data = data.round(0)

# 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()
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=['grey', (0.99,0.55,0.38), (0.85,0.37,0.01)])
# cmap = LinearSegmentedColormap.from_list(name='day3', colors=['grey', (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())

# 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_intens = data_count['counts'].tolist()
color_map = cm.ScalarMappable(cmap=cmap)
color_map.set_array(color_intens)
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)