"""Newton-Raphson residuals and Jacobians for liquid and gel phase
equilibria."""

import numpy as np
from si_units import MOL, KELVIN

from .interface_feos import get_feos_npt_state
from .elastic_contribution import (
    flory_rehner_dad_v,
    miao_dad_v,
    phantom_dad_v,
)


_MAX_NEWTON_ITERS = 30
_MAX_BACKTRACK = 10


def _elastic_dad_v(params, partial_vol_polymer, vol_frac_polymer):
    elasticity = params["elasticity"]
    if elasticity == "Flory":
        return flory_rehner_dad_v(
            partial_molar_vol_pol=partial_vol_polymer,
            elasticity_c=params["C"],
            vol_frac_pol_0=params["ref_volume_frac"],
            vol_frac_pol=vol_frac_polymer,
        )
    if elasticity == "Miao":
        return miao_dad_v(
            partial_molar_vol_pol=partial_vol_polymer,
            elasticity_c=params["C"],
            vol_frac_pol_0=params["ref_volume_frac"],
            vol_frac_pol=vol_frac_polymer,
            vol_frac_pol_max=params["max_volume_poly_frac"],
        )
    if elasticity == "Phantom":
        return phantom_dad_v(
            partial_molar_vol_pol=partial_vol_polymer,
            elasticity_c=params["C"],
            vol_frac_pol_0=params["ref_volume_frac"],
            vol_frac_pol=vol_frac_polymer,
        )
    raise ValueError(f"Unknown elasticity model: {elasticity}")


# pylint: disable=too-many-locals


def estimate_elastic_fugacity_derivatives(
    temperature, pressure, params, state, dev=0.01, return_press=False
):
    """Estimate elastic contribution and its derivatives vs ln(x_j)."""
    rho_polymer = state.partial_density[0]
    partial_vol_polymer = state.partial_molar_volume()[0]
    partial_vol_species = state.partial_molar_volume()
    vol_frac_polymer = partial_vol_polymer * rho_polymer

    ref_moles = state.moles

    d_a_elast_dv = _elastic_dad_v(
        params, partial_vol_polymer, vol_frac_polymer
    )
    f_elast_ref = partial_vol_species * d_a_elast_dv

    d_fi_d_lnxj = np.zeros((len(f_elast_ref), len(f_elast_ref)), dtype=float)

    for idx_j, moles in enumerate(ref_moles):
        new_moles = (ref_moles / MOL).copy() * MOL
        new_moles[idx_j] = moles * (1 + dev)
        total_moles = np.sum(new_moles / MOL) * MOL
        x_lst = np.array([m / total_moles for m in new_moles])

        dev_state = get_feos_npt_state(
            temperature, pressure, params, xi_lst=x_lst
        )

        dev_rho_polymer = dev_state.partial_density[0]
        dev_partial_vol_polymer = dev_state.partial_molar_volume()[0]
        dev_partial_vol_species = dev_state.partial_molar_volume()
        dev_vol_frac_polymer = dev_partial_vol_polymer * dev_rho_polymer

        dev_d_a_elast_dv = _elastic_dad_v(
            params, dev_partial_vol_polymer, dev_vol_frac_polymer
        )

        dev_f_elast = dev_partial_vol_species * dev_d_a_elast_dv

        d_fi_d_lnxj[:, idx_j] = ref_moles[idx_j] * (
            (dev_f_elast - f_elast_ref) / (new_moles[idx_j] - ref_moles[idx_j])
        )

    if return_press:
        return (
            f_elast_ref,
            d_fi_d_lnxj,
            partial_vol_species,
            partial_vol_polymer,
        )

    return f_elast_ref, d_fi_d_lnxj


# pylint: disable=too-many-locals


def get_residuals_and_jacobian(y_vec, spec_idx, spec_val, pressure, params):
    """Return residual vector and Jacobian for Newton-Raphson solver."""
    no_components = params["number_components"]

    log_gel_i = y_vec[:no_components]
    x_gel_i = np.exp(log_gel_i)
    log_gel_ii = y_vec[no_components : 2 * no_components]
    x_gel_ii = np.exp(log_gel_ii)
    lnki_values = y_vec[2 * no_components : -1]
    temperature = np.exp(y_vec[-1])

    mix_state_i = get_feos_npt_state(
        temperature, pressure, params, xi_lst=x_gel_i
    )

    mix_state_ii = get_feos_npt_state(
        temperature, pressure, params, xi_lst=x_gel_ii
    )

    ln_phi_i = mix_state_i.ln_phi()
    dln_phi_dnj_i = mix_state_i.dln_phi_dnj()
    dln_phi_dt_i = mix_state_i.dln_phi_dt()
    nt_i = mix_state_i.total_moles

    ln_phi_ii = mix_state_ii.ln_phi()
    dln_phi_dnj_ii = mix_state_ii.dln_phi_dnj()
    dln_phi_dt_ii = mix_state_ii.dln_phi_dt()
    nt_ii = mix_state_ii.total_moles

    residuals = np.zeros_like(y_vec)

    residuals[:no_components] = lnki_values - ln_phi_ii + ln_phi_i

    residuals[no_components : 2 * no_components] = (
        lnki_values - log_gel_i + log_gel_ii
    )
    residuals[2 * no_components] = np.sum(x_gel_i) - 1
    residuals[2 * no_components + 1] = np.sum(x_gel_ii) - 1
    residuals[2 * no_components + 2 :] = y_vec[spec_idx] - spec_val

    jacobian = np.zeros((len(y_vec), len(y_vec)))

    # (5.37): J_{i,j}
    for i in range(no_components):
        for j in range(no_components):
            jacobian[i, j] = x_gel_i[j] * (nt_i * dln_phi_dnj_i)[i, j]

    # (5.38): J_{i,NC+j}
    for i in range(no_components):
        for j in range(no_components):
            jacobian[i, no_components + j] = (
                -1 * x_gel_ii[j] * (nt_ii * dln_phi_dnj_ii)[i, j]
            )

    # (5.39): J_{i,2NC+i}
    for i in range(no_components):
        jacobian[i, 2 * no_components + i] = 1

    # (5.40): J_{NC+i, i}
    for i in range(no_components):
        jacobian[no_components + i, i] = -1

    # (5.41): J_{NC+i,NC+i}
    for i in range(no_components):
        jacobian[no_components + i, no_components + i] = 1

    # (5.42): J_{2NC+i,NC+i}
    for i in range(no_components):
        jacobian[no_components + i, 2 * no_components + i] = 1

    # (5.43): J_{2NC+1,j}
    for j in range(no_components):
        jacobian[2 * no_components, j] = x_gel_i[j]

    # (5.44): J_{2NC+2,NC+j}
    for j in range(no_components):
        jacobian[2 * no_components + 1, no_components + j] = x_gel_ii[j]

    # (5.45): J_{2NC+3, -1} (specification_row)
    for idx, s_index in enumerate(spec_idx):
        jacobian[2 * no_components + 2 + idx, s_index] = 1

    # (5.46): temperature derivatives
    dln_phi_dt_diff = dln_phi_dt_i - dln_phi_dt_ii
    for i in range(no_components):
        jacobian[i, -1] = temperature * KELVIN * dln_phi_dt_diff[i]

    if params["gel"]:
        f_elast_num, jacobian_elast_num = (
            estimate_elastic_fugacity_derivatives(
                temperature, pressure, params, mix_state_ii
            )
        )

        f_elast = np.zeros_like(y_vec)
        f_elast[:no_components] = f_elast_num

        jacobian_elast = np.zeros((len(y_vec), len(y_vec)))

        # (5.38): J_{i,NC+j} elasticity only important for gel phase (II)
        for i in range(no_components):
            for j in range(no_components):
                jacobian_elast[i, no_components + j] = jacobian_elast_num[i, j]

        residuals = residuals - f_elast
        jacobian = jacobian - jacobian_elast

        # delete row for polymer fugacities and polymer partition coeff
        residuals[0] = 0.0
        residuals[no_components] = 0.0

    return residuals, jacobian


# pylint: disable=too-many-arguments,too-many-positional-arguments


def newton_raphson(
    y_init, spec_idx, spec_val, pressure, params, tol=1e-4, verbosity=False
):
    """Solve the coupled equilibrium equations with backtracking."""
    y_vec = y_init.copy()

    residuals, jacobian = get_residuals_and_jacobian(
        y_vec=y_vec,
        spec_idx=spec_idx,
        spec_val=spec_val,
        pressure=pressure,
        params=params,
    )

    norm_f = np.linalg.norm(residuals)

    for i in range(_MAX_NEWTON_ITERS):
        if norm_f < tol:
            break

        delta_y = np.linalg.solve(jacobian, residuals)

        for j in range(_MAX_BACKTRACK):
            candidate_y = y_vec - delta_y / (j + 1)

            f_candidate, jacobian_candidate = get_residuals_and_jacobian(
                y_vec=candidate_y,
                spec_idx=spec_idx,
                spec_val=spec_val,
                pressure=pressure,
                params=params,
            )

            norm_f_candidate = np.linalg.norm(f_candidate)

            if norm_f_candidate < norm_f:
                break

        y_vec = candidate_y
        jacobian = jacobian_candidate
        residuals = f_candidate
        norm_f = norm_f_candidate

    if (i >= _MAX_NEWTON_ITERS - 1) or (norm_f > tol):
        raise RuntimeError(
            "Newton method did not converge within set number of iterations"
        )

    # get sensitivity
    jacobian_inv = np.linalg.inv(jacobian)
    df_ds = np.zeros_like(residuals, dtype=float)

    # tracing element: first specification in lst
    no_components = params["number_components"]
    df_ds[2 * no_components + 2] = -1

    dy_ds = -1 * jacobian_inv.dot(df_ds)

    if verbosity:
        print(f"T={np.exp(spec_val[0])}; Iter {i}; Final f_norm {norm_f}")

    return y_vec, dy_ds, (i + 1), norm_f
