"""
Simulation tests.
"""

from carousel.core import (
    logging, data_sources, outputs, formulas, calculations, simulations, models,
    UREG
)
from carousel.contrib.readers import ArgumentReader
from carousel.tests import PROJ_PATH
import numpy as np
import os
import sympy

LOGGER = logging.getLogger(__name__)


def test_make_sim_metaclass():
    """
    Test setting the simulation parameter file as class attributes versus
    specifying the simulation parameter file in the model parameter file.

    :return: simulation
    """

    class SimTest1(simulations.Simulation):
        sim_file = 'Tuscon.json'
        sim_path = os.path.join(PROJ_PATH, 'simulations', 'Standalone')

    sim_test1 = SimTest1()
    return sim_test1


class PythagorasData(data_sources.DataSource):
    data_cache_enabled = False
    data_reader = ArgumentReader
    a = {'units': 'cm', 'argpos': 0}
    b = {'units': 'cm', 'argpos': 2}
    a_unc = {'units': 'cm', 'argpos': 1}
    b_unc = {'units': 'cm', 'argpos': 3}

    def __prepare_data__(self):
        keys = self.parameters.keys()
        for k in keys:
            if k.endswith('_unc'):
                unc = self.data.pop(k)
                self.data_source.pop(k)
                kunc = k[:-4]
                v = self.data[kunc]
                if not unc.dimensionless:
                    unc = (unc / v)
                # raises dimensionality error if not dimensionless
                self.uncertainty[kunc] = {kunc: unc.to(UREG.percent)}
            else:
                self.isconstant[k] = True


class PythagorasOutput(outputs.Output):
    c = {'units': 'cm', 'isconstant': True}


def f_hypotenuse(a, b):
    a, b = np.atleast_1d(a), np.atleast_1d(b)
    return np.sqrt(a * a + b * b).reshape(1, -1)


class PythagorasFormula(formulas.Formula):
    module = 'carousel.tests.test_sim'
    formulas = {
        'f_hypotenuse': {
            'args': ['a', 'b'],
            'units': [('=A', ), ('=A', '=A')],
            'isconstant': []
        }
    }


class PythagorasCalc(calculations.Calc):
    static = [{
        'formula': 'f_hypotenuse',
        'args': {'data': {'a': 'a', 'b': 'b'}},
        'returns': ['c']
    }]


class PythagorasSim(simulations.Simulation):
    ID = 'Pythagorean Theorem'
    commands = ['start', 'load', 'run', 'pause']
    path = '~/Carousel_Tests'
    thresholds = None
    interval = [1, 'hour']
    sim_length = [0, 'hour']
    write_frequency = 1
    write_fields = {'data': ['a', 'b'], 'outputs': ['c']}
    display_frequency = 1
    display_fields = {'data': ['a', 'b'], 'outputs': ['c']}


class PythagorasModel(models.Model):
    modelpath = os.path.dirname(__file__)
    data = [PythagorasData]
    outputs = [PythagorasOutput]
    formulas = [PythagorasFormula]
    calculations = [PythagorasCalc]
    simulations = [PythagorasSim]


def test_call_sim_with_args():
    a, a_unc, b, b_unc = 3.0, 0.1, 4.0, 0.1
    c = f_hypotenuse(a, b)
    m1 = PythagorasModel()
    data = {'PythagorasData': {'a': a, 'b': b, 'a_unc': a_unc, 'b_unc': b_unc}}
    m1.command('run', data=data)
    assert m1.registries['outputs']['c'].m == c
    assert m1.registries['outputs']['c'].u == UREG.cm
    x, y = sympy.symbols('x, y')
    z = sympy.sqrt(x * x + y * y)
    fx = sympy.lambdify((x, y), z.diff(x))
    fy = sympy.lambdify((x, y), z.diff(y))
    dz = np.sqrt(fx(a, b) ** 2 * a_unc ** 2 + fy(a, b) ** 2 * b_unc ** 2)
    c_unc = c * np.sqrt(m1.registries['outputs'].variance['c']['c'])
    LOGGER.debug('uncertainty in c is %g', c_unc)
    assert np.isclose(dz, np.array(c_unc))
    c_unc = c * m1.registries['outputs'].uncertainty['c']['c']
    assert np.isclose(dz, np.array(c_unc))
    return m1


if __name__ == '__main__':
    m = test_call_sim_with_args()

@@ -140,16 +140,121 @@
  UREG%0A)%0A
+from carousel.core.data_sources import DataParameter%0Afrom carousel.core.formulas import FormulaParameter%0A
 from car
@@ -297,16 +297,16 @@
 tReader%0A
-
 from car
@@ -956,16 +956,32 @@
     a = 
+DataParameter(**
 %7B'units'
@@ -1000,16 +1000,17 @@
 pos': 0%7D
+)
 %0A    b =
@@ -1010,16 +1010,32 @@
     b = 
+DataParameter(**
 %7B'units'
@@ -1054,16 +1054,17 @@
 pos': 2%7D
+)
 %0A    a_u
@@ -1064,24 +1064,40 @@
     a_unc = 
+DataParameter(**
 %7B'units': 'c
@@ -1112,16 +1112,17 @@
 pos': 1%7D
+)
 %0A    b_u
@@ -1122,24 +1122,40 @@
     b_unc = 
+DataParameter(**
 %7B'units': 'c
@@ -1170,16 +1170,17 @@
 pos': 3%7D
+)
 %0A%0A    de
@@ -1981,95 +1981,49 @@
     
-module = 'carousel.tests.test_sim'%0A    formulas = %7B%0A        'f_hypotenuse': %7B%0A 
+f_hypotenuse = FormulaParameter(%0A
         
    '
@@ -2018,27 +2018,21 @@
         
-   '
 args
-': 
+=
 %5B'a', 'b
@@ -2047,21 +2047,14 @@
     
-  
-  '
 units
-': 
+=
 %5B('=
@@ -2083,21 +2083,16 @@
         
-   
- '
 isconsta
@@ -2097,29 +2097,77 @@
 tant
-': 
+=
 %5B%5D%0A    
+)%0A%0A
     
-%7D%0A    %7D
+class Meta:%0A        module = 'carousel.tests.test_sim'
 %0A%0A%0Ac

update test_sim with Parameters