from yt.mods import *
import matplotlib.pyplot as plt
import yt.data_objects.derived_quantities
import os.path
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import MaxNLocator
import numpy as np


# Creates 2D PDF of Velocity Anomaly from Keplerian, for all R, and restricting gas to > RhoCrit

resolution = 2.5 #au
Rdisk = 100.*1.496e13
sink = 4.*resolution


if len(sys.argv) == 2:
    last = int(sys.argv[1])
else:
    # Gets the last data file
    for i in range(10000,0,-1):
        last = i;
        if not os.path.exists('./data.%04d.3d.hdf5' %i): continue
        else: break


pf = load('data.%04d.3d.hdf5' %last)

data = pf.h.all_data()
masses = data['particle_mass']
if masses.size == 0:
    c = [0.,0.,0.]
    mass = 0.
else:
    index = masses.argmax()
    c = [data['particle_position_' + direction][index] for direction in ['x', 'y', 'z']]
    mass = masses[index]

sp = pf.h.sphere(c, 50./pf['au'])
L = sp.quantities['AngularMomentumVector']() # Orient disk with L pointing North.


# Set Field Parameters:

gamma = 5./3.
mu  = 2.33
mp = 1.67262171e-24
kb = 1.3806505e-16
AU = 1.49597871e13 # cm
year = 3.15569e7 # s
sigma = 5.670373e-5 #
G = 6.6726e-8 #
Msolar = 1.9891e33 # g


def _ThermalEnergy(field, data):
    return (data['energy-density']
          - 0.5*(data['X-momentum']**2+data['Y-momentum']**2+data['Z-momentum']**2)/data['density'])
add_field("ThermalEnergy", function=_ThermalEnergy)

def _Temperature(field, data):
    import yt.utilities.physical_constants as physc
    Eth = (data['energy-density'] - 
          0.5*(data['X-momentum']**2+data['Y-momentum']**2+data['Z-momentum']**2)/data['density'])
    return (gamma-1.)*Eth*mu*mp/kb/data['density']
add_field("Temperature", function=_Temperature)

def _SoundSpeed(field, data):
    return na.sqrt(gamma*kb/(mp*mu)*data['Temperature'])
add_field("SoundSpeed", function=_SoundSpeed)


def _CylinderRadius(field, data):
    center = data.get_field_parameter("center")
    norm = data.get_field_parameter("normal")
    coords = na.array([data['x'] - center[0],
                       data['y'] - center[1],
                       data['z'] - center[2]])
    CrossR = na.cross(norm,coords,axis=0)
    return na.sqrt(na.sum(na.square(CrossR),0))
def _ConvertCylinderRadiusCGS(data):
    return data.convert("cm")
add_field("CylinderRadius", function=_CylinderRadius,validators=[ValidateSpatial(0),ValidateParameter("center")],
          convert_function = _ConvertCylinderRadiusCGS,units=r"\rm{cm}")


def _Height(field, data):
    cylrad = data['CylinderRadius']
    rad = data['Radius']
    h = na.sqrt(na.square(rad)-na.square(cylrad))
    return h
def _ConvertHeightCGS(data):
    return data.convert("cm")
add_field("Height", function=_Height,validators=[ValidateSpatial(0),ValidateParameter("center")],
          convert_function = _ConvertHeightCGS,units=r"\rm{cm}")


def _VectorKeplerianDiff(field, data):
    masses = data.get_field_parameter("sink")
    ##### Eventually Add Sphere mass here too #######
    c = data.get_field_parameter("center")
    L = -data.get_field_parameter("normal")

    r_vec = na.array([data['x'] - c[0],
                      data['y'] - c[1],
                      data['z'] - c[2]])

    v_vec = na.array([data['x-velocity'],
                      data['y-velocity'],
                      data['z-velocity']])

    r_u = r_vec/na.sqrt(na.sum(na.square(r_vec),0))
    tan_u = na.cross(L,r_u,axis=0)
    kepvelocity = na.sqrt(G*masses.max()/data['CylinderRadius'])*tan_u
    
    v_anomaly = v_vec-kepvelocity
    
    return na.sqrt(na.sum(na.square(v_anomaly),0))/data["SoundSpeed"]
add_field("VectorKeplerianDiff", function=_VectorKeplerianDiff,validators=[ValidateParameter("center")])


def _RhoCrit(field,data):
    tauCrit = 1.
    SigmaCrit = 2.*tauCrit/(4.8*(data['Temperature']/300.)**0.8);
    rhoSG = na.pi*G*SigmaCrit**2./(2.*kb/(mu*mp)*20.)
    # NSG H = na.sqrt(pi/2) Cs/Omega so that integrating becomes Sigma = 2 rho H
    rhoNSG = na.sqrt(G*mass*SigmaCrit**2/(2.*np.pi*kb/(mu*mp)*20.*data['CylinderRadius']**3.))
    return (rhoSG + rhoNSG)/2.
add_field("RhoCrit", function=_RhoCrit)


def _CellVolumeRestrict(field,data):
    msink = data.get_field_parameter("sink")
    Volume = data["CellVolume"]
    ScaleH = data["SoundSpeed"]*na.sqrt(data["CylinderRadius"]**3./(G*msink))
    mask = na.all([(data["Density"] < data["RhoCrit"]),(data["Height"]>ScaleH)],axis=0)
    Volume[mask] = 0.
    return Volume / data['Radius']**2.   ###### Scales by 1/R^2 to remove volume bias of sampling
add_field("CellVolumeRestrict", function=_CellVolumeRestrict)

def _CellMassRestrict(field,data):
    msink = data.get_field_parameter("sink")
    Mass = data["CellMass"]
    ScaleH = data["SoundSpeed"]*na.sqrt(data["CylinderRadius"]**3./(G*msink))
    mask = na.all([(data["Density"] < data["RhoCrit"]),(data["Height"]>ScaleH)],axis=0)
    Mass[mask] = 0.
    return Mass
add_field("CellMassRestrict", function=_CellMassRestrict)

def _CylinderRadiusAU(field, data):
    return data['CylinderRadius']/1.49597871e13
add_field("CylinderRadiusAU", function=_CylinderRadiusAU,validators=[ValidateParameter("center")],units=r"\rm{AU}")

if is_root():
    try: os.stat('PDF')
    except: os.mkdir('PDF')
    os.chdir('./PDF')

t = pf.current_time*pf['years']/1000
r_bins = int(Rdisk/(resolution*AU))
v_bins = 100


# Volume PDF

sp = pf.h.sphere(c,Rdisk)
sp.set_field_parameter('angular_momentum_vector', L)
sp.set_field_parameter('center',c)
sp.set_field_parameter('normal',L)
sp.set_field_parameter('sink',mass)

sp.pf.field_info['CylinderRadiusAU'].take_log = False
sp.pf.field_info['VectorKeplerianDiff'].take_log = False

plot = PhasePlot(sp,"CylinderRadiusAU","VectorKeplerianDiff","CellVolumeRestrict",weight_field=None,fractional=True)
#plot.x_title='Radius [AU]'
#plot.y_title='$\widetilde{\mathbf{V}}$ [Mach]'
#plot.z_title='Volume Fraction'
#plot.set_title("CellVolumeRestrict", 'PDF of Disk Turbulence at %.4s kyr' %(t))
plot.save('./TurbulenceVolume.2D.%04d.png' %(i))