"""
Unit test for Valve.

This script includes unit tests for the CircAdapt unit test framework.
Additionally, this page uses the same functions to render a visualisation
of the unittest for the documentation.
"""

import circadapt
import circadapt.model
import circadapt.model.benchmark

from circadapt.model.benchmark import Valve
import matplotlib.pyplot as plt
import numpy as np

import scipy.optimize

import pytest
from _settings import test_dts, test_solver, settings

# %% analytical solution
def analytical_solution(t_true, model):
    """
    Analytical solution of Valve.

    Parameters
    ----------
    t_true : array
        timesteps to caluclate solution.
    model : circadapt.model.benchmark.Valve
        Model to get parameters from.

    Returns
    -------
    array
        analytical solution of flow.

    """
    model_options = model.model_options

    LenValve = model.get('Model.V.l')
    A_open = model.get('Model.V.A_open')
    A_leak = model.get('Model.V.A_leak')

    A_prox = 2*A_open
    A_dist = 2*A_open

    rho_b = 1050

    # A from implementation
    a_open = 1 / (1.5 * (LenValve / A_open +
                         0.5/np.sqrt(A_prox) +
                         0.5/np.sqrt(A_dist)))
    a_leak = 1 / (1.5 * (LenValve / A_leak +
                         0.5/np.sqrt(A_prox) +
                         0.5/np.sqrt(A_dist)))

    b_open = 0.5 * (A_open**-2 - (A_prox)**-2)
    b_leak = 0.5 * (A_leak**-2 - (A_prox)**-2)

    c_open = 1e3/rho_b
    c_leak = -1e3/rho_b

    def flow(t, q0, a, b, c):
        if q0 < - np.sqrt(np.abs(c_leak/b_leak)):
            C = (q0)**2 * b / c - 1
            aux = c/b*(1+C*np.exp(2*a*np.sqrt(b*np.abs(c))*t/(rho_b*2)))
            q = np.sign(aux)*np.sqrt(np.abs(aux))
            return q
        if q0 > np.sqrt(np.abs(c/b)):
            C = (q0)**2 * b / c - 1
            aux = c/b*(1+C*np.exp(-2*a*np.sqrt(b*np.abs(c))*t))
            q = np.sign(aux)*np.sqrt(np.abs(aux))
            return q

        if c < 0 and q0 > 0:
            C = (np.sign(c) * np.arctan(q0 / (np.sqrt(np.abs(c)/b))) /
                 np.sqrt(b*np.abs(c)))
            q = - np.sqrt(np.abs(c)/b) * np.tan(
                    np.sqrt(b*np.abs(c)) * (a*t+C)
                )

            # as soon as q<0, equation changes
            # q=0 when a*t+C=0
            # t=-C/a
            t0 = -C/a
            q[t > t0] = flow(t[t > t0] - t0, 0, a, b, c)

            return q

        C = (np.sign(c) *
             np.arctanh(q0 / (np.sqrt(np.abs(c)/b))) /
             np.sqrt(b*np.abs(c)))
        return np.sign(c)*np.sqrt(np.abs(c)/b) * np.tanh(
                np.sqrt(b*np.abs(c)) * (a*t+C)
            )

    # get parameters from model_options
    t_on = model_options.get('t_on')
    t_off = model_options.get('t_off')
    q0 = model_options.get('q0')

    # phase 0: Decay
    q_true = flow(t_true, q0, a_open, b_open, c_leak)

    # phase 1: negative flow with A_eff = A_leak
    t_root = scipy.optimize.root(
        lambda x: flow(x, q0, a_open, b_open, c_leak),
        0).x[0]
    q_true[t_true >= t_root] = flow(
        t_true[t_true >= t_root]-t_root,
        0,
        a_leak, b_leak, c_leak)

    # phase 2: increasing flow with A_eff = A_open
    q_true[t_true >= t_on] = flow(
        t_true[t_true >= t_on] - t_on,
        q_true[np.argmax(t_true >= t_on)],
        a_open, b_open, c_open)

    # phase 3 == phase 0: decreasing flow with A_eff = A_open
    q_true[t_true >= t_off] = flow(
        t_true[t_true >= t_off] - t_off,
        q_true[np.argmax(t_true >= t_off)],
        a_open, b_open, c_leak)

    # phase 4 == phase 1: negative flow with A_eff = A_leak
    t_root = scipy.optimize.root(lambda x: flow(
        x,
        q0,
        a_open, b_open, c_leak), 0.75).x[0]
    t_root = (t_off +
              np.arctan(q_true[np.argmax(t_true >= t_off)] /
                        (np.sqrt(np.abs(c_leak)/b_open))) /
              np.sqrt(b_open*np.abs(c_leak)) / a_open)
    q_true[t_true >= t_root] = flow(
        t_true[t_true >= t_root] - t_root,
        0,
        a_leak, b_leak, c_leak,
        )

    return {'t': t_true, 'q': q_true*1e3}


# %% simulate
def create_model(dt, solver):
    # get solver name and order from string
    solver_name = solver[:]
    order = None
    if solver_name[-1].isnumeric() and solver_name[-3:-1] == '_o':
        order = int(solver_name[-1])
        solver_name = solver_name[:-3]

    # create model and set solver properties
    model = Valve(solver_name)
    if order is not None:
        model.set('Solver.order', order)
    model.set('Solver.dt', dt)
    model.set('Solver.dt_export', dt)

    return model


def get_model_signals(model):
    signals = {'t': model['Solver']['t']}
    signals['q'] = model['Valve']['q'][:, 0] * 1e3

    return signals


def calculate_errors(dt, solver, signal):
    try:
        model = create_model(dt, solver)
        model.run(1)
        t_true = model['Solver']['t']
        true_signals = analytical_solution(t_true, model)
        model_signals = get_model_signals(model)
        return np.sqrt(np.mean(
            ((true_signals[signal] - model_signals[signal]) /
             np.max(np.abs(true_signals[signal]))
             )**2))
    except circadapt.error.ModelCrashed:
        return np.nan


# %% pytest
@pytest.mark.parametrize("dt", test_dts, ids=lambda d: f"dt={d*1e3}ms")
@pytest.mark.parametrize("solver", test_solver, ids=lambda d: f"solver={d}")
def test_valve(dt, solver):  
    """
    Test elements:
        """

    # which signal to compare
    signal_to_test = 'q'

    # run tests
    e = 1e-1
    assert np.abs(calculate_errors(dt, solver, signal_to_test)) < e
    

# %% main
if __name__ == "__main__":
    # plot settings
    plot_dt = np.array([1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1])
    plot_solvers = test_solver

    # get options
    model = create_model(5e-4, 'forward_euler')
    signal = 'q'

    # get all errors
    error = np.ndarray((len(plot_solvers), len(plot_dt)))
    for i_solver, solver in enumerate(plot_solvers):
        for i_dt, dt in enumerate(plot_dt):
            print(f'\rsimulate data solver {i_solver} with Delta t = {dt}    ',
                  end='\r')
            error[i_solver, i_dt] = calculate_errors(dt, solver, signal)
    print('Simulated all data                                                ')
    print(error)

    # %% make plot
    t_true = np.linspace(0, model['General']['t_cycle'], 10000)
    true_signals = analytical_solution(t_true, model)
    model = create_model(1e-4, plot_solvers[3])
    model.run(1)
    model_signals = get_model_signals(model)

    # Create a new figure
    plt.figure(1, clear=True, figsize=(8, 8))

    # Plot signal
    ax1 = plt.subplot(2, 2, 1)
    # plt.plot(true_signals['t'],
    #          true_signals['q'],
    #          label='True Signals', ls='--')
    plt.plot(model_signals['t'],
             model_signals['q'],
             label='Model Signals', ls='-', lw=1)
    plt.title('Flow', fontsize=14)
    plt.ylabel('$q [L/s]$', fontsize=12)
    plt.xlabel('$t [ms]$', fontsize=12)

    # Plot L2
    ax2 = plt.subplot(2, 2, 3)
    for i in range(error.shape[0]):
        ax2.plot(plot_dt*1e3, error[i, :], c=settings['solver_color'][plot_solvers[i]])
    plt.yscale('log')
    plt.xscale('log')

    # Move the legend outside the plot to the right
    # plt.legend(plot_solvers, bbox_to_anchor=(1.05, 1), loc='upper left')
    ax2.legend(plot_solvers, loc='lower left', bbox_to_anchor=(0.35, 0.0, 0.5, 0.5))

    plt.title('Mean L2 Error', fontsize=14)
    plt.xlabel('${\\Delta}t [ms]$', fontsize=12)

    # Super title for the entire figure
    plt.suptitle('Verification of Valve', fontsize=18, weight='bold')

    # Remove right and top spines of the first subplot
    ax1.spines['right'].set_visible(False)
    ax1.spines['top'].set_visible(False)
    ax2.spines['right'].set_visible(False)
    ax2.spines['top'].set_visible(False)

    # Adjust layout
    # plt.tight_layout()
    plt.subplots_adjust(top=0.92,
        bottom=0.079,
        left=0.15,
        right=0.9,
        hspace=0.285,
        wspace=0.2)


    # Show the plot
    plt.show()