from typing import List
from typing import Optional
from typing import Sequence
from typing import Tuple

import numpy
from PyMca5.PyMcaMath.linalg import lstsq


def fit_profiles(
    base_profiles: numpy.ndarray,
    profiles: numpy.ndarray,
    positive_constraint: Sequence[int] = (),
    weight: bool = True,
    profiles_stdev: Optional[numpy.ndarray] = None,
) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
    """Solve linear system of equations `A.p = b` for several profiles.
    Some parameters might be have `>=0` constraints. This implementation
    use the `lstsq` solver from PyMca.

    :param base_profiles: `A` with shape ``(num_channels, num_parameters)``
    :param profiles: `b` with shape ``(num_channels, num_profiles)``
    :param positive_constraint: indices of parameters that should be constraint to `>=0`.
    :param weight: use weights (all ones otherwise).
    :param profiles_stdev: shape ``(num_channels,)`` or ``(num_channels, num_profiles)``
    :returns: parameters `p` with shape ``(num_parameters, num_profiles)``,
              standard deviations on `p` with shape ``(num_parameters, num_profiles)``,
              number of free parameters with shape ``(num_profiles,)``
    """
    num_channels, num_profiles = profiles.shape
    if not weight:
        profiles_stdev = None

    if profiles_stdev is not None:
        if profiles_stdev.ndim == 1:
            profiles_stdev = profiles_stdev[:, numpy.newaxis]
        if profiles_stdev.shape not in (
            (num_channels, num_profiles),
            (num_channels, 1),
        ):
            raise ValueError("Profile and profile STDEV shape is not compatible")

    if len(base_profiles) != num_channels:
        raise ValueError(
            "Profiles and base-profiles must have the same number of observations"
        )

    # SVD or matrix inversion
    matrix_inversion = weight and profiles_stdev is None
    svd = not matrix_inversion

    result = lstsq(
        base_profiles,
        profiles,
        digested_output=True,
        weight=weight,
        sigma_b=profiles_stdev,
        svd=svd,
    )

    parameters = result["parameters"]
    uncertainties = result["uncertainties"]
    num_parameters, _ = parameters.shape  # shape (num_parameters, num_profiles)
    num_free_parameters = numpy.array([num_parameters] * num_profiles)
    if not positive_constraint:
        return parameters, uncertainties, num_free_parameters

    nrefit_max = 2 * len(positive_constraint) + 1
    for refit_count in range(1, nrefit_max + 2):
        # Detect negative values for parameters with a positive constraint
        refit_parameters_indices, profiles_to_fit, negative_list = _refit_info(
            parameters, positive_constraint
        )

        # Re-fit profiles without the parameters that are negative so the parameters
        # that are negative are effectively forced to 0.
        num_profiles_refit = profiles_to_fit.sum() if profiles_to_fit is not None else 0
        num_profiles_min = int(num_profiles / 400)
        do_refit = (
            num_profiles_refit  # no negative parameters
            and num_profiles_refit >= num_profiles_min  # not with refitting few spectra
            and refit_count <= nrefit_max  # probably infinite loop
        )
        if do_refit:
            result = lstsq(
                base_profiles[:, refit_parameters_indices],
                profiles[:, profiles_to_fit],
                digested_output=True,
                weight=weight,
                sigma_b=profiles_stdev,
                svd=svd,
            )

            # Replace re-fitted parameters
            for fit_param_idx, parametersi, uncertaintiesi in zip(
                refit_parameters_indices, result["parameters"], result["uncertainties"]
            ):
                parameters[fit_param_idx, profiles_to_fit] = parametersi
                uncertainties[fit_param_idx, profiles_to_fit] = uncertaintiesi

        # Fix negative parameters to zero
        if negative_list:
            num_free_parameters[:] = num_parameters
        for negative_param_idx, profile_negative_mask in negative_list:
            parameters[negative_param_idx, profile_negative_mask] = 0
            uncertainties[negative_param_idx, profile_negative_mask] = 0
            num_free_parameters[profile_negative_mask] -= 1

        if not do_refit:
            break

    return parameters, uncertainties, num_free_parameters


def _refit_info(
    parameters: numpy.ndarray, positive_constraint: Sequence[int]
) -> Tuple[List[int], Optional[numpy.ndarray], List[Tuple[int, numpy.ndarray]]]:
    """
    :param parameters: shape ``(num_parameters, num_profiles)``
    :param positive_constraint: indices of parameters that should be constraint to `>=0`.
    :returns: indices of free parameters, profile indices to be re-fitted, negative profiles (mask) per parameter
    """
    negative_list: List[Tuple[int, numpy.ndarray]] = []
    # parameter index, profiles (mask) where the parameter is negative

    negative_mask = parameters[positive_constraint] < 0
    for negative_param_idx, profile_negative_mask in zip(
        positive_constraint, negative_mask
    ):
        if profile_negative_mask.any():
            negative_list.append((negative_param_idx, profile_negative_mask))

    if not negative_list:
        return list(positive_constraint), None, negative_list

    negative_list.sort(reverse=True)

    fixed_at_zero_indices = [negative_list[0][0]]
    _, profiles_to_fit = negative_list[0]

    for negative_param_idx, profile_negative_mask in negative_list[1:]:
        combined_mask = profiles_to_fit & profile_negative_mask
        if combined_mask.any():
            fixed_at_zero_indices.append(negative_param_idx)
            profiles_to_fit = combined_mask

    refit_parameters_indices = [
        i for i in range(len(parameters)) if i not in fixed_at_zero_indices
    ]

    return refit_parameters_indices, profiles_to_fit, negative_list
