# -*- 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
from carbatpy.models.components.flowmachine_comp import FlowMachine

# 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 that 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
        Overrides to the base configuration (e.g., for optimization).
        Passed via ``**kwargs``.
    verbose : bool, optional
        If True, print additional information. Passed via ``**kwargs``.
        Default is False.
    plotting : bool, optional
        If True, create a combined T-H_dot plot. Passed via ``**kwargs``.
        Default is False.
    
    Returns
    -------
    dict
        Result dictionary with the following keys:
        
        - ``"COP"`` (float): Coefficient of performance of the heat pump.
        - ``"output"`` (dict): Component outputs and the configuration used 
          (accessible via ``"config"`` key).
        - ``"warnings"`` (dict): Warnings per component and additional items 
          such as ``"pressure_ratio"``.
        - ``"figure"`` (matplotlib.figure.Figure or None): Figure object if 
          ``plotting=True``, otherwise None.
        - ``"axes"`` (matplotlib.axes.Axes or None): Axes object if 
          ``plotting=True``, otherwise None.
        - ``"costs"`` (float or dict): Aggregate cost metric from components, 
          computed via ``cb.orc_comp.all_costs()``.
    
    Raises
    ------
    KeyError
        If the compressor power is missing in the configuration 
        (``process.fixed.compressor.power``).
    ValueError
        If input parameters are outside valid ranges.
    
    See Also
    --------
    compute_orc_carnot_battery : Compute ORC cycle
    load_config : Load configuration from YAML file
    
    Notes
    -----
    The coefficient of performance (COP) is calculated as:
    
    .. math::
        
        COP = \\frac{\\dot{Q}_{cond}}{W_{comp}}
    
    where :math:`\\dot{Q}_{cond}` is the condenser heat output and 
    :math:`W_{comp}` is the compressor work input.
    
    Component Sequence
    ^^^^^^^^^^^^^^^^^^
    The heat pump cycle follows this sequence:
    
    1. **Start**: Initialize working fluid state
    2. **Compressor**: Compress refrigerant to high pressure
    3. **Condenser**: Reject heat to hot storage tank
    4. **Throttle**: Expand refrigerant to low pressure (isenthalpic)
    5. **Evaporator**: Absorb heat from cold storage tank
    
    Configuration Priority
    ^^^^^^^^^^^^^^^^^^^^^^
    Configurations are merged in the following order (later overrides earlier):
    
    1. ``CB_DEFAULTS`` (base defaults)
    2. Configuration from ``dir_name`` (file or dict)
    3. User-provided ``config`` parameter (via kwargs)
    
    Examples
    --------
    Basic usage with configuration file:
    
    >>> result = heat_pump("config/hp_default.yaml", verbose=True)
    >>> print(f"COP: {result['COP']:.2f}")
    COP: 3.45
    
    With configuration dict and plotting:
    
    >>> config_dict = {
    ...     "working_fluid": {"name": "R134a"},
    ...     "process": {"fixed": {"compressor": {"power": 5000}}}
    ... }
    >>> result = heat_pump(config_dict, plotting=True)
    >>> result['figure'].savefig('heat_pump_cycle.png')
    
    With configuration overrides for optimization:
    
    >>> result = heat_pump(
    ...     "config/hp_default.yaml",
    ...     config={
    ...         "components": {
    ...             "compressor": {"eta_is": 0.85},
    ...             "evaporator": {"delta_T": 5.0}
    ...         }
    ...     }
    ... )
    >>> print(f"Optimized COP: {result['COP']:.2f}")
    Optimized COP: 3.52
    
    Access component outputs:
    
    >>> result = heat_pump("config.yaml")
    >>> compressor_output = result['output']['compressor']
    >>> print(f"Discharge temp: {compressor_output['T_out']:.1f} K")
    Discharge temp: 323.5 K
    """

    cb.utils.io_utils.reset_fluid_registry()

    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 = 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}}
    )
    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}}
    )
    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 = {
        "compressor": {"calc_type": "const_eta"},
        # "compressor": {"calc_type": "turbo * RadialCompressor"},
        # "compressor": {"calc_type": "process map CentrifugalCompressor"},
        # "compressor": {"calc_type": "process map CentrifugalCompressorGeared"},
        # "compressor": {"calc_type": "process map ScrewCompressor"},
        "cold_storage": {"temp_low": 250.0},
        "working_fluid": {"p_high": 20e5},
    }

    res_act = heat_pump(dir_name_out, config=conf_m, 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']}")
    plt.show()
