"""Prepare gel parameters, trace binary gel seeds, and continue them
into ternary binodals."""

from __future__ import annotations

import numpy as np

from equi_gel.newton_raphson_tracing import (
    add_component_to_y_vec,
    build_initial_guess_binary,
    interpret_y_vector,
    trace_parameter,
)
from equi_gel.parameters import (
    add_mixing_association,
    get_subset_of_parameters,
    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_gel(
    base_params: dict,
    *,
    t_high: float,
    t_low: float,
    pressure: float = 1.13,
    y_crosslinker: float | None = None,
    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,
    ternary_step_size_init: float = 1e-3,
    ternary_newton_tol: float = 1e-3,
    ternary_max_step_fraction: float = 0.01,
    stuck_progress_tol: float | None = None,
    verbosity: bool = False,
    y_init: np.ndarray | None = None,
) -> dict:
    """
    Trace PNIPAAm-water-acetic acid gel binodal:
    1) binary trace in T (components 1-2)
    2) ternary trace in composition (trace_index=2)
    Returns the interpreted binodal dict.
    """
    if y_crosslinker is None:
        y_crosslinker = base_params["Y_CROSSLINKER"]

    params = prepare_gel_params(base_params, y_crosslinker=y_crosslinker)
    params_binary = get_subset_of_parameters(params.copy(), [1, 2])

    if y_init is None:
        y_init = build_initial_guess_binary(t_high, pressure, params_binary)

        y_lst, *_ = trace_parameter(
            trace_index=-1,
            y_init=y_init,
            trace_start=t_high,
            trace_target=t_low,
            pressure=pressure,
            params=params_binary,
            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,
        )

        # sanity: did we actually reach near t_low?
        if abs(np.exp(y_lst[-1][-1]) - t_low) > 3.0:
            raise RuntimeError(
                f"T-trace did not reach target. "
                f"Got T={np.exp(y_lst[-1][-1]):.3f}, target={t_low:.3f}"
            )

        y_init_ternary = add_component_to_y_vec(y_lst[-1].copy(), 1e-5, 1e-5)

    else:
        y_init_ternary = add_component_to_y_vec(y_init.copy(), 1e-5, 1e-5)

    kwargs = {}
    if stuck_progress_tol is not None:
        kwargs["stuck_progress_tol"] = stuck_progress_tol

    y_lst, *_ = trace_parameter(
        trace_index=2,
        y_init=y_init_ternary,
        trace_start=1e-5,
        trace_target=1.0 - 1e-5,
        fixed_index=[-1],
        pressure=pressure,
        params=params,
        step_size_init=ternary_step_size_init,
        newton_tol=ternary_newton_tol,
        critical_point_tol=trace_critical_tol,
        max_step_fraction=ternary_max_step_fraction,
        verbosity=verbosity,
        **kwargs,
    )

    return interpret_y_vector(y_lst, pressure, params)
