"""Simulate behavior of virtual source algorithms.

The simulation recreates an observer-cape, the virtual Source and a virtual target
- input = hdf5-file with a harvest-recording
- output = optional as hdf5-file

The output file can be analyzed and plotted with shepherds tool suite.
"""

from contextlib import ExitStack
from pathlib import Path
from typing import Optional

import numpy as np
from tqdm import tqdm

from shepherd_core.data_models.base.calibration import CalibrationEmulator
from shepherd_core.data_models.content.virtual_source import VirtualSourceConfig
from shepherd_core.logger import logger
from shepherd_core.reader import Reader
from shepherd_core.writer import Writer

from .target_model import TargetABC
from .virtual_source_model import VirtualSourceModel


def simulate_source(
    config: VirtualSourceConfig,
    target: TargetABC,
    path_input: Path,
    path_output: Optional[Path] = None,
    *,
    monitor_internals: bool = False,
) -> float:
    """Simulate behavior of virtual source algorithms.

    FN returns the consumed energy of the target.
    """
    stack = ExitStack()
    file_inp = Reader(path_input, verbose=False)
    stack.enter_context(file_inp)
    cal_emu = CalibrationEmulator()
    cal_inp = file_inp.get_calibration_data()

    if path_output:
        file_out = Writer(
            path_output, cal_data=cal_emu, mode="emulator", verbose=False, force_overwrite=True
        )
        stack.enter_context(file_out)
        file_out.store_hostname("emu_sim_" + config.name)
        file_out.store_config(config.model_dump())
        cal_out = file_out.get_calibration_data()

    src = VirtualSourceModel(
        config,
        cal_emu,
        dtype_in=file_inp.get_datatype(),
        log_intermediate=False,
        window_size=file_inp.get_window_samples(),
        voltage_step_V=file_inp.get_voltage_step(),
    )
    i_out_nA = 0
    e_out_Ws = 0.0
    if monitor_internals and path_output:
        stats_sample = 0
        stats_internal = np.empty((round(file_inp.runtime_s * file_inp.samplerate_sps), 11))
        try:
            # keep dependencies low
            from matplotlib import pyplot as plt
        except ImportError:
            logger.warning("Matplotlib not installed, plotting of internals disabled")
            stats_internal = None
    else:
        stats_internal = None

    for _t, v_inp, i_inp in tqdm(
        file_inp.read(is_raw=True), total=file_inp.chunks_n, desc="Chunk", leave=False
    ):
        v_uV = 1e6 * cal_inp.voltage.raw_to_si(v_inp)
        i_nA = 1e9 * cal_inp.current.raw_to_si(i_inp)

        for _n in range(len(_t)):
            v_uV[_n] = src.iterate_sampling(
                V_inp_uV=int(v_uV[_n]),
                I_inp_nA=int(i_nA[_n]),
                I_out_nA=i_out_nA,
            )
            i_out_nA = target.step(int(v_uV[_n]), pwr_good=src.cnv.get_power_good())
            i_nA[_n] = i_out_nA

            if stats_internal is not None:
                stats_internal[stats_sample] = [
                    _t[_n] * 1e-9,  # s
                    src.hrv.voltage_hold * 1e-6,
                    src.cnv.V_input_request_uV * 1e-6,  # V
                    src.hrv.voltage_set_uV * 1e-6,
                    src.cnv.V_mid_uV * 1e-6,
                    src.hrv.current_hold * 1e-6,  # mA
                    src.hrv.current_delta * 1e-6,
                    i_out_nA * 1e-6,
                    src.cnv.P_inp_fW * 1e-12,  # mW
                    src.cnv.P_out_fW * 1e-12,
                    src.cnv.get_power_good(),
                ]
                stats_sample += 1

        e_out_Ws += (v_uV * i_nA).sum() * 1e-15 * file_inp.sample_interval_s
        if path_output:
            v_out = cal_out.voltage.si_to_raw(1e-6 * v_uV)
            i_out = cal_out.current.si_to_raw(1e-9 * i_nA)
            file_out.append_iv_data_raw(_t, v_out, i_out)

    stack.close()

    if stats_internal is not None:
        stats_internal = stats_internal[:stats_sample, :]
        fig, axs = plt.subplots(4, 1, sharex="all", figsize=(20, 4 * 6), layout="tight")
        fig.suptitle(f"VSrc-Sim with {config.name}, Inp={path_input.name}, E={e_out_Ws} Ws")
        axs[0].set_ylabel("Voltages [V]")
        axs[0].plot(stats_internal[:, 0], stats_internal[:, 1:5])
        axs[0].legend(["V_cv_hold", "V_inp_Req", "V_cv_set", "V_cap"], loc="upper right")

        axs[1].set_ylabel("Current [mA]")
        axs[1].plot(stats_internal[:, 0], stats_internal[:, 5:8])
        axs[1].legend(["C_cv_hold", "C_cv_delta", "C_out"], loc="upper right")

        axs[2].set_ylabel("Power [mW]")
        axs[2].plot(stats_internal[:, 0:1], stats_internal[:, 8:10])
        axs[2].legend(["P_inp", "P_out"], loc="upper right")

        axs[3].set_ylabel("PwrGood [n]")
        axs[3].plot(stats_internal[:, 0], stats_internal[:, 10])
        axs[3].legend(["PwrGood"], loc="upper right")

        axs[3].set_xlabel("Runtime [s]")

        for ax in axs:
            # deactivates offset-creation for ax-ticks
            ax.get_yaxis().get_major_formatter().set_useOffset(False)
            ax.get_xaxis().get_major_formatter().set_useOffset(False)

        plt.savefig(path_output.with_suffix(".png"))
        plt.close(fig)
        plt.clf()

    return e_out_Ws
