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

# # simulation

# In[1]:


import S1DBedloadSolver as S1Dbed
import numpy as np


# In[2]:


get_ipython().run_cell_magic('time', '', "length = 10000.0\ndx = 100.0\nimax = int(length/dx) + 1\ndt = 10.0\ntotalTime = 5000.1*3600.0\noutTimeStep = 200.0*3600.0\nRunUpTime = 3.0 * 3600.0\nhini = 1.0\nmanning  = 0.03\nib = 1.0/200.0\nib2 = 1.0/1000.0\noutputfilename = '1D_case2.json'\n\n# grain diameter classification\nscreenclass = np.array( [4.0], dtype=float )/1000\ndsize = np.array( [4.0/1000.0], dtype=float )\n\n# percentage of grain size under exchange layer\ndratioStandard1 = np.full_like(dsize, 1.0/len(dsize), dtype=float)\ndratioStandard = np.full( (imax, len(dsize) ), dratioStandard1, dtype=float)\n\n# initial percentage of grain size in exchange layer\ndratio = np.copy(dratioStandard)\n\n# thickness of exchange layer \nhExlayer = dsize[-1]\n\n# Initial & Boundary condition\nB = np.full(imax, 1.0, dtype=float)\nA = hini*B\nQ = ib**0.5*(hini)**(5.0/3.0)/manning*B #normal flow\nzb = np.zeros(imax)\nfor i in range(1,imax):\n    zb[i] = zb[i-1] + ib2*dx if i < 50 else zb[i-1] + ib*dx\n#     zb[i] = zb[i-1] + ib*dx # if i < 50 else zb[i-1] + ib*dx\n    \nzb = zb[::-1]\n\ndAb = np.zeros(imax)\nQup = Q[0]\n\ndef Adown(time, Q, dzb, ib):\n    return ( manning**2*Q**2/ib/B[-1]**2 )**0.3 * B[-1]\n\ndef Qup(time):\n    return Q[0]\n\nS1Dbed.bedvariation(\ndx,dt,manning,totalTime,outTimeStep,RunUpTime\n,dsize ,dratioStandard ,dratio ,hExlayer ,A ,Q ,B ,zb ,dAb ,Qup ,Adown\n,outputfilename, screenclass\n )\n")


# # json to NetCDF

# In[3]:


import xarray as xr
import json
import numpy as np


# In[4]:


outputfile = '1D_case2.json'
data = json.load( open(outputfile, 'r') )

cond = data['condition']
d = data['output']

time = [ p['time']/3600 for p in d ]
x = [ p['distance'] for p in d[0]['profile'] ]
screenclass =  np.array( cond['screenclass'] )*1000
zb = np.array( cond['elevation'] )


# In[5]:


dep = []
dzb = []
for dr in d :
    A = np.array( [ dd['A']  for dd in dr['profile'] ] )
    dAb = np.array( [ dd['dAb'] for dd in dr['profile'] ] )
    dep.append( zb + A + dAb)
    dzb.append( zb + dAb)
    
dep = np.array(dep)
dzb = np.array(dzb)


# In[6]:


ds = xr.Dataset({'H': (['t','x'], dep), 'zb': (['t','x'], dzb) }, coords={'x': x , 't': time})
out = ds.to_netcdf('case2.nc')

del out, ds


# # figure

# In[7]:


import xarray as xr
import hvplot.xarray


# In[8]:


dsin =xr.open_dataset('case2.nc')


# In[9]:


dini = dsin.isel(t=0)
dini['H_ini'] = dini['H']
dini['zb_ini'] = dini['zb']
dini = dini.drop(['H','zb'])
fini = dini.hvplot()
f = dsin.hvplot(groupby='t',x='x')
f * fini


# In[10]:


del dsin, f, fini