import yt
import numpy as np
from yt.visualization.api import Streamlines
import matplotlib.pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# Load the dataset
ds = yt.load("IsolatedGalaxy/galaxy0030/galaxy0030")

# Define c: the center of the box, N: the number of streamlines,
# scale: the spatial scale of the streamlines relative to the boxsize,
# and then pos: the random positions of the streamlines.
c = ds.arr([0.5]*3, 'code_length')
N = 30
scale = ds.quan(15, 'kpc').in_units('code_length')  # 15 kpc in code units
pos_dx = np.random.random((N,3))*scale-scale/2.
pos = c+pos_dx

# Create the streamlines from these positions with the velocity fields as the
# fields to be traced
streamlines = Streamlines(ds, pos, 'velocity_x', 'velocity_y', 'velocity_z', length=1.0)
streamlines.integrate_through_volume()

# Create a 3D matplotlib figure for visualizing the streamlines
fig=pl.figure()
ax = Axes3D(fig)

# Trace the streamlines through the volume of the 3D figure
for stream in streamlines.streamlines:
    stream = stream[np.all(stream != 0.0, axis=1)]

    # Make the colors of each stream vary continuously from blue to red
    # from low-x to high-x of the stream start position (each color is R, G, B)
    # can omit and just set streamline colors to a fixed color
    x_start_pos = ds.arr(stream[0,0], 'code_length')
    x_start_pos -= ds.arr(0.5, 'code_length')
    x_start_pos /= scale
    x_start_pos += 0.5
    color = np.array([x_start_pos, 0, 1-x_start_pos])

    # Plot the stream in 3D
    ax.plot3D(stream[:,0], stream[:,1], stream[:,2], alpha=0.3, color=color)

# Create a sphere object centered on the highest density point in the simulation
# with radius = 1 Mpc
sphere = ds.sphere("max", (1.0, "Mpc"))

# Identify the isodensity surface in this sphere with density = 1e-24 g/cm^3
surface = ds.surface(sphere, "density", 1e-24)

# Color this isodensity surface according to the log of the temperature field
colors = yt.apply_colormap(np.log10(surface["temperature"]), cmap_name="hot")

# Render this surface
p3dc = Poly3DCollection(surface.triangles, linewidth=0.0)
colors = colors[0,:,:]/255.  # scale to [0,1]
colors[:,3] = 0.3            # alpha = 0.3
p3dc.set_facecolors(colors)
ax.add_collection(p3dc)

# Save the figure
pl.savefig('streamlines_isocontour.png')
