"""Build feos states and expose thermodynamic properties used by the
tracing routines."""

from feos import (  # pylint: disable=no-name-in-module
    PureRecord,
    Identifier,
    BinaryRecord,
    Parameters,
    EquationOfState,
    State,
    PhaseEquilibrium,
)

import numpy as np
from si_units import PASCAL, MOL, KILO, METER, BAR, KELVIN

KIJ_REF_T = 298.15


def _build_pure_records(temperature, params):
    """Build PureRecord list for all components at temperature."""
    pure_params = []

    for idx in range(params["number_components"]):
        pure_params.append(
            PureRecord(
                identifier=Identifier(
                    name=params[f"comp{idx + 1}_name"],
                ),
                molarweight=params[f"comp{idx + 1}_molarweight"],
                m=params[f"comp{idx + 1}_m"],
                sigma=params[f"comp{idx + 1}_sigma"]
                + params[f"comp{idx + 1}_sigma_fac1"]
                * np.exp(params[f"comp{idx + 1}_sigma_exp_fac1"] * temperature)
                + params[f"comp{idx + 1}_sigma_fac2"]
                * np.exp(
                    params[f"comp{idx + 1}_sigma_exp_fac2"] * temperature
                ),
                epsilon_k=params[f"comp{idx + 1}_epsilon_k"],
                association_sites=[
                    {
                        "na": params[f"comp{idx + 1}_na"],
                        "nb": params[f"comp{idx + 1}_nb"],
                        "kappa_ab": params[f"comp{idx + 1}_kappa_ab"],
                        "epsilon_k_ab": params[f"comp{idx + 1}_epsilon_k_ab"],
                    }
                ],
            )
        )

    return pure_params


def _build_binary_records(temperature, params):
    """Build BinaryRecord list for all component pairs at temperature."""
    binary_records_lst = []

    for idx in range(params["number_components"] - 1):
        for idx_ii in range(idx + 1, params["number_components"]):
            k_ij = (
                params[f"comp{idx + 1}{idx_ii + 1}_kij_intercept"]
                + (temperature - KIJ_REF_T)
                * params[f"comp{idx + 1}{idx_ii + 1}_kij_slope"]
            )

            # check for association in mixing rule
            if params[
                f"comp{idx + 1}{idx_ii + 1}_epsilon_k_ab"
            ] and not np.isnan(
                params[f"comp{idx + 1}{idx_ii + 1}_epsilon_k_ab"]
            ):
                binary_records_lst.append(
                    BinaryRecord(
                        Identifier(name=params[f"comp{idx + 1}_name"]),
                        Identifier(name=params[f"comp{idx_ii + 1}_name"]),
                        k_ij=k_ij,
                        association_sites=[
                            {
                                "kappa_ab": params[
                                    f"comp{idx + 1}{idx_ii + 1}_kappa_ab"
                                ],
                                "epsilon_k_ab": params[
                                    f"comp{idx + 1}{idx_ii + 1}_epsilon_k_ab"
                                ],
                            }
                        ],
                    )
                )
            else:
                binary_records_lst.append(
                    BinaryRecord(
                        Identifier(name=params[f"comp{idx + 1}_name"]),
                        Identifier(name=params[f"comp{idx_ii + 1}_name"]),
                        k_ij=k_ij,
                    )
                )

    return binary_records_lst


def _weightfracs_to_molefracs(wi_lst, molar_weights):
    """Convert weight fractions to mole fractions."""
    molar_weights = np.asarray(molar_weights)
    molar_weights_inv = 1 / molar_weights
    weightfracs = np.asarray(wi_lst)
    return (
        weightfracs
        * molar_weights_inv
        / np.sum(weightfracs * molar_weights_inv)
    )


def get_feos_npt_state(
    temperature, pressure, params, xi_lst=None, wi_lst=None
):
    """Create a feos State at T/P with composition in mole or weight fractions.

    Args:
        temperature: Temperature in Kelvin.
        pressure: Pressure in bar.
        params: Parameter dict with compX_ fields and binary interaction params
        xi_lst: Optional mole fraction list.
        wi_lst: Optional weight fraction list; overrides xi_lst if provided.
    """
    pure_params = _build_pure_records(temperature, params)

    molar_weights = [
        params[f"comp{idx + 1}_molarweight"]
        for idx in range(params["number_components"])
    ]

    binary_records_lst = _build_binary_records(temperature, params)

    # if weightfracs are given, transform into molfracs
    if wi_lst is not None:
        xi_lst = _weightfracs_to_molefracs(wi_lst, molar_weights)

    feos_params = Parameters.from_records(
        pure_params,
        binary_records=binary_records_lst,
    )

    eos = EquationOfState.pcsaft(feos_params)

    return State(
        eos,
        temperature=temperature * KELVIN,
        pressure=pressure * BAR,
        molefracs=np.array(xi_lst),
        density_initialization="liquid",
    )


def get_feos_pure_prop(temperature, pressure, params):
    """
    Return density (mol/m^3) and vapor pressure (Pa) for a pure component.
    """
    pure_params = _build_pure_records(temperature, params)

    feos_params = Parameters.from_records(
        pure_params,
    )

    eos = EquationOfState.pcsaft(feos_params)

    return (
        State(
            eos,
            temperature=temperature * KELVIN,
            pressure=pressure * BAR,
            density_initialization="liquid",
        ).density
        / (KILO * MOL / (METER**3))
    ), (PhaseEquilibrium.vapor_pressure(eos, temperature * KELVIN)[0] / PASCAL)


def is_density_iteration_sane(
    temperature, pressure, params, trials=100, eps=0.10
):
    """Basic sanity check for liquid-state convergence across compositions."""
    w_lst = np.linspace(0, 1.0, trials)

    pure_vols = np.array(
        [
            1
            / get_feos_npt_state(
                temperature, pressure, params, xi_lst=[1.0, 0.0]
            ).density,
            1
            / get_feos_npt_state(
                temperature, pressure, params, xi_lst=[0.0, 1.0]
            ).density,
        ]
    )

    for w in w_lst:
        mix_state = get_feos_npt_state(
            temperature, pressure, params, wi_lst=[w, 1.0 - w]
        )

        vol_id = np.sum(np.array(mix_state.molefracs) * pure_vols)

        if (vol_id - 1 / mix_state.density) / (1 / mix_state.density) > eps:
            return False

    return True
