import numpy as np
from halotools.empirical_models import PrebuiltSubhaloModelFactory
model = PrebuiltSubhaloModelFactory('smhm_binary_sfr')
from halotools.sim_manager import CachedHaloCatalog
halocat = CachedHaloCatalog(simname = 'zheng07', redshift = 0, halo_finder = 'rockstar')
model.populate_mock(halocat)

sample_mask = model.mock.galaxy_table['stellar_mass'] > 1e10
gals = model.mock.galaxy_table[sample_mask]

#Compute total stellar mass
from halotools.utils import group_member_generator
gals.sort('halo_hostid')
grouping_key = 'halo_hostid'
requested_columns = ['stellar_mass']
group_gen = group_member_generator(gals, grouping_key, requested_columns)

total_stellar_mass = np.zeros(len(gals))
for first, last, member_props in group_gen:
    stellar_mass_of_members = member_props[0]
    total_stellar_mass[first:last] = sum(stellar_mass_of_members)

gals['halo_total_stellar_mass'] = total_stellar_mass

#Compute host halo mass
gals.sort(['halo_hostid', 'halo_upid'])
grouping_key = 'halo_hostid'
requested_columns = ['halo_mvir']
group_gen = group_member_generator(gals, grouping_key, requested_columns)

host_mass = np.zeros(len(gals))
for first, last, member_props in group_gen:
    mvir_members = member_props[0]
    mvir_host = mvir_members[0]
    host_mass[first:last] = mvir_host

gals['halo_mhost'] = host_mass

#Compare total vs host mass and plot
from halotools.mock_observables import mean_y_vs_x

bins = np.logspace(12, 15, 25)
result = mean_y_vs_x(gals['halo_mhost'].data,
                     gals['halo_total_stellar_mass'].data,
                     bins = bins,
                     error_estimator = 'variance')

host_mass, mean_stellar_mass, mean_stellar_mass_err = result
print host_mass, mean_stellar_mass, mean_stellar_mass_err

from seaborn import plt

plt.errorbar(host_mass, mean_stellar_mass, yerr=mean_stellar_mass_err,
             fmt = "none", ecolor='gray')
plt.plot(host_mass, mean_stellar_mass, 'D', color='k')

plt.loglog()
plt.xticks(size=18)
plt.yticks(size=18)
plt.xlabel(r'$M_{\rm halo}/M_{\odot}$', fontsize=25)
plt.ylabel(r'$\langle M_{\ast}^{\rm tot}/M_{\odot}\rangle$', fontsize=25)
plt.ylim(ymax=5e12)


#Compute velocities
from halotools import mock_observables
x = halocat.halo_table['halo_x']
y = halocat.halo_table['halo_y']
z = halocat.halo_table['halo_z']
sample = np.vstack((x,y,z)).T
vx = halocat.halo_table['halo_vx']
vy = halocat.halo_table['halo_vy']
vz = halocat.halo_table['halo_vz']
velocities = np.vstack((vx,vy,vz)).T