# -*- coding: utf-8 -*-
"""
Heat pump with two two-tank storages (component-based formulation, comp.py).

Created on Sun Jan  4 11:06:10 2026

Author: atakan
University of Duisburg-Essen, Germany

Within the Priority Programme: "Carnot Batteries: Inverse Design from
Markets to Molecules" (SPP 2403)
https://www.uni-due.de/spp2403/
https://git.uni-due.de/spp-2403/residuals_weather_storage
"""
import copy
from types import SimpleNamespace
import numpy as np
import matplotlib.pyplot as plt
import carbatpy as cb
# import pandas as pd  # only needed in the optional optimization example below


def deep_merge(dst: dict, src: dict) -> dict:
    """
    Recursively merge two nested dictionaries.

    - Values from src are written into dst.
    - For nested dicts, merging is done recursively.
    - Non-dict values overwrite the value in dst.

    Notes
    -----
    - dst is modified in-place.
    - src is not modified.

    Parameters
    ----------
    dst : dict
        Target dictionary; will be updated in-place.
    src : dict
        Source dictionary; its values are merged into dst.

    Returns
    -------
    dict
        The updated dst dictionary (same reference).
    """
    for k, v in src.items():
        if isinstance(v, dict) and isinstance(dst.get(k), dict):
            deep_merge(dst[k], v)
        else:
            dst[k] = v
    return dst


def resolve_config(dir_or_dict, override=None) -> dict:
    """
    Create an isolated configuration and apply optional overrides.

    - If dir_or_dict is a path, load the YAML configuration.
    - If dir_or_dict is a dict, use a deep copy of it.
    - Overrides are deep-merged into the base configuration.
    - Always returns a new, isolated structure (no references to global defaults).

    Parameters
    ----------
    dir_or_dict : str or dict
        Path to a YAML file or a configuration dictionary.
    override : dict, optional
        Dictionary with parameters to override (e.g., for optimization).

    Returns
    -------
    dict
        Fully isolated and merged configuration structure.
    """
    if isinstance(dir_or_dict, dict):
        base = copy.deepcopy(dir_or_dict)
    else:
        # read_config gibt direkt zurück wenn dict, daher hier deepcopy
        loaded = cb.utils.io_utils.read_config(dir_or_dict)
        base = copy.deepcopy(loaded)
    
    if override:
        deep_merge(base, copy.deepcopy(override))
    return base


def heat_pump(dir_name, **kwargs):
    """
    Compute a heat pump with two two-tank storages (component-based formulation).

    The function builds an isolated, merged configuration from a base (file or dict),
    computes the component sequence (start, compressor, condenser, throttle, evaporator),
    and returns COP, outputs, warnings, and optionally a combined plot.

    Key properties
    - No in-place modifications to CB_DEFAULTS or externally provided configurations.
    - Configurations are handled in isolation (deepcopy/recursive merges).
    - Matplotlib Figure/Axes are consistently passed to all plot calls so that the
      entire cycle is drawn on the same axes.

    Parameters
    ----------
    dir_name : str or dict
        Path to a YAML configuration file or a configuration dict.
    config : dict, optional (via kwargs)
        Overrides to the base configuration (e.g., for optimization).
    verbose : bool, optional (via kwargs)
        If True, print additional information.
    plotting : bool, optional (via kwargs)
        If True, create a combined T-H_dot plot.

    Returns
    -------
    dict
        Result dictionary with keys:
        - "COP" : float
            Coefficient of performance of the heat pump.
        - "output" : dict
            Component outputs and the configuration used ("config").
        - "warnings" : dict
            Warnings per component (and additional items, e.g., pressure_ratio).
        - "figure" : matplotlib.figure.Figure or None
            Figure object, if plotting=True.
        - "axes" : matplotlib.axes.Axes or None
            Axes object, if plotting=True.
        - "costs" : any
            Aggregate cost metric from components (from cb.orc_comp.all_costs).

    Raises
    ------
    KeyError
        If the compressor power is missing in the configuration
        (process.fixed.compressor.power).
    """
    new_config = kwargs.get("config", None)
    verbose = kwargs.get("verbose", False)
    plotting = kwargs.get("plotting", False)

    fig = None
    ax = None
    warnings = {}
    outputs = {}

    def add_w_o(comp):
        # collect warnings/outputs per component
        warnings[comp.name] = copy.deepcopy(comp.warning)
        outputs[comp.name] = copy.deepcopy(comp.output)

    # Build isolated, merged configuration
    config = resolve_config(dir_name, new_config)

    # Preliminary start state (mass flow will be set after the compressor)
    start0 = cb.comp.Start("start", config, m_dot=10e-3, verbose=verbose)

    # Compressor power from configuration
    try:
        power = config["process"]["fixed"]["compressor"]["power"]
    except Exception:
        raise KeyError("Configuration is missing process.fixed.compressor.power.")

    # Compressor
    compressor = cb.comp.FlowMachine("compressor", config, verbose=verbose)
    p_high = config["working_fluid"]["p_high"]
    p_ratio = p_high / start0.output["state_out"]["working_fluid"][1]
    if p_ratio < 1.0:
        warnings["pressure_ratio"] = SimpleNamespace(
            value=10.0 / p_ratio,
            message="Pressure ratio is wrong"
        )

    compressor.calculate(
        start0.output["state_out"],
        {"working_fluid": [600, p_high, 5e5]},
        run_param={"power": power}
    )

    # Working-fluid mass flow
    m_dot_w = compressor.output["m_dot"]["working_fluid"]
    m_dot_param = {"m_dot": {"working_fluid": m_dot_w}}

    # Start state with correct mass flow
    start = cb.comp.Start("start", config, m_dot=m_dot_w, verbose=verbose)
    add_w_o(start)
    add_w_o(compressor)

    # Condenser
    condenser = cb.comp.StaticHeatExchanger("condenser", config, verbose=verbose)
    inp_cond, _ = condenser.set_in_out({"working_fluid": compressor.output["state_out"]["working_fluid"]})
    condenser.calculate(in_states=inp_cond, run_param={"m_dot": {"working_fluid": m_dot_w}}, verbose=False)
    volumes_c = condenser.calculate_volume()
    if verbose:
        print(f"Storage Volumes condenser: {volumes_c}")
    add_w_o(condenser)

    # Throttle
    throttle = cb.comp.Throttle("throttle", config, verbose=verbose)
    throttle.calculate(condenser.output["state_out"], compressor.output["state_in"], run_param=m_dot_param)
    add_w_o(throttle)

    # Evaporator
    evaporator = cb.comp.StaticHeatExchanger("evaporator", config, verbose=verbose)
    inp_ev_wf, _ = evaporator.set_in_out({"working_fluid": throttle.output["state_out"]["working_fluid"]})
    inp_ev_sec, out_ev_sec = evaporator.set_in_out(start.output["state_in"], False)
    evaporator.calculate(inp_ev_wf, out_ev_sec, run_param={"m_dot": {"working_fluid": m_dot_w}}, verbose=False)
    volumes_e = evaporator.calculate_volume()
    if verbose:
        print(f"Storage Volumes evaporator: {volumes_e}")
    add_w_o(evaporator)

    # Performance figure
    cop = np.abs(condenser.output["q_dot"] / power)
    if verbose:
        print(f"COP: {cop:.4f}")

    costs = cb.orc_comp.all_costs([compressor, condenser, throttle, evaporator],
                                  verbose)

    # Combined plot (ensure all elements share the same figure/axes)
    if plotting:
        fig, ax = plt.subplots(1)
        # Copy plot defaults (plain values), then attach fig/ax
        plot_info = copy.deepcopy(cb.CB_DEFAULTS["Components"]["Plot"])
        plot_info["ax"] = ax
        plot_info["fig"] = fig
        plot_info["x-shift"] = [0, 0]

        # Start point: use shallow copy to preserve fig/ax references
        start_plot_info = plot_info.copy()
        start_plot_info.update({
            "label": ["start", ""],
            "col": ["ok", "bv"],
            "direction": 1
        })
        shift, direct = start.plot(start_plot_info)

        # Subsequent components on the same axes
        plot_info["x-shift"] = shift
        plot_info["direction"] = direct

        plot_info.update({
            "label": [compressor.name, ""],
            "col": ["-r", "bv-"]
        })
        shift, direct = compressor.plot(plot_info)

        plot_info["x-shift"] = shift
        plot_info["direction"] = direct
        plot_info.update({
            "label": [condenser.name, ""],
            "col": [":r", "rv-"]
        })
        shift, direct = condenser.plot(plot_info)

        plot_info["x-shift"] = shift
        plot_info["direction"] = direct
        plot_info.update({
            "label": [throttle.name, ""],
            "col": ["-b", "bv-"]
        })
        shift, direct = throttle.plot(plot_info)

        plot_info["x-shift"] = shift
        plot_info["direction"] = direct
        plot_info.update({
            "label": [evaporator.name, ""],
            "col": [":b", "bv-"]
        })
        evaporator.plot(plot_info)

    # Collect results
    outputs["config"] = copy.deepcopy(config)
    results = {
        "COP": cop,
        "output": outputs,
        "warnings": warnings,
        "figure": fig,
        "axes": ax,
        "costs": costs
    }
    return results

if __name__ == "__main__":
    dir_name_out = cb.CB_DEFAULTS["General"]["CB_DATA"] + "\\io-hp-data.yaml"

    conf_m = {
        "cold_storage": {"temp_low": 250.0},
        "working_fluid": {"p_high": 35e5}
        }

    res_act = heat_pump(dir_name_out, verbose=True, plotting=True)
    if any(ns.value != 0 for ns in res_act["warnings"].values()):
        print(f"Check warnings, at least one deviates from 0!\n{res_act['warnings']}")
