from typing import Iterable, Union

import numpy as np

from cca_zoo.utils import _process_parameter

from ._base import BaseIterative, BaseLoop
from ._deflation import DeflationMixin


class SCCA_ADMM(BaseIterative, DeflationMixin):
    r"""
    Fits a sparse CCA model by alternating ADMM for two or more views.

    .. math::

        w_{opt}=\underset{w}{\mathrm{argmax}}\{\sum_i\sum_{j\neq i} \|X_iw_i-X_jw_j\|^2 + \|w_i\|_1\}\\

        \text{subject to:}

        w_i^TX_i^TX_iw_i=1

    Parameters
    ----------
    latent_dimensions : int, default=1
        Number of latent dimensions to use in the model.
    copy_data : bool, default=True
        Whether to copy the data or overwrite it.
    random_state : int, default=None
        Random seed for initialisation.
    deflation : str, default="cca"
        Deflation method to use. Options are "cca" and "pls".
    tau : float or list of floats, default=None
        Regularisation parameter. If a single float is given, the same value is used for all views.
        If a list of floats is given, the values are used for each view.
    mu : float or list of floats, default=None
        Regularisation parameter. If a single float is given, the same value is used for all views.
        If a list of floats is given, the values are used for each view.
    lam : float or list of floats, default=None
        Regularisation parameter. If a single float is given, the same value is used for all views.
        If a list of floats is given, the values are used for each view.
    eta : float or list of floats, default=None
        Regularisation parameter. If a single float is given, the same value is used for all views.
        If a list of floats is given, the values are used for each view.
    tol : float, default=1e-9
        Tolerance for convergence.

    References
    ----------
    Suo, Xiaotong, et al. "Sparse canonical correlation analysis." arXiv preprint arXiv:1705.10865 (2017).

    Examples
    --------
    >>> from cca_zoo.linear import SCCA_ADMM
    >>> import numpy as np
    >>> rng=np.random.RandomState(0)
    >>> X1 = rng.random((10,5))
    >>> X2 = rng.random((10,5))
    >>> model = SCCA_ADMM(random_state=0,tau=[1e-1,1e-1])
    >>> model.fit((X1,X2)).score((X1,X2))
    array([0.84348183])
    """

    def __init__(
        self,
        latent_dimensions: int = 1,
        copy_data=True,
        random_state=None,
        deflation="cca",
        tau: Union[Iterable[float], float] = None,
        mu: Union[Iterable[float], float] = None,
        lam: Union[Iterable[float], float] = None,
        eta: Union[Iterable[float], float] = None,
        initialization: Union[str, callable] = "pls",
        tol: float = 1e-3,
    ):
        self.tau = tau
        self.mu = mu
        self.lam = lam
        self.eta = eta
        super().__init__(
            latent_dimensions=latent_dimensions,
            copy_data=copy_data,
            initialization=initialization,
            tol=tol,
            random_state=random_state,
            deflation=deflation,
        )

    def _check_params(self):
        self.tau = _process_parameter("tau", self.tau, 0, self.n_views_)
        self.lam = _process_parameter("lam", self.lam, 1, self.n_views_)
        self.eta = _process_parameter("eta", self.eta, 0, self.n_views_)

    def _get_module(self, weights=None, k=None):
        return SCCA_ADMM_PL(
            weights=weights,
            k=k,
            tau=self.tau,
            lam=self.lam,
            eta=self.eta,
            n_samples_=self.n_samples_,
            n_views_=self.n_views_,
        )


class SCCA_ADMM_PL(BaseLoop):
    def __init__(
        self,
        weights,
        k=None,
        n_samples_=None,
        n_views_=None,
        tau=None,
        eta=None,
        lam=None,
        mu=None,
    ):
        super().__init__(weights=weights, k=k)
        self.eta = [np.ones(n_samples_) * eta for eta in eta]
        self.z = [np.ones(n_samples_)] * n_views_
        self.mu = mu

    def training_step(self, batch, batch_idx):
        views = batch["views"]
        scores = np.stack(self(views))
        for view_index, view in enumerate(views):
            targets = np.ma.array(scores, mask=False)
            targets.mask[view_index] = True
            gradient = views[view_index].T @ targets.sum(axis=0).filled()
            mu = self.mu[view_index]
            lam = self.lam[view_index]
            N = views[view_index].shape[0]
            unnorm_z = []
            norm_eta = []
            norm_weights = []
            norm_proj = []
            for _ in range(self.max_iter):
                # We multiply 'tau' by N in order to make regularisation match across the different sparse cca methods
                self.weights[view_index] = self._prox_mu_f(
                    self.weights[view_index]
                    - mu
                    / lam
                    * views[view_index].T
                    @ (
                        views[view_index] @ self.weights[view_index]
                        - self.z[view_index]
                        + self.eta[view_index]
                    ),
                    mu,
                    gradient,
                    N * self.tau[view_index],
                )
                unnorm_z.append(
                    np.linalg.norm(
                        views[view_index] @ self.weights[view_index]
                        + self.eta[view_index]
                    )
                )
                self.z[view_index] = self._prox_lam_g(
                    views[view_index] @ self.weights[view_index] + self.eta[view_index]
                )
                self.eta[view_index] = (
                    self.eta[view_index]
                    + views[view_index] @ self.weights[view_index]
                    - self.z[view_index]
                )
                norm_eta.append(np.linalg.norm(self.eta[view_index]))
                norm_proj.append(
                    np.linalg.norm(views[view_index] @ self.weights[view_index])
                )
                norm_weights.append(np.linalg.norm(self.weights[view_index], 1))

    def _prox_mu_f(self, x, mu, c, tau):
        u_update = x.copy()
        mask_1 = x + (mu * c) > mu * tau
        # if mask_1.sum()>0:
        u_update[mask_1] = x[mask_1] + mu * (c[mask_1] - tau)
        mask_2 = x + (mu * c) < -mu * tau
        # if mask_2.sum() > 0:
        u_update[mask_2] = x[mask_2] + mu * (c[mask_2] + tau)
        mask_3 = ~(mask_1 | mask_2)
        u_update[mask_3] = 0
        return u_update

    def _prox_lam_g(self, x):
        norm = np.linalg.norm(x)
        if norm < 1:
            return x
        else:
            return x / max(1, norm)
