"""Various integration tests with multimode systems
"""
from __future__ import annotations
import pytest
import numpy
# import pandas

from cosapp.base import Port, System
from cosapp.drivers import EulerExplicit, NonLinearSolver
from cosapp.recorders import DataFrameRecorder
from typing import List


class ElecPort(Port):
    def setup(self):
        self.add_variable('V', 0.0, unit='V')
        self.add_variable('I', 1.0, unit='A')


class Resistor(System):
    """Resistor component

    Attributes:
        R : float
            Resistance in Ohms
    """
    def setup(self, R=0):
        self.add_input(ElecPort, 'elec_in')
        self.add_output(ElecPort, 'elec_out')

        self.add_inward('R', abs(float(R)), unit='ohm', desc='Internal resistance')
        self.add_outward('deltaV')

    def compute(self):
        self.elec_out.I = self.elec_in.I
        self.elec_out.V = self.elec_in.V - (self.elec_out.I * self.R)
        self.deltaV = self.elec_in.V - self.elec_out.V


class Node(System):
    """System representing an electric circuit node with
    an arbitrary number of incoming and outgoing branches.

    Nodes provide an off-design problem ensuring potential equality
    and global current balance (see 'Unknowns' and 'Equations' below).

    Constructor arguments:
    ----------------------
    - n_in [int], optional: Number of incoming branches. Defaults to 1.
    - n_out [int], optional: Number of outgoing branches. Defaults to 1.

    Properties:
    -----------
    - n_in [int]: Number of incoming branches.
    - n_out [int]: Number of outgoing branches.
    - incoming: Tuple containing all `ElecPort` inputs.
    - outgoing: Tuple containing all `ElecPort` outputs.

    Unknowns:
    ---------
    - n_out current fractions (one per outgoing branch), if n_out > 1.

    Equations:
    ----------
    - (n_in - 1) potential equality conditions for incoming branches.
    - 1 total current balance equation, if n_out > 1.
    """
    def setup(self, n_in=1, n_out=1):
        """Node constructor.

        Arguments:
        -----------
        - n_in [int], optional: Number of incoming branches. Defaults to 1.
        - n_out [int], optional: Number of outgoing branches. Defaults to 1.
        """
        self.add_property('n_in', int(n_in))
        self.add_property('n_out', int(n_out))

        if min(self.n_in, self.n_out) < 1:
            raise ValueError("Node needs at least one incoming and one outgoing branch")

        self.add_property('incoming',
            tuple(
                self.add_input(ElecPort, f"elec_in{i}")
                for i in range(self.n_in)
            )
        )
        self.add_property('outgoing',
            tuple(
                self.add_output(ElecPort, f"elec_out{i}")
                for i in range(self.n_out)
            )
        )

        if self.n_out > 1:  # unnecessary otherwise
            self.add_inward('I_frac',
                value = numpy.full(self.n_out, 1.0 / self.n_out),
                desc = f"Current fractions distributed to outgoing branches",
                limits = (0, 1),
            )
            self.add_unknown('I_frac', lower_bound=0, upper_bound=1)
            self.add_equation('sum(I_frac) == 1', name='Current balance')

        for i in range(1, self.n_in):   # case where node is 'joiner'
            self.add_equation(f'elec_in{i}.V == elec_in0.V')

        self.add_outward('V', 0., unit='V', desc='Actual node voltage')
        self.add_outward('sum_I_in', 0., unit='A', desc='Sum of all incoming currents')
        self.add_outward('sum_I_out', 0., unit='A', desc='Sum of all outgoing currents')

    def compute(self):
        # Sum of incoming currents
        self.sum_I_in = I = sum(port.I for port in self.incoming)

        # Output voltage
        self.V = V = numpy.mean([port.V for port in self.incoming])

        # Current distribution
        try:
            I_frac = self.I_frac
        except AttributeError:
            I_frac = [1]
        for j, port in enumerate(self.outgoing):
            port.V = V
            port.I = I * I_frac[j]

        self.sum_I_out = I * sum(I_frac)

    @classmethod
    def make(cls, parent, name, incoming: List[ElecPort], outgoing: List[ElecPort], pulling=None) -> Node:
        """Factory method making appropriate connections with parent system"""
        node = cls(name, n_in=max(len(incoming), 1), n_out=max(len(outgoing), 1))
        parent.add_child(node, pulling=pulling)
        
        for branch_elec, node_elec in zip(incoming, node.incoming):
            parent.connect(branch_elec, node_elec)
        
        for branch_elec, node_elec in zip(outgoing, node.outgoing):
            parent.connect(branch_elec, node_elec)

        return node


class TwoWayCircuit(System):
    def setup(self):
        self.add_outward_modevar('upbranch', True)
        self.reconfig()

        self.add_outward('Requiv', 0.0)
        self.add_outward('deltaV', 0.0)
        self.add_event('switch', trigger='elec_in.V == elec_out.V')
    
    def compute(self):
        self.deltaV = self.elec_in.V - self.elec_out.V
        I = self.elec_in.I
        self.Requiv = self.deltaV / I if abs(I) > 0 else numpy.nan

    def transition(self):
        if self.switch.present:
            self.upbranch = not self.upbranch
            self.reconfig()

    def reconfig(self):
        for name in self.children:
            self.pop_child(name)
        pulled_ports = ['elec_in', 'elec_out']
        if self.upbranch:
            child = self.add_child(Resistor("R1", R=100), pulling=pulled_ports)
        else:
            child = self.add_child(Resistor("R2", R=500), pulling=pulled_ports)
        self.elec_out.V = child.elec_out.V = 0.0


class TwoWayCircuitWithEq(System):
    def setup(self):
        self.add_outward_modevar('upbranch', True)
        self.add_child(Resistor("R0", R=100), pulling=['elec_in', 'elec_out'])
        self.reconfig()

        self.add_outward('Requiv', 0.0)
        self.add_outward('deltaV', 0.0)
        self.add_event('switch', trigger='elec_in.V == elec_out.V')
    
    def compute(self):
        self.deltaV = self.elec_in.V - self.elec_out.V
        I = self.elec_in.I
        self.Requiv = self.deltaV / I if I != 0 else numpy.nan

    def transition(self):
        if self.switch.present:
            self.upbranch = not self.upbranch
            self.reconfig()

    def reconfig(self):
        problem = self.problem
        if not self.upbranch and problem.is_empty():
            self.add_unknown("R0.R").add_equation("R0.R == 500")
        else:
            problem.clear()
            self.R0.R = 100


@pytest.fixture
def case_TwoWayCircuit():
    circuit = TwoWayCircuit('circuit')
    driver = circuit.add_driver(
        EulerExplicit(time_interval=[0, 1], dt=0.1)
    )
    driver.add_recorder(
        DataFrameRecorder(includes=['elec_in.*', 'elec_out.V', 'Requiv', 'deltaV', 'upbranch']),
        period = 0.1,
    )
    solver = driver.add_child(NonLinearSolver('solver', tol=1e-9))
    solver.add_unknown('elec_in.I').add_equation("elec_out.V == 0")
    return circuit, driver


def test_TwoWayCircuit(case_TwoWayCircuit):
    circuit, driver = case_TwoWayCircuit
    omega = 6
    driver.set_scenario(
        values = {
            'elec_in.V': f"cos({omega} * t)",
        }
    )
    circuit.run_drivers()

    df = driver.recorder.export_data()
    # print("", df.drop(['Section', 'Status', 'Error code'], axis=1), sep="\n")

    for i, row in df.iterrows():
        I = row['elec_in.I']
        R = row['Requiv']
        context = f"row #{i}, I = {I}, {list(circuit.exec_order)}"
        if I > 1e-12:
            assert R == pytest.approx(100), context
        elif I < -1e-12:
            assert R == pytest.approx(500), context
    
    assert [record.time for record in driver.recorded_events] == pytest.approx(
        [(2 * k + 1) * 0.5 * numpy.pi / omega for k in range(2)]
    )


@pytest.fixture
def case_TwoWayCircuitWithEq():
    circuit = TwoWayCircuitWithEq('circuit')
    driver = circuit.add_driver(
        EulerExplicit(time_interval=[0, 1], dt=0.1)
    )
    driver.add_recorder(
        DataFrameRecorder(includes=['elec_in.*', 'elec_out.V', 'Requiv', 'deltaV', 'upbranch']),
        period = 0.1,
    )
    solver = driver.add_child(NonLinearSolver('solver', tol=1e-9))
    solver.add_unknown('elec_in.I').add_equation("elec_out.V == 0")
    return circuit, driver


def test_TwoWayCircuitWithEq(case_TwoWayCircuitWithEq):
    circuit, driver = case_TwoWayCircuitWithEq
    omega = 6
    driver.set_scenario(
        values = {
            'elec_in.V': f"cos({omega} * t)",
        }
    )
    circuit.run_drivers()

    df = driver.recorder.export_data()
    # print("", df.drop(['Section', 'Status', 'Error code'], axis=1), sep="\n")

    for i, row in df.iterrows():
        I = row['elec_in.I']
        R = row['Requiv']
        context = f"row #{i}, I = {I}"
        if I > 1e-12:
            assert R == pytest.approx(100), context
        elif I < -1e-12:
            assert R == pytest.approx(500), context
    
    assert [record.time for record in driver.recorded_events] == pytest.approx(
        [(2 * k + 1) * 0.5 * numpy.pi / omega for k in range(2)]
    )


def test_MultimodeSystem_mode_manager():
    """Test a pattern in which a mode manager transfers
    a mode variable to a sibling multimode system.

    The multimode system is expected to be in synch with
    the mode manager at each transition.
    """
    class ModeManager(System):
        def setup(self):
            self.add_inward('x', 0.0)
            self.add_event('pif', trigger="x > 1")
            self.add_event('paf', trigger="x < 1")
            self.add_event('zap', trigger="x > 2")
            self.add_outward_modevar('mode', init=0, dtype=int)

        def transition(self) -> None:
            if self.pif.present:
                self.mode = 1
            if self.paf.present:
                self.mode = 0
            if self.zap.present:
                self.mode = 2
    
    class MultimodeSystem(System):
        def setup(self):
            self.add_inward_modevar("mode", value=2)
            self.add_outward_modevar("state", init="get_state(mode)")

        @staticmethod
        def get_state(mode) -> float:
            if mode == 0:
                state = "A"
            elif mode == 1:
                state = "B"
            elif mode == 2:
                state = "C"
            else:
                raise ValueError
            return state

        def transition(self) -> None:
            self.state = self.get_state(self.mode)

    class TopSystem(System):
        def setup(self) -> None:
            self.add_child(ModeManager('manager'))
            self.add_child(MultimodeSystem('foo'))

            self.connect(self.manager, self.foo, 'mode')
    
    s = TopSystem('s')
    driver = s.add_driver(EulerExplicit(dt=0.1, time_interval=(0, 1)))
    driver.add_recorder(DataFrameRecorder(), period=driver.dt)
    driver.set_scenario(
        values={
            'manager.x': '3 * t',  # mode changes @ t=1/3 and 2/3
        }
    )

    s.run_drivers()

    data = driver.recorder.export_data()
    assert all(data['foo.mode'] == data['manager.mode'])
    assert all(list(map(s.foo.get_state, data['foo.mode'])) == data['foo.state'])


def test_MultimodeSystem_close_events():
    """Test event occurence within the time step directly following another event"""
    class TwoEventSystem(System):
        def setup(self) -> None:
            self.add_event('foo')
            self.add_event('bar')
    
    s = TwoEventSystem('s')

    driver = s.add_driver(EulerExplicit(dt=0.1, time_interval=(0, 0.3)))

    # Set triggers in two different time steps
    s.foo.trigger = "t == 0.12"
    s.bar.trigger = "t == 0.23"
    s.run_drivers()

    assert len(driver.recorded_events) == 2
    assert len(driver.recorded_events[0].events) == 1
    assert driver.recorded_events[0].time == 0.12
    assert driver.recorded_events[0].events[0] is s.foo 
    assert driver.recorded_events[1].time == 0.23
    assert len(driver.recorded_events[1].events) == 1
    assert driver.recorded_events[1].events[0] is s.bar 

    # Set both triggers within the same time step
    s.foo.trigger = "t == 0.12"
    s.bar.trigger = "t == 0.13"
    s.run_drivers()

    assert len(driver.recorded_events) == 2
    assert len(driver.recorded_events[0].events) == 1
    assert driver.recorded_events[0].time == 0.12
    assert driver.recorded_events[0].events[0] is s.foo
    assert driver.recorded_events[1].time == 0.13
    assert len(driver.recorded_events[1].events) == 1
    assert driver.recorded_events[1].events[0] is s.bar


def test_MultimodeSystem_event_init():
    """Test checking that the system is up-to-date when primary events are initialized.
    """
    class Kinematics(System):
        def setup(self):
            self.add_inward('v', numpy.zeros(2))
            self.add_transient('x', der='v')

    class Wall(System):
        def setup(self):
            self.add_inward('x', numpy.zeros(2), desc='Point on the wall')
            self.add_inward('n_dir', numpy.r_[1., 1.], desc='Normal direction')
            self.add_outward('n', numpy.r_[0., 0.], desc='Unit normal vector')

        def compute(self) -> None:
            self.n = self.n_dir / numpy.linalg.norm(self.n_dir)
        
        def distance(self, point) -> float:
            """Signed distance to wall"""
            return self.n.dot(numpy.asarray(point) - self.x)

    class Assembly(System):
        def setup(self):
            self.add_child(Kinematics('kinematics'), pulling=['x', 'v'])
            self.add_child(Wall('wall'))
            self.add_event('rebound', trigger='wall.distance(x) == 0')

        def transition(self) -> None:
            if self.rebound.present:
                n = self.wall.n
                v = self.v
                v -= (2 * v.dot(n)) * n

    s = Assembly('s')

    driver = s.add_driver(EulerExplicit('driver', time_interval=(0, 1), dt=1.0))
    driver.set_scenario(
        init={
            'x': [0., 0.],
        },
        values={
            'v': [1., 0.],
            'wall.x': [0.25, 0.],
            'wall.n_dir': [-1., 0.],
        },
        stop=s.rebound,
    )
    # Initialize the wall normal with an **incorrect** value.
    # If the system is not up-to-date, the actual distance to the wall
    # will be not be correctly initialized, and the rebound will be missed.
    s.wall.n[:] = [1., 0.]

    s.run_drivers()

    assert len(driver.recorded_events) == 1
    record = driver.recorded_events[-1]
    assert len(record.events) == 2
    assert record.events[0] is s.rebound
    assert record.events[1] is driver.scenario.stop
    assert record.time == pytest.approx(0.25, rel=1e-15)


def test_MultimodeSystem_transition_order():
    """Check mode initialisation and transition
    across a multimode system tree.
    """
    class ModeManager(System):
        def setup(self):
            self.add_inward('x', 0.0)
            self.add_event('pif', trigger="x > 1")
            self.add_event('paf', trigger="x < 1")
            self.add_outward_modevar('mode', init="0 if x < 1 else 1", dtype=int)

        def transition(self) -> None:
            if self.pif.present:
                self.mode = 1
            if self.paf.present:
                self.mode = 0
    
    class MultimodeSystem(System):
        def setup(self, shift=10):
            self.add_property('shift', shift)
            self.add_inward_modevar("m_in", value=2, dtype=int)
            self.add_outward_modevar("m_out", init=f"m_in + {shift}", dtype=int)

        def transition(self) -> None:
            self.m_out = self.m_in + self.shift
    
    class MultimodeAssembly(System):
        def setup(self):
            a = self.add_child(MultimodeSystem('a', shift=10), pulling='m_in')
            b = self.add_child(MultimodeSystem('b', shift=20), pulling='m_in')
            c = self.add_child(MultimodeSystem('c', shift=30), pulling='m_out')

            self.connect(b, c, {'m_out': 'm_in'})

    class TopSystem(System):
        def setup(self) -> None:
            self.add_child(ModeManager('manager'))
            self.add_child(MultimodeAssembly('foo'))

            self.connect(self.manager, self.foo, {'mode': 'm_in'})
    
    s = TopSystem('s')
    driver = s.add_driver(EulerExplicit(dt=0.1, time_interval=(0, 1)))
    driver.add_recorder(DataFrameRecorder(excludes='*.shift'), period=driver.dt)
    driver.set_scenario(
        values={
            "manager.x": "4 * t if t < 0.5 else 4 * (1 - t)",  # mode changes @ t=1/4 and 3/4
        }
    )

    s.run_drivers()

    data = driver.recorder.export_data()
    # data = data.drop(['Section', 'Status', 'Error code'], axis=1)
    # import pandas
    # pandas.set_option('display.width', 200)
    # pandas.set_option('display.max_columns', 200)
    # print("\n", data)

    expected_modes = numpy.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0])

    assert numpy.array_equal(data['foo.m_in'], expected_modes)
    assert numpy.array_equal(data['foo.m_in'], data['manager.mode'])
    assert numpy.array_equal(data['foo.a.m_in'], data['foo.m_in'])     # pulling
    assert numpy.array_equal(data['foo.b.m_in'], data['foo.m_in'])     # pulling
    assert numpy.array_equal(data['foo.c.m_in'], data['foo.b.m_out'])  # sibling connection
    assert numpy.array_equal(data['foo.c.m_out'], numpy.asarray(data['foo.c.m_in']) + s.foo.c.shift)
    assert numpy.array_equal(data['foo.a.m_out'], expected_modes + s.foo.a.shift)
    assert numpy.array_equal(data['foo.b.m_out'], expected_modes + s.foo.b.shift)
    assert numpy.array_equal(data['foo.c.m_out'], expected_modes + s.foo.b.shift + s.foo.c.shift)


def test_MultimodeSystem_loop_recomposition():
    """Check multimode system containing new loops after transition.
    """
    class System1(System):
        def setup(self):
            self.add_inward('x', 1.0)
            self.add_inward('y', 1.0)
            self.add_outward('z', 0.0)

        def compute(self):
            self.z = self.x + self.y

    class System2(System):
        def setup(self):
            self.add_inward('k', 1.0)
            self.add_inward('x', 1.0)
            self.add_outward('y', 0.0)

        def compute(self):
            self.y = self.k + self.x**2

    class MultimodeAssembly(System):
        def setup(self):
            a = self.add_child(System1('a'))
            b = self.add_child(System2('b'))
            c = self.add_child(System2('c'))
            
            self.add_event('click')

            self.connect(a, b, {'z': 'x', 'y': 'y'})  # loop
            self.connect(b, c, {'y': 'x'})

        def transition(self):
            if self.click.present:
                self.connect(self.a, self.c, {'x': 'y'})
    
    s = MultimodeAssembly('s')
    s.add_driver(NonLinearSolver('solver'))

    s.a.x = -2.0
    s.b.k = 1.0
    s.c.k = -2.0
    s.run_drivers()
    # print("", s.drivers['solver'].problem, sep="\n")
    assert s.drivers['solver'].results.success
    assert s.drivers['solver'].problem.shape == (1, 1)

    s.drivers.clear()
    driver = s.add_driver(EulerExplicit(dt=1.0, time_interval=(0, 1)))
    driver.add_child(NonLinearSolver('solver'))
    # driver.add_recorder(DataFrameRecorder(), period=1.0)

    s.click.trigger = "t == 0.4"

    try:
        s.run_drivers()
    except:
        raise
    finally:
        # data = driver.recorder.export_data()
        # data = data.drop(['Section', 'Status', 'Error code'], axis=1)
        # pandas.set_option('display.width', 200)
        # pandas.set_option('display.max_columns', 200)
        # print("", driver.solver.problem, "", data, sep="\n")
        problem = driver.solver.problem
        assert problem.shape == (2, 2)


def test_MultimodeSystem_event_init():
    """Test checking that the system is up-to-date when primary events are initialized.
    """
    class Kinematics(System):
        def setup(self):
            self.add_inward('v', numpy.zeros(2))
            self.add_transient('x', der='v')

    class Wall(System):
        def setup(self):
            self.add_inward('x', numpy.zeros(2), desc='Point on the wall')
            self.add_inward('n_dir', numpy.r_[1., 1.], desc='Normal direction')
            self.add_outward('n', numpy.r_[0., 0.], desc='Unit normal vector')

        def compute(self) -> None:
            self.n = self.n_dir / numpy.linalg.norm(self.n_dir)
        
        def distance(self, point) -> float:
            """Signed distance to wall"""
            return self.n.dot(numpy.asarray(point) - self.x)

    class Assembly(System):
        def setup(self):
            self.add_child(Kinematics('kinematics'), pulling=['x', 'v'])
            self.add_child(Wall('wall'))
            self.add_event('rebound', trigger='wall.distance(x) == 0')

        def transition(self) -> None:
            if self.rebound.present:
                n = self.wall.n
                v = self.v
                v -= (2 * v.dot(n)) * n

    s = Assembly('s')

    driver = s.add_driver(EulerExplicit('driver', time_interval=(0, 1), dt=1.0))
    driver.set_scenario(
        init={
            'x': [0., 0.],
        },
        values={
            'v': [1., 0.],
            'wall.x': [0.25, 0.],
            'wall.n_dir': [-1., 0.],
        },
        stop=s.rebound,
    )
    # Initialize the wall normal with an **incorrect** value.
    # If the system is not up-to-date, the actual distance to the wall
    # will be not be correctly initialized, and the rebound will be missed.
    s.wall.n[:] = [1., 0.]

    s.run_drivers()

    assert len(driver.recorded_events) == 1
    record = driver.recorded_events[-1]
    assert len(record.events) == 2
    assert record.events[0] is s.rebound
    assert record.events[1] is driver.scenario.stop
    assert record.time == pytest.approx(0.25, rel=1e-15)
