"""Prepare PNIPAAm-water gel parameters and trace binary gel branches
in temperature."""

from __future__ import annotations


from equi_gel.newton_raphson_tracing import (
    build_initial_guess_binary,
    interpret_y_vector,
    trace_parameter,
)
from equi_gel.parameters import (
    add_mixing_association,
    scale_component_molarweight,
)


def prepare_gel_params(params_in: dict, y_crosslinker: float) -> dict:
    """Return a derived, ready-to-trace parameter dict for gel calculations."""
    params = params_in.copy()

    nf = params["NETWORK_FUNCTIONALITY"]
    c = (nf - 2.0) / nf

    m_desired = (2.0 / nf) * (1.0 - y_crosslinker) / y_crosslinker
    v_max = params["V_MAX_FAC"] * 0.94 * m_desired**2

    params["C"] = c * params["C_FAC"]
    params["max_volume_poly_frac"] = 1.0 / v_max

    wrong_m = params["comp1_m"]
    wrong_mw = params["comp1_molarweight"]
    molar_weight = (m_desired / wrong_m) * wrong_mw

    params = scale_component_molarweight(
        params, 1, molar_weight, in_place=False
    )
    params = add_mixing_association(params)
    return params


def trace_binodal_binary_gel(
    base_params: dict,
    *,
    t_high: float,
    t_low: float,
    pressure: float = 1.13,
    y_crosslinker: float | None = None,
    # tracing controls
    trace_step_size_init: float = 1e-8,
    trace_newton_tol: float = 1e-6,
    trace_critical_tol: float = 1e-5,
    trace_max_step_fraction: float = 1.0,
    verbosity: bool = False,
):
    """
    Trace the binary PNIPAAm-water gel binodal in temperature (trace_index=-1).
    Optionally detect VPT and trace third phase curve.

    Returns:
      (binodal_dict, third_curve_dict)
    """
    if y_crosslinker is None:
        y_crosslinker = base_params["y_crosslinker"]

    params = prepare_gel_params(
        base_params,
        y_crosslinker=y_crosslinker,
    )

    y_init = build_initial_guess_binary(t_high, pressure, params)

    y_lst, *_ = trace_parameter(
        trace_index=-1,
        y_init=y_init,
        trace_start=t_high,
        trace_target=t_low,
        pressure=pressure,
        params=params,
        step_size_init=trace_step_size_init,
        newton_tol=trace_newton_tol,
        critical_point_tol=trace_critical_tol,
        max_step_fraction=trace_max_step_fraction,
        verbosity=verbosity,
    )

    binodal_dict = interpret_y_vector(y_lst, pressure, params)

    return binodal_dict, y_lst
