"""Distance metrics between experimental data and binary or ternary
binodals."""

import numpy as np

SQRT3 = np.sqrt(3.0)


def ternary_pct(data, key, idx_order):
    """
    Return (N, 3) mass fractions in percent for a specific component order.
    """
    return 100 * np.array([np.take(v, idx_order) for v in data[key]])


def binary_pct(data, key, idx_order, t_scale=1.0, which_idx=0):
    """Return (N, 2) array for binary data: [T/T_scale, 100 * w_i]."""
    w = np.array([np.take(v, idx_order) for v in data[key]], dtype=float)
    w_pct = 100.0 * w[:, which_idx]
    temperature = np.asarray(data["temperature_K"], dtype=float)
    return np.column_stack([temperature / t_scale, w_pct])


def ternary_to_xy(w_frac):
    """Project ternary composition (N, 3) into 2D Cartesian coords."""
    _, b, c = w_frac[:, 0], w_frac[:, 1], w_frac[:, 2]
    x = b + 0.5 * c
    y = (SQRT3 / 2.0) * c
    return np.column_stack([x, y])


def resample_polyline(points, n=120):
    """Resample a polyline to n points by arclength."""
    if len(points) < 2:
        return points.copy()

    seg = np.linalg.norm(np.diff(points, axis=0), axis=1)
    s = np.concatenate([[0.0], np.cumsum(seg)])
    if s[-1] == 0:
        return points[:1].copy()

    s_new = np.linspace(0.0, s[-1], n)
    return np.column_stack(
        [np.interp(s_new, s, points[:, k]) for k in range(points.shape[1])]
    )


def resample_binodal(binodal_pct, n=120):
    """Resample binodal data in percent and return both pct and XY points."""

    resampled = resample_polyline(binodal_pct, n=n)

    if binodal_pct.shape[1] == 3:
        xy_resampled = ternary_to_xy(resampled)
    else:
        xy_resampled = resampled

    return resampled, xy_resampled


def min_dist2_points_to_polyline(points_xy, curve_xy):
    """Return minimal squared distance from each point to a polyline."""
    if len(curve_xy) < 2:
        return np.full(points_xy.shape[0], np.inf)

    a = curve_xy[:-1]
    b = curve_xy[1:]
    ab = b - a
    inv_ab2 = 1.0 / np.sum(ab * ab, axis=1)

    p = points_xy[:, None, :]
    ap = p - a[None, :, :]

    t = np.sum(ap * ab[None, :, :], axis=2) * inv_ab2[None, :]
    t = np.clip(t, 0.0, 1.0)
    proj = a[None, :, :] + t[:, :, None] * ab[None, :, :]

    d2 = np.sum((p - proj) ** 2, axis=2)
    return np.min(d2, axis=1)


def _collect_exp_points_ternary(dict_exp_lst, permutation_list):
    exp_points = []
    for dict_idx, current_dict in enumerate(dict_exp_lst):
        comp_idx = permutation_list[1][dict_idx]
        phase_i = ternary_pct(current_dict, "w_I_lst", comp_idx)
        exp_points.append(phase_i)
        if current_dict["type"] == "tie_line":
            phase_ii = ternary_pct(current_dict, "w_II_lst", comp_idx)
            exp_points.append(phase_ii)
    return np.vstack(exp_points)


def _collect_exp_points_binary(
    dict_exp_lst, permutation_list, t_scale, which_idx
):
    exp_points = []
    for dict_idx, current_dict in enumerate(dict_exp_lst):
        comp_idx = permutation_list[1][dict_idx]
        phase_i = binary_pct(
            current_dict, "w_I_lst", comp_idx, t_scale, which_idx=which_idx
        )
        exp_points.append(phase_i)
        if current_dict["type"] == "tie_line":
            phase_ii = binary_pct(
                current_dict,
                "w_II_lst",
                comp_idx,
                t_scale,
                which_idx=which_idx,
            )
            exp_points.append(phase_ii)
    return np.vstack(exp_points)


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


def get_minimal_distance(
    dict_binodal_calc,
    dict_exp_lst,
    permutation_list,
    binodal_sampling=20,
    t_scale=1.0,
    which_idx_binary=0,
):
    """
    Compute mean squared distance from experimental points to binodal curves.
    """
    no_components = dict_binodal_calc["params"]["number_components"]
    comp_idx_calc = permutation_list[0][0]

    if no_components == 3:
        phase_i_calc = ternary_pct(dict_binodal_calc, "w_I_lst", comp_idx_calc)
        phase_ii_calc = ternary_pct(
            dict_binodal_calc, "w_II_lst", comp_idx_calc
        )
        exp_pct = _collect_exp_points_ternary(dict_exp_lst, permutation_list)
        exp_xy = ternary_to_xy(exp_pct)
    elif no_components == 2:
        phase_i_calc = binary_pct(
            dict_binodal_calc,
            "w_I_lst",
            comp_idx_calc,
            t_scale,
            which_idx=which_idx_binary,
        )
        phase_ii_calc = binary_pct(
            dict_binodal_calc,
            "w_II_lst",
            comp_idx_calc,
            t_scale,
            which_idx=which_idx_binary,
        )
        exp_pct = _collect_exp_points_binary(
            dict_exp_lst, permutation_list, t_scale, which_idx_binary
        )
        exp_xy = exp_pct
    else:
        raise ValueError("Only binary or ternary systems are supported")

    phase_i_res, phase_i_xy = resample_binodal(
        phase_i_calc, n=binodal_sampling
    )
    phase_ii_res, phase_ii_xy = resample_binodal(
        phase_ii_calc, n=binodal_sampling
    )

    dist2_i = min_dist2_points_to_polyline(exp_xy, phase_i_xy)
    dist2_ii = min_dist2_points_to_polyline(exp_xy, phase_ii_xy)
    avg_dist_sq = np.mean(np.minimum(dist2_i, dist2_ii))

    return avg_dist_sq, phase_i_res, phase_ii_res


def get_minimal_distance_swelling(
    dict_binodal_calc,
    dict_exp,
):
    """
    Compute mean squared distance from experimental points to binodal curves.
    """
    no_components = dict_binodal_calc["params"]["number_components"]

    if no_components == 3:
        dos_exp = np.asarray(dict_exp["DOS_g_per_g"])
        w_3_exp = np.asarray(dict_exp["w_3_g_per_g"])

        dos_calc = np.asarray(
            [1 / w[0] for w in dict_binodal_calc["w_II_lst"]]
        )
        w_3_calc = np.asarray([w[2] for w in dict_binodal_calc["w_I_lst"]])

        # check for the boundaries, no extrapolation if calc got stuck
        if (max(w_3_exp) - max(w_3_calc)) > 1e-3:
            return 10000 * (max(w_3_exp) - max(w_3_calc))

        # check for the boundaries, no extrapolation if calc got stuck
        if -1 * (min(w_3_exp) - min(w_3_calc)) > 1e-3:
            return 10000 * (-1 * (min(w_3_exp) - min(w_3_calc)))

        # sort by w_3_calc (np.interp needs ascending x)
        idx = np.argsort(w_3_calc)
        x = w_3_calc[idx]
        y = dos_calc[idx]

        # interpolate DOS at experimental w_3 values
        dos_calc_at_exp = np.interp(w_3_exp, x, y)

    elif no_components == 2:

        dos_exp = np.asarray(dict_exp["DOS_g_per_g"])
        t_exp = np.asarray(dict_exp["temperature_K"])

        dos_calc = np.asarray(
            [1.0 / w[0] for w in dict_binodal_calc["w_II_lst"]]
        )
        t_calc = np.asarray(dict_binodal_calc["temperature_K"])

        # sort by w_3_calc (np.interp needs ascending x)
        idx = np.argsort(t_calc)
        x = t_calc[idx]
        y = dos_calc[idx]

        # interpolate DOS at experimental w_3 values
        dos_calc_at_exp = np.interp(t_exp, x, y)

    else:
        raise ValueError("Only binary and ternary systems are supported")

    r = (dos_calc_at_exp - dos_exp) / dos_exp
    return float(np.sqrt(np.mean(r * r)))
