import carbatpy as cb
import numpy as np

from carbatpy.models.components.comp import FlowDevice

from carbatpy.optional.spp_machines import Compressor, Turbine, require_spp_machines
from carbatpy.optional.process_maps import (
    InverseDesign as PMCompressor,
    require_process_maps,
)
from carbatpy.optional.expander_liege import ExpanderSE, require_expander_liege


class FlowMachine(FlowDevice):
    """FlowMachine means Compressor, Expander, and Pump."""

    def initialize(self, **kwargs):
        super().initialize(**kwargs)
        try:
            if cb.CB_DEFAULTS["Fluid_Defaults"]["PROPS"] == "REFPROP":
                self.TS = cb.CB_DEFAULTS["Fluid_Defaults"]["TRANS_STRING"]

            elif cb.CB_DEFAULTS["Fluid_Defaults"]["PROPS"] == "TREND":

                if len(self.composition) == 1:
                    self.TS = cb.CB_DEFAULTS["Fluid_Defaults"]["TRANS_TREND"]
                elif len(self.composition) > 1:
                    self.TS = cb.CB_DEFAULTS["Fluid_Defaults"]["TRANS_TREND_MIX"]
                else:
                    raise Exception(
                        "Length of composition array not correct: ",
                        len(self.composition),
                    )

            else:
                raise Exception(
                    "cb.CB_FLUID_DEFAULT['PROPS'] unknown. Try REFPROP or TREND"
                )
        except Exception as e:
            raise Exception("Error: ", e)

    def calculate(self, in_states, out_states, run_param=None, **kwargs):
        """
        Calculates a Flowmachine. Available models:
            - const_eta: compressor (Pump) / expander
            - turbo * RadialCompressor: from spp machines
            - turbo * RadialTurbine: from spp machines
            - scroll * Expander: from Liège

        For more information on the models see documentation.

        Results written in self.output

        Parameters
        ----------
        in_states : dictionary
            keys: fluidnmae (str)
            index: state entering, as defined in THERMO_STRING T,p,h,v,s.
        out_states : dictionary
            keys: fluidnmae (str)
            index: output state only the pressure is used out_state[1].
        run_params : dict, optional
            if further parameters must be passed for the calculation.
            - power : power given, mass flow rate calculated
        **kwargs : TYPE
            DESCRIPTION.

        Raises
        ------
        Exception
            DESCRIPTION.

        Returns
        -------
        None.

        """
        if self.verbose:
            self.logger.info(f"Calculating {self.name}")

        fluid_act = self.inputs["act_fluids"][self.fl_name]

        self.p_out = out_states[self.fl_name][1]

        state_in = fluid_act.set_state(
            [in_states[self.fl_name][1], in_states[self.fl_name][2]],
            "PH",
            output="FluidState",
            wanted=self.TS,
        )

        self.composition = np.array(state_in.total_z)
        self.p_in = state_in.pressure
        self.temp_in = state_in.temperature
        self.v_in = state_in.sp_volume
        self.h_in = state_in.enthalpy
        self.s_in = state_in.entropy

        state_out_isentropic = fluid_act.set_state(
            [self.s_in, self.p_out], "SP", output="FluidState"
        )

        self.diff_enthalpy_s = state_out_isentropic.enthalpy - self.h_in

        # Expander or Compressor via pressure ratio
        self.expander = self.p_in > self.p_out

        # Check for run_param
        implemented = {"power", "m_dot"}
        fixed = list(run_param.keys())[0]
        if fixed not in implemented:
            raise NotImplementedError(
                f"""FlowMachine:
                                              {self.inputs['parameters']['fixed']}
                                              not implemented!"""
            )

        if self.calc_type == "const_eta":
            eff_is = float(self.inputs["parameters"]["eta_s"])
            work_specific, state_out = self._calculate_state_out_const_eta(
                fluid_act,
                self.expander,
                eff_is,
            )

            if fixed == "power":
                power = run_param["power"]
                m_dot = power / work_specific

            elif fixed == "m_dot":
                m_dot = run_param["m_dot"][self.fl_name]
                power = m_dot * work_specific

            n_stages = None

        elif self.calc_type == "turbo * RadialCompressor":
            require_spp_machines()
            compressor = Compressor()

            if fixed == "power":
                power = float(run_param["power"])

                # Volume flow is unknown: iteration with start of eta from config
                eta_guess = float(self.inputs["parameters"]["eta_s"])

                m_dot, work_specific, compressor, state_out = self._iterate_fixed_power(
                    self.calc_type, power, fluid_act, eta_guess
                )

                power = m_dot * work_specific
                eff_is = float(compressor.IsentropicEfficiency[-1])
                n_stages = compressor.BasicParameters.StageNumber

                if self.verbose:
                    self.logger.info(f"Compressor Power: {power}\n")

            elif fixed == "m_dot":
                m_dot = float(run_param["m_dot"][self.fl_name])
                v_dot_in = m_dot * self.v_in

                # Radial Compressor DesignTool
                compressor = Compressor()
                state_out, _ = compressor.designCompressor(
                    self.working_fluid,
                    self.composition,
                    self.p_in,
                    self.temp_in,
                    self.p_out,
                    v_dot_in,
                    GearedCompressor=False,
                    EfficiencyConvergence=1e-2,
                    MaxIterations=20,
                    Plot=False,
                )

                work_specific = state_out.enthalpy - self.h_in
                m_dot = compressor.MassFlow
                power = m_dot * work_specific
                eff_is = float(compressor.IsentropicEfficiency[-1])
                n_stages = compressor.BasicParameters.StageNumber

        elif self.calc_type == "turbo * RadialTurbine":
            require_spp_machines()
            if fixed == "power":
                raise Exception(
                    "Power case not implemented for calc_type:" "turbo * RadialTurbine"
                )

            elif fixed == "m_dot":
                m_dot = run_param["m_dot"][self.fl_name]

                # Radial Turbine Design Tool
                turbine = Turbine()
                state_out, _ = turbine.designTurbine(
                    self.working_fluid,
                    self.composition,
                    state_in,
                    self.p_out,
                    m_dot,
                    Plot=False,
                )

                eff_is = turbine.IsentropicEfficiency[-1]
                work_specific = state_out.enthalpy - self.h_in
                power = m_dot * work_specific
                n_stages = turbine.BasicParameters.StageNumber

        elif self.calc_type == "scroll * Expander":
            require_expander_liege
            if fixed == "power":
                raise Exception(
                    "Power case not implemented for calc_type:" "turbo * RadialTurbine"
                )

            elif fixed == "m_dot":
                m_dot = run_param["m_dot"][self.fl_name]

                # Semi Empirical scroll Expander
                expander = ExpanderSE()
                inputs = {
                    "P_su": self.p_in,
                    "h_su": self.h_in,
                    "P_ex": self.p_out,
                    "m_dot": m_dot,
                    "T_amb": 288.15,
                    "fluid": self.working_fluid,
                    "comp": self.composition,
                }

                default_values = {
                    "AU_amb": 0,  #  Heat transfer to the ambient neglected
                    "AU_su_n": 0,  # Heat transfer to the suction side neglected
                    "AU_ex_n": 0,  # Heat transfer to the exhaust side neglected
                    "d_su1": 0,  # Pressure drop at the suction side neglected
                    "m_dot_n": 1,  # Nominal mass flow rate
                    "A_leak": 1e-12,  # Very small leakage area to neglect the leakage effect
                    "W_dot_loss_0": 0,  # Idle losses neglected
                    "alpha": 0,  # Proportionality rate for mechanical losses neglected
                    "C_loss": 0,  # Torque losses neglected
                    "rv_in": 2.5,  # Default volume ratio
                    "V_s": 0.001,  # Default swept volume
                    "mode": "P_M",  # Default mode
                }

                expander.set_inputs(**inputs)
                expander.set_parameters(**default_values)

                # Solve the expander component
                expander.solve()

                state_out = fluid_act.set_state(
                    [expander.h_ex, self.p_out],
                    "HP",
                    output="FluidState",
                    wanted=self.TS,
                )

                if self.verbose:
                    expander.print_setup()
                    expander.print_results()

                eff_is = expander.epsilon_is
                work_specific = state_out.enthalpy - self.h_in
                power = m_dot * work_specific
                n_stages = None

        elif "process map" in self.calc_type.lower():
            require_process_maps
            ct = self.calc_type.lower()

            if "screw" in ct:
                pm_machine = "screw"

            elif "geared" in ct or "igc" in ct:
                pm_machine = "centrifugal_IGC"

            else:
                pm_machine = "centrifugal_single"

            if fixed == "power":
                power = float(run_param["power"])

                # Volume flow is unknown: iteration with start of eta from config
                eta_guess = float(self.inputs["parameters"]["eta_s"])

                m_dot, work_specific, compressor, state_out = self._iterate_fixed_power(
                    self.calc_type, power, fluid_act, eta_guess, pm_machine=pm_machine
                )

                power = m_dot * work_specific
                eff_is = float(compressor.eta_is)
                n_stages = None

                if self.verbose:
                    self.logger.info(f"Compressor Power: {power}\n")

        else:
            raise NotImplementedError(f"calc_type '{self.calc_type}' not implemented")

        if eff_is < 0:
            raise Exception(f"Isentropic efficiency is less than 0: {eff_is}")

        fluid_act.set_state([state_out.enthalpy, self.p_out], "HP", output="FluidState")

        self.inputs["act_fluids"]["working_fluid"] = fluid_act

        self.output.update(
            {
                "state_in": {self.fl_name: np.array(state_in.state)},
                "state_out": {self.fl_name: np.array(state_out.state)},
                "isentropic_efficiency": eff_is,
                "n_stages": n_stages,
                "work_specific": work_specific,
                "power": power,
                "m_dot": {self.fl_name: m_dot},
                "warning": self.warning,
            }
        )
        self.entropy_production()
        self.all_data = {self.fl_name: np.array([state_in, state_out])}

    def _calculate_state_out_const_eta(self, fluid_act, expander: bool, eta_s: float):

        if expander:
            _work_specific = self.diff_enthalpy_s * eta_s
        else:
            _work_specific = self.diff_enthalpy_s / eta_s

        _state_out = fluid_act.set_state(
            [self.h_in + _work_specific, self.p_out],
            "HP",
            output="FluidState",
            wanted=self.TS,
        )

        return _work_specific, _state_out

    def _refprop_mix(self, wf, x):
        f = [c.strip() for c in wf.split("*") if c.strip()]
        s = sum(x)
        return "REFPROP::" + "&".join(f"{c}[{xi/s}]" for c, xi in zip(f, x) if xi)

    def _iterate_fixed_power(
        self,
        calc_type,
        power,
        fluid_act,
        eta_guess,
        tol=1e-5,
        maxiter=10,
        pm_machine: str | None = None,
    ):

        work_guess, _ = self._calculate_state_out_const_eta(
            fluid_act,
            expander=self.expander,
            eta_s=eta_guess,
        )
        m_dot_iter = power / work_guess

        if calc_type == "turbo * RadialCompressor":
            require_spp_machines()
            compressor = Compressor()

            v_dot_iter = m_dot_iter * self.v_in
            for _ in range(maxiter):
                st_out, _ = compressor.designCompressor(
                    self.working_fluid,
                    self.composition,
                    self.p_in,
                    self.temp_in,
                    self.p_out,
                    v_dot_iter,
                    GearedCompressor=False,
                    EfficiencyConvergence=1e-2,
                    MaxIterations=20,
                    Plot=False,
                )
                w_iter = st_out.enthalpy - self.h_in

                P_iter = m_dot_iter * w_iter

                if self.verbose:
                    self.logger.info(f"Compressor Power per Iter: {P_iter}\n")

                if abs(1 - power / P_iter) < tol:
                    break

                m_dot_iter = power / w_iter
                v_dot_iter = m_dot_iter * self.v_in

            return m_dot_iter, w_iter, compressor, st_out

        elif "process map" in calc_type.lower():
            fluid = self._refprop_mix(self.working_fluid, self.composition)
            print(pm_machine)
            for _ in range(maxiter):
                comp = PMCompressor(
                    fluid,
                    self.p_in,
                    self.temp_in,
                    self.p_out / self.p_in,
                    m_dot_iter,
                    machine=pm_machine,
                )

                w_iter = self.diff_enthalpy_s / comp.eta_is

                P_iter = m_dot_iter * w_iter
                if abs(1 - power / P_iter) < tol:
                    break

                m_dot_iter = power / w_iter

            h_out = w_iter + self.h_in

            st_out = fluid_act.set_state(
                [h_out, self.p_out], "HP", output="FluidState", wanted=self.TS
            )

            return m_dot_iter, w_iter, comp, st_out

        else:
            raise NotImplementedError(
                f"Iteration for calc_type '{calc_type}' not implemented"
            )
