#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
from astropy.cosmology import Planck15

def miyamoto_nagai_params_from_medians(m_d, m_z):
    m_d_pred = lambda x: exp(1.43475163 + (1.04827148*log(x)-1.09023112)*(arctan(log(x)/2.30939056)/pi+0.5))
    m_z_pred = lambda x: 5.77340E-01*x
    b_over_a = scipy.optimize.brentq(lambda x: m_d_pred(x)/m_z_pred(x)-m_d/m_z, 1/64, 64)
    a1 = m_d/m_d_pred(b_over_a)
    a2 = m_z/m_z_pred(b_over_a)
    a = 0.5*(a1 + a2)
    b = a*b_over_a
    return a, b

# Simumation parameters
h0 = 0.6774
snap, z = loadtxt('snapshots.dat', unpack=True)
snap = snap.astype(int)
a = 1/(1+z)
t = Planck15.age(z).value

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

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

def get_half_mass_radius(m, r):
    i = argsort(r)
    r_sorted = r[i]
    m_sorted = m[i]
    m_cum = cumsum(m_sorted)
    j = searchsorted(m_cum, m_cum[-1]/2)
    return r_sorted[j] # close enough

def get_transformation(m, X, V=None):
    V_center = None
    if not V is None: X_center, V_center = def_dc(m, X, V)
    else: X_center = def_dc(m, X)
    X_shifted = X - X_center
    r = linalg.norm(X_shifted, axis=1)
    rh = get_half_mass_radius(m, r)
    mask = r < 2*rh
    Q = ellipsoids.quadrupole_tensor(*X_shifted[mask].T, m[mask])
    eigenvalues, eigenvectors = linalg.eig(Q)
    R = ellipsoids.rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues)
    return X_center, V_center, R, mask

f = h5py.File(file_name, 'r')

# First get the global coordinate system
i, snapshot = len(snap)-1, snap[-1]
m = f[str(snapshot)][particle_types['stars']]['Masses'][...]
X = f[str(snapshot)][particle_types['stars']]['Coordinates'][...] * a[i] / h0
X_center_glob, _, R_glob, _ = get_transformation(m, X)

for i in range(len(snap)):
    snapshot = snap[0] + i

    m_dm  = f[str(snapshot)][particle_types['dm']]['Masses'][...]
    X_dm  = f[str(snapshot)][particle_types['dm']]['Coordinates'][...] * a[i] / h0
    m_gas = f[str(snapshot)][particle_types['gas']]['Masses'][...]
    X_gas = f[str(snapshot)][particle_types['gas']]['Coordinates'][...] * a[i] / h0
    m = append(m_dm, m_gas)
    X = vstack([X_dm, X_gas])

    X = (R_glob @ (X - X_center_glob).T).T
    X_center_halo = def_dc(m, X)
    r = linalg.norm(X-X_center_halo, axis=1)
    rh = get_half_mass_radius(m, r)
    b_halo_plummer = 0.76642093654*rh
    M_halo_plummer = sum(m)

    j = r.argsort()
    rr = r[j]
    m_cumulative = cumsum(m[j])/sum(m)
    nfw_cumulative_mass = lambda rho_0, b, r: 4*pi*rho_0*b**3*(log(1+r/b) - r/(b+r))
    cost = lambda args: sum((nfw_cumulative_mass(*args, rr) - m_cumulative)**2)
    res = scipy.optimize.minimize(cost, [1e-3, 30.], method='Nelder-Mead')
    rho_0_nfw, b_nfw = res.x
    rho_0_nfw *= sum(m)
    lsq_nfw = res.fun/len(rr)

    m = f[str(snapshot)][particle_types['stars']]['Masses'][...]
    X = f[str(snapshot)][particle_types['stars']]['Coordinates'][...] * a[i] / h0
    V = f[str(snapshot)][particle_types['stars']]['Velocities'][...] * sqrt(a[i])
    X = (R_glob @ (X - X_center_glob).T).T
    V = (R_glob @ V.T).T

    X_center_stars, V_center_stars, R, mask = get_transformation(m, X, V)
    direction = R[2,:]
    if direction[2] < 0: direction = -direction
    theta_inertia = arccos(direction[2]/linalg.norm(direction))
    phi_inertia   = arctan2(direction[1], direction[0])
    L = cross(X-X_center_stars, V-V_center_stars)
    L = sum(L[mask], axis=0)
    phi_L   = arctan2(L[1], L[0])
    theta_L = arccos(L[2]/linalg.norm(L))

    X_new = (R @ (X - X_center_stars).T).T
    x, y, z = X_new.T
    m_z = median(abs(z[mask]))
    m_d = median(sqrt(x[mask]**2+y[mask]**2))
    a_mn, b_mn = miyamoto_nagai_params_from_medians(m_d, m_z)
    M_disk = sum(m)
    print('%d %.8E   %15.8E %15.8E %15.8E   %15.8E %15.8E   %15.8E %15.8E   %.8E %.8E %.8E   %15.8E %15.8E %15.8E   %.8E %.8E   %.8E %.8E %.8E' % (snapshot, t[i], *X_center_stars, phi_inertia, theta_inertia, phi_L, theta_L, M_disk, a_mn, b_mn, *X_center_halo, M_halo_plummer, b_halo_plummer, rho_0_nfw, b_nfw, lsq_nfw))

f.close()