# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Load Packages

# %matplotlib inline 
from ifis_tools import database_tools as db 
from ifis_tools import asynch_manager as am 
from ifis_tools import auxiliar as aux
from wmf import wmf
import pandas as pd 
import numpy as np 
import os 
import pylab as pl
from string import Template
from param_ident import core 
from aclayer import core as ac

names = {'south': {'path': 'data/for_hlm/south_skunk/'},
        'turkey':{'path':'data/for_hlm/turkey/'}}

# # Define Watershed setups
#
# In this step we will produce the setups of the watersheds making variations of $v_0$ and $\lambda_1$ for different sizes of sub-watersheds.

# +
#Get the watershed
south = am.hlmModel(linkid=272678, ExtraParams=['h_order'])

#Temporal rvr to get the topology of the watershed
south.write_rvr('data/for_hlm/south_skunk/south_skunk.rvr')
topo = south.topo.copy()
topo.set_index('link_id', inplace = True)
idx = south.Table.index.intersection(topo.index)
south.Table['dest'] = topo['id'].loc[idx]

# +
#Get the watershed
turkey = am.hlmModel(linkid=434514, ExtraParams=['h_order'])

#Temporal rvr to get the topology of the watershed
turkey.write_rvr('data/for_hlm/turkey/turkey.rvr')
topo = turkey.topo.copy()
topo.set_index('link_id', inplace = True)
idx = turkey.Table.index.intersection(topo.index)
turkey.Table['dest'] = topo['id'].loc[idx]
# -

wat = {'south': south, 'turkey': turkey}

# ## Obtain sub-watersheds at different prun-levels
#
# Here we include levels from 7 to 4, then we will produce random values of $v_0$ and $\lambda_1$ using a cascade approach.

# obtain the sub-watersheds for each one (this takes time try not to run it, instead read the data)
core.Get_sub_watersheds(wat['south'], 7, 3)
core.Get_sub_watersheds(wat['turkey'], 7, 3)
# Save to csv files to avoid doing this again.
wat['south'].Table.to_csv('data/for_hlm/south_skunk/watershed_table.csv')
wat['turkey'].Table.to_csv('data/for_hlm/turkey/watershed_table.csv')

wat['south'].Table = pd.read_csv('data/for_hlm/south_skunk/watershed_table.csv', index_col=0)
wat['turkey'].Table = pd.read_csv('data/for_hlm/turkey/watershed_table.csv', index_col=0)


# ## Distribute $v_0$ and $\lambda_1$ using the cascade
#

# Synchronic production of the virtual variables $v_0$ and $\lambda_1$.

def cascade_values_sync(Table, param, levels, seed_values, norm_dev = 0.15):
    '''Generates a virtual setup of a watershed for a given set of values (seed_values)
    the result is updates on the cu.Table DataFrame. Contrasted with the async version, this 
    version creates values that depends on the previous level.
    Params:
        - level: the minimum orton order to distribute the seed_values.
        - param: name of the parameter in the table.
        - seed_values: a list with the initial values for the biggest horton order.
        - norm_dev: deviation aprameter to produce the random cascade.
    Resutls:
        - cu.Table[param]'''
    p = param + '7'
    Table[p] = 0.0    
    for sub, value in zip(Table['prun_7'].unique(), seed_values):
        Table.at[Table['prun_7'] == sub, p] = value

    Factors = {}
    #Iterate for each level
    for level in np.arange(6,levels-1,-1):
        
        #Initial value
        name = 'prun_'+str(level)
        p = param + str(level)
        Table[p] = Table[param + str(level+1)].copy()
        fac = []
        
        #Iterate through the sub-watersheds of each level `
        for subw in Table[name].unique():
            factor = np.random.normal(1, norm_dev)            
            Table.loc[Table[name] == subw, p] *= factor            
            fac.append(factor)
        Factors.update({str(level): fac})
    return Factors


for name, seed in zip(['turkey','south'], [[0.15,0.2,0.18,0.25,0.3,0.35,0.4], [0.15,0.3,0.22]]):
    f_vo = cascade_values_sync(wat[name].Table, 'vo', 4,  [0.15,0.3,0.22], 0.2)
for name, seed in zip(['turkey','south'], [[0.19,0.22,0.2,0.22,0.2,0.23,0.3], [0.18,0.23,0.25]]):
    f_l1 = cascade_values_sync(wat[name].Table, 'la', 4, seed, 0.2)

x = a['vo4'] * (a['Acum'] / a['Acum'].sum())
x.sum()

x = a['la4'] * (a['Acum'] / a['Acum'].sum())
x.sum()

a['la4'].mean()

# ### Plot with the obtained distribution at each case.

wat['south'].Table = pd.read_csv('data/for_hlm/south_skunk/water_virtual.csv', index_col = 0)
wat['turkey'].Table = pd.read_csv('data/for_hlm/turkey/water_virtual.csv', index_col = 0)

# + hide_input=true
fig  = pl.figure(figsize=(20,12))
bins = np.arange(0.1,0.5,0.05)

###################### Vo

ax = fig.add_subplot(221)
for n,alpha in zip(['vo6','vo5','vo4'],[0.85,0.6,0.4]):
    h,b = ac.get_histogram(wat['south'].Table[n], bins = bins, inf = 0.1, sup = 0.5)
    pl.bar(b[:-1], 100 * h / h.sum(), 0.05, align='edge', edgecolor = 'k', lw = 3, alpha = alpha, 
          label = 'Order ' + n[-1])
ax.tick_params(labelsize = 23)
ax.grid()
ax.set_ylim(0,65)
ax.set_title('a) $v_0$ for South skunk', size = 24)
ax.set_ylabel('Fraction [$\%$]', size = 23)
ax.legend(loc = 0, fontsize = 'xx-large')
ax.set_xticklabels([])

ax = fig.add_subplot(222)
for n,alpha in zip(['vo6','vo5','vo4'],[0.85,0.6,0.4]):
    h,b = ac.get_histogram(wat['turkey'].Table[n], bins = bins, inf = 0.1, sup = 0.5)
    pl.bar(b[:-1], 100 * h / h.sum(), 0.05, align='edge', edgecolor = 'k', lw = 3, alpha = alpha, 
          label = 'Order ' + n[-1])
ax.tick_params(labelsize = 23)
ax.grid()
ax.set_ylim(0,65)
ax.set_title('b) $v_0$ for Turkey', size = 24)
ax.set_yticklabels([])
ax.set_xticklabels([])


###################### LAMBDA

ax = fig.add_subplot(223)
for n,alpha in zip(['la6','la5','la4'],[0.85,0.6,0.4]):
    h,b = ac.get_histogram(wat['south'].Table[n], bins = bins, inf = 0.1, sup = 0.5)
    pl.bar(b[:-1], 100 * h / h.sum(), 0.05, align='edge', edgecolor = 'k', lw = 3, alpha = alpha, 
          label = 'Order ' + n[-1])
ax.tick_params(labelsize = 23)
ax.grid()
ax.set_ylim(0,65)
ax.set_title('c) $\lambda_1$ for South skunk', size = 24)
ax.set_ylabel('Fraction [$\%$]', size = 23)

ax = fig.add_subplot(224)
for n,alpha in zip(['la6','la5','la4'],[0.85,0.6,0.4]):
    h,b = ac.get_histogram(wat['turkey'].Table[n], bins = bins, inf = 0.1, sup = 0.5)
    pl.bar(b[:-1], 100 * h / h.sum(), 0.05, align='edge', edgecolor = 'k', lw = 3, alpha = alpha, 
          label = 'Order ' + n[-1])
ax.tick_params(labelsize = 23)
ax.grid()
ax.set_ylim(0,65)
ax.set_title('d) $\lambda_1$ for Turkey', size = 24)
ax.set_yticklabels([])
pl.savefig('figures/Vo_and_L1_distribution_levels_654.png', bbox_inches = 'tight')
# -

# # Write Waterhsed data for the HLM model to run 

# Creates the control points for both watersheds

#Creates a sav file of the control points.
for name in wat.keys():
    for k in [4,5,6]:
        f = open(names[name]['path'] + 'control_vr'+str(k)+'.sav','w')
        for i in wat[name].Table['prun_'+str(k)].unique():
            f.write('%d\n' % wat[name].Table.loc[wat[name].Table['prun_'+str(k)] == i, 'Acum'].idxmax())
        f.close()

# Setup the global parameters for both cases 

# +
#Set the global parameters for both watersheds.

Params = {'v_r': {'value': 0.02,'format': '%.4f'},
    'a_r': {'value': 1.67,'format': '%.2f'},
    'a': {'value': 1.5e-6,'format': '%.2e'},
    'b': {'value': 17,'format': '%.2f'},
    'c': {'value': 5.5e-7,'format': '%.2e'}, # 0.001
    'd': {'value': 35,'format': '%.2f'}, # 0.03
    'm1': {'value': 2.045e-6,'format': '%.2e'},
    'ki_fac': {'value': 0.02,'format': '%.4f'},
    'TopDepth': {'value': 0.1,'format': '%.2f'},
    'NoFlow': {'value': 1.48,'format': '%.2f'},
    'Td': {'value': 999,'format': '%.2f'}, # 1.54
    'Beta': {'value': 1.67,'format': '%.2f'},
    'lambda1': {'value': 0.18,'format': '%.3f'},
    'lambda2': {'value': -0.1,'format': '%.2f'},
    'channel_vo': {'value': 0.25,'format': '%.3f'},} 
Order = ['v_r','a_r','a','b','c','d','m1','ki_fac',
    'TopDepth','NoFlow','Td','Beta','lambda1', 'lambda2','channel_vo']
for name in wat.keys():
    wat[name].set_parameters(DictP=Params, DicOrder=Order)
# -

# Setup the virtually distributed parameters of both cases.

#setupo the distributed virtual parameters of both cases for the different levels.
for name in wat.keys():
    for i in [4,5,6]:
        # setup the parameters for each watershed.
        wat[name].Table.at[wat[name].Table.index, 'channel_vo'] = wat[name].Table['vo'+str(i)].values
        wat[name].Table.at[wat[name].Table.index, 'lambda1'] = wat[name].Table['la'+str(i)].values
        
        #Write the param file to run HLM
        wat[name].write_prm(names[name]['path']+'Tparam_o'+str(i)+'.prm', 
            extraNames=Order,
            extraFormats=[Params[k]['format'] for k in Order])
        

# Write the global files.

for name in wat.keys():
    path = names[name]['path']    
    for y in range(2002,2019):
        #Writes the global files and run files    
        for k in [4,5,6]:
            #Set the name of the global file
            global_name = 'global_'+str(k)+'_'+str(y)+'.gbl'
            #Writes the global file
            wat[name].write_Global(path + global_name,model_uid=608,
                date1=str(y)+'-03-01 01:00', 
                date2=str(y)+'-12-31 01:00',
                rvrFile='/Users/nicolas/Parameter_identifiability/'+path+'topology.rvr', 
                rvrLink='',
                prmFile='/Users/nicolas/Parameter_identifiability/'+path+'param_o'+str(k)+'.prm',
                initialFile='/Users/nicolas/Parameter_identifiability/data/for_hlm/initial.uini',
                initialType = 1,
                rainPath = '/Dedicated/IFC/data_bin/stage4/st4_'+str(y)+'/iowa_',
                controlFile='/Users/nicolas/Parameter_identifiability/'+path+'control_vr5.sav',
                nComponents = 1,
                Components = [0],
                datResults='out_'+str(k)+'_'+str(y)+'.dat',
                evpFile='/Users/nicolas/Parameter_identifiability/data/for_hlm/evap.mon')
            #Updates the dictionary with the parameters to setup the run
            Proc={global_name:{'nproc': 14, 'secondplane': False}}
            #Writes the runfile with all the comands
            wat[name].write_runfile(path + 'run_'+str(k)+'_'+str(y)+'.sh', 
                Proc, jobName='o'+str(k)+'_'+str(y), nCores= 14) 

# ## Run the prun level 4 setup to save snapshots to then run just relevant events 

for name in wat.keys():
    path = names[name]['path']    
    for y in range(2002,2019):
        #Writes the global files and run files    
        for k in [4,]:
            #Set the name of the global file
            global_name = 'global_'+str(k)+'_'+str(y)+'.gbl'
            #Writes the global file
            wat[name].write_Global(path + global_name,model_uid=608,
                date1=str(y)+'-03-01 01:00', 
                date2=str(y)+'-12-31 01:00',
                rvrFile='/Users/nicolas/Parameter_identifiability/'+path+'topology.rvr', 
                rvrLink='',
                prmFile='/Users/nicolas/Parameter_identifiability/'+path+'param_o'+str(k)+'.prm',
                initialFile='/Users/nicolas/Parameter_identifiability/data/for_hlm/initial.uini',
                initialType = 1,
                rainPath = '/Dedicated/IFC/data_bin/stage4/st4_'+str(y)+'/iowa_',
                controlFile='/Users/nicolas/Parameter_identifiability/'+path+'control_vr5.sav',
                nComponents = 1,
                Components = [0],
                datResults='out_'+str(k)+'_'+str(y)+'.dat',
                snapType = 4,
                snapTime = 1440,
                snapPath = '/Users/nicolas/Parameter_identifiability/data/for_hlm/south_skunk/snapshots/p4.h5',
                evpFile='/Users/nicolas/Parameter_identifiability/data/for_hlm/evap.mon')
            #Updates the dictionary with the parameters to setup the run
            Proc={global_name:{'nproc': 14, 'secondplane': False}}
            #Writes the runfile with all the comands
            wat[name].write_runfile(path + 'run_'+str(k)+'_'+str(y)+'.sh', 
                Proc, jobName='o'+str(k)+'_'+str(y), nCores= 14) 

# # Random watersheds of order 5.
#
# Here we generate 10 random watersheds that distribute the parameters at the prun level 5. 

y = 815

path = names['south']['path']

years = [2008,2010,2013,2016,2018]

# +
#setupo the distributed virtual parameters of both cases for the different levels.
for name in wat.keys():
    for i in range(1,2):
        # Create random parameters for the prun 5 level 
        f_vo = cascade_values_sync(wat['south'].Table, 'vo', 5, [0.15,0.3,0.22], 0.2)
        f_l1 = cascade_values_sync(wat['south'].Table, 'la', 5, [0.18,0.23,0.25], 0.2)
        
        # setup the parameters for each watershed.
        wat['south'].Table.at[wat['south'].Table.index, 'channel_vo'] = wat['south'].Table['vo5'].values
        wat['south'].Table.at[wat['south'].Table.index, 'lambda1'] = wat['south'].Table['la5'].values
        
        #Write the param file to run HLM
        wat['south'].write_prm(names['south']['path']+'param_s'+str(i)+'.prm', 
            extraNames=Order,
            extraFormats=[Params[k]['format'] for k in Order])
        
        #Writes the global and the runfile for each experiment
        #for y in range(2002,2019):            
        #Set the name of the global file
        global_name = 'global_s'+str(i)+'_'+str(y)+'.gbl'
        #Writes the global file
        wat[name].write_Global(path + global_name,model_uid=608,
            date1='2016'+'-03-01 01:00', 
            date2='2016'+'-12-31 01:00',
            rvrFile='/Users/nicolas/Parameter_identifiability/'+path+'topology.rvr', 
            rvrLink='',
            prmFile='/Users/nicolas/Parameter_identifiability/'+path+'param_s'+str(i)+'.prm',
            initialFile='/Users/nicolas/Parameter_identifiability/data/for_hlm/initial.uini',
            initialType = 1,
            rainPath = '/Dedicated/IFC/data_bin/stage4/st4_'+str(y)+'/iowa_',
            controlFile='/Users/nicolas/Parameter_identifiability/'+path+'control_vr5.sav',
            nComponents = 1,
            Components = [0],
            datResults='out_'+str(i)+'_'+str(y)+'.h5',
            datType = 5,
            evpFile='/Users/nicolas/Parameter_identifiability/data/for_hlm/evap.mon')
        #Updates the dictionary with the parameters to setup the run
        Proc={global_name:{'nproc': 14, 'secondplane': False}}
        #Writes the runfile with all the comands
        wat[name].write_runfile(path + 'run_s'+str(i)+'_'+str(y)+'.sh', 
            Proc, jobName='o'+str(i)+'_'+str(y), nCores= 14) 

        
# -

q = pd.read_hdf('data/for_hlm/south_skunk/out_1_815.h5',)

import h5py
import h5

h = h5py.File('data/for_hlm/south_skunk/out_1_815.h5','r')

h.keys()

h.keys()

g = h['/']

o = h.get['outputs']

o.value



StarcoderdataJupyterScriptsDedupFiltered