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/Walk2.csv'
HEATMAP_PATH = "Assets/Data_image/heatmap2.png"

# Generate only if obstacles change
def set_obstacles(maximum):
    """
    set_obstacles creates a DataFrame which marks the positions of the obstacles. 
    The obstacles are marked with the negative maximum and the file is placed at the path.
    The negative maximum is used so that the cmap can assign the correct colors to the values.
    Since no other negative numbers exist and the next highest number is 0, the negative values are painted 'grey' 
    and the other [0-maximum] are painted according to the color gradient between the two specified colors.
    
    :param maximum: Is the maximum to the associated positions file
    """
    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] = -maximum
    positions.to_pickle(OBSTACLE_PATH)

def set_patches(plt):
    """
    set_patches uses the positions of the obstacles to draw a black border around them. 
    
    :param plt: The plot, where the figure is placed, so it can be modified by this function
    """
    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.1 Get maximum value of positions data
max1 = pd.DataFrame(positions).max().max() # Gets the maximum of each column and than the maximum of the maximums

# 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(max1)
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)