"""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

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_config import VirtualSourceConfig
from shepherd_core.logger import log
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: Path | None = None,
    *,
    monitor_internals: bool = False,
) -> float:
    """Simulate behavior of virtual source algorithms.

    FN returns the consumed energy of the target.
    """
    with ExitStack() as stack:
        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  # noqa: PLC0415
            except ImportError:
                log.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)

    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
