#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import yt
import os
# yt.toggle_interactivity()
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

#Create a derived field, the gas pressure.
def _pressure(field, data):
    
    gamma_eso=5.0/3.0
    
    return (gamma_eso - 1.0) * data["PartType0","Density"] * data["PartType0","InternalEnergy"]

def _Temperature(field, data):
    mhyd=1.67e-24 #mass of hydrogen in grams
    boltz=1.38e-16 #boltzmann constant in erg/K
    mu=0.5 #mean molecular weight
    gamma=5/3
    factor= (mu*(mhyd*(gamma-1)))/boltz
    
    return (data["PartType0", "InternalEnergy"])*factor*1.907e+15


###############################################################################
###############################################################################
fname = 'snapshot_110.hdf5'

unit_base2 = {'UnitLength_in_cm'         : 6.957e+10, #R_sol
                'UnitMass_in_g'            : 1.989e+33, #M_sol
                'UnitVelocity_in_cm_per_s' : 4.367e+7}

bbox_lim = 3.0 # R_sol

bbox = [[-bbox_lim,bbox_lim],
        [-bbox_lim,bbox_lim],
        [-bbox_lim,bbox_lim]]
 
ds = yt.load(fname, unit_base=unit_base2, bounding_box = bbox)
ds.index
ad = ds.all_data()

# adding a derived pressure field
ds.add_field(('gas', 'pressure'), function=_pressure, units="g/(cm*s**2)", 
              display_name='pressure', sampling_type='particle', force_override=True)

# adding a derived Temperature
ds.add_field(('gas', 'Temperature'), function=_Temperature, units="code_specific_energy", 
              display_name='Temperature', sampling_type='particle', force_override=True)


plot = yt.LinePlot(ds, ("PartType0", "Density"), (0.0, 0.0, 0.0), (0.0, 1.0, 0.0), 1000)