from ._analytic_rotation import target_rotation
from ._gpa_rotation import (
    GPA,
    CF_objective,
    ff_partial_target,
    ff_target,
    oblimin_objective,
    orthomax_objective,
    rotateA,
    vgQ_partial_target,
    vgQ_target,
)

__all__ = []


def rotate_factors(A, method, *method_args, **algorithm_kwargs):
    r"""
    Subroutine for orthogonal and oblique rotation of the matrix :math:`A`.
    For orthogonal rotations :math:`A` is rotated to :math:`L` according to

    .. math::

        L =  AT,

    where :math:`T` is an orthogonal matrix. And, for oblique rotations
    :math:`A` is rotated to :math:`L` according to

    .. math::

        L =  A(T^*)^{-1},

    where :math:`T` is a normal matrix.

    Parameters
    ----------
    A : numpy matrix (default None)
        non rotated factors
    method : str
        should be one of the methods listed below
    method_args : list
        additional arguments that should be provided with each method
    algorithm_kwargs : dictionary
        algorithm : str (default gpa)
            should be one of:

            * 'gpa': a numerical method
            * 'gpa_der_free': a derivative free numerical method
            * 'analytic' : an analytic method

        Depending on the algorithm, there are algorithm specific keyword
        arguments. For the gpa and gpa_der_free, the following
        keyword arguments are available:

        max_tries : int (default 501)
            maximum number of iterations

        tol : float
            stop criterion, algorithm stops if Frobenius norm of gradient is
            smaller then tol

        For analytic, the supported arguments depend on the method, see above.

        See the lower level functions for more details.

    Returns
    -------
    The tuple :math:`(L,T)`

    Notes
    -----
    What follows is a list of available methods. Depending on the method
    additional argument are required and different algorithms
    are available. The algorithm_kwargs are additional keyword arguments
    passed to the selected algorithm (see the parameters section).
    Unless stated otherwise, only the gpa and
    gpa_der_free algorithm are available.

    Below,

        * :math:`L` is a :math:`p\times k` matrix;
        * :math:`N` is :math:`k\times k` matrix with zeros on the diagonal and ones
          elsewhere;
        * :math:`M` is :math:`p\times p` matrix with zeros on the diagonal and ones
          elsewhere;
        * :math:`C` is a :math:`p\times p` matrix with elements equal to
          :math:`1/p`;
        * :math:`(X,Y)=\operatorname{Tr}(X^*Y)` is the Frobenius norm;
        * :math:`\circ` is the element-wise product or Hadamard product.

    oblimin : orthogonal or oblique rotation that minimizes
        .. math::
            \phi(L) = \frac{1}{4}(L\circ L,(I-\gamma C)(L\circ L)N).

        For orthogonal rotations:

        * :math:`\gamma=0` corresponds to quartimax,
        * :math:`\gamma=\frac{1}{2}` corresponds to biquartimax,
        * :math:`\gamma=1` corresponds to varimax,
        * :math:`\gamma=\frac{1}{p}` corresponds to equamax.

        For oblique rotations rotations:

        * :math:`\gamma=0` corresponds to quartimin,
        * :math:`\gamma=\frac{1}{2}` corresponds to biquartimin.

        method_args:

        gamma : float
            oblimin family parameter
        rotation_method : str
            should be one of {orthogonal, oblique}

    orthomax : orthogonal rotation that minimizes

        .. math::
            \phi(L) = -\frac{1}{4}(L\circ L,(I-\gamma C)(L\circ L)),

        where :math:`0\leq\gamma\leq1`. The orthomax family is equivalent to
        the oblimin family (when restricted to orthogonal rotations).
        Furthermore,

        * :math:`\gamma=0` corresponds to quartimax,
        * :math:`\gamma=\frac{1}{2}` corresponds to biquartimax,
        * :math:`\gamma=1` corresponds to varimax,
        * :math:`\gamma=\frac{1}{p}` corresponds to equamax.

        method_args:

        gamma : float (between 0 and 1)
            orthomax family parameter

    CF : Crawford-Ferguson family for orthogonal and oblique rotation which
    minimizes:

        .. math::

            \phi(L) =\frac{1-\kappa}{4} (L\circ L,(L\circ L)N)
                     -\frac{1}{4}(L\circ L,M(L\circ L)),

        where :math:`0\leq\kappa\leq1`. For orthogonal rotations the oblimin
        (and orthomax) family of rotations is equivalent to the
        Crawford-Ferguson family.
        To be more precise:

        * :math:`\kappa=0` corresponds to quartimax,
        * :math:`\kappa=\frac{1}{p}` corresponds to varimax,
        * :math:`\kappa=\frac{k-1}{p+k-2}` corresponds to parsimax,
        * :math:`\kappa=1` corresponds to factor parsimony.

        method_args:

        kappa : float (between 0 and 1)
            Crawford-Ferguson family parameter
        rotation_method : str
            should be one of {orthogonal, oblique}

    quartimax : orthogonal rotation method
        minimizes the orthomax objective with :math:`\gamma=0`

    biquartimax : orthogonal rotation method
        minimizes the orthomax objective with :math:`\gamma=\frac{1}{2}`

    varimax : orthogonal rotation method
        minimizes the orthomax objective with :math:`\gamma=1`

    equamax : orthogonal rotation method
        minimizes the orthomax objective with :math:`\gamma=\frac{1}{p}`

    parsimax : orthogonal rotation method
        minimizes the Crawford-Ferguson family objective with
        :math:`\kappa=\frac{k-1}{p+k-2}`

    parsimony : orthogonal rotation method
        minimizes the Crawford-Ferguson family objective with :math:`\kappa=1`

    quartimin : oblique rotation method that minimizes
        minimizes the oblimin objective with :math:`\gamma=0`

    quartimin : oblique rotation method that minimizes
        minimizes the oblimin objective with :math:`\gamma=\frac{1}{2}`

    target : orthogonal or oblique rotation that rotates towards a target

    matrix : math:`H` by minimizing the objective

        .. math::

            \phi(L) =\frac{1}{2}\|L-H\|^2.

        method_args:

        H : numpy matrix
            target matrix
        rotation_method : str
            should be one of {orthogonal, oblique}

        For orthogonal rotations the algorithm can be set to analytic in which
        case the following keyword arguments are available:

        full_rank : bool (default False)
            if set to true full rank is assumed

    partial_target : orthogonal (default) or oblique rotation that partially
    rotates towards a target matrix :math:`H` by minimizing the objective:

        .. math::

            \phi(L) =\frac{1}{2}\|W\circ(L-H)\|^2.

        method_args:

        H : numpy matrix
            target matrix
        W : numpy matrix (default matrix with equal weight one for all entries)
            matrix with weights, entries can either be one or zero

    Examples
    --------
    >>> A = np.random.randn(8,2)
    >>> L, T = rotate_factors(A,'varimax')
    >>> np.allclose(L,A.dot(T))
    >>> L, T = rotate_factors(A,'orthomax',0.5)
    >>> np.allclose(L,A.dot(T))
    >>> L, T = rotate_factors(A,'quartimin',0.5)
    >>> np.allclose(L,A.dot(np.linalg.inv(T.T)))
    """
    if "algorithm" in algorithm_kwargs:
        algorithm = algorithm_kwargs["algorithm"]
        algorithm_kwargs.pop("algorithm")
    else:
        algorithm = "gpa"
    assert (
        "rotation_method" not in algorithm_kwargs
    ), "rotation_method cannot be provided as keyword argument"
    T = L = None
    p, k = A.shape
    # set ff or vgQ to appropriate objective function, compute solution using
    # recursion or analytically compute solution

    if method == "orthomax":
        assert len(method_args) == 1, (
            f"Only {method} family parameter should be provided"
        )
        rotation_method = "orthogonal"
        gamma = method_args[0]
        if algorithm == "gpa":

            def vgQ(L=None, A=None, T=None):
                return orthomax_objective(
                    L=L, A=A, T=T, gamma=gamma, return_gradient=True
                )
            ff = None

        elif algorithm == "gpa_der_free":

            def ff(L=None, A=None, T=None):
                return orthomax_objective(
                    L=L, A=A, T=T, gamma=gamma, return_gradient=False
                )
            vgQ = None

        else:
            raise ValueError(
                f"Algorithm {algorithm} is not possible for {method} rotation."
            )
    elif method == "oblimin":
        assert len(method_args) == 2, (
            "Both %s family parameter and "
            "rotation_method should be "
            "provided" % method
        )
        rotation_method = method_args[1]
        assert rotation_method in [
            "orthogonal",
            "oblique",
        ], "rotation_method should be one of {orthogonal, oblique}"
        gamma = method_args[0]
        if algorithm == "gpa":

            def vgQ(L=None, A=None, T=None):
                return oblimin_objective(
                    L=L, A=A, T=T, gamma=gamma, return_gradient=True
                )
            ff = None

        elif algorithm == "gpa_der_free":

            def ff(L=None, A=None, T=None):
                return oblimin_objective(
                    L=L,
                    A=A,
                    T=T,
                    gamma=gamma,
                    rotation_method=rotation_method,
                    return_gradient=False,
                )
            vgQ = None

        else:
            raise ValueError(
                "Algorithm %s is not possible for %s rotation" % (algorithm, method)
            )
    elif method == "CF":
        assert len(method_args) == 2, (
            "Both %s family parameter and "
            "rotation_method should be provided" % method
        )
        rotation_method = method_args[1]
        assert rotation_method in [
            "orthogonal",
            "oblique",
        ], "rotation_method should be one of {orthogonal, oblique}"
        kappa = method_args[0]
        if algorithm == "gpa":

            def vgQ(L=None, A=None, T=None):
                return CF_objective(
                    L=L,
                    A=A,
                    T=T,
                    kappa=kappa,
                    rotation_method=rotation_method,
                    return_gradient=True,
                )
            ff = None

        elif algorithm == "gpa_der_free":

            def ff(L=None, A=None, T=None):
                return CF_objective(
                    L=L,
                    A=A,
                    T=T,
                    kappa=kappa,
                    rotation_method=rotation_method,
                    return_gradient=False,
                )
            vgQ = None

        else:
            raise ValueError(
                "Algorithm %s is not possible for %s rotation" % (algorithm, method)
            )
    elif method == "quartimax":
        return rotate_factors(A, "orthomax", 0, **algorithm_kwargs)
    elif method == "biquartimax":
        return rotate_factors(A, "orthomax", 0.5, **algorithm_kwargs)
    elif method == "varimax":
        return rotate_factors(A, "orthomax", 1, **algorithm_kwargs)
    elif method == "equamax":
        return rotate_factors(A, "orthomax", 1 / p, **algorithm_kwargs)
    elif method == "parsimax":
        return rotate_factors(
            A, "CF", (k - 1) / (p + k - 2), "orthogonal", **algorithm_kwargs
        )
    elif method == "parsimony":
        return rotate_factors(A, "CF", 1, "orthogonal", **algorithm_kwargs)
    elif method == "quartimin":
        return rotate_factors(A, "oblimin", 0, "oblique", **algorithm_kwargs)
    elif method == "biquartimin":
        return rotate_factors(A, "oblimin", 0.5, "oblique", **algorithm_kwargs)
    elif method == "target":
        assert len(method_args) == 2, (
            "only the rotation target and orthogonal/oblique should be provide"
            " for %s rotation" % method
        )
        H = method_args[0]
        rotation_method = method_args[1]
        assert rotation_method in [
            "orthogonal",
            "oblique",
        ], "rotation_method should be one of {orthogonal, oblique}"
        if algorithm == "gpa":

            def vgQ(L=None, A=None, T=None):
                return vgQ_target(H, L=L, A=A, T=T, rotation_method=rotation_method)
            ff = None

        elif algorithm == "gpa_der_free":

            def ff(L=None, A=None, T=None):
                return ff_target(H, L=L, A=A, T=T, rotation_method=rotation_method)
            vgQ = None

        elif algorithm == "analytic":
            assert rotation_method == "orthogonal", (
                "For analytic %s rotation only orthogonal rotation is supported"
            )
            T = target_rotation(A, H, **algorithm_kwargs)
        else:
            raise ValueError(
                f"Algorithm {algorithm} is not possible for {method} rotation"
            )
    elif method == "partial_target":
        assert len(method_args) == 2, (
            "2 additional arguments are expected for %s rotation" % method
        )
        H = method_args[0]
        W = method_args[1]
        rotation_method = "orthogonal"
        if algorithm == "gpa":

            def vgQ(L=None, A=None, T=None):
                return vgQ_partial_target(H, W=W, L=L, A=A, T=T)
            ff = None
        elif algorithm == "gpa_der_free":

            def ff(L=None, A=None, T=None):
                return ff_partial_target(H, W=W, L=L, A=A, T=T)
            vgQ = None
        else:
            raise ValueError(
                "Algorithm %s is not possible for %s rotation" % (algorithm, method)
            )
    else:
        raise ValueError("Invalid method")
    # compute L and T if not already done

    if T is None:
        L, phi, T, table = GPA(
            A, vgQ=vgQ, ff=ff, rotation_method=rotation_method, **algorithm_kwargs
        )
    if L is None:
        assert T is not None, "Cannot compute L without T"
        L = rotateA(A, T, rotation_method=rotation_method)
    return L, T
