"""
base_class_two_phase.py

Base class for two-phase flow correlations.
Handles shared logic:
  - Fluid-state initialisation (liquid / vapour phases, mixtures)
  - Surface tension
  - Rouhani–Axelsson void fraction
  - Log-mean void fraction (Deng Eq. 23)
  - Biberg stratified angle approximation
  - Wojtan/Kattan dimensionless geometry
"""

import numpy as np
import carbatpy as cb

g = 9.80665  # m/s²


class TwoPhaseBase:
    """
    Base class for two-phase correlations.

    Parameters
    ----------
    state  : FluidState  – two-phase mixture state
    fluid  : carbatpy Fluid object
    D      : float       – inner tube diameter [m]
    G      : float       – mass flux [kg/(m² s)]
    """

    def __init__(self, state, fluid, D, G):
        self.state = state
        self.fluid = fluid
        self.D = D
        self.G = G

        self.TRANS = self._get_trans_string()

        if self.fluid.props == "TREND":
            self.fluid.set_state([state.pressure, state.enthalpy], "PH")
        elif self.fluid.props == "REFPROP":    
            self.fluid.set_state([state.pressure, state.quality], "PQ")

        if fluid.no_compounds > 1:
            self._set_liq_vap_phase_for_mixture()
            # For zeotropic mixtures state_l / state_v are at T_mix,
            # so dT_GL = 0 there.  Use bubble/dew endpoints instead.
            state_bubble = fluid.set_state([self.state.pressure, 0], "PQ",
                                           wanted=self.TRANS, output="FluidState")
            state_dew    = fluid.set_state([self.state.pressure, 1], "PQ",
                                           wanted=self.TRANS, output="FluidState")
            self.dH_GL = state_dew.enthalpy    - state_bubble.enthalpy
            self.dT_GL = state_dew.temperature - state_bubble.temperature
        else:
            self._init_pure_states()
            self._calc_sigma()
            self.dH_GL = self.state_v.enthalpy    - self.state_l.enthalpy
            self.dT_GL = self.state_v.temperature - self.state_l.temperature

    # ------------------------------------------------------------------
    # Fluid-state helpers
    # ------------------------------------------------------------------

    def _init_pure_states(self):
        """
        Set self.state_l and self.state_v for a pure fluid.
        """
        self.state_l = self.fluid.set_state([self.state.pressure, 0], "PQ",
                                            wanted=self.TRANS, output="FluidState")
        self.state_v = self.fluid.set_state([self.state.pressure, 1], "PQ",
                                            wanted=self.TRANS, output="FluidState")

    def _get_trans_string(self):
        DEFAULTS = cb.CB_DEFAULTS
        if self.fluid.props == 'TREND':
            return DEFAULTS["Fluid_Defaults"]['TRANS_TREND_MIX']
        return DEFAULTS["Fluid_Defaults"]['TRANS_STRING']

    def _set_liq_vap_phase_for_mixture(self):
        overall_comp = self.fluid.composition
        comp_l = self.fluid.properties.liq_x
        comp_v = self.fluid.properties.vap_y

        try:
            self.fluid.set_composition(comp_l)
        except Exception:
            raise ValueError("comp_l or comp_v == None")

        self.state_l = self.fluid.set_state([self.state.temperature, 0], "TQ",
                                            wanted=self.TRANS, output="FluidState")
        self._calc_sigma()

        self.fluid.set_composition(comp_v)
        self.state_v = self.fluid.set_state([self.state.temperature, 1], "TQ",
                                            wanted=self.TRANS, output="FluidState")
        self.fluid.set_composition(overall_comp)

    def _calc_sigma(self):
        if self.fluid.props == 'TREND':
            thermo = cb.CB_DEFAULTS["Fluid_Defaults"]['THERMO_TREND']
            sigma = self.fluid.set_state([self.state.temperature, 0], "TQ",
                                         wanted=thermo + ['ST'])[-1]
            if sigma <= 0:
                print("Warning: surface tension not correctly calculated.")
        else:
            sigma = self.fluid.set_state([self.state.temperature, 0], "TQ",
                                         wanted="STN")[0]
        self.sigma = sigma

    # ------------------------------------------------------------------
    # Void fraction
    # ------------------------------------------------------------------

    def _ra_void_fraction(self, x):
        """
        Rouhani–Axelsson (Steiner version) void fraction – Wojtan Eq. (4).
        """
        term1 = (1.0 + 0.12 * (1.0 - x)) * (
            x * self.state_v.sp_volume + (1.0 - x) * self.state_l.sp_volume)
        term2 = (1.18 * (1.0 - x) *
                 (g * self.sigma * (1 / self.state_l.sp_volume
                                    - 1 / self.state_v.sp_volume)) ** 0.25
                 / (self.G * self.state_l.sp_volume ** -0.5))
        return (x * self.state_v.sp_volume) / (term1 + term2)

    def _calc_alpha(self):
        """
        Compute Rouhani–Axelsson void fraction (alpha_ra) and log-mean void
        fraction (alpha) following Deng Eq. (23).
        Stores: self.alpha_ra, self.alpha
        """
        alpha_h = (1 + ((1 - self.state.quality) / self.state.quality) *
                   (self.state_v.sp_volume ** -1
                    / self.state_l.sp_volume ** -1)) ** -1
        self.alpha_ra = self._ra_void_fraction(self.state.quality)
        self.alpha = ((alpha_h - self.alpha_ra)
                      / np.log(alpha_h / self.alpha_ra))  # Eq. 23

    # ------------------------------------------------------------------
    # Geometry helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _biberg_theta_strat(epsilon):
        """
        Biberg approximation for θ_strat(ε) – Wojtan Eq. (9).
        Returns θ_strat in [rad], range (0, 2π).
        """
        e = float(min(max(epsilon, 1e-12), 1.0 - 1e-12))
        term1 = np.pi * (1.0 - e)
        term2 = (1.5 * np.pi) ** (1.0 / 3.0) * (
            1.0 - 2.0 * (1.0 - e)
            + (1.0 - e) ** (1.0 / 3.0)
            - e ** (1.0 / 3.0))
        term3 = (0.005 * (1.0 - e) * e * (1.0 - 2.0 * (1.0 - e))
                 * (1.0 + 4.0 * ((1.0 - e) ** 2 + e ** 2)))
        theta = 2.0 * np.pi - 2.0 * (term1 + term2 - term3)
        return float(min(max(theta, 1e-9), 2.0 * np.pi - 1e-9))

    def _wojtan_geom_from_e_theta(self, epsilon, theta_strat):
        """
        Dimensionless geometric quantities (Wojtan/Kattan).
        ALD = A_L/D²,  AVD = A_V/D²,  hLD = (1−cos((2π−θ)/2))/2,
        PiD = sin((2π−θ)/2).
        """
        A_total = 0.25 * np.pi * self.D ** 2
        AL = A_total * (1.0 - epsilon)
        AV = A_total * epsilon
        ALD = AL / self.D ** 2
        AVD = AV / self.D ** 2
        hLD = 0.5 * (1.0 - np.cos((2 * np.pi - theta_strat) / 2))
        PiD = np.sin((2 * np.pi - theta_strat) / 2)
        return {"A": A_total, "A_L": AL, "A_V": AV,
                "ALD": ALD, "AVD": AVD,
                "hLD": float(hLD), "PiD": float(PiD)}
