"""
Author:
    Dr. Thor Tepper-García

About:
    Visualise the potential created with save_6dpytreegrav.py
    Data structure:
    x | y | z | vx | vy | vz | mass | Specigic potential
    Coordinates in kpc, velocities in km/s, masses in Msunm potential in[ km/s]^2

Run:
    $> py39 plot_pot.py

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



range = 10.
data_range = ([-range,range],[-range,range])
nbins = 128
pot_scale = 1e4

data_dir = './data/'

# Load data
filen = f'{data_dir}/hbd_7_hyrdisc_sink_4_comp5_1_3_6dpytreegrav_skip10_snap00101.npy'

print(f'Loading data {filen}')
phase_space = np.load(filen)
print('Done.')

# Retrieve coordinates and potential
x = phase_space[:,0]
y = phase_space[:,1]
z = phase_space[:,2]
pot = phase_space[:,7] / pot_scale

# Calculate projected historgams
histo_xy = \
        stats.binned_statistic_2d(x, y, pot, \
        statistic='median', \
        bins=nbins, range=data_range)
img_data_xy = histo_xy.statistic.T

histo_xz = \
        stats.binned_statistic_2d(x, z, pot, \
        statistic='median', \
        bins=nbins, range=data_range)
img_data_xz = histo_xz.statistic.T

# Visualise
fig, axs = plt.subplots(1,2, figsize=(9,3.5))
_extent = (*data_range[0], *data_range[1])

ax = axs[0]
ax.set_xlabel('$x$ (kpc)')
ax.set_ylabel('$y$ (kpc)')
img = ax.imshow(img_data_xy,
    cmap='tab20c',
    extent=_extent, origin='lower', aspect='equal')
ax = axs[1]
ax.set_xlabel('$x$ (kpc)')
ax.set_ylabel('$z$ (kpc)')
img = ax.imshow(img_data_xz,
    cmap='tab20c',
    extent=_extent, origin='lower', aspect='equal')
plt.subplots_adjust(bottom=0.15)
# add colour bar
pos1 = ax.get_position()  # Get current position
plt.subplots_adjust(right=0.96*pos1.x1)
pos1 = ax.get_position()  # Get current position
sm = plt.cm.ScalarMappable(cmap='tab20c',
    norm=plt.Normalize(vmin=None, vmax=None))
pos2 = [pos1.x1, pos1.y0, 0.02, pos1.height]
cbar = fig.add_axes(pos2) # (left, bottom, width, length)
plt.colorbar(sm, cax=cbar, label=f'Potential ({pot_scale:.0e} km$^2$s$^{-2}$)')

plt.show()
plt.close()