"""Set of helper utilities to work with single component system"""

from __future__ import annotations

from typing import TYPE_CHECKING, cast

import numpy as np

from .lnpiseries import lnPiCollection
from .utils import RootResultDict, array_to_scalar, rootresults_to_rootresultdict

if TYPE_CHECKING:
    from typing import Any, Mapping, Sequence

    import xarray as xr
    from numpy.typing import ArrayLike

    from ._typing import MyNDArray
    from .lnpidata import lnPiMasked
    from .segment import BuildPhasesBase


def tag_phases_singlecomp(x: Sequence[lnPiMasked]) -> Sequence[int] | MyNDArray:
    """
    Function to tag phases with a unique id.

    This is for analyzing a single component system.
    If multiple phases passed (length of input > 1), then
    sort by ``argmax = x.local_argmax()``, i.e., the location of the maxima.
    Otherwise, assign `phase_id = 0` if ``argmax < Nmax // 2`` where ``Nmax`` is
    the maximum number of particles in lnPi.

    Parameters
    ----------
    x : sequence of lnPiMasked
        lnPi objects to be tagged.

    Returns
    -------
    phase_id : sequence of int
        Phase code for each element in ``x``.
    """
    if len(x) > 2:
        msg = "bad tag function"
        raise ValueError(msg)

    argmax0 = np.array([xx.local_argmax()[0] for xx in x])
    if len(x) == 2:
        return np.argsort(argmax0)

    return np.where(argmax0 <= x[0].shape[0] / 2, 0, 1)


def get_lnz_min(
    target: float,
    collection: lnPiCollection,
    build_phases: BuildPhasesBase,
    phase_id: int | None = 0,
    component: int | None = None,
    dlnz: float = 0.5,
    dfac: float = 1.0,
    ref: lnPiMasked | None = None,
    build_kws: Mapping[str, Any] | None = None,
    ntry: int = 20,
    solve_kws: Mapping[str, Any] | None = None,
) -> tuple[lnPiCollection, RootResultDict]:
    """
    Solve for target density

    Parameters
    ----------
    target : float
        target density of specified component
    collection : CollectionPhases
        initial guess to work from.  This is assumed to be sorted in ascending `lnz` order
    build_phases : callable
        function to build new phases objects
    phase_id : int, default=0
        phase id tag to consider.  If None, consider min over 'phases'
    component : int, optional
        component to consider.
        If component is None, use build_phases.index.
        If component=='None', then consider total density
    dlnz : float, default=0.5
        how to change the chemical potential if `collection` doesn't already
        bracket solution
    ref : lnPiMasked, optional
        optional lnPiMasked object to pass to build_phases
    build_kws : dictionary, optional
        optional arguments to `build_phases`
    ntry : int, default=20
        number of times to attempted finding left/right bracketing objects
    solve_kws : dictionary, optional
        optional arguments to `scipy.optimize.brentq`

    Returns
    -------
    phases : Phases object as solution lnz
        solution_parameters : optional
    info : :class:`RootResultDict`
    """
    from scipy.optimize import brentq

    if phase_id is not None and phase_id not in collection.index.get_level_values(
        "phase"
    ):
        msg = f"no phase {phase_id}"
        raise ValueError(msg)

    lnz_idx = build_phases.index
    if isinstance(component, str) and component.lower() == "none":

        def getter_comp(x: lnPiCollection) -> xr.DataArray:
            return x.xge.dens_tot

    else:
        component = component or lnz_idx

        def getter_comp(x: lnPiCollection) -> xr.DataArray:
            return x.xge.dens.sel(component=component)

    if phase_id is None:

        def getter(x: lnPiCollection) -> xr.DataArray:
            return getter_comp(x).min("phase")

    else:

        def getter(x: lnPiCollection) -> xr.DataArray:
            return getter_comp(x).sel(phase=phase_id)

    s = getter(collection).to_series().dropna()
    # builder
    if build_kws is None:
        build_kws = {}

    # left bracket
    left = None
    ss = s[s < target]
    if len(ss) > 0:
        left = collection.mloc[ss.index[[-1]]]
    else:
        # TODO(wpk): Need to make collection work better.  Might need to just have a single class for everything...
        # reveal_type(collection.mloc[s.index[[0]]])
        new_lnz = collection.mloc[s.index[[0]]]._get_lnz(lnz_idx)
        dlnz_left = dlnz
        for _i in range(ntry):
            new_lnz -= dlnz_left
            p = build_phases(new_lnz, ref=ref, **build_kws)
            if phase_id in p._get_level("phase") and getter(p).to_numpy() < target:
                left = p
                break
            dlnz_left *= dfac

    if left is None:
        msg = "could not find left bounds"
        raise RuntimeError(msg)

    # right bracket
    right = None
    ss = s[s > target]
    if len(ss) > 0:
        right = collection.mloc[ss.index[[0]]]
    else:
        new_lnz = collection.mloc[s.index[[-1]]]._get_lnz(lnz_idx)
        dlnz_right = dlnz

        for _i in range(ntry):
            new_lnz += dlnz_right
            p = build_phases(new_lnz, ref=ref, **build_kws)

            if phase_id not in p._get_level("phase"):
                # went to far
                new_lnz -= dlnz_right
                # reset to half dlnz
                dlnz_right = dlnz_right * 0.5
            elif getter(p).to_numpy() > target:
                right = p
                break
            else:
                dlnz_right *= dfac

    if right is None:
        msg = "could not find right bounds"
        raise RuntimeError(msg)

    def f(x: float) -> float:
        lnz_new = x
        p = build_phases(lnz_new, ref=ref, **build_kws)
        f.lnpi = p  # type: ignore[attr-defined]
        return array_to_scalar(getter(p).values) - target

    a, b = sorted([x._get_lnz(lnz_idx) for x in [left, right]])

    xx, r = brentq(f, a, b, full_output=True, **(solve_kws or {}))

    return f.lnpi, rootresults_to_rootresultdict(r, residual=f(xx))  # type: ignore[attr-defined]


def get_lnz_max(
    edge_distance_min: int,
    build_phases: BuildPhasesBase,
    ref: lnPiMasked | None = None,
    lnz_start: float | None = None,
    collection: lnPiCollection | None = None,
    dlnz: float = 0.5,
    dfac: float = 1.0,
    threshold_abs: float = 1e-4,
    ntry: int = 50,
    build_kws: Mapping[str, Any] | None = None,
) -> tuple[lnPiCollection, dict[str, int | float]]:
    """Find max lnz by bisection"""

    build_kws = build_kws or {}
    ref = ref or build_phases.phase_creator.ref

    if ref is None:
        msg = "must specify `ref` or build_phases must have access to reference lnPiMasked object"
        raise ValueError(msg)

    lnz_idx = build_phases.index
    lnz_start = cast(float, lnz_start or ref.lnz[lnz_idx])

    # need left/right bounds
    # left is greatest lnz point with edge_distance > edge_distance_min
    # right is least lnz point with edge_distance < edge_distance_min
    left: lnPiCollection | None = None
    right: lnPiCollection | None = None
    lnz_left: float | None = None
    lnz_right: float | None = None
    n_left: int | None = None
    n_right: int | None = None

    def getter(p: lnPiCollection) -> xr.DataArray:
        v = p.xge.edge_distance(ref)
        if not p._xarray_unstack:
            v = v.unstack(p._concat_dim)  # noqa: PD010
        return v.min("phase")

    # if have collection, try to use it
    if collection is not None:
        # see if bound contained in collection

        s = getter(collection).to_series()

        # left
        ss = s[s > edge_distance_min]
        if len(ss) > 0:
            left = collection.mloc[ss.index[[-1]]]
        else:
            lnz_left = collection.zloc[[0]]._get_lnz(lnz_idx)

        # right
        ss = s[s < edge_distance_min]
        if len(ss) > 0:
            right = collection.mloc[ss.index[[0]]]
        else:
            lnz_right = collection.zloc[[-1]]._get_lnz(lnz_idx)

    # if left not set, try to find it
    if left is None:
        if lnz_left is None:
            lnz_left = lnz_start
        dlnz_loc = dlnz
        for i in range(ntry):
            lnz_left -= dlnz_loc
            p = build_phases(lnz_left, ref=ref, **build_kws)
            if getter(p).to_numpy() >= edge_distance_min:
                left = p
                n_left = i
                break
            # grow dlnz
            dlnz_loc *= dfac

    if left is None:
        msg = "could not find left bound"
        raise RuntimeError(msg)

    # if right not set, try to find it
    if right is None:
        if lnz_right is None:
            lnz_right = lnz_start

        dlnz_loc = dlnz
        for i in range(ntry):
            lnz_right += dlnz_loc
            p = build_phases(lnz_right, ref=ref, **build_kws)

            if getter(p).to_numpy() < edge_distance_min:
                right = p
                n_right = i
                break
            # grow dlnz
            dlnz_loc *= dfac

    if right is None:
        msg = "could not find right bound"
        raise RuntimeError(msg)

    # not do bisection
    bracket = [left, right]
    values = [getter(x).to_numpy() for x in bracket]

    for i in range(ntry):  # noqa: B007
        lnz = [x._get_lnz(lnz_idx) for x in bracket]
        delta = np.abs(lnz[1] - lnz[0])
        if delta < threshold_abs:
            break
        lnz_mid = 0.5 * (lnz[0] + lnz[1])
        mid = build_phases(lnz_mid, ref=ref, **build_kws)
        y_mid = getter(mid).to_numpy()

        index = 0 if y_mid >= edge_distance_min else 1

        bracket[index] = mid
        values[index] = y_mid

    left, right = bracket

    info = {"ntry": i, "ntry_left": n_left, "ntry_right": n_right, "precision": delta}
    return left, info


def build_grid(
    x: ArrayLike | MyNDArray | None = None,
    dx: float | None = None,
    x_range: Sequence[float] | MyNDArray | None = None,
    x0: float | None = None,
    offsets: Sequence[float] | MyNDArray | None = None,
    even_grid: bool = False,
    digits: int | None = None,
    unique: bool = True,
    outlier: bool = False,
) -> MyNDArray:
    """
    Build a grid of values


    Parameters
    ----------
    x : array, optional
        if passes, use this as the base array
    dx : float, optional
        delta x. See below for uses
    x_range : array of length 2, optional
        if x not specified, make an array with
        `x=np.arange(x_range[0], x_range[1] + dx/2, dx)`
    x0 : float, optional
        see offsets
    offsets : array of length 2, optional
        if x or x_range not specified, set x_range = x0 + offsets
    even_grid : bool, default=True
        if True, make sure all values are multiples of dx
    digits : int, optional
        if passes, round output to this many digits
    outlier : bool, default=False
        if outlier is True, allow values outsize x

    Returns
    -------
    out : array of values
    """

    if x is None:
        if dx is None:
            msg = "Must specify dx"
            raise ValueError(msg)

        if x_range is not None:
            x_range = np.asarray(x_range)
        elif offsets is not None and x0 is not None:
            x_range = x0 + np.asarray(offsets)
        else:
            msg = "must specify x_range or x0 and offsets"
            raise ValueError(msg)

        if len(x_range) != 2:
            raise ValueError

        # for consistency with previous code.
        if even_grid:
            new_range = np.round(x_range / dx) * dx
            if new_range[0] < x_range[0]:
                new_range[0] += dx
            if new_range[1] > x_range[1]:
                new_range[1] -= dx
            x_range = new_range
        x = np.arange(x_range[0], x_range[1] + dx * 0.5, dx)  # type: ignore[index]
    else:
        x = np.asarray(x)

    if not outlier:
        lb, ub = x.min(), x.max()

    if even_grid:
        if dx is None:
            raise TypeError
        x = np.round(x / dx) * dx

    if digits is not None:
        x = np.round(x, digits)  # type: ignore[arg-type]

    if unique:
        x = np.unique(x)  # type: ignore[arg-type]

    if not outlier:
        x = x[(lb <= x) & (x <= ub)]  # type: ignore[index]

    return x  # type: ignore[return-value]


def limited_collection(
    build_phases: BuildPhasesBase,
    dlnz: float,
    lnz_range: Sequence[float] | MyNDArray | None = None,
    offsets: list[int] | None = None,
    digits: int | None = None,
    even_grid: bool = True,
    course_step: int | None = None,
    edge_distance_min: int | None = None,
    dens_min: float | None = None,
    lnz_min_kws: Mapping[str, Any] | None = None,
    lnz_max_kws: Mapping[str, Any] | None = None,
    ref: lnPiMasked | None = None,
    build_kws: Mapping[str, Any] | None = None,
    # build_stability_kws: Mapping[str, Any] | None = None,
    nmax: int | None = None,
    xarray_output: bool = True,
    collection_kws: Mapping[str, Any] | None = None,
    limit_course: bool = False,
) -> tuple[lnPiCollection, lnPiCollection]:
    """
    Build a lnPiCollection over a range of lnz values

    Parameters
    ----------
    build_phases :
    """

    if lnz_range is None:
        ref = ref or build_phases.phase_creator.ref
        if ref is None:
            msg = "Must pass in ref or build_phases must have access to reference lnPiMasked object"
            raise ValueError(msg)
        x0 = ref.lnz[build_phases.index]
    else:
        x0 = None

    # TODO: update tests to get rid of outlier=True  # noqa: TD002
    lnzs = build_grid(
        dx=dlnz,
        x_range=lnz_range,
        x0=x0,
        offsets=offsets,
        even_grid=even_grid,
        digits=digits,
        outlier=True,
    )

    if course_step is None:
        course_step = max(len(lnzs) // 40, 1)
    if collection_kws is None:
        collection_kws = {}

    # limit lnz
    c_course = None
    lnz_min, lnz_max = lnzs[0], lnzs[-1]
    if edge_distance_min is not None or dens_min is not None:
        c_course = lnPiCollection.from_builder(
            lnzs[::course_step],
            build_phases=build_phases,
            build_kws=build_kws,
            nmax=nmax,
            xarray_output=xarray_output,
            **collection_kws,
        )

        if dens_min is not None:
            try:
                if lnz_min_kws is None:
                    lnz_min_kws = {}
                p_min, o = get_lnz_min(
                    dens_min, c_course, build_phases, build_kws=build_kws, **lnz_min_kws
                )
                if o["converged"]:
                    lnz_min = p_min.iloc[0].lnz[build_phases.index] - dlnz
            except Exception:  # noqa: BLE001, S110
                pass

        if edge_distance_min is not None:
            try:
                if lnz_max_kws is None:
                    lnz_max_kws = {}
                p_max, _ = get_lnz_max(
                    edge_distance_min,
                    build_phases,
                    build_kws=build_kws,
                    collection=c_course,
                    **lnz_max_kws,
                )
                lnz_max = p_max.iloc[0].lnz[build_phases.index]
            except Exception:  # noqa: BLE001,S110
                pass

    lnzs = lnzs[(lnzs >= lnz_min) & (lnzs <= lnz_max)]
    c = lnPiCollection.from_builder(
        lnzs,
        build_phases=build_phases,
        build_kws=build_kws,
        nmax=nmax,
        xarray_output=xarray_output,
        **collection_kws,
    )

    if c_course is None:
        c_course = c
    elif limit_course:
        slnz = f"lnz_{build_phases.index}"
        q = f"{slnz} >= {lnz_min} and {slnz} <= {lnz_max}"
        c_course = c_course.query(q)

    return c_course, c
