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

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

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

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

[ideal]
Vc = engine.pace * 1

[mb]
dot(V) = (hh.I1 + mm.I2 + tt.I3) / 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 * (mb.V - 0.05 [V])
    in [pA]

[mm]
dot_C = beta * O - alpha * C
dot(C) = dot_C
alpha = 0.3 [1/s]
    in [1/s]
beta = 0.4 [1/s]
    in [1/s]
O = 1 - C
I2 = 2 [pS] * O * (mb.V + 0.02 [V])
    in [pA]

[tt]
I3 = 4 [pS] * mm.O * (mb.V + 0.02 [V])
tTest = 0
t_Test = 1
t_test = 2
"""

units_output = """
/*
This file was generated by Myokit.
*/

/* Input parameters */
in_i { }

/* Engine: pace */
/* E.g.
  -80 * (1 - sigmoid((t-100)*5000))
  -120 * (sigmoid((t-100)*5000) - sigmoid((t-200)*5000))
*/
enginePace { 0.0 } /* engine.pace */

/* Constants: hh */
hhXInf { 0.8 } /* hh.x.inf */
hhXTau { 3.0 } /* hh.x.tau [s] */
hhYAlpha { 0.1 } /* hh.y.alpha [S/F] */
hhYBeta { 0.2 } /* hh.y.beta [S/F] */

/* Constants: mm */
mmAlpha { 0.3 } /* mm.alpha [S/F] */
mmBeta { 0.4 } /* mm.beta [S/F] */

/* Constants: tt */
ttTTest1 { 0.0 } /* tt.tTest */
ttTTest2 { 1.0 } /* tt.t_Test */
ttTTest3 { 2.0 } /* tt.t_test */

/* Constants: mb */
mbC { 20.0 } /* mb.C [pF] */

/* Initial conditions */
u_i {
  mbV = -0.08 * 1000 [1 (0.001)], /* mb.V [mV] */
  hhX = 0.1, /* hh.x */
  hhY = 0.9, /* hh.y */
  mmC = 0.9, /* mm.C */
}

/* Variables: hh */
hhI1 { 3.0 * hhX * hhY * (mbV / 1000.0 - 0.05) * 0.05 } /* hh.I1 [A/F] */

/* Variables: mm */
mmO { 1.0 - mmC } /* mm.O */
mmI2 { 2.0 * mmO * (mbV / 1000.0 + 0.02) * 0.05 } /* mm.I2 [A/F] */
mmDotC { mmBeta * mmO - mmAlpha * mmC } /* mm.dot_C */

/* Variables: ideal */
idealVc { enginePace * 1.0 } /* ideal.Vc */

/* Variables: tt */
ttI3 { 4.0 * mmO * (mbV / 1000.0 + 0.02) } /* tt.I3 */

/* Solve */
F_i {
  (hhI1 / 0.05 + mmI2 / 0.05 + ttI3) / mbC * 1000.0 / 1000.0,
  (hhXInf - hhX) / hhXTau / 1000.0,
  (hhYAlpha * (1.0 - hhY) - hhYBeta * hhY) / 1000.0,
  mmDotC / 1000.0,
}

/* Output */
out_i {
  hhX,
  hhY,
  mbV,
  mmC,
}
"""


class DiffSLExporterTest(unittest.TestCase):
    """Tests DiffSL export."""

    def test_diffsl_exporter(self):
        # Tests exporting a model

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

            e = myokit.formats.diffsl.DiffSLExporter()

            # Test with simple model
            e.model(path, model)

            # Test with protocol set
            with self.assertRaisesRegex(ValueError, 'input protocol'):
                e.model(path, model, protocol='-80 + 120*heaviside(t-10)')

            # 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')
            with self.assertRaisesRegex(myokit.ExportError, 'valid model'):
                e.model(path, model)

    def test_explicit_time_dependence(self):
        # Tests that explicit time dependence (dot(y) = -y * t) is supported
        m_td = myokit.parse_model("""
[[model]]
c.y = 1.0

[engine]
time = 0
    bind time

[c]
dot(y) = -y * engine.time
""")
        e = myokit.formats.diffsl.DiffSLExporter()
        with TemporaryDirectory() as d:
            path = d.path('diffsl.model')
            e.model(path, m_td)
            with open(path) as f:
                content = f.read()
        self.assertRegex(content, r'F_i\s*\{\s*-cY\s*\*\s*t,\s*\}')

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

        # Export model
        m = myokit.parse_model(units_model)
        e = myokit.formats.diffsl.DiffSLExporter()
        with TemporaryDirectory() as d:
            path = d.path('diffsl.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('mb.V').set_rhs('hh.I1 + mm.I2')
        m.get('mb').remove_variable(m.get('mb.C'))
        with TemporaryDirectory() as d:
            path = d.path('diffsl.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('diffsl.model')
            with WarningCollector() as c:
                e.model(path, m)
            self.assertIn('Unable to convert engine.time', c.text())

    def test_diffsl_exporter_fetching(self):
        # Tests getting an DiffSL exporter via the 'exporter' interface

        e = myokit.formats.exporter('diffsl')
        self.assertIsInstance(e, myokit.formats.diffsl.DiffSLExporter)

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

    def test_non_cardiac_model(self):
        # Tests exporting a non-cardiac model without membrane potential
        # Regression test for issue #1185

        # Create a simple pharmacokinetic model
        pk_model = myokit.parse_model("""
            [[model]]
            pk.A1 = 100

            [environment]
            t = 0 bind time
                in [s]

            [pk]
            dot(A1) = -CL * C1
                in [mg]
            V1 = 10 [L]
                in [L]
            C1 = A1 / V1
                in [mg/L]
            CL = 0.5 [L/s]
                in [L/s]
            """)

        # Export should not crash
        e = myokit.formats.diffsl.DiffSLExporter()
        with TemporaryDirectory() as d:
            path = d.path('pk_model.diffsl')
            e.model(path, pk_model)

            # Verify the file was created
            with open(path, 'r') as f:
                content = f.read()

            # Basic sanity checks
            self.assertIn('pkA1', content)
            self.assertIn('F_i', content)
            self.assertNotIn('G_i', content)
            self.assertNotIn('dudt_i', content)

    def test_inputs_and_outputs(self):
        # Tests custom inputs and outputs parameters

        # Load example model
        model = myokit.load_model('example')

        e = myokit.formats.diffsl.DiffSLExporter()
        with TemporaryDirectory() as d:
            # Test with custom inputs
            path = d.path('test_inputs.diffsl')
            inputs = [model.get('membrane.C')]
            e.model(path, model, inputs=inputs)
            with open(path, 'r') as f:
                content = f.read()
            self.assertIn('in_i {\n  membraneC = ', content)
            # Verify input is not in constants section
            self.assertNotIn('membraneC {', content)

            # Test with custom outputs
            path = d.path('test_outputs.diffsl')
            outputs = [model.get('membrane.V'), model.get('ina.m')]
            e.model(path, model, outputs=outputs)
            with open(path, 'r') as f:
                content = f.read()
            # Check that outputs are in specified order (not alphabetical)
            lines = content.split('\n')
            out_section = []
            in_out = False
            for line in lines:
                if 'out_i {' in line:
                    in_out = True
                    continue
                if in_out:
                    if '}' in line:
                        break
                    out_section.append(line.strip())
            self.assertEqual(len(out_section), 2)
            self.assertIn('membraneV', out_section[0])
            self.assertIn('inaM', out_section[1])

            # Test with both custom inputs and outputs
            path = d.path('test_both.diffsl')
            e.model(path, model, inputs=inputs, outputs=outputs)
            with open(path, 'r') as f:
                content = f.read()
            self.assertIn('in_i {\n  membraneC = ', content)
            # Verify outputs are present in out_i block in correct order
            match = re.search(r'out_i\s*\{([^}]*)\}', content, re.DOTALL)
            self.assertIsNotNone(match)
            out_content = match.group(1)
            self.assertIn('membraneV', out_content)
            self.assertIn('inaM', out_content)
            # Verify order: membraneV should appear before inaM
            self.assertLess(
                out_content.index('membraneV'), out_content.index('inaM')
            )

            # Test with engine.pace as an input
            path = d.path('test_pace_input.diffsl')
            e.model(path, model, inputs=[model.get('engine.pace')])
            with open(path, 'r') as f:
                content = f.read()
            # pace should appear in in_i block
            match = re.search(r'in_i\s*\{([^}]*)\}', content, re.DOTALL)
            self.assertIsNotNone(match)
            self.assertIn('enginePace', match.group(1))
            # pace should not also appear in a separate variable block
            self.assertNotIn('/* Engine: pace */', content)

    def test_inputs_outputs_validation(self):
        # Tests validation of inputs and outputs parameters

        # Load example model
        model = myokit.load_model('example')

        e = myokit.formats.diffsl.DiffSLExporter()
        with TemporaryDirectory() as d:
            path = d.path('test.diffsl')

            # Test 1: State variable cannot be input
            with self.assertRaisesRegex(
                myokit.ExportError,
                'Input variable membrane.V must be a constant'
            ):
                e.model(path, model, inputs=[model.get('membrane.V')])

            # Test 2: Intermediate variable (non-const, non-state) cannot
            # be input
            with self.assertRaisesRegex(
                myokit.ExportError,
                'Input variable ina.m.alpha must be a constant'
            ):
                e.model(path, model, inputs=[model.get('ina.m.alpha')])

            # Test 3: Constant cannot be output
            with self.assertRaisesRegex(
                myokit.ExportError,
                'Output variable membrane.C must be time-varying'
            ):
                e.model(path, model, outputs=[model.get('membrane.C')])

            # Test 4: Multiple invalid inputs
            with self.assertRaisesRegex(
                myokit.ExportError, 'Input variable .* must be a constant'
            ):
                e.model(
                    path,
                    model,
                    inputs=[
                        model.get('membrane.C'),  # Valid
                        model.get('membrane.V'),  # Invalid (state)
                    ]
                )

            # Test 5: Multiple invalid outputs
            with self.assertRaisesRegex(
                myokit.ExportError,
                'Output variable .* must be time-varying'
            ):
                e.model(
                    path,
                    model,
                    outputs=[
                        model.get('membrane.V'),  # Valid (state)
                        model.get('ina.gNa'),     # Invalid (constant)
                    ]
                )

            # Test 6: Input from different model
            other_model = myokit.load_model('example')
            with self.assertRaisesRegex(
                myokit.ExportError,
                'Input variable .* does not belong to the model being '
                'exported'
            ):
                e.model(path, model, inputs=[other_model.get('membrane.C')])

            # Test 7: Output from different model
            with self.assertRaisesRegex(
                myokit.ExportError,
                'Output variable .* does not belong to the model being '
                'exported'
            ):
                e.model(path, model, outputs=[other_model.get('membrane.V')])

    def test_get_state_index(self):
        # Tests the get_state_index introspection method

        # Load example model
        model = myokit.load_model('example')

        e = myokit.formats.diffsl.DiffSLExporter()
        with TemporaryDirectory() as d:
            path = d.path('test.diffsl')

            # Test 1: Calling before model() raises error
            with self.assertRaisesRegex(
                RuntimeError, 'can only be called after model'
            ):
                e.get_state_index(model.get('membrane.V'))

            # Export the model
            e.model(path, model)

            # Test 2: Get index for state variables
            states = list(model.states())
            for i, state in enumerate(states):
                idx = e.get_state_index(state)
                self.assertEqual(idx, i)

            # Test 3: Non-state variable returns None
            const_var = model.get('membrane.C')
            self.assertIsNone(e.get_state_index(const_var))

            # Test 4: Another non-state variable returns None
            inter_var = model.get('ina.m.alpha')
            self.assertIsNone(e.get_state_index(inter_var))


class DiffSLExpressionWriterTest(myokit.tests.ExpressionWriterTestCase):
    """Test conversion to DiffSL syntax."""

    _name = 'diffsl'
    _target = myokit.formats.diffsl.DiffSLExpressionWriter

    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)')

        # No expm in DiffSL
        self.eq(Minus(Exp(Number(2)), Number(1)), 'exp(2.0) - 1.0')
        self.eq(Minus(Number(1), Exp(Number(3))), '1.0 - exp(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), '(log(a) / log(10.0))')
        self.eq(Log(a, b), '(log(a) / log(b))')

    def test_supported_functions(self):
        a = self.a

        self.eq(Abs(a), 'abs(a)')
        self.eq(Cos(a), 'cos(a)')
        self.eq(Exp(a), 'exp(a)')
        self.eq(Log(a), 'log(a)')
        self.eq(Sin(a), 'sin(a)')
        self.eq(Sqrt(a), 'sqrt(a)')
        self.eq(Tan(a), 'tan(a)')

    def test_unsupported_functions(self):
        a = self.a

        with WarningCollector() as wc:
            self.eq(ACos(a), 'acos(a)')
        self.assertIn('Unsupported', wc.text())

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

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

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

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

    def test_conditional_operators(self):
        a, b, c, d = self.abcd

        self.eq(Equal(a, b), 'heaviside(a - b) * heaviside(b - a)')

        self.eq(Less(a, b), '(1 - heaviside(a - b))')

        self.eq(LessEqual(a, b), 'heaviside(b - a)')

        self.eq(More(a, b), '(1 - heaviside(b - a))')

        self.eq(MoreEqual(a, b), 'heaviside(a - b)')

        self.eq(NotEqual(a, b), '(1 - heaviside(a - b) * heaviside(b - a))')

        self.eq(
            Not(NotEqual(a, b)),
            '(1 - (1 - heaviside(a - b) * heaviside(b - a)))',
        )

        self.eq(
            Not(Not(Equal(a, b))),
            '(1 - (1 - heaviside(a - b) * heaviside(b - a)))',
        )

        self.eq(
            And(Equal(a, b), NotEqual(c, d)),
            'heaviside(a - b) * heaviside(b - a)'
            ' * (1 - heaviside(c - d) * heaviside(d - c))',
        )

        self.eq(
            Or(More(d, c), MoreEqual(b, a)),
            '(1 - (1 - (1 - heaviside(c - d))) * (1 - heaviside(b - a)))',
        )

        self.eq(
            Or(Less(d, c), LessEqual(b, a)),
            '(1 - (1 - (1 - heaviside(d - c))) * (1 - heaviside(a - b)))',
        )

        self.eq(
            Not(Or(Equal(Number(1), Number(2)), Equal(Number(3), Number(4)))),
            '(1 - (1 - (1 - heaviside(1.0 - 2.0) * heaviside(2.0 - 1.0)) '
            '* (1 - heaviside(3.0 - 4.0) * heaviside(4.0 - 3.0))))',
        )

        self.eq(
            Not(Less(Number(1), Number(2))), '(1 - (1 - heaviside(1.0 - 2.0)))'
        )

    def test_if_expressions(self):
        a, b, c, d = self.abcd

        self.eq(
            If(Equal(a, b), c, d),
            '(heaviside(a - b) * heaviside(b - a) * c'
            ' + (1 - heaviside(a - b) * heaviside(b - a)) * d)',
        )

        self.eq(
            If(Equal(a, b), c, Number(0)),
            '(heaviside(a - b) * heaviside(b - a) * c'
            ' + (1 - heaviside(a - b) * heaviside(b - a)) * 0.0)',
        )

        self.eq(
            If(Equal(a, b), Number(0), d),
            '(heaviside(a - b) * heaviside(b - a) * 0.0'
            ' + (1 - heaviside(a - b) * heaviside(b - a)) * d)',
        )

        self.eq(
            If(Equal(a, b), c, Number(1)),
            '(heaviside(a - b) * heaviside(b - a) * c'
            ' + (1 - heaviside(a - b) * heaviside(b - a)) * 1.0)',
        )

        self.eq(
            If(Equal(a, b), Number(1), d),
            '(heaviside(a - b) * heaviside(b - a) * 1.0'
            ' + (1 - heaviside(a - b) * heaviside(b - a)) * d)',
        )

        self.eq(
            If(NotEqual(a, b), c, d),
            '((1 - heaviside(a - b) * heaviside(b - a)) * c'
            ' + (1 - (1 - heaviside(a - b) * heaviside(b - a))) * d)',
        )

        self.eq(
            If(More(a, b), c, d),
            '((1 - heaviside(b - a)) * c + (1 - (1 - heaviside(b - a))) * d)',
        )

        self.eq(
            If(MoreEqual(a, b), c, d),
            '(heaviside(a - b) * c + (1 - heaviside(a - b)) * d)',
        )

        self.eq(
            If(Less(a, b), c, d),
            '((1 - heaviside(a - b)) * c + (1 - (1 - heaviside(a - b))) * d)',
        )

        self.eq(
            If(LessEqual(a, b), c, d),
            '(heaviside(b - a) * c + (1 - heaviside(b - a)) * d)',
        )

    def test_piecewise_expressions(self):
        a, b, c, d = self.abcd

        self.eq(Piecewise(Equal(a, b), c, d), self.w.ex(If(Equal(a, b), c, d)))

        self.eq(
            Piecewise(NotEqual(a, b), c, d),
            self.w.ex(
                If(NotEqual(a, b), c, d),
            ),
        )

        self.eq(
            Piecewise(More(a, b), c, d),
            self.w.ex(
                If(More(a, b), c, d),
            ),
        )

        self.eq(
            Piecewise(MoreEqual(a, b), c, d),
            self.w.ex(If(MoreEqual(a, b), c, d)),
        )

        self.eq(Piecewise(Less(a, b), c, d), self.w.ex(If(Less(a, b), c, d)))

        self.eq(
            Piecewise(LessEqual(a, b), c, d),
            self.w.ex(
                If(LessEqual(a, b), c, d),
            ),
        )

        self.eq(
            Piecewise(Equal(a, b), c, Equal(a, d), Number(3), Number(4)),
            self.w.ex(
                If(Equal(a, b), c, If(Equal(a, d), Number(3), Number(4)))
            ),
        )

        self.eq(
            Piecewise(Less(a, b), Number(0), Less(c, d), Number(0), Number(5)),
            '((1 - heaviside(a - b)) * 0.0 '
            '+ (1 - (1 - heaviside(a - b))) * ((1 - heaviside(c - d)) * 0.0 '
            '+ (1 - (1 - heaviside(c - d))) * 5.0))',
        )

    def test_heaviside_numerical(self):
        """Test generated heaviside expressions with numerical values"""

        def heaviside(x):
            return 1 if x >= 0 else 0

        values = itertools.product(
            [-10e9, -1, -1e-9, 0, 1e-9, 1, 10e9], repeat=4
        )

        for a, b, c, d in values:
            # a == b
            result = int(a == b)
            expr = self.w.ex(Equal(Number(a), Number(b)))
            self.assertEqual(eval(expr), result)

            # a < b
            result = int(a < b)
            expr = self.w.ex(Less(Number(a), Number(b)))
            self.assertEqual(eval(expr), result)

            # a <= b
            result = int(a <= b)
            expr = self.w.ex(LessEqual(Number(a), Number(b)))
            self.assertEqual(eval(expr), result)

            # a > b
            result = int(a > b)
            expr = self.w.ex(More(Number(a), Number(b)))
            self.assertEqual(eval(expr), result)

            # a >= b
            result = int(a >= b)
            expr = self.w.ex(MoreEqual(Number(a), Number(b)))
            self.assertEqual(eval(expr), result)

            # a != b
            result = int(a != b)
            expr = self.w.ex(NotEqual(Number(a), Number(b)))
            self.assertEqual(eval(expr), result)

            # not(a != b)
            result = int(not (a != b))
            expr = self.w.ex(Not(NotEqual(Number(a), Number(b))))
            self.assertEqual(eval(expr), result)

            # not(not(a == b))
            result = int(not (not (a == b)))
            expr = self.w.ex(Not(Not(Equal(Number(a), Number(b)))))
            self.assertEqual(eval(expr), result)

            # (a == b) and (c != d)
            result = int((a == b) and (c != d))
            expr = self.w.ex(
                And(
                    Equal(Number(a), Number(b)), NotEqual(Number(c), Number(d))
                )
            )
            self.assertEqual(eval(expr), result)

            # (d > c) or (b >= a)
            result = int((d > c) or (b >= a))
            expr = self.w.ex(
                Or(More(Number(d), Number(c)), MoreEqual(Number(b), Number(a)))
            )
            self.assertEqual(eval(expr), result)

            # (d < c) or (b <= a)
            result = int((d < c) or (b <= a))
            expr = self.w.ex(
                Or(Less(Number(d), Number(c)), LessEqual(Number(b), Number(a)))
            )
            self.assertEqual(eval(expr), result)

            # (a == b) or (c == d)
            result = int((a == b) or (c == d))
            expr = self.w.ex(
                Or(Equal(Number(a), Number(b)), Equal(Number(c), Number(d)))
            )
            self.assertEqual(eval(expr), result)

            # not(a < b)
            result = int(not (a < b))
            expr = self.w.ex(Not(Less(Number(a), Number(b))))
            self.assertEqual(eval(expr), result)

            # if(a > b, c, d)
            result = c if (a > b) else d
            expr = self.w.ex(
                If(More(Number(a), Number(b)), Number(c), Number(d))
            )
            self.assertEqual(eval(expr), result)

            # if(a >= b, c, d)
            result = c if (a >= b) else d
            expr = self.w.ex(
                If(MoreEqual(Number(a), Number(b)), Number(c), Number(d))
            )
            self.assertEqual(eval(expr), result)

            # if(a < b, c, d)
            result = c if (a < b) else d
            expr = self.w.ex(
                If(Less(Number(a), Number(b)), Number(c), Number(d))
            )
            self.assertEqual(eval(expr), result)

            # if(a <= b, c, d)
            result = c if (a <= b) else d
            expr = self.w.ex(
                If(LessEqual(Number(a), Number(b)), Number(c), Number(d))
            )
            self.assertEqual(eval(expr), result)

            # piecewise(a > b, c, d)
            result = c if (a > b) else d
            expr = self.w.ex(
                Piecewise(More(Number(a), Number(b)), Number(c), Number(d))
            )
            self.assertEqual(eval(expr), result)

            # piecewise(a >= b, c, d)
            result = c if (a >= b) else d
            expr = self.w.ex(
                Piecewise(
                    MoreEqual(Number(a), Number(b)), Number(c), Number(d)
                )
            )
            self.assertEqual(eval(expr), result)

            # piecewise(a < b, c, d)
            result = c if (a < b) else d
            expr = self.w.ex(
                Piecewise(Less(Number(a), Number(b)), Number(c), Number(d))
            )
            self.assertEqual(eval(expr), result)

            # piecewise(a <= b, c, d)
            result = c if (a <= b) else d
            expr = self.w.ex(
                Piecewise(
                    LessEqual(Number(a), Number(b)), Number(c), Number(d)
                )
            )
            self.assertEqual(eval(expr), result)

            # piecewise(a == b, c, a == d, 3, 4)
            result = c if (a == b) else (3 if (a == d) else 4)
            expr = self.w.ex(
                Piecewise(
                    Equal(Number(a), Number(b)),
                    Number(c),
                    Equal(Number(a), Number(d)),
                    Number(3),
                    Number(4),
                )
            )
            self.assertEqual(eval(expr), result)

            # piecewise(a < b, 0, c < d, 0, 5)
            result = 0 if (a < b) else (0 if (c < d) else 5)
            expr = self.w.ex(
                Piecewise(
                    Less(Number(a), Number(b)),
                    Number(0),
                    Less(Number(c), Number(d)),
                    Number(0),
                    Number(5),
                ),
            )
            self.assertEqual(eval(expr), result)


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