from copy import copy

import json

import glob

import numpy as np

import os

import re

import shutil

import unittest

from difflib import Differ

from numpy.testing import assert_array_almost_equal

from prms_python import (

modify_params, Parameters, Scenario, ScenarioSeries, Simulation,

SimulationSeries, Data

)

OPJ = os.path.join

class TestSimulationSeries(unittest.TestCase):

'''

Expect that a SimulationSeries

'''

def setUp(self):

self.test_model_data_dir = os.path.join(

'test', 'data', 'models', 'lbcd'

)

self.simulation_dir = os.path.join(self.test_model_data_dir, 'tmp_sim')

def tearDown(self):

for g in glob.glob(OPJ(self.simulation_dir, '*')):

shutil.rmtree(g)

def test_simulation_series(self):

tdd = self.test_model_data_dir

data = Data(OPJ(tdd, 'data'))

base_parameters = Parameters(OPJ(tdd, 'parameters'))

def _copy_mod(base_parameters, val):

ret = copy(base_parameters)

ret['dday_intcp'][:] = val

return ret

parameters_gen = (

(val, _copy_mod(base_parameters, val))

for val in (-50, -40, -30)

)

control_path = OPJ(tdd, 'control')

series = SimulationSeries(

Simulation.from_data(

data, parameters, control_path,

self.simulation_dir + str(val)

)

for val, parameters in parameters_gen

)

outputs = list(series.run().outputs_iter())

self.assertEqual(len(series), 3)

self.assertEqual(len(outputs), 3)

for out in outputs:

sdir = out['simulation_dir']

assert os.path.isdir(out['simulation_dir'])

assert os.path.exists(OPJ(sdir, 'outputs', 'statvar.dat'))

assert os.path.exists(OPJ(sdir, 'inputs', 'data'))

assert os.path.exists(OPJ(sdir, 'inputs', 'parameters'))

shutil.rmtree(sdir)

class TestSimulation(unittest.TestCase):

"""

Simulations should take a base directory and return a simulation directory

"""

def setUp(self):

self.test_data_dir = os.path.join('test', 'data')

self.test_model_data_dir = os.path.join(

'test', 'data', 'models', 'lbcd'

)

self.simulation_dir = os.path.join(self.test_data_dir, 'tmp_sim')

def tearDown(self):

if os.path.exists(self.simulation_dir):

shutil.rmtree(self.simulation_dir)

def test_simulation_no_simdir(self):

"Simulation should run and write outputs to input directory when simulation_dir is not specified"

s = Simulation(self.test_model_data_dir)

s.run()

g = [

os.path.basename(f)

for f in glob.glob(os.path.join(self.test_model_data_dir, '*'))

]

self.assertIn('prms_ic.out', g)

self.assertIn('prms.out', g)

self.assertIn('statvar.dat', g)

def test_simulation_w_simdir(self):

"Simulation should create sim dir with inputs and outputs directory when simulation_dir is specified"

s = Simulation(self.test_model_data_dir, self.simulation_dir)

s.run()

gs = [

os.path.basename(f)

for f in glob.glob(os.path.join(self.simulation_dir, '*'))

]

self.assertIn('inputs', gs)

self.assertIn('outputs', gs)

assert_valid_input_dir(

self, os.path.join(self.simulation_dir, 'inputs')

)

assert_valid_output_dir(

self, os.path.join(self.simulation_dir, 'outputs')

)

def test_simulation_from_data(self):

"""

Use @classmethod from_data to build a simulation from Parameters and Data instances

"""

tdd = self.test_model_data_dir

data = Data(OPJ(tdd, 'data'))

parameters = Parameters(OPJ(tdd, 'parameters'))

ctrl = OPJ(tdd, 'control')

test_dir = OPJ(tdd, 'test-sim-dir')

if os.path.isdir(test_dir):

shutil.rmtree(test_dir)

s = Simulation.from_data(data, parameters, ctrl, test_dir)

s.run()

g = [

os.path.basename(f)

for f in glob.glob(OPJ(test_dir, 'outputs', '*'))

]

self.assertIn('prms_ic.out', g)

self.assertIn('prms.out', g)

self.assertIn('statvar.dat', g)

# clean up

shutil.rmtree(test_dir)

class TestScenario(unittest.TestCase):

def setUp(self):

self.test_data_dir = os.path.join('test', 'data')

self.test_model_data_dir = os.path.join(

'test', 'data', 'models', 'lbcd'

)

self.scenario_dir = os.path.join(self.test_data_dir, 'tmp_scenario')

def tearDown(self):

if os.path.exists(self.scenario_dir):

shutil.rmtree(self.scenario_dir)

def test_create_scenario(self):

"""

a simulation setup should create a simulation directory with correct scenario data

"""

s = Scenario(

self.test_model_data_dir, self.scenario_dir,

title='Scenario Uno', description='test scenario for prms_python'

)

param_mods = {

'snow_adj': lambda x: 1.1*x,

'rad_trncf': lambda x: 0.9*x

}

s.build(param_mod_funs=param_mods)

assert_valid_input_dir(self, self.scenario_dir) # os.path.join(self.scenario_dir, 'inputs'))

s.run()

assert_valid_input_dir(

self, os.path.join(self.scenario_dir, 'inputs')

)

assert_valid_output_dir(

self, os.path.join(self.scenario_dir, 'outputs')

)

md_json_path = os.path.join(self.scenario_dir, 'metadata.json')

assert os.path.isfile(md_json_path)

md_json = json.loads(open(md_json_path).read())

assert md_json['title'] == 'Scenario Uno'

assert md_json['description'] == 'test scenario for prms_python'

assert 'start_datetime' in md_json

assert 'end_datetime' in md_json

assert 'mod_funs_dict' in md_json

p_base = Parameters(

os.path.join(self.test_model_data_dir, 'parameters')

)

p_scen = Parameters(

os.path.join(self.scenario_dir, 'inputs', 'parameters')

)

assert_array_almost_equal(p_base['snow_adj']*1.1, p_scen['snow_adj'])

assert_array_almost_equal(p_base['rad_trncf']*0.9, p_scen['rad_trncf'])

class TestScenarios(unittest.TestCase):

def setUp(self):

self.test_data_dir = os.path.join('test', 'data')

self.test_model_data_dir = os.path.join(

self.test_data_dir, 'models', 'lbcd'

)

self.scenarios_dir = os.path.join(self.test_data_dir, 'tmp_scenarios')

def tearDown(self):

if os.path.exists(self.scenarios_dir):

shutil.rmtree(self.scenarios_dir)

def test_scenario_series(self):

"create_many_simulations should create many simulation directories and correct data"

s = ScenarioSeries(

self.test_model_data_dir, self.scenarios_dir,

title='scenario series uno',

description='''

Each scenario is given a title with the schema

'"<param_name1>":<scale_value1>|"<param_name2>":<scale_value2>', which can be

easily parsed into a Python dictionary later. We use the pipe instead of

comma for easier visual inspection, though this is entirely up to the

user/developer. The scenario consists of parameters scaled for each of the

specified parameters.

'''

)

def _scale_param(val):

def scale_by_val(x):

return x*val

return scale_by_val

scale_arange = np.arange(0.7, 0.9, 0.1)

series_funs = [

{

'title': '"rad_trncf":{0:.1f}|"snow_adj":{0:.1f}'.format(val),

'rad_trncf': _scale_param(val),

'snow_adj': _scale_param(val)

}

for val in scale_arange

]

s.build(series_funs)

g_series = glob.glob(os.path.join(self.scenarios_dir, '*'))

# this should be 5 because of series_metadata.json and inputs_dir

assert len(g_series) == 5, g_series

uuid_pattern = re.compile(

r'^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'

)

uuid_dir_count = 0

non_uuid_dir_count = 0

series_dirs = []

found_metadata = False

for d in g_series:

# XXX TODO this is in flux. at this point maybe all will be dirs

if os.path.isdir(d):

# at this point the files are in the base scenario dir;

# they will be moved later when the scenario is run

assert_valid_input_dir(self, os.path.join(d))

if re.match(uuid_pattern, os.path.basename(d)):

uuid_dir_count += 1

assert re.match(uuid_pattern, os.path.basename(d)), os.path.basename(d)

else:

non_uuid_dir_count += 1

else:

found_metadata = True

series_dirs.append(d)

assert uuid_dir_count == 3

assert non_uuid_dir_count == 1

assert found_metadata

s.run()

g_series = glob.glob(os.path.join(self.scenarios_dir, '*'))

series_md_path = os.path.join(

self.scenarios_dir, 'series_metadata.json'

)

self.assertIn(series_md_path, g_series)

series_md = json.loads(open(series_md_path).read())

dir_titles = series_md['uuid_title_map']

titles = dir_titles.values()

assert '"rad_trncf":0.7|"snow_adj":0.7' in titles

assert '"rad_trncf":0.8|"snow_adj":0.8' in titles

assert '"rad_trncf":0.9|"snow_adj":0.9' in titles

p_base = Parameters(

os.path.join(self.test_model_data_dir, 'parameters')

)

for g in g_series:

if os.path.isdir(g) and 'input' not in g:

assert_valid_output_dir(self, os.path.join(g, 'outputs'))

md = json.loads(open(os.path.join(g, 'metadata.json')).read())

title = md['title']

scale_vals = eval('{' + title.replace('|', ',') + '}')

p_scen = Parameters(os.path.join(g, 'inputs', 'parameters'))

rad_base = p_base['rad_trncf']

rad_scen = p_scen['rad_trncf']

snow_base = p_base['snow_adj']

snow_scen = p_scen['snow_adj']

assert_array_almost_equal(

rad_base*scale_vals['rad_trncf'], rad_scen)

assert_array_almost_equal(

snow_base*scale_vals['snow_adj'], snow_scen)

def assert_valid_input_dir(test_case, d):

gs = [

os.path.basename(f)

for f in glob.glob(os.path.join(d, '*'))

]

test_case.assertIn('control', gs, d)

test_case.assertIn('parameters', gs, d)

test_case.assertIn('data', gs, d)

def assert_valid_output_dir(test_case, d):

go = [

os.path.basename(f)

for f in

glob.glob(os.path.join(d, '*'))

]

test_case.assertIn('prms_ic.out', go)

test_case.assertIn('prms.out', go)

test_case.assertIn('statvar.dat', go)

class TestParameters(unittest.TestCase):

def setUp(self):

self.test_data_dir = os.path.join('test', 'data')

self.test_param = os.path.join(self.test_data_dir, 'parameters')

self.temp_dir = os.path.join('test', 'data', 'tmp')

if os.path.isdir(self.temp_dir):

shutil.rmtree(self.temp_dir)

os.mkdir(self.temp_dir)

def tearDown(self):

if os.path.isdir(self.temp_dir):

shutil.rmtree(self.temp_dir)

def test_faithful_copy(self):

"""

Parameters class should load and write identical file

"""

test_copy = os.path.join(self.temp_dir, 'param_copy_test')

p = Parameters(self.test_param)

p.write(test_copy)

expected = open(self.test_param)

generated = open(test_copy)

# ignore first two lines ignore now date written on Parameters.write

expected = expected.readlines()[2:]

generated = generated.readlines()[2:]

assert expected == generated

def test_modify_params(self):

"parameter and data file modification access should work as expected"

# test parameters mod with 1d rad_trncf

params_out = os.path.join(self.temp_dir, 'modify_params_out')

mod_d = {'rad_trncf': lambda x: 1.1*x}

modify_params(self.test_param, params_out, mod_d)

generated = open(params_out).readlines()[2:]

expected = open(

os.path.join(

self.test_data_dir, 'expected_rad_trncf_1.1'

)

).readlines()[2:]

assert expected == generated

# test parameters mod with 2d snow_adj

params_out = os.path.join(self.temp_dir, 'modify_params_out')

mod_d = {'snow_adj': lambda x: 1.1*x}

modify_params(self.test_param, params_out, mod_d)

generated = open(params_out).readlines()[2:]

expected = open(

os.path.join(

self.test_data_dir, 'expected_snow_adj_1.1'

)

).readlines()[2:]

for a, b in zip(expected, generated):

try:

a = float(a)

b = float(b)

self.assertAlmostEqual(a, b, places=5)

except ValueError:

assert a == b

def colored_string_diff(s1, s2):

""" Writes differences between strings s1 and s2 """

d = Differ()

diff = d.compare(s1.splitlines(), s2.splitlines())

diffList = [el for el in diff

if el[0] != ' ' and el[0] != '?']

for l in diffList:

if l[0] == '+':

print(bcolors.GREEN + '+' + l[1:] + bcolors.ENDC)

elif l[0] == '-':

print(bcolors.RED + '-' + l[1:] + bcolors.ENDC)

else:

assert False, 'Error, diffList entry must start with + or -'

class bcolors:

HEADER = '\033[95m'

BLUE = '\033[94m'

GREEN = '\033[92m'

YELLOW = '\033[93m'

RED = '\033[91m'

ENDC = '\033[0m'