#!/usr/bin/env python
# coding: utf-8

# ## using obspy + yt to extract properties along a seismic ray path 
# 
# 

# In[1]:


import yt
import xarray as xr
import numpy as np
import pandas as pd 
import matplotlib.pyplot as plt
from obspy.taup import TauPyModel
from obspy import UTCDateTime
from obspy.clients.fdsn import Client
from pygeodesy.sphericalNvector import LatLon
import cartopy.crs as ccrs


# ## load the data into yt

# In[2]:


ddir = "/home/chavlin/hdd/data/yt_data/IRIS_models/"
modelfile='GYPSUMS_kmps.nc'
xr_ds = xr.open_dataset(ddir+modelfile)
bbox = np.array([
    (xr_ds[dim].values[:].min(),xr_ds[dim].values[:].max()) for dim in xr_ds.vs.dims
])
bbox[1][0]=-89.99
bbox[1][1]=89.99
bbox[2][0]=0.001
bbox[2][1]=360 - 0.001

data = {'vs':xr_ds.vs.values}
ds = yt.load_uniform_grid(data,data['vs'].shape,1.0,bbox = bbox, 
                          geometry=("internal_geographic",
                                    ['depth','latitude','longitude']
                                   )
                         )


# In[3]:


target_depth = 150
center = ds.domain_center
p = yt.SlicePlot(ds,"depth",'vs',center=[target_depth,center[1],center[2]]) 
p.set_cmap('vs', 'dusk_r')
p.set_log('vs', False)
 
p.set_mpl_projection(('Mollweide', (), {'central_longitude':-110}))#{'central_longitude':-100, 'central_latitude':40}))
#p.set_mpl_projection(('Mollweide', (), {'central_longitude':0}))#{'central_longitude':-100, 'central_latitude':40}))

p._setup_plots()

p.plots['vs'].axes.coastlines()
p.show()




# ## use obspy + pygeodesy to get event info and generate a ray path

# In[4]:


client = Client()
cat = client.get_events(eventid=3279407) # Tohoku
cat.events


# In[5]:


event = cat.events[0]
depth_km = event.origins[0].depth / 100
lat = event.origins[0].latitude
lon = event.origins[0].longitude
[depth_km,lat,lon]


# In[6]:


stat = client.get_stations(network='HV',station='AIN')[0][0]
stat


# In[7]:


eq_start = LatLon(lat,lon)
station_loc = LatLon(stat.latitude,stat.longitude)
gc_dist_km = eq_start.distanceTo(station_loc) / 1000
gc_dist_deg = gc_dist_km /111. # approx 111 km /deg
gc_dist_deg


# In[8]:


model = TauPyModel('iasp91')
arrivals = model.get_ray_paths(source_depth_in_km=depth_km,
                               distance_in_degree=gc_dist_deg,
                               phase_list=["P"])


# In[9]:


arrivals.plot_rays()


# ## extracting data along the path 
# 
# pull out the depth and surface projection (lat/lon) along the path:

# In[10]:


path = arrivals[0].path
latlondepths = []

path_frac = np.linspace(0,1,len(path))
for frac,pathpt in zip(path_frac,path):
    ilat,ilon = eq_start.intermediateTo(station_loc,frac).latlon    
    if ilon < 0: 
        ilon = ilon + 360
    latlondepths.append([ilat,ilon,pathpt[3]])


# extract a quantity from the yt dataset

# In[11]:


path_vs = []
for lld in latlondepths:     
    pt = ds.point(lld[::-1])
    path_vs.append(pt['vs'])

path_vs = np.array(path_vs)
latlondepths = np.array(latlondepths)


# plot the surface projection of the ray

# In[12]:


target_depth = 50
center = ds.domain_center
p = yt.SlicePlot(ds,"depth",'vs',center=[target_depth,center[1],center[2]]) 
p.set_cmap('vs', 'dusk_r')
p.set_log('vs', False)
 
p.set_mpl_projection(('NearsidePerspective', (), 
                      {'central_longitude':190.,'central_latitude':20.}))

p._setup_plots()

p.plots['vs'].axes.coastlines()
p.plots['vs'].axes.scatter(latlondepths[:,1],latlondepths[:,0],
           transform=ccrs.PlateCarree(),color=[0,.8,0.,.5])
p.show()


# In[13]:


plt.plot(path_frac*gc_dist_deg,latlondepths[:,2])
plt.xlabel("along path distance (degrees)")
plt.ylabel("depth along path")
plt.show()


# In[14]:


plt.plot(path_frac*gc_dist_deg,path_vs)
plt.xlabel("along path distance (degrees)")
plt.ylabel("vs along path")
plt.show()


# In[15]:


plt.plot(latlondepths[:,2][path_frac<=0.5],path_vs[path_frac<=0.5],label='GyPSuM down')
plt.plot(latlondepths[:,2][path_frac>0.5],path_vs[path_frac>0.5],label='GyPSuM up')
plt.ylim([3,4.8])
plt.xlim([0,800])
plt.legend(loc='lower right')


# In[16]:


# compare to the 1d model that generated the raypath.... 

df_iasp91 = pd.read_csv('/home/chavlin/hdd/data/yt_data/IRIS_refModels/IASP91_IDV.csv',header=1)
vs_iasp = np.interp(latlondepths[:,2], df_iasp91['Depth[unit="km"]'], df_iasp91['Vs[unit="km/s"]'])


# In[17]:


plt.plot(path_frac*gc_dist_deg,path_vs,label='GyPSuM')
plt.plot(path_frac*gc_dist_deg,vs_iasp,label='iasp91')
plt.xlabel("along path distance (degrees)")
plt.ylabel("vs along path")
plt.legend(loc='lower center')
plt.show()


# In[18]:


plt.plot(latlondepths[:,2],path_vs,label='GyPSuM')
plt.plot(latlondepths[:,2],vs_iasp,label='iasp91')
plt.xlabel("depth along path ")
plt.ylabel("vs along path")
plt.legend(loc='lower center')
plt.show()


# In[ ]: