# -*- coding: utf-8 -*-
"""
Created on Sun Jan  4 11:28:38 2026

@author: atakan
Universität Duisburg-Essen, Germany

In the framework of 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


def orc(dir_name, cop, q_dot_high, **kwargs):
    """
    Compute an Organic Rankine Cycle (ORC) for a Carnot battery.

    The ORC uses an isolated configuration (from file or dict). The evaporator heat
    input q_dot_high and the COP of the preceding charging heat pump are prescribed;
    the working fluid mass flow is determined accordingly. Components: pump, evaporator,
    expander, condenser to ambient (condenser_sur), and main condenser.

    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.
    cop : float
        COP of the preceding charging heat pump.
    q_dot_high : float
        Prescribed heat flow rate to the evaporator [W].
    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 cycle plot.

    Returns
    -------
    dict
        Result dictionary with keys:
        - "eta_th" : float
            Thermal efficiency (net electric power / evaporator heat input).
        - "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 required configuration entries are missing (e.g., working_fluid.p_high).
    """
    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)

    # Isolated, merged configuration
    config = cb.hp_comp.resolve_config(dir_name, new_config)

    # Initial start (mass flow will be set later)
    start0 = cb.comp.Start("start", config, m_dot=10e-3, verbose=verbose)
    p_low = config["working_fluid"]["p_low"]
    p_high = config["working_fluid"]["p_high"]

    # Prescribed heat flows (sign convention: negative for heat rejected)
    run_p_evap = {"q_dot": q_dot_high}
    q_dot_loss = -q_dot_high * (1.0 - 1.0 / cop)  # heat to cold storage (negative)
    run_p_cond = {"q_dot": q_dot_loss}

    # Pump (first pass with m_dot=1 to set states)
    pump = cb.comp.FlowMachine("pump", config, verbose=verbose)
    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"
        )

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

    # Evaporator: determines actual m_dot via q_dot_high
    evaporator = cb.comp.StaticHeatExchanger("evaporator", config, verbose=verbose)
    inp_evap, _ = evaporator.set_in_out({"working_fluid": pump.output["state_out"]["working_fluid"]})
    evaporator.calculate(inp_evap, run_param=run_p_evap, verbose=verbose)
    volumes_e = evaporator.calculate_volume()
    add_w_o(evaporator)
    if verbose:
        print(f"Storage Volumes evaporator: {volumes_e}")

    # Correct start with determined m_dot
    m_dot_w = evaporator.output["m_dot"]["working_fluid"]
    start = cb.comp.Start("start", config, m_dot=m_dot_w, verbose=verbose)
    run_p_machine = {"m_dot": {"working_fluid": m_dot_w}}

    # Recompute pump with correct m_dot (if needed by your component implementation)
    pump.calculate(
        start.output["state_out"],
        {"working_fluid": [600, p_high, 5e5]},
        run_param=run_p_machine,
        m_dot=m_dot_w
    )
    add_w_o(start)
    add_w_o(pump)

    # Expander
    expander = cb.comp.FlowMachine("expander", config, verbose=verbose)
    expander.calculate(
        evaporator.output["state_out"],
        {"working_fluid": [600, p_low, 5e5]},
        m_dot=m_dot_w,
        run_param=run_p_machine,
        verbose=verbose
    )
    add_w_o(expander)

    # Net power
    power_tot = expander.output["power"] + pump.output["power"]

    # Condenser to surroundings (heat not going into storage)
    q_dot_sur = -(q_dot_high + q_dot_loss + power_tot)  # choose sign so run_param is positive
    run_p_cond_sur = {"q_dot": -q_dot_sur, "m_dot": {"working_fluid": m_dot_w}}

    condenser_sur = cb.comp.StaticHeatExchanger("condenser_sur", config, verbose=verbose)
    inp_cs, _ = condenser_sur.set_in_out({"working_fluid": expander.output["state_out"]["working_fluid"]})
    condenser_sur.calculate(inp_cs, run_param=run_p_cond_sur, verbose=verbose)
    add_w_o(condenser_sur)

    # Main condenser (to cold storage)
    run_p_cond.update({"q_dot": -q_dot_loss, "m_dot": {"working_fluid": m_dot_w}})
    condenser = cb.comp.StaticHeatExchanger("condenser", config, verbose=verbose)
    inp_c, _ = condenser.set_in_out({"working_fluid": condenser_sur.output["state_out"]["working_fluid"]})
    sec_in, sec_out = condenser.set_in_out({"working_fluid": start.output["state_in"]["working_fluid"]}, False)
    condenser.calculate(inp_c, sec_out, run_param=run_p_cond, verbose=verbose)
    volumes_c = condenser.calculate_volume()
    add_w_o(condenser)
    if verbose:
        print(f"Storage Volumes condenser: {volumes_c}")

    # Efficiency and costs
    eta_th = np.abs(power_tot) / np.abs(evaporator.output["q_dot"])
    costs = cb.orc_comp.all_costs([pump, evaporator, expander, condenser_sur, condenser],
                                  verbose)

    if verbose:
        print(f"eta_th: {eta_th:.4f}")

    # Plot: shared Figure/Axes, do not deepcopy fig/ax
    if plotting:
        fig, ax = plt.subplots(1)
        # Copy plot defaults (pure 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 on same axes: use shallow copy so fig/ax references are preserved
        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": [pump.name, ""], "col": ["-r", "bv-"]})
        shift, direct = pump.plot(plot_info)

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

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

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

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

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


def all_costs(components_all, verbose):
    cost_total = 0
    cost_dict ={}

    for proc in components_all:
        cost = proc.estimate_costs()
        cost_total += cost
        cost_dict[proc.name] = cost
        if 'Storage_costs' in proc.output.keys():
            for key, val in proc.output['Storage_costs'].items():
                cost_dict[proc.name+"-storage-" + key] = val
                cost_total +=val
                
        if verbose:
            print(f'{proc.name} --Exergy destr.rate: {proc.output["exergy_destruction_rate"]:.2e} W, costs: {proc.cost:.2e}')
            if 'Storage_costs' in proc.output.keys():
                print(f"\tStorage costs: {proc.output['Storage_costs']}")
    return {"total_costs":cost_total, "cost_distribution":cost_dict}


if __name__ == "__main__":
    dir_name_m = cb.CB_DEFAULTS["General"]["CB_DATA"]+"\\io-orc-data.yaml"
    conf_m = {"cold_storage": {"temp_low": 250.},
              "working_fluid": {"p_low": 7e5,
                                "p_high": 11e5,
                                'species': 'ethane * propane * pentane * hexane',
                                'fractions': [0.25,
                                              0.55,
                                              0.15,
                                              0.05]}}
    res_act = orc(dir_name_m, 3, 1000, config=conf_m,
                                         plotting =True,
                                         verbose=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']}")
    