from __future__ import annotations

import bisect
import math
import numbers
import warnings
from typing import TYPE_CHECKING

from ...compat import np, scipy
from ...errors import UndefinedBehavior
from ...util import UnitsContainer, infer_base_unit

if TYPE_CHECKING:
    from ..._typing import UnitLike
    from .quantity import PlainQuantity


def _get_reduced_units(
    quantity: PlainQuantity, units: UnitsContainer
) -> UnitsContainer:
    # loop through individual units and compare to each other unit
    # can we do better than a nested loop here?
    for unit1, exp in units.items():
        # make sure it wasn't already reduced to zero exponent on prior pass
        if unit1 not in units:
            continue
        for unit2 in units:
            # get exponent after reduction
            exp = units[unit1]
            if unit1 != unit2:
                power = quantity._REGISTRY._get_dimensionality_ratio(unit1, unit2)
                if power:
                    units = units.add(unit2, exp / power).remove([unit1])
                    break
    return units


def ito_reduced_units(quantity: PlainQuantity) -> None:
    """Return PlainQuantity scaled in place to reduced units, i.e. one unit per
    dimension. This will not reduce compound units (e.g., 'J/kg' will not
    be reduced to m**2/s**2), nor can it make use of contexts at this time.
    """

    # shortcuts in case we're dimensionless or only a single unit
    if quantity.dimensionless:
        return quantity.ito({})
    if len(quantity._units) == 1:
        return None

    units = quantity._units.copy()
    new_units = _get_reduced_units(quantity, units)

    return quantity.ito(new_units)


def to_reduced_units(
    quantity: PlainQuantity,
) -> PlainQuantity:
    """Return PlainQuantity scaled in place to reduced units, i.e. one unit per
    dimension. This will not reduce compound units (intentionally), nor
    can it make use of contexts at this time.
    """

    # shortcuts in case we're dimensionless or only a single unit
    if quantity.dimensionless:
        return quantity.to({})
    if len(quantity._units) == 1:
        return quantity

    units = quantity._units.copy()
    new_units = _get_reduced_units(quantity, units)

    return quantity.to(new_units)


def to_compact(
    quantity: PlainQuantity, unit: UnitsContainer | None = None
) -> PlainQuantity:
    """ "Return PlainQuantity rescaled to compact, human-readable units.

    To get output in terms of a different unit, use the unit parameter.


    Examples
    --------

    >>> import pint
    >>> ureg = pint.UnitRegistry()
    >>> (200e-9 * ureg.s).to_compact()
    <Quantity(200.0, 'nanosecond')>
    >>> (1e-2 * ureg("kg m/s^2")).to_compact("N")
    <Quantity(10.0, 'millinewton')>
    """

    if not isinstance(quantity.magnitude, numbers.Number) and not hasattr(
        quantity.magnitude, "nominal_value"
    ):
        warnings.warn(
            "to_compact applied to non numerical types has an undefined behavior.",
            UndefinedBehavior,
            stacklevel=2,
        )
        return quantity

    qm = (
        quantity.magnitude
        if not hasattr(quantity.magnitude, "nominal_value")
        else quantity.magnitude.nominal_value
    )
    if quantity.unitless or qm == 0 or math.isnan(qm) or math.isinf(qm):
        return quantity

    SI_prefixes: dict[int, str] = {}
    for prefix in quantity._REGISTRY._prefixes.values():
        try:
            scale = prefix.converter.scale
            # Kludgy way to check if this is an SI prefix
            log10_scale = int(math.log10(scale))
            if log10_scale == math.log10(scale):
                SI_prefixes[log10_scale] = prefix.name
        except Exception:
            SI_prefixes[0] = ""

    SI_prefixes_list = sorted(SI_prefixes.items())
    SI_powers = [item[0] for item in SI_prefixes_list]
    SI_bases = [item[1] for item in SI_prefixes_list]

    if unit is None:
        unit = infer_base_unit(quantity, registry=quantity._REGISTRY)
    else:
        unit = infer_base_unit(quantity.__class__(1, unit), registry=quantity._REGISTRY)

    q_base = quantity.to(unit)

    magnitude = q_base.magnitude
    # Support uncertainties
    if hasattr(magnitude, "nominal_value"):
        magnitude = magnitude.nominal_value

    units = list(q_base._units.items())
    units_numerator = [a for a in units if a[1] > 0]

    if len(units_numerator) > 0:
        unit_str, unit_power = units_numerator[0]
    else:
        unit_str, unit_power = units[0]

    if unit_power > 0:
        power = math.floor(math.log10(abs(magnitude)) / float(unit_power) / 3) * 3
    else:
        power = math.ceil(math.log10(abs(magnitude)) / float(unit_power) / 3) * 3

    index = bisect.bisect_left(SI_powers, power)

    if index >= len(SI_bases):
        index = -1

    prefix_str = SI_bases[index]

    new_unit_str = prefix_str + unit_str
    new_unit_container = q_base._units.rename(unit_str, new_unit_str)

    return quantity.to(new_unit_container)


def _get_unprefixed(quantity: PlainQuantity) -> PlainQuantity:
    units = list(quantity._units.items())
    new_units = []

    base_unit_map = {}

    for unit in quantity._REGISTRY._base_units:
        base_unit_map[unit] = str(quantity._REGISTRY.get_base_units(unit)[1]._units)

    # TODO: This code is very similar to util.infer_base_unit.
    # Look at combining the two functions.
    for unit in units:
        candidates = quantity._REGISTRY.parse_unit_name(unit[0])
        _, unprefix_unit, _ = candidates[0]
        # the unprefix_unit could be g rather than kg.
        # check whether the base_unit for the unprefix_unit also maps to g
        if unprefix_unit in base_unit_map:
            unprefix_unit = base_unit_map[unprefix_unit]
        new_units.append((unprefix_unit, unit[1]))
    new_unit_container = UnitsContainer(new_units)
    return new_unit_container


def to_unprefixed(quantity: PlainQuantity) -> PlainQuantity:
    units = _get_unprefixed(quantity)
    return quantity.to(units)


def ito_unprefixed(quantity: PlainQuantity) -> PlainQuantity:
    units = _get_unprefixed(quantity)
    return quantity.ito(units)


def to_preferred(
    quantity: PlainQuantity, preferred_units: list[UnitLike] | None = None
) -> PlainQuantity:
    """Return Quantity converted to a unit composed of the preferred units.

    Examples
    --------

    >>> import pint
    >>> ureg = pint.UnitRegistry()
    >>> (1 * ureg.acre).to_preferred([ureg.meters])
    <Quantity(4046.87261, 'meter ** 2')>
    >>> (1 * (ureg.force_pound * ureg.m)).to_preferred([ureg.W])
    <Quantity(4.44822162, 'watt * second')>
    """

    units = _get_preferred(quantity, preferred_units)
    return quantity.to(units)


def ito_preferred(
    quantity: PlainQuantity, preferred_units: list[UnitLike] | None = None
) -> PlainQuantity:
    """Return Quantity converted to a unit composed of the preferred units.

    Examples
    --------

    >>> import pint
    >>> ureg = pint.UnitRegistry()
    >>> (1 * ureg.acre).to_preferred([ureg.meters])
    <Quantity(4046.87261, 'meter ** 2')>
    >>> (1 * (ureg.force_pound * ureg.m)).to_preferred([ureg.W])
    <Quantity(4.44822162, 'watt * second')>
    """

    units = _get_preferred(quantity, preferred_units)
    return quantity.ito(units)


def _get_preferred(
    quantity: PlainQuantity, preferred_units: list[UnitLike] | None = None
) -> PlainQuantity:
    if preferred_units is None:
        preferred_units = quantity._REGISTRY.default_preferred_units

    if not quantity.dimensionality:
        return quantity._units.copy()

    # The optimizer isn't perfect, and will sometimes miss obvious solutions.
    # This sub-algorithm is less powerful, but always finds the very simple solutions.
    def find_simple():
        best_ratio = None
        best_unit = None
        self_dims = sorted(quantity.dimensionality)
        self_exps = [quantity.dimensionality[d] for d in self_dims]
        s_exps_head, *s_exps_tail = self_exps
        n = len(s_exps_tail)
        for preferred_unit in preferred_units:
            dims = sorted(preferred_unit.dimensionality)
            if dims == self_dims:
                p_exps_head, *p_exps_tail = (
                    preferred_unit.dimensionality[d] for d in dims
                )
                if all(
                    s_exps_tail[i] * p_exps_head == p_exps_tail[i] ** s_exps_head
                    for i in range(n)
                ):
                    ratio = p_exps_head / s_exps_head
                    ratio = max(ratio, 1 / ratio)
                    if best_ratio is None or ratio < best_ratio:
                        best_ratio = ratio
                        best_unit = preferred_unit ** (s_exps_head / p_exps_head)
        return best_unit

    simple = find_simple()
    if simple is not None:
        return simple

    # For each dimension (e.g. T(ime), L(ength), M(ass)), assign a default base unit from
    # the collection of base units

    unit_selections = {
        base_unit.dimensionality: base_unit
        for base_unit in map(quantity._REGISTRY.Unit, quantity._REGISTRY._base_units)
    }

    # Override the default unit of each dimension with the 1D-units used in this Quantity
    unit_selections.update(
        {
            unit.dimensionality: unit
            for unit in map(quantity._REGISTRY.Unit, quantity._units.keys())
        }
    )

    # Determine the preferred unit for each dimensionality from the preferred_units
    # (A prefered unit doesn't have to be only one dimensional, e.g. Watts)
    preferred_dims = {
        preferred_unit.dimensionality: preferred_unit
        for preferred_unit in map(quantity._REGISTRY.Unit, preferred_units)
    }

    # Combine the defaults and preferred, favoring the preferred
    unit_selections.update(preferred_dims)

    # This algorithm has poor asymptotic time complexity, so first reduce the considered
    # dimensions and units to only those that are useful to the problem

    # The dimensions (without powers) of this Quantity
    dimension_set = set(quantity.dimensionality)

    # Getting zero exponents in dimensions not in dimension_set can be facilitated
    # by units that interact with that dimension and one or more dimension_set members.
    # For example MT^1 * LT^-1 lets you get MLT^0 when T is not in dimension_set.
    # For each candidate unit that interacts with a dimension_set member, add the
    # candidate unit's other dimensions to dimension_set, and repeat until no more
    # dimensions are selected.

    discovery_done = False
    while not discovery_done:
        discovery_done = True
        for d in unit_selections:
            unit_dimensions = set(d)
            intersection = unit_dimensions.intersection(dimension_set)
            if 0 < len(intersection) < len(unit_dimensions):
                # there are dimensions in this unit that are in dimension set
                # and others that are not in dimension set
                dimension_set = dimension_set.union(unit_dimensions)
                discovery_done = False
                break

    # filter out dimensions and their unit selections that don't interact with any
    # dimension_set members
    unit_selections = {
        dimensionality: unit
        for dimensionality, unit in unit_selections.items()
        if set(dimensionality).intersection(dimension_set)
    }

    # update preferred_units with the selected units that were originally preferred
    preferred_units = list(
        {u for d, u in unit_selections.items() if d in preferred_dims}
    )
    preferred_units.sort(key=str)  # for determinism

    # and unpreferred_units are the selected units that weren't originally preferred
    unpreferred_units = list(
        {u for d, u in unit_selections.items() if d not in preferred_dims}
    )
    unpreferred_units.sort(key=str)  # for determinism

    # for indexability
    dimensions = list(dimension_set)
    dimensions.sort()  # for determinism

    # the powers for each elemet of dimensions (the list) for this Quantity
    dimensionality = [quantity.dimensionality[dimension] for dimension in dimensions]

    # Now that the input data is minimized, setup the optimization problem

    # use scipy.optimize.milp to select units from preferred units

    # prefer units that interact with fewer dimensions
    cost = [len(p.dimensionality) for p in preferred_units]

    # set the cost for non preferred units to a higher number
    bias = (
        max(map(abs, dimensionality)) * max((1, *cost)) * 10
    )  # arbitrary, just needs to be larger
    cost.extend([bias] * len(unpreferred_units))

    all_units = preferred_units + unpreferred_units
    num_units = len(all_units)
    num_dims = len(dimensions)

    # where [u_1 ... u_N] are powers of N candidate units (vars)
    # introduce auxiliary variables [a_1 ... a_N] to handle absolute values
    # such that |u_i| <= a_i
    # variables: x = [u_1, ..., u_N, a_1, ..., a_N]
    # objective: min c @ x
    # where c = [0, ..., 0, cost_1, ..., cost_i]
    c = np.concatenate([np.zeros(num_units), np.array(cost)])

    # u: (-inf, inf), a: (0, inf)
    bounds = scipy.optimize.Bounds(
        lb=np.concatenate([-np.inf * np.ones(num_units), np.zeros(num_units)]),
        ub=np.inf * np.ones(2 * num_units),
    )

    # constraint: 1 means integer
    integrality = np.ones(2 * num_units, dtype=np.intp)

    # constraint: |u_i| <= a_i
    E = np.eye(num_units)
    constraints = [
        # 1)  u_i - a_i <= 0
        scipy.optimize.LinearConstraint(A=np.hstack([E, -E]), lb=-np.inf, ub=0),
        # 2) -u_i - a_i <= 0
        scipy.optimize.LinearConstraint(A=np.hstack([-E, -E]), lb=-np.inf, ub=0),
    ]

    # constraint: [D, 0] @ [u, a] = dimensionality
    if num_dims > 0:
        D = np.array(
            [[unit.dimensionality[d] for d in dimensions] for unit in all_units]
        ).transpose()

        A_eq = np.hstack([D, np.zeros((num_dims, num_units))])
        b_eq = np.array(dimensionality)
        constraints.append(scipy.optimize.LinearConstraint(A=A_eq, lb=b_eq, ub=b_eq))

    # run the mips minimizer and extract the result if successful
    res = scipy.optimize.milp(
        c, constraints=constraints, integrality=integrality, bounds=bounds
    )

    if res.success and res.x is not None:
        exponents = np.round(res.x[:num_units]).astype(int)

        result_unit = quantity._REGISTRY.Unit("")
        for unit, exponent in zip(all_units, exponents, strict=True):
            if exponent != 0:
                result_unit *= unit**exponent

        return result_unit

    # for whatever reason, a solution wasn't found
    # return the original quantity
    return quantity._units.copy()
