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

# In[19]:


get_ipython().run_line_magic('matplotlib', 'inline')
from matplotlib import pylab
import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from salishsea_tools import tidetools

from salishsea_tools import (nc_tools,viz_tools)

import matplotlib.cm as cm

from matplotlib import animation

from numpy import *
from pylab import *
import matplotlib.patches as patches


# In[2]:


result = nc.Dataset('/ocean/vdo/MEOPAR/ariane-runs/20160701_20160704_4d/ariane_trajectories_qualitative.nc')
grid = nc.Dataset('/data/vdo/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')


# In[59]:


result.variables['final_x'][:]


# In[50]:


lont=result.variables['traj_lon']
latt=result.variables['traj_lat']


# In[60]:


result.variables['init_x'][:]


# In[61]:


result.variables['init_y'][:]


# In[62]:


result.variables['init_z'][:]


# In[63]:


result.variables['final_z'][:]


# In[64]:


result.variables['final_y'][:]


# In[65]:


result.variables['init_t'][:]


# In[66]:


result.variables['final_t'][:]


# In[124]:


lont[1:,0]


# In[123]:


n = np.arange(1)

# 2D
fig, ax = plt.subplots(1,1,figsize=(10,6.6))

for N,c in zip(n,['red']):
    ax.scatter(lont[1:,N],latt[1:,N],color=c)
    ax.scatter(lont[0,N],latt[0,N],color='0.30',marker='s')

viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-124.9,-124.6])
ax.set_ylim([49.2,49.8])

ax.set_title('Surface particles')
ax.set_xlabel('lon')
ax.set_ylabel('lat')


# In[2]:


3*27


# In[4]:


import xarray as xr


# In[5]:


grid2 = xr.open_dataset('/data/vdo/MEOPAR/NEMO-forcing/grid/bathymetry_201702.nc')


# In[74]:


fig,ax = plt.subplots(1,1,figsize=(20,16))
grid2.Bathymetry.plot()
viz_tools.set_aspect(ax)
ax.plot(125,599,'r*', markersize=15)
ax.plot(126.89,605.537, 'y*', markersize=15)
ax.set_ylim((550,650))
ax.set_xlim((100,150))


# In[6]:


test = nc.Dataset('/ocean/vdo/MEOPAR/ariane-runs/test/ariane_trajectories_qualitative.nc')


# In[79]:


test.variables['final_y'][:]


# In[80]:


fig,ax = plt.subplots(1,1,figsize=(20,16))
grid2.Bathymetry.plot()
viz_tools.set_aspect(ax)
ax.plot(125,599,'r*', markersize=15)
ax.plot(120.45,620.097, 'y*', markersize=15)
ax.set_ylim((550,650))
ax.set_xlim((100,150))


# In[81]:


test2 = nc.Dataset('/ocean/vdo/MEOPAR/ariane-runs/test2/ariane_trajectories_qualitative.nc')


# In[83]:


test2.variables['final_y'][:]


# In[5]:


2694*744


# In[6]:


168*407


# In[136]:


fig,ax = plt.subplots(1,1,figsize=(20,16))
grid2.Bathymetry.plot()
viz_tools.set_aspect(ax)
ax.plot(125,599,'r*', markersize=15)
ax.plot(124.18,608.815, 'y*', markersize=15)
ax.set_ylim((590,635))
ax.set_xlim((115,135))
plt.grid('on')


# In[99]:


grid.variables['Bathymetry'][:].shape


# In[20]:


mesh_mask=nc.Dataset('/data/vdo/MEOPAR/NEMO-forcing/grid/mesh_mask201702.nc')
mesh_mask.variables['tmask'][0,0,600,slice(120,130)]


# In[145]:


for x in np.arange(120,133):
        for y in np.arange(599,630):
            for z in range(1,40): 
                if (mesh_mask.variables['tmask'][0,z,y,x] == 1) and (mesh_mask.variables['tmask'][0,z+2,y,x] == 1) and (mesh_mask.variables['tmask'][0,z,y+1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x+1] == 1) and (mesh_mask.variables['tmask'][0,z,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x-1] == 1):
                    print(str(x) +' ' + str(y) +' '+ str(z)+ ' 0.5 1.0')


# In[239]:


with open('initial_positions.txt', 'a') as f:
    for x in np.arange(120,133):
        for y in np.arange(598,608):
            for z in range(2,40):
                if ((mesh_mask.variables['tmask'][0,z,y,x] == 1) and (mesh_mask.variables['tmask'][0,z+2,y,x] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y+1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x+1] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x-1] == 1)
                and (mesh_mask.variables['tmask'][0,z+2,y,x+1] == 1) and (mesh_mask.variables['tmask'][0,z+2,y,x-1] == 1)
               and (mesh_mask.variables['tmask'][0,z+2,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z+2,y+1,x] == 1)
               and (mesh_mask.variables['tmask'][0,z+2,y-1,x-1] == 1) and (mesh_mask.variables['tmask'][0,z+2,y-1,x+1] == 1)
               and (mesh_mask.variables['tmask'][0,z+2,y+1,x-1] == 1) and (mesh_mask.variables['tmask'][0,z+2,y+1,x+1] == 1)):
                    f.write((str(x) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))


# In[235]:


n


# In[192]:


fig,ax = plt.subplots(1,1)
mesh = ax.pcolormesh(mesh_mask.variables['tmask'][0,:,620,:])
ax.set_ylim((30,0))
ax.set_xlim((118,133))
fig.colorbar(mesh,ax=ax)
ax.plot(120,12,'r*')
plt.grid('on')


# In[165]:


mesh_mask.variables['tmask'][0,:,614,120]


# In[223]:


test = nc.Dataset('/ocean/vdo/MEOPAR/ariane-runs/test/ariane_trajectories_qualitative.nc')


# In[3]:


with open('initial_positions.txt', 'a') as f:
    for x in np.arange(120,130):
        for y in np.arange(735,750):
            for z in range(2,40):
                for t in range(1,30):
                    if ((mesh_mask.variables['tmask'][0,z,y,x] == 1) and (mesh_mask.variables['tmask'][0,z+2,y,x] == 1) 
                    and (mesh_mask.variables['tmask'][0,z,y+1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x+1] == 1) 
                    and (mesh_mask.variables['tmask'][0,z,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x-1] == 1)
                    and (mesh_mask.variables['tmask'][0,z+2,y,x+1] == 1) and (mesh_mask.variables['tmask'][0,z+2,y,x-1] == 1)
                   and (mesh_mask.variables['tmask'][0,z+2,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z+2,y+1,x] == 1)
                   and (mesh_mask.variables['tmask'][0,z+2,y-1,x-1] == 1) and (mesh_mask.variables['tmask'][0,z+2,y-1,x+1] == 1)
                   and (mesh_mask.variables['tmask'][0,z+2,y+1,x-1] == 1) and (mesh_mask.variables['tmask'][0,z+2,y+1,x+1] == 1)):
                        f.write((str(x) +' ' + str(y) +' '+ str(z+0.5)+' '+ str(t)+' 1.0\n'))


# In[8]:


Jan04 = nc.Dataset('/ocean/vdo/MEOPAR/completed-runs/SalishSeaLake/LakeJan0.4/SalishSea_1h_20170101_20170107_grid_T.nc')


# In[9]:


Jan04.variables['deptht'][:]


# In[16]:


Jan04.variables['deptht'][20:]


# In[38]:


with open('initial_positions2.txt', 'a') as f:
    for x in np.arange(118,133):
        for y in np.arange(608,659):
            for z in range(21,30):
                if ((mesh_mask.variables['tmask'][0,z,y,x] == 1) and (mesh_mask.variables['tmask'][0,z+1,y,x] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y+1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x+1] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x-1] == 1)
                and (mesh_mask.variables['tmask'][0,z+1,y,x+1] == 1) and (mesh_mask.variables['tmask'][0,z+1,y,x-1] == 1)
               and (mesh_mask.variables['tmask'][0,z+1,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z+1,y+1,x] == 1)
               and (mesh_mask.variables['tmask'][0,z+1,y-1,x-1] == 1) and (mesh_mask.variables['tmask'][0,z+1,y-1,x+1] == 1)
               and (mesh_mask.variables['tmask'][0,z+1,y+1,x-1] == 1) and (mesh_mask.variables['tmask'][0,z+1,y+1,x+1] == 1)):
                    f.write((str(x) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.25) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.5) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.75) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x) +' ' + str(y+0.5) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))


# In[21]:


grid.variables['Bathymetry'][608:659, 118:133].max()


# In[ ]:


for x in np.arange(118,133):
    for y in np.arange(608,659):
        if grid.variables['Bathymetry'][y,x] >= 30:


# In[30]:


July = nc.Dataset('/results/SalishSea/hindcast/01jul17/SalishSea_1h_20170701_20170701_grid_T.nc')


# In[33]:


July.variables['deptht'][:][21:]


# In[40]:


with open('initial_positions3.txt', 'a') as f:
    for x in np.arange(118,133):
        for y in np.arange(608,659):
            for z in range(21,32):
                if ((mesh_mask.variables['tmask'][0,z,y,x] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y+1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x+1] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x-1] == 1)):
                    f.write((str(x) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.25) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.5) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.75) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x) +' ' + str(y+0.5) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))


# In[22]:


grid = nc.Dataset('/data/vdo/MEOPAR/NEMO-forcing/grid/bathymetry_201702.nc')


# In[23]:


fig, ax = plt.subplots(1,1, figsize=(16,12))
bathy = ax.pcolormesh(grid.variables['Bathymetry'][:])
ax.set_ylim(590,670)
ax.set_xlim(115,135)
viz_tools.set_aspect(ax)
ax.add_patch(patches.Rectangle((118,598), 13, 18, fill=False, linewidth=3))


# In[24]:


with open('initial_positions2.txt', 'a') as f:
    for x in np.arange(118,132):
        for y in np.arange(598,617):
            for z in range(21,26):
                if ((mesh_mask.variables['tmask'][0,z,y,x] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y+1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x+1] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x-1] == 1)):
                    f.write((str(x) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.25) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.5) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.75) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x) +' ' + str(y+0.5) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))


# In[ ]:


with open('initial_positions2.txt', 'a') as f:
    for x in np.arange(118,132):
        for y in np.arange(598,617):
            for z in range(21,26):
                if ((mesh_mask.variables['tmask'][0,z,y,x] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y+1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x+1] == 1) 
                and (mesh_mask.variables['tmask'][0,z,y-1,x] == 1) and (mesh_mask.variables['tmask'][0,z,y,x-1] == 1)):
                    f.write((str(x) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.25) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.5) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.75) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x) +' ' + str(y+0.5) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))


# In[25]:


list(range(24,25))


# In[27]:


with open('initial_positions2.txt', 'a') as f:
    for x in np.arange(118,132):
        for y in np.arange(598,617):
            for z in range(24,26):
                if ((mesh_mask.variables['tmask'][0,z,y,x]) == 1):
                    f.write((str(x) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.25) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.5) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x+0.75) +' ' + str(y) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))
                    f.write((str(x) +' ' + str(y+0.5) +' '+ str(z+0.5)+ ' 0.5 1.0\n'))


# In[ ]: