#!/usr/bin/env python3
#
# Tests the EasyML module.
#
# This file is part of Myokit.
# See http://myokit.org for copyright, sharing, and licensing details.
#
import os
import unittest

import myokit
import myokit.formats
import myokit.formats.easyml
import myokit.tests

from myokit import (
    Number, PrefixPlus, PrefixMinus, Plus, Minus,
    Multiply, Divide, Quotient, Remainder, Power, Sqrt,
    Exp, Log, Log10, Sin, Cos, Tan, ASin, ACos, ATan, Floor, Ceil, Abs,
    Not, And, Or, Equal, NotEqual, More, Less, MoreEqual, LessEqual,
    If, Piecewise,
)

from myokit.tests import TemporaryDirectory, WarningCollector, DIR_DATA


# Model that requires unit conversion
units_model = """
[[model]]
membrane.V = -0.08
hh.x = 0.1
hh.y = 0.9
mm.C = 0.9

[engine]
time = 0 [s]
    in [s]
    bind time

[membrane]
dot(V) = (hh.I1 + mm.I2) / C
    in [V]
C = 20 [pF]
    in [pF]

[hh]
dot(x) = (inf - x) / tau
    inf = 0.8
    tau = 3 [s]
        in [s]
dot(y) = alpha * (1 - y) - beta * y
    alpha = 0.1 [1/s]
        in [1/s]
    beta = 0.2 [1/s]
        in [1/s]
I1 = 3 [pS] * x * y * (membrane.V - 0.05 [V])
    in [pA]

[mm]
dot(C) = beta * O - alpha * C
alpha = 0.3 [1/s]
    in [1/s]
beta = 0.4 [1/s]
    in [1/s]
O = 1 - C
I2 = 2 [pS] * O * (membrane.V + 0.02 [V])
    in [pA]
"""
units_output = """
/*
This file was generated by Myokit.
*/

V; .nodal(); .external(Vm);
Iion; .nodal(); .external();

V_init = -80.0;
x_init = 0.1;
y_init = 0.9;
C_init = 0.9;

// hh
I1 = 3.0 * x * y * (V / 1000.0 - 0.05) * 0.05;
x_inf = 0.8;
tau_x = 3.0 * 1000.0;
alpha_y = 0.1 * 0.001;
beta_y = 0.2 * 0.001;

// mm
diff_C = (mm_beta * O - mm_alpha * C) * 0.001;
I2 = 2.0 * O * (V / 1000.0 + 0.02) * 0.05;
O = 1.0 - C;
mm_alpha = 0.3;
mm_beta = 0.4;

// Sum of currents
Iion = I1 + I2;

// Markov model
group {
  C;
}.method(markov_be);

// Trace all currents and state variables
group {
  I1;
  I2;
  V;
  x;
  y;
  C;
}.trace();

// Parameters
group {
  mm_alpha;
  mm_beta;
}.param();
"""


class EasyMLExporterTest(unittest.TestCase):
    """ Tests EasyML export. """

    def test_easyml_exporter(self):
        # Tests exporting a model

        model = myokit.load_model('example')
        with TemporaryDirectory() as d:
            path = d.path('easy.model')

            # Test with simple model
            e = myokit.formats.easyml.EasyMLExporter()
            e.model(path, model)

            # Test with extra bound variables
            model.get('membrane.C').set_binding('hello')
            e.model(path, model)

            # Test without V being a state variable
            v = model.get('membrane.V')
            v.demote()
            v.set_rhs(3)
            e.model(path, model)

            # Test with invalid model
            v.set_rhs('2 * V')
            self.assertRaisesRegex(
                myokit.ExportError, 'valid model', e.model, path, model)

    def test_easyml_exporter_static(self):
        # Tests exporting a model (with HH and markov states) and compares
        # against reference output.

        # Export model
        m = myokit.load_model(os.path.join(DIR_DATA, 'decker-2009.mmt'))
        e = myokit.formats.easyml.EasyMLExporter()
        with TemporaryDirectory() as d:
            path = d.path('decker.model')
            e.model(path, m)
            with open(path, 'r') as f:
                observed = f.readlines()

        # Load expected output
        with open(os.path.join(DIR_DATA, 'decker.model'), 'r') as f:
            expected = f.readlines()

        # Compare (line by line, for readable output)
        for ob, ex in zip(observed, expected):
            self.assertEqual(ob, ex)
        self.assertEqual(len(observed), len(expected))

    def test_unit_conversion(self):
        # Tests exporting a model that requires unit conversion

        # Export model
        m = myokit.parse_model(units_model)
        e = myokit.formats.easyml.EasyMLExporter()
        with TemporaryDirectory() as d:
            path = d.path('easy.model')
            e.model(path, m)
            with open(path, 'r') as f:
                observed = f.read().strip().splitlines()

        # Get expected output
        expected = units_output.strip().splitlines()

        # Compare (line by line, for readable output)
        for ob, ex in zip(observed, expected):
            self.assertEqual(ob, ex)
        self.assertEqual(len(observed), len(expected))

        # Test warnings are raised if conversion fails
        m.get('membrane.V').set_rhs('hh.I1 + mm.I2')
        m.get('membrane').remove_variable(m.get('membrane.C'))
        with TemporaryDirectory() as d:
            path = d.path('easy.model')
            with WarningCollector() as c:
                e.model(path, m)
            self.assertIn('Unable to convert hh.I1', c.text())
            self.assertIn('Unable to convert mm.I2', c.text())

        m.get('engine.time').set_unit(myokit.units.cm)
        with TemporaryDirectory() as d:
            path = d.path('easy.model')
            with WarningCollector() as c:
                e.model(path, m)
            self.assertIn('Unable to convert time units [cm]', c.text())

    def test_export_reused_variable(self):
        # Tests exporting when an `inf` or other special variable is used twice

        # Create model re-using tau and inf
        m = myokit.parse_model(
            """
            [[model]]
            m.V = -80
            c.x = 0.1
            c.y = 0.1

            [m]
            time = 0 bind time
            i_ion = c.I
            dot(V) = -i_ion

            [c]
            inf = 0.5
            tau = 3
            dot(x) = (inf - x) / tau
            dot(y) = (inf - y) / tau
            I = x * y * (m.V - 50)
            """)

        # Export, and read back in
        e = myokit.formats.easyml.EasyMLExporter()
        with TemporaryDirectory() as d:
            path = d.path('easy.model')
            e.model(path, m)
            with open(path, 'r') as f:
                x = f.read()

        self.assertIn('x_inf =', x)
        self.assertIn('y_inf =', x)
        self.assertIn('tau_x =', x)
        self.assertIn('tau_y =', x)

    def test_easyml_exporter_fetching(self):
        # Tests getting an EasyML exporter via the 'exporter' interface

        e = myokit.formats.exporter('easyml')
        self.assertIsInstance(e, myokit.formats.easyml.EasyMLExporter)

    def test_capability_reporting(self):
        # Tests if the correct capabilities are reported
        e = myokit.formats.easyml.EasyMLExporter()
        self.assertTrue(e.supports_model())


class EasyMLExpressionWriterTest(myokit.tests.ExpressionWriterTestCase):
    """ Test conversion to EasyML syntax. """
    _name = 'easyml'
    _target = myokit.formats.easyml.EasyMLExpressionWriter

    def test_number(self):
        self.eq(Number(1), '1.0')
        self.eq(Number(-2), '-2.0')
        self.eq(Number(13, 'mV'), '13.0')

    def test_name(self):
        # Inherited from CBasedExpressionWriter
        self.eq(self.a, 'a')
        w = self._target()
        w.set_lhs_function(lambda v: v.var().qname().upper())
        self.assertEqual(w.ex(self.a), 'COMP.A')

    def test_derivative(self):
        # Inherited from CBasedExpressionWriter
        self.eq(myokit.Derivative(self.a), 'dot(a)')

    def test_partial_derivative(self):
        e = myokit.PartialDerivative(self.a, self.b)
        self.assertRaisesRegex(NotImplementedError, 'Partial', self.w.ex, e)

    def test_initial_value(self):
        e = myokit.InitialValue(self.a)
        self.assertRaisesRegex(NotImplementedError, 'Initial', self.w.ex, e)

    def test_prefix_plus_minus(self):
        # Inherited from CBasedExpressionWriter
        p = Number(11, 'kV')
        a, b, c = self.abc
        self.eq(PrefixPlus(p), '+11.0')
        self.eq(PrefixPlus(PrefixPlus(PrefixPlus(p))), '+(+(+11.0))')
        self.eq(Divide(PrefixPlus(Plus(a, b)), c), '+(a + b) / c')
        self.eq(PrefixMinus(p), '-11.0')
        self.eq(PrefixMinus(PrefixMinus(p)), '-(-11.0)')
        self.eq(PrefixMinus(Number(-1)), '-(-1.0)')
        self.eq(PrefixMinus(Minus(a, b)), '-(a - b)')
        self.eq(Multiply(PrefixMinus(Plus(b, a)), c), '-(b + a) * c')
        self.eq(PrefixMinus(Divide(b, a)), '-(b / a)')

    def test_plus_minus(self):
        a, b, c = self.abc
        self.eq(Plus(a, b), 'a + b')
        self.eq(Plus(Plus(a, b), c), 'a + b + c')
        self.eq(Plus(a, Plus(b, c)), 'a + (b + c)')

        self.eq(Minus(a, b), 'a - b')
        self.eq(Minus(Minus(a, b), c), 'a - b - c')
        self.eq(Minus(a, Minus(b, c)), 'a - (b - c)')

        self.eq(Minus(a, b), 'a - b')
        self.eq(Plus(Minus(a, b), c), 'a - b + c')
        self.eq(Minus(a, Plus(b, c)), 'a - (b + c)')
        self.eq(Minus(Plus(a, b), c), 'a + b - c')
        self.eq(Minus(a, Plus(b, c)), 'a - (b + c)')

        # Substitution of expm for 1 - exp() and exp() - 1
        self.eq(Minus(Exp(Number(2)), Number(1)), 'expm1(2.0)')
        self.eq(Minus(Number(1), Exp(Number(3))), '-expm1(3.0)')

    def test_multiply_divide(self):
        # Inherited from CBasedExpressionWriter
        a, b, c = self.abc
        self.eq(Multiply(a, b), 'a * b')
        self.eq(Multiply(Multiply(a, b), c), 'a * b * c')
        self.eq(Multiply(a, Multiply(b, c)), 'a * (b * c)')
        self.eq(Divide(a, b), 'a / b')
        self.eq(Divide(Divide(a, b), c), 'a / b / c')
        self.eq(Divide(a, Divide(b, c)), 'a / (b / c)')

    def test_quotient(self):
        # Inherited from CBasedExpressionWriter
        a, b, c = self.abc
        with WarningCollector():
            self.eq(Quotient(a, b), 'floor(a / b)')
            self.eq(Quotient(Plus(a, c), b), 'floor((a + c) / b)')
            self.eq(Quotient(Divide(a, c), b), 'floor(a / c / b)')
            self.eq(Quotient(a, Divide(b, c)), 'floor(a / (b / c))')
            self.eq(Multiply(Quotient(a, b), c), 'floor(a / b) * c')
            self.eq(Multiply(c, Quotient(a, b)), 'c * (floor(a / b))')

    def test_remainder(self):
        # Inherited from CBasedExpressionWriter
        a, b, c = self.abc
        with WarningCollector():
            self.eq(Remainder(a, b), '(a - b * floor(a / b))')
            self.eq(Remainder(Plus(a, c), b),
                    '(a + c - b * floor((a + c) / b))')
            self.eq(Multiply(Remainder(a, b), c), '(a - b * floor(a / b)) * c')
            self.eq(Divide(c, Remainder(b, a)), 'c / ((b - a * floor(b / a)))')

    def test_power(self):
        # Inherited from CBasedExpressionWriter
        a, b, c = self.abc
        self.eq(Power(a, b), 'pow(a, b)')
        self.eq(Power(Power(a, b), c), 'pow(pow(a, b), c)')
        self.eq(Power(a, Power(b, c)), 'pow(a, pow(b, c))')

    def test_log(self):
        # Inherited from CBasedExpressionWriter
        a, b = self.ab
        self.eq(Log(a), 'log(a)')
        self.eq(Log10(a), 'log10(a)')
        self.eq(Log(a, b), '(log(a) / log(b))')

    def test_functions(self):
        a, b = self.ab

        self.eq(Sqrt(a), 'sqrt(a)')
        self.eq(Exp(a), 'exp(a)')
        self.eq(Cos(a), 'cos(a)')
        self.eq(ACos(a), 'acos(a)')
        self.eq(Abs(a), 'fabs(a)')

        with WarningCollector() as c:
            self.eq(Sin(a), 'sin(a)')
        self.assertIn('Unsupported', c.text())

        with WarningCollector() as c:
            self.eq(Tan(a), 'tan(a)')
        self.assertIn('Unsupported', c.text())

        with WarningCollector() as c:
            self.eq(ASin(a), 'asin(a)')
        self.assertIn('Unsupported', c.text())

        with WarningCollector() as c:
            self.eq(ATan(a), 'atan(a)')
        self.assertIn('Unsupported', c.text())

        with WarningCollector() as c:
            self.eq(Floor(a), 'floor(a)')
        self.assertIn('Unsupported', c.text())

        with WarningCollector() as c:
            self.eq(Ceil(a), 'ceil(a)')
        self.assertIn('Unsupported', c.text())

    def test_conditions(self):
        # Inherited from CBasedExpressionWriter
        a, b, c, d = self.abcd
        self.eq(Not(NotEqual(a, b)), '(!(a != b))')
        self.eq(And(Equal(a, b), NotEqual(c, d)), '((a == b) && (c != d))')
        self.eq(Or(More(d, c), MoreEqual(b, a)), '((d > c) || (b >= a))')
        self.eq(Or(Less(d, c), LessEqual(b, a)), '((d < c) || (b <= a))')
        self.eq(
            Not(Or(Equal(Number(1), Number(2)), Equal(Number(3), Number(4)))),
            '(!((1.0 == 2.0) || (3.0 == 4.0)))')
        self.eq(Not(Less(Number(1), Number(2))), '(!(1.0 < 2.0))')

    def test_conditionals(self):
        # Inherited from CBasedExpressionWriter
        a, b, c, d = self.abcd
        self.eq(If(Equal(a, b), d, c), '((a == b) ? d : c)')
        self.eq(Piecewise(NotEqual(d, c), b, a), '((d != c) ? b : a)')
        self.eq(Piecewise(Equal(a, b), c, Equal(a, d), Number(3), Number(4)),
                '((a == b) ? c : ((a == d) ? 3.0 : 4.0))')


if __name__ == '__main__':
    unittest.main()
