# Minimum working example

import h5py
import yt
import numpy as np
from yt.visualization.volume_rendering.transfer_function_helper import TransferFunctionHelper
from yt.visualization.volume_rendering.api import Scene, VolumeSource

# Load data using uniform grid
ds = h5py.File('/Volumes/triton/Data/ModelChallenge/R2349/rhybrid.h5')
data = {k:v[:] for k,v in ds.items() if 'number_density' in k}
bbox = np.array([[-4,4],[-4,4],[-4,4]])
ds = yt.load_uniform_grid(data, data['H_p1_number_density'].shape, bbox=bbox)

# Create slice plots to see that the data has been correctly loaded
p = yt.SlicePlot(ds, 'y',['H_p1_number_density', 'O2_p1_number_density'])
p.set_zlim('H_p1_number_density', 1e-4,1000)
p.set_zlim('O2_p1_number_density', 1e-4,1000)
p.show()

# Start volume rendering------------------

# Initiatize scene and camera, try to make the vr look like sliceplots
sc = Scene()

cam = sc.add_camera(ds, lens_type='perspective')
cam.position = ds.arr([0,8,0], 'code_length')
cam.focus = ds.arr([0,0,0], 'code_length')
cam.north_vector = [1,0,0]


# Create H+ data source, make a good transfer function and add it to the scene
source = VolumeSource(ds.all_data(), 'H_p1_number_density')

tfh = TransferFunctionHelper(ds)
tfh.set_field('H_p1_number_density')
tfh.set_log(True)
tfh.set_bounds((0.8,50))
tfh.build_transfer_function()
tfh.grey_opacity = True
tfh.tf.add_layers(10, colormap='viridis')

source.set_transfer_function(tfh.tf)
sc.add_source(source)


# Create an O2+ data source, make a transfer function, and add it
source = VolumeSource(ds.all_data(), 'O2_p1_number_density')

tfh = TransferFunctionHelper(ds)
tfh.set_field('O2_p1_number_density')
tfh.set_log(True)
tfh.grey_opacity = True
tfh.set_bounds((0.01,300))
tfh.build_transfer_function()
tfh.tf.add_layers(10, colormap='viridis')
tfh.tf.map_to_colormap(-2.0,10.0, colormap='viridis', scale=0.4)
source.set_transfer_function(tfh.tf)
sc.add_source(source)


# Render/show scene
sc.render()
sc.show(sigma_clip=6.0)