"""
Example of a custom component definition
2023, January 2
"""

import numpy as np
import fonsim as fons
import fonsim.constants.norm as cnorm


class Container(fons.Component):
    """
    A Container object is a (by default, empty) container or tank.
    It has one terminal named 'a'.
    It has one state named 'mass',
    which represents the mass of the fluid inside the container.

    :param label: label
    :param fluid: fluid, must be compressible
    :param volume: volume of the container in m^3.
    """
    def __init__(self, label=None, fluid=None, volume=None):
        super().__init__(label)

        # Define terminals and states
        self.set_terminals(
            fons.Terminal('a',
                          [fons.Variable('pressure', 'across',  label='p',
                                         initial_value=cnorm.pressure_atmospheric,
                                         range=(0, np.inf)),
                           fons.Variable('massflow', 'through', label='mf')]))
        self.set_states(fons.Variable('mass', 'local', label='m',
                                      initial_value=volume*fluid.rho_stp))

        # Define behaviour
        @self.auto_state
        def update_state(dt, mf, m):
            m_new = m + mf*dt
            return {'m': m_new}
        self.update_state = update_state

        @self.auto
        def evaluate(p, m):
            residual = m*cnorm.pressure_atmospheric - volume*fluid.rho_stp*p
            return np.array([residual], dtype=float)
        self.evaluate = evaluate


class Container_bis(Container):
    def __init__(self, label=None, fluid=None, volume=None):
        super().__init__(label=label, fluid=fluid, volume=volume)

        # With derivative specified
        @self.auto_state
        def update_state(dt, mf, m):
            jacobian = {}
            m_new = m + dt * mf
            jacobian['m/p'] = 0
            jacobian['m/mf'] = dt
            return {'m': m_new}, jacobian
        self.update_state = update_state

        @self.auto
        def evaluate(p, m):
            mass_stp = volume*fluid.rho_stp
            values, jacobian = np.zeros([1], dtype=float), [{}]
            values[0] = m*cnorm.pressure_atmospheric - mass_stp*p
            jacobian[0]['m'] = cnorm.pressure_atmospheric
            jacobian[0]['p'] = -mass_stp
            return values, jacobian
        self.evaluate = evaluate


# Run a small demo with this component to show that it works
import matplotlib.pyplot as plt

# Pressure reference
waves = [(0.0, 0.2), (0.5, 0.0), (1.0, 0.2), (1.5, 0.0)]
wave_function = fons.wave.Custom(waves)*1e5 + fons.pressure_atmospheric
# Compressible fluid
fluid = fons.air
# System
system = fons.System()
system.add(fons.PressureSource('source', pressure=wave_function))
system.add(fons.CircularTube('tube', fluid=fluid, diameter=2e-3, length=0.4))
system.add(Container('container', fluid=fluid, volume=250e-6))
system.connect('tube', 'source')
system.connect('tube', 'container')

# Create and run simulation
sim = fons.Simulation(system, duration=2, step_max=1e-3)
sim.run()

# Plot simulation results
fig, axs = plt.subplots(3, sharex=True, tight_layout=True)
fons.plot(axs[0], sim, 'pressure', 'bar', ('source', 'container'))
fons.plot_state(axs[1], sim, 'mass', 'g', ['container'])
fons.plot(axs[2], sim, 'massflow', 'g/s', ['tube'])
for a in axs: a.legend()
plt.show()