"""Trace binary and ternary phase boundaries with predictor-corrector
Newton steps."""

import time
import numpy as np

from .alternating_tangents import find_binodal_by_alternating_tangents
from .interface_feos import get_feos_npt_state
from .newton_raphson import newton_raphson


_MIN_STEP_SIZE = 1e-25
_MAX_RESCALES = 40


def build_initial_guess_binary(
    temperature, pressure, params, reverse_order=False
):
    """Build initial Newton-Raphson vector for a binary system.

    Uses alternating tangents to estimate tie-line compositions.
    """
    original_number_components = params["number_components"]
    params["number_components"] = 2
    tie_line = find_binodal_by_alternating_tangents(
        temperature, pressure, params
    )
    tie_line = sorted(tie_line, reverse=reverse_order)

    x_guess = [
        get_feos_npt_state(
            temperature, pressure, params, wi_lst=[w, 1.0 - w]
        ).molefracs
        for w in tie_line
    ]
    k_init = x_guess[0] / x_guess[1]

    y_init = np.concatenate(
        [
            np.clip(np.log(x_guess[0]), -300, 300),
            np.clip(np.log(x_guess[1]), -300, 300),
            np.log(k_init),
            [np.log(temperature)],
        ]
    )
    params["number_components"] = original_number_components

    return y_init


# pylint: disable=too-many-arguments,too-many-positional-arguments,
# pylint: disable=too-many-locals,too-many-branches,too-many-statements


def trace_parameter(
    trace_index,
    y_init,
    trace_start,
    trace_target,
    pressure,
    params,
    fixed_index=None,
    step_size_init=0.5,
    desired_steps=5,
    newton_tol=1e-3,
    critical_point_tol=1e-2,
    stuck_progress_tol=0.0005,
    max_step_fraction=0.01,
    verbosity=False,
):
    """
    Trace a variable by stepping its log value and solving with Newton-Raphson.
    """
    if fixed_index is None:
        fixed_index = []

    tracing_sign = -1 if trace_start > trace_target else 1

    spec_idx = [trace_index] + fixed_index
    spec_val = [np.log(trace_start)] + [y_init[val] for val in fixed_index]

    y_vec_lst = []
    spec_val_lst = []
    spec_idx_lst = []
    step_sizes = []
    iter_amounts = []
    f_norm_lst = []

    step_size = step_size_init
    spec_val_old = -1
    no_components = int((len(y_init) - 1) / 3)

    rescale_counter = 0
    final_step = False

    y_init_old = y_init.copy()
    dy_ds_old = np.zeros_like(y_init, dtype=float)

    trace_current = trace_start

    start = time.perf_counter()

    while (tracing_sign * (trace_target - trace_current) > 0) or final_step:
        spec_val[0] = np.log(trace_current)

        try:
            y_init, dy_ds, n_iter, f_norm = newton_raphson(
                y_init=y_init,
                spec_idx=spec_idx,
                spec_val=spec_val,
                pressure=pressure,
                params=params,
                tol=newton_tol,
            )

            y_vec_lst.append(y_init.copy())
            spec_val_lst.append(spec_val.copy())
            spec_idx_lst.append(spec_idx.copy())
            step_sizes.append(step_size)
            iter_amounts.append(n_iter)
            f_norm_lst.append(f_norm)

            max_linear_step = max_step_fraction  # existing linear cap
            max_growth = 1.3

            # --- adaptive step control with linear-space cap ---

            step_size_adaptive = desired_steps / n_iter * step_size
            step_size_adaptive = min(
                step_size_adaptive, step_size * max_growth
            )

            # linear-space cap converted to log-space
            proposed_linear = trace_current + tracing_sign * max_linear_step

            if trace_current > 0 and proposed_linear > 0:
                max_log_step = abs(
                    np.log(proposed_linear) - np.log(trace_current)
                )
                step_size = min(step_size_adaptive, max_log_step)
            else:
                step_size = step_size_adaptive

            y_init_old = y_init.copy()
            dy_ds_old = dy_ds.copy()
            spec_val_old = spec_val[0]
            rescale_counter = 0

            if (
                np.linalg.norm(
                    np.exp(y_init[:no_components])
                    - np.exp(y_init[no_components : 2 * no_components])
                )
                < critical_point_tol
            ):
                if verbosity:
                    print("critical point identified")
                y_vec_lst.pop()
                spec_val_lst.pop()
                spec_idx_lst.pop()
                step_sizes.pop()
                iter_amounts.pop()
                f_norm_lst.pop()
                break

            if final_step:
                break

        except np.linalg.LinAlgError as exc:
            if verbosity:
                print("Singular matrix encountered - stopping calculation.")
                print(exc)
            break

        except RuntimeError as exc:
            if "Newton method did not converge" in str(exc):
                if verbosity:
                    print(exc)
                    print("Rescaling step size.")
                final_step = False
                step_size *= 0.5
                rescale_counter += 1
            elif "Invalid state in validate: density" in str(exc):
                if verbosity:
                    print("Newton failed with invalid state.")
                    print(exc)
                break
            else:
                raise

        step_size = max(step_size, _MIN_STEP_SIZE)
        spec_val[0] = spec_val_old + tracing_sign * step_size
        trace_current = np.exp(spec_val[0])

        if (tracing_sign * (trace_target - trace_current)) < 0:
            step_size = abs(np.log(trace_target) - spec_val_old)
            spec_val[0] = np.log(trace_target)
            trace_current = np.exp(spec_val[0])
            final_step = True

        if rescale_counter >= _MAX_RESCALES:
            raise RuntimeError("Tracing method got stuck by rescaling")

        if len(step_sizes) > 40:
            if np.sum(step_sizes[-39:]) < 40 * 1e-24:
                if verbosity:
                    print(
                        "Tracing method not moving forward:"
                        + " step size vanishing"
                    )
                break
            if (
                np.exp(y_vec_lst[-39][spec_idx_lst[-39][0]]) - trace_current
            ) / (trace_start - trace_target) < stuck_progress_tol:
                if verbosity:
                    print(
                        "Tracing method not moving forward:"
                        + " no tracing progress"
                    )
                break

        y_init = y_init_old + tracing_sign * dy_ds_old * step_size

    end = time.perf_counter()
    if verbosity:
        print(f"Execution time: {end - start:.6f} seconds")

    return (
        y_vec_lst,
        spec_val_lst,
        spec_idx_lst,
        step_sizes,
        iter_amounts,
        f_norm_lst,
    )


def add_component_to_y_vec(y_vec, x_i=1e-5, x_ii=1e-5):
    """Append a component to the Newton vector with default log fractions."""
    no_components = int((len(y_vec) - 1) / 3)

    y_vec = np.concatenate(
        [
            y_vec[:no_components],
            [np.log(x_i)],
            y_vec[no_components : 2 * no_components],
            [np.log(x_ii)],
            y_vec[2 * no_components : 3 * no_components],
            [np.log(x_i / x_ii)],
            [y_vec[-1]],
        ]
    )
    return y_vec


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


def interpret_y_vector(
    y_vec,
    pressure,
    params,
    step_sizes=None,
    iter_amounts=None,
    f_norm_lst=None,
):
    """Interpret a Newton-Raphson vector into compositions and states."""
    y = np.asarray(y_vec, dtype=float)

    was_single = y.ndim == 1
    if was_single:
        y = y[None, :]

    result = {"type": "calculation"}

    if len(y[0]) == 7:
        no_components = 2
    elif len(y[0]) == 10:
        no_components = 3
    elif len(y[0]) == 13:
        no_components = 4
    else:
        raise ValueError("Unexpected vector length for y_vec")

    params["number_components"] = no_components

    for i in range(no_components):
        result[f"comp{i+1}_name"] = params[f"comp{i+1}_name"]

    x_i_all, x_ii_all = [], []
    w_i_all, w_ii_all = [], []
    temperature_all = []

    for row in y:
        x_i = np.exp(row[:no_components])
        x_ii = np.exp(row[no_components : 2 * no_components])
        temperature = np.exp(row[-1])

        mix_state_i = get_feos_npt_state(
            temperature, pressure, params, xi_lst=x_i
        )
        mix_state_ii = get_feos_npt_state(
            temperature, pressure, params, xi_lst=x_ii
        )

        x_i_all.append(x_i)
        x_ii_all.append(x_ii)
        w_i_all.append(mix_state_i.massfracs())
        w_ii_all.append(mix_state_ii.massfracs())
        temperature_all.append(temperature)

    if was_single:
        result |= {
            "x_I_lst": x_i_all[0],
            "x_II_lst": x_ii_all[0],
            "w_I_lst": w_i_all[0],
            "w_II_lst": w_ii_all[0],
            "temperature_K": temperature_all[0],
        }
        if f_norm_lst:
            result |= {"f_norm": f_norm_lst[0]}
        if step_sizes:
            result |= {"step_sizes": step_sizes[0]}
        if iter_amounts:
            result |= {"iter_amounts": iter_amounts[0]}
    else:
        result |= {
            "x_I_lst": x_i_all,
            "x_II_lst": x_ii_all,
            "w_I_lst": w_i_all,
            "w_II_lst": w_ii_all,
            "temperature_K": temperature_all,
        }
        if f_norm_lst:
            result |= {"f_norm": f_norm_lst}
        if step_sizes:
            result |= {"step_sizes": step_sizes}
        if iter_amounts:
            result |= {"iter_amounts": iter_amounts}

    result["params"] = params

    return result
