#pylint: disable=W0401,W0614,W0622
#All pylint codes: http://pylint.pycqa.org/en/latest/technical_reference/features.html

from pylab import *
import h5py
from density_center import def_dc
import ellipsoids
import scipy.optimize

# Simumation parameters
h0 = 0.6774

# Halo parameters
file_name = 'data/subhalo_411321.hdf5'
# centers_file_name = 'data/centers_411321.hdf5'

### Snapshot parameters ###
snapshot = 99
a        = 1.0 # Should read this value from somewhere!

# Read the centre from separate file
# DISABLED: we calculate on our own.
# f = h5py.File(centers_file_name, 'r')
# X_center = f[str(snapshot)]['Coordinates'][...]
# V_center = f[str(snapshot)]['Velocities'][...]
# f.close()

# Dictionary of particle types
particle_types = {}
particle_types['gas']   = '0'
particle_types['dm']    = '1'
particle_types['stars'] = '4'
particle_types['bhs']   = '5'

# Read stars
particle_type = 'stars'
with h5py.File(file_name, 'r') as f:
    m = f[str(snapshot)][particle_types[particle_type]]['Masses'][...]
    X = f[str(snapshot)][particle_types[particle_type]]['Coordinates'][...] * a / h0
    V = f[str(snapshot)][particle_types[particle_type]]['Velocities'][...] * sqrt(a)
M_tot = sum(m)

# Calculate density centre and shift appropriately
X_center_new, V_center_new = def_dc(m, X, V)
X -= X_center_new
V -= V_center_new

# Rotate such that the short axis is z and medium axis is y
r = linalg.norm(X, axis=1)
rh = median(r)
mask = r < 2*rh
Q = ellipsoids.quadrupole_tensor(*X[mask].T, m[mask])
eigenvalues, eigenvectors = np.linalg.eig(Q)
R = ellipsoids.rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues)
X_new = (R @ X.T).T
x, y, z = X_new.T

# d is the axial distance
d = sqrt(x**2 + y**2)

# Cread a two-dimensional grid
# The grid sizes don't _have to_ be equal in both directions, but there is some
# logic in keeping them the same.
n_grid_d, n_grid_z = 16, 16
d_max = 2*median(d)
z_max = 2*median(abs(z))
d_grid = linspace(0, d_max, n_grid_d+1)
z_grid = linspace(0, z_max, n_grid_z+1)

# Calculates the mass in each (d,z)-cell
# We fold negative z-values assuming symmetry with respect to the xy-plane
values, _, _ = histogram2d(d, abs(z), bins=[d_grid, z_grid], weights=m)
values = values.T # Needed because how histogram2d workds

# Calculate the volume of each (d,z)-cell (cylinder subtraction)
d_edges, z_edges = meshgrid(d_grid, z_grid, indexing='xy')
volumes = pi*(d_edges[1:,1:]**2 - d_edges[1:,:-1]**2)*(z_edges[1:,1:] - z_edges[:-1,1:])

# Finally we have the density as a function of d and z. The normalization is to
# make the numbers easier to work with for the minimization routine.
rho_measured_normalized = values/volumes/M_tot

# Define Plummer and Miyamoto-Nagai density
rho_plummer = lambda r, M, b: (3*M/(4*pi*b**3))*(1+(r/b)**2)**(-2.5)
rho_mn = lambda d, z, M, a, b: \
    (b**2 * M / (4*pi)) * \
    (a*d**2 + (a + 3*sqrt(z**2+b**2))*(a+sqrt(z**2+b**2))**2) / \
    ((d**2 + (a+sqrt(z**2+b**2))**2)**2.5 * (z**2+b**2)**1.5)

# The minimization procedure
means = lambda arr: .5*(arr[:-1]+arr[1:]) # small helper function
# Define a grid of the centre of each cell from the histogram we created earlier
dd, zz = meshgrid(means(d_grid), means(z_grid), indexing='xy')
def cost(args):
    a_mn, b_mn = args
    rho = rho_mn(dd, zz, 1, a_mn, b_mn)
    square_diff = (rho - rho_measured_normalized)**2
    return sum(square_diff)
minimization_result = \
    scipy.optimize.minimize(cost, [median(abs(z)), median(d)], method='Nelder-Mead', tol=1e-6, options={'maxiter':5000})
a_mn, b_mn = minimization_result.x
####
if False:
    print('!!!Setting parameters artificially!!!')
    a_mn, b_mn = 2.98, 1.61
####
rho_best_fit = rho_mn(dd, zz, 1, a_mn, b_mn)
print(f'a_mn = {a_mn:.4f} kpc   b_mn = {b_mn:.4f} kpc')
print(f'M_mn = {M_tot:.2e} MSun')

# Compare with Matteo's results
with h5py.File(file_name, 'r') as f:
    particle_id = f[str(snapshot)][particle_types[particle_type]]['ParticleIDs'][...]
    matteo_id = f[str(snapshot)][particle_types[particle_type]]['IDs_truncated'][...]
    matteo_component_tag = f[str(snapshot)][particle_types[particle_type]]['Component'][...]

# Find the array indices in the original arrays that appear in Matteo's list
i = particle_id.argsort()
matteo_disk_ids = matteo_id[matteo_component_tag==1]
i_matteo_disk = i[searchsorted(particle_id[i], matteo_disk_ids)]
matteo_bulge_ids = matteo_id[matteo_component_tag==2]
i_matteo_bulge = i[searchsorted(particle_id[i], matteo_bulge_ids)]

# Print Matteo's results
M_matteo_bulge = sum(m[i_matteo_bulge])
M_matteo_disk = sum(m[i_matteo_disk])
rh_matteo_bulge = median(r[i_matteo_bulge])
print('=== Matteo\'s results ===')
print(f'M_bulge = {M_matteo_bulge:.2e} MSun   M_disk = {M_matteo_disk:.2e} MSun')
print(f'rh_bulge = {rh_matteo_bulge:.4f} kpc (equivalent Plummer radius: {0.76642*rh_matteo_bulge:.4f}) kpc')

# Plot the density as a function of d for three values of z.
rho = lambda d, z: rho_mn(d, z, 1, a_mn, b_mn)
for c, i in enumerate([0, 3, 6, 9]):
    semilogy(dd[i,:], M_tot*rho_measured_normalized[i,:], c=f'C{c}', ls='-', label=f'h={1000*zz[i,0]:.0f} pc')
    semilogy(dd[i,:], M_tot*rho(dd[i,:], zz[i,0]), c=f'C{c}', ls='--', alpha=0.5)

    legend()
    xlabel('x [kpc]')
    ylabel(r'$\rho\ [\rm M_\odot\ kpc^{-3}]$')

savefig('subhalo_411321_stellar_fit.png')
show()