import numpy
import pylab
from pymses import RamsesOutput
from pymses.utils.regions import Cylinder
from pymses.analysis import sample_points, bin_cylindrical
from pymses.utils import constants as C

# Galactic cylinder parameters
gal_center = [ 0.567811, 0.586055, 0.559156 ]        # in box units
gal_radius = 0.00024132905460547268                  # in box units
gal_thickn = 0.00010238202316595811                  # in box units
gal_normal = [ -0.172935, 0.977948, -0.117099 ]      # Norm = 1

# RamsesOutput
ro = RamsesOutput("/data/Aquarius/output", 193)

# Prepare to read the density field only
source = ro.amr_source(["rho"])

# Cylinder region
cyl = Cylinder(gal_center, gal_normal, gal_radius, gal_thickn)

# AMR density field point sampling
points = cyl.random_points(1.0e6) # 1M sampling points
point_dset = sample_points(source, points)
rho_weight_func = lambda dset: dset["rho"]
r_bins = numpy.linspace(0.0, gal_radius, 200)

# Profile computation
rho_profile = bin_cylindrical(point_dset, gal_center, gal_normal,
		rho_weight_func, r_bins, divide_by_counts=True)
	
# Plot
# Geometrical midpoint of the bins
length = ro.info["unit_length"].express(C.kpc)
bins_centers = (r_bins[1:]+r_bins[:-1])/2. * length
dens = ro.info["unit_density"].express(C.H_cc)
pylab.semilogy(bins_centers, rho_profile * dens, 'r-')
pylab.xlabel("r (kpc)")
pylab.ylabel(r"$\rho$ (H/cc)")
pylab.title("Cylindrical density profile")