import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from scipy.stats import gaussian_kde

def TrajHM(x, y, name):
    # Perform kernel density estimation (KDE) on a grid
    x = np.array(x)
    y = np.array(y)
    #xmin, xmax = x.min() - 1, x.max() + 1  # Define x range with padding
    #ymin, ymax = y.min() - 1, y.max() + 1  # Define y range with padding
    xmin, xmax = -2, 10
    ymin, ymax = -2, 10
    xx, yy = np.mgrid[xmin:xmax:100j, ymin:ymax:100j]  # Create a mesh grid (100x100 resolution)
    #print('xx', xx)
    #print('yy', yy)

    positions = np.vstack([xx.ravel(), yy.ravel()])  # Flatten the grid for KDE evaluation
    #print(positions[0][1])
    kde = gaussian_kde(np.vstack([x, y]), bw_method=0.5)  # Fit KDE to trajectory data
    density = kde(positions).reshape(xx.shape)  # Compute density on the grid
    # Normalize density values to range [0,1]
    density = (density - density.min()) / (density.max() - density.min())
    # density+=0.1

    # Create custom colormap that sets lowest density to white
    custom_cmap = LinearSegmentedColormap.from_list("custom_cmap", ["white", "red", "orange", "yellow"], N=256)

    # Plot smooth heatmap
    plt.figure(figsize=(8, 8))
    #plt.imshow(density.T, origin='lower', extent=[0, 10, 0, 10], cmap=custom_cmap, aspect='auto')
    plt.scatter(x[-1], y[-1], s=100, facecolors='none', edgecolors='blue')
    plt.imshow(density.T, origin='lower', extent=[xmin, xmax, ymin, ymax], cmap=custom_cmap, aspect='auto')
    plt.colorbar(label="Density")  # Add color bar
    plt.xlabel("X-axis")
    plt.ylabel("Y-axis")
    plt.title("Algorithm Trajectory Heatmap")
    plt.savefig(name)
    plt.close()