"""Plots the Antennae using Kernel Density Estimation to map the density"""

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

from analysis.colormaps import red_map
from analysis import utils

def plotKDE(i, theta, phi, view):
    """Plot the Antennae galaxies using kernel density estimation

    Parameters:
        i (int): Timestep at which to plot them
        theta (float): polar angle
        phi (float): azimuthal angle
        view (float): rotation along the line of sight
    """
    data = utils.loadData('antennae', 13500)
    r_vec = data['r_vec'][::]
    
    # Rotate to view from (theta, phi, view) viewpoint
    M1 = np.array([[1, 0, 0],
        [0, np.cos(theta), np.sin(theta)],
        [0, -np.sin(theta), np.cos(theta)]])
    M2 = np.array([[np.cos(phi), np.sin(phi), 0],
        [-np.sin(phi), np.cos(phi), 0],
        [0, 0, 1]])
    M3 = np.array([[np.cos(view), np.sin(view), 0],
        [-np.sin(view), np.cos(view), 0],
        [0, 0, 1]])
    M = np.matmul(M2, np.matmul(M1, M3))
    r_vec = np.tensordot(r_vec, M, axes=[1,0])
    
    # Plotting
    plt.figure(figsize=(4,5), dpi=400)
    # Perform Kernel Density Estimation to colour by density
    xy = np.vstack([r_vec[:,0],r_vec[:,1]])
    c = gaussian_kde(xy)(xy)
    # Densest points are plotted last
    idx = c.argsort()
    x, y, c = r_vec[:,0][idx], r_vec[:,1][idx], c[idx]

    plt.scatter(x, y, c=c, s=3, edgecolor='', cmap=red_map)
    plt.axis('square')
    plt.show()
    

plotKDE(135, 1.86, 4.10, 5.10)