import UnityEngine as ue
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.colors import LinearSegmentedColormap 
import pandas as pd

WIDTH = int(70)
HEIGHT = int(35)
OBSTACLE_PATH = "Assets/Data_image/obstacle.pkl"
POSITION_PATH = ue.Application.dataPath + '/Data_position/Walk1.csv'
HEATMAP_PATH = "Assets/Data_image/heatmap1.png"

# Generate only if obstacles change
def set_obstacles():
    global list_obstacle_length
    positions = pd.DataFrame(np.zeros((HEIGHT, WIDTH)))
    obstacles = ue.Object.FindObjectsOfType(ue.GameObject)
    for obstacle in obstacles:
        if(obstacle.layer == 15 or obstacle.layer == 12):
            startWidth = int(obstacle.transform.position.x - obstacle.transform.localScale.x / 2)
            endWidth = int(obstacle.transform.position.x + obstacle.transform.localScale.x / 2)
            startHeight = int(obstacle.transform.position.z - obstacle.transform.localScale.z / 2)
            endHeight = int(obstacle.transform.position.z + obstacle.transform.localScale.z / 2)
            for currentW in range(startWidth, endWidth, 1):
                for currentH in range(startHeight, endHeight, 1):
                    positions[currentW][currentH] = -500
    positions.to_pickle(OBSTACLE_PATH)

def set_patches(plt):
    obstacles = ue.Object.FindObjectsOfType(ue.GameObject)
    for obstacle in obstacles:
        if(obstacle.layer == 15):
            startWidth = int(obstacle.transform.position.x - obstacle.transform.localScale.x / 2)
            endWidth = int(obstacle.transform.position.x + obstacle.transform.localScale.x / 2)
            startHeight = int(obstacle.transform.position.z - obstacle.transform.localScale.z / 2)
            endHeight = int(obstacle.transform.position.z + obstacle.transform.localScale.z / 2)
            # plt.gca().add_patch(
            plt.add_patch(
                patches.Rectangle(
                    (startWidth, startHeight),
                    endWidth - startWidth,
                    endHeight - startHeight,
                    fill=False,
                    color='black'
                )
            )

# 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. Create wide-form DataFrame for generating heatmap
positions = data_count.loc[:,:].reset_index().pivot(index='Position z', columns='Position x', values='counts')

# 4. Fill missing values
positions.fillna(0, inplace=True)

# 5. reindex DataFrame (70,35) size of Surface; (70, 35) first x width then z height
positions = positions.reindex_axis(range(0, HEIGHT), axis=0, fill_value=0)
positions = positions.reindex_axis(range(0, WIDTH), axis=1, fill_value=0)

# 6. Get obstacles (obstacles, market stalls) and paste specific value in positions
# Therefore save pkl file in folder and read from it afterwards
set_obstacles()
positions_heatmap = pd.read_pickle(OBSTACLE_PATH)

# 6.1 Merge positions data with obstacles data
positions_heatmap.where(positions_heatmap != 0, positions, inplace=True)

# 6.2 Debug Output 
# positions_heatmap = pd.DataFrame(positions_heatmap)
# positions_heatmap.to_html('Assets/Data_image/positions_heatmap.html')

# 7. Plot the heatmap
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2)
ax4.remove()
# cmap = LinearSegmentedColormap.from_list(name='greenToRed', colors=['grey', 'limegreen', 'chartreuse', 'yellow', 'darkorange', 'red'])
cmap19 = LinearSegmentedColormap.from_list(name='2019', colors=['grey', (0.40,0.76,0.65), (0.11,0.62,0.47)])
cmap20 = LinearSegmentedColormap.from_list(name='2020', colors=['grey', (0.99,0.55,0.38), (0.85,0.37,0.01)])
cmap21 = LinearSegmentedColormap.from_list(name='2021', colors=['grey', (0.55,0.63,0.80), (0.46,0.44,0.70)])
heatmap = sns.heatmap(positions_heatmap, cmap=cmap19, cbar=False, square=True, yticklabels=False, xticklabels=False, ax=ax1)
heatmap.invert_yaxis()
heatmap.set_title(cmap19.name)
heatmap2 = sns.heatmap(positions_heatmap, cmap=cmap20, cbar=False, square=True, yticklabels=False, xticklabels=False, ax=ax2)
heatmap2.invert_yaxis()
heatmap2.set_title(cmap20.name)
heatmap3 = sns.heatmap(positions_heatmap, cmap=cmap21, cbar=False, square=True, yticklabels=False, xticklabels=False, ax=ax3)
heatmap3.invert_yaxis()
heatmap3.set_title(cmap21.name)
plt.xlabel('')
plt.ylabel('')

# 9. Mark the Market stalls
set_patches(ax1)
set_patches(ax2)
set_patches(ax3)

# 9.1 
plt.show()

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