import click

import io

import itertools

import json

import matplotlib.pyplot as plt

import numpy as np

import os

import re

from datetime import datetime

from matplotlib.backends.backend_pdf import PdfPages

from ..scenario import ScenarioSeries

from ..util import load_data_file, load_statvar, nash_sutcliffe

@click.group()

def prmspy():

"access PRMS-Python functionality from the command line"

click.echo('\n*** Welcome to PRMS-Python! ***\n')

@prmspy.command()

@click.argument('base_data_dir', nargs=1)

@click.option('--params', '-p', nargs=1, type=str, multiple=True,

help='list of parameters to edit')

@click.option('--scale_vals', '-s', nargs=1, type=str, multiple=True,

help='list of scaling values; must be divisible by len(params), '

'e.g. "[0.8, 0.9, 1.0, 1.1, 1.2]"'

)

@click.option('--output-dir', '-o', nargs=1, type=str,

help='directory where scenario data should be written')

@click.option('--title', '-t', nargs=1, type=str,

help='title for this parameter scaling experiment')

@click.option('--description', '-d', nargs=1, type=str,

help='description for this parameter scaling experiment')

@click.option('--run-prms', is_flag=True,

help='run PRMS after creating scenario input data')

@click.option('--prms-exec', '-e', default='prms',

help='PRMS executable to be used, eg prmsV4')

@click.option('--nproc', '-n', default=None, type=int,

help='number of processors to use')

@click.option('--analyze-output', is_flag=True,

help='create analysis of model output in the form of a '

'"Nash-Sutcliffe Matrix"')

def param_scale_sim(base_data_dir,

params,

scale_vals,

output_dir,

title,

description,

run_prms,

prms_exec,

nproc,

analyze_output):

'Provide params and scaling values; run PRMS scenarios'

# scale_vals must be divisible by the number of params

# valid_input = len(scale_vals) % len(params) == 0

print(scale_vals)

scale_vals = [eval(el) for el in scale_vals]

# if not valid_input:

# raise IOError('The length of scale_vals list is not divisible '

# 'by the length of the parameter list')

if not output_dir:

output_dir =\

re.sub(r'\/|\\\\', '-', base_data_dir) + \

'-modified-params-' + datetime.now().isoformat()

os.mkdir(output_dir)

s = ScenarioSeries(

base_data_dir, output_dir, title=title, description=description

)

def _scale_fun(val):

def scale(x):

return x * val

return scale

# assign scaling values to parameter names

# dimensionality correctness guaranteed from valid_input check

pval_len = len(scale_vals)/len(params)

def _slice(idx):

return slice(pval_len*idx, pval_len*(idx+1))

param_val_dict = {

param: zip(itertools.repeat(param), scale_vals[idx])

for idx, param in enumerate(params)

}

# param_val_dict = {

# param: zip(itertools.repeat(param, pval_len), scale_vals[_slice(idx)])

# for idx, param in enumerate(params)

# }

# for every combination we need to iterate over every value of

# each param we have

combinations = itertools.product(*param_val_dict.values())

# now build the scenario_list with re-callable titles

scenario_list = []

for combo in combinations:

scenario_def = {

'title': _build_title(combo)

}

scenario_def.update(

{

param: _scale_fun(val)

for param, val in combo

}

)

scenario_list.append(scenario_def)

s.build(scenario_list)

if run_prms:

s.run(prms_exec=prms_exec, nproc=nproc)

if analyze_output:

print('\n** Sorry, analyze_output has not yet been implemented! **\n')

@prmspy.command()

@click.argument('data_dir', nargs=1)

@click.argument('output_pdf_path', nargs=1)

def nash_sutcliffe_matrix(data_dir, output_pdf_path):

'Save a PDF of the Nash-Sutcliffe created from <data_dir> to <output_pdf_path>'

observed = load_data_file(

os.path.join(data_dir, 'base_inputs', 'data')

).runoff_1

series_metadata = json.loads(

io.open(os.path.join(data_dir, 'series_metadata.json')).read()

)

modeled_flows = {

title: load_statvar(

os.path.join(data_dir, uu, 'outputs', 'statvar.dat')

).basin_cfs_1

for uu, title in series_metadata['uuid_title_map'].items()

}

titles = list(modeled_flows.keys())

params = [pair.split(':')[0] for pair in titles[0].split('|')]

if len(params) > 2:

print('This command is only supported for two covarying parameters!')

exit(1)

p1_vals = sorted({

float(pair[0].split(':')[1])

for pair in map(lambda t: t.split('|'), titles)

})

p2_vals = sorted({

float(pair[1].split(':')[1])

for pair in map(lambda t: t.split('|'), titles)

})

print(p1_vals)

nash_sutcliffe_mat = np.zeros((len(p1_vals), len(p2_vals)))

for p1_idx, p1_val in enumerate(p1_vals):

for p2_idx, p2_val in enumerate(p2_vals):

coord = (p1_idx, p2_idx)

title = '{p1}:{p1val:.3f}|{p2}:{p2val:.3f}'.format(

p1=params[0], p2=params[1], p1val=p1_val, p2val=p2_val

)

modeled = modeled_flows[title]

nash_sutcliffe_mat[coord] = nash_sutcliffe(observed, modeled)

with PdfPages(output_pdf_path) as pdf:

fig, ax = plt.subplots()

cax = ax.matshow(nash_sutcliffe_mat, cmap='viridis')

ytix = p1_vals

xtix = p2_vals

plt.yticks(range(len(ytix)), ytix)

plt.xticks(range(len(xtix)), xtix)

ax.xaxis.set_ticks_position('bottom')

plt.ylabel(params[0].replace('"', ''))

plt.xlabel(params[1].replace('"', ''))

for i, j in itertools.product(range(len(ytix)), range(len(xtix))):

plt.text(j, i, '%.2f' % nash_sutcliffe_mat[i, j],

horizontalalignment='center',

color='w'

if nash_sutcliffe_mat[i, j] <

np.mean(nash_sutcliffe_mat.flatten())

else 'k')

plt.title('Nash-Sutcliffe Matrix')

plt.grid(b=False)

fig.colorbar(cax)

pdf.savefig()

plt.close()

def _build_title(combo):

'''

Given a list of tuples with parameter/scale_val pairs, build the title

for the parameterization. For example,

>>> combo = [('snow_adj', 0.8), ('rad_trncf', 0.9), ('jh_coef', 1.1)]

>>> assert _build_title(combo) ==\

'"snow_adj":0.800|"rad_trncf":0.900|"jh_coef":1.100'

'''

return '|'.join('"{0}":{1:.3f}'.format(param, val) for param, val in combo)