import struct
import matplotlib.cm as cm
import numpy as np
import os
import sys
import yt
import yt.analysis_modules.sphgr.api as sphgr
import yt.analysis_modules.sphgr.group_funcs as gf
from readgadget import *

SCALE  = 5.
Npoints = 100000

snap   = sys.argv[1]
output = sys.argv[2]

basedir  = os.path.dirname(os.path.abspath(snap))
snapname = os.path.basename(snap)

# load data
ds   = yt.load(snap)
obj  = sphgr.load_sphgr_data('%s/sphgr_%s' % (basedir,snapname), ds)

#glist = obj.global_particle_lists.halo_glist[::100]
glist  = obj.global_particle_lists.halo_glist
indexes = np.where(glist > -1)[0]

#if len(indexes) > Npoints:
#    indexes = indexes[::int(float(len(indexes))/float(Npoints))]

# center and rotate galaxy
pos = obj.particle_data['gpos']/obj.boxsize  - 0.5 #- (obj.boxsize/2.)
hsml= obj.particle_data['ghsml']/obj.boxsize
rho = np.log10(obj.particle_data['grho'].to('g/cm**3')/1.67e-24)

pos *= SCALE
hsml*= SCALE

pos = pos[indexes]
hsml= hsml[indexes]
rho = rho[indexes]
print len(indexes)


#indexes = np.where(rho > -4)[0]
#pos = pos[indexes]
#hsml = hsml[indexes]
#rho = rho[indexes]

# map densities onto a color map (rgb)
rgb = cm.ScalarMappable(cmap=cm.viridis).to_rgba(rho)
rgb = rgb[:,0:3]

starpos = obj.particle_data['spos']/obj.boxsize - 0.5
sage    = obj.particle_data['sage']

starrgb = cm.ScalarMappable(cmap=cm.autumn).to_rgba(sage)
starrgb = starrgb[:,0:3]
starhsml = np.full(len(sage),np.median(hsml))

starpos *= SCALE
starhsml*= SCALE

pos = np.concatenate((pos,starpos))
hsml = np.concatenate((hsml, starhsml))
rgb = np.concatenate((rgb,starrgb))

print len(hsml)

#from pylab import *
#plot(pos[:,0],pos[:,1],'o',ms=1)
#plot(rho,t,'o',ms=1)
#show()

#"""
# write binary data
f = open(output,'wb')
f.write(struct.pack('<I',len(pos)))
for data in [pos,rgb,hsml]:
    f.write(data.astype('f').tostring())
f.write(struct.pack('<I',len(pos)))
f.close()
#"""