# -*- coding: utf-8 -*-

"""
@author: Dussap Bastien
This .py contains the KernelQuantifier with and without RFF
Code inspired by QuaPy : https://github.com/HLT-ISTI/QuaPy
"""

import numpy as np
import torch

from typing import Union
from cvxopt import matrix, solvers


from .rff import bandwidthRFF, available_kernel_rff, select_kernel_rff
from .data import LabelledCollection
from .base import BaseQuantifier
from .utils import choose_device
from .kernel import GaussianKernel, bandwidth_kernel, available_kernel


class KernelQuantifierRFF(BaseQuantifier):
    """
    KernelQuantifier with RFF.
    """

    def __init__(self, kernel_type: str = "gaussian", seed: int = 123):
        """
        kernel_type [str] : "gaussian".
        device [torch.device] : device.
        seed [float] : seed, defaults is 123.
        """

        self.kernel_type = kernel_type
        self.seed = seed

        assert self.kernel_type in available_kernel_rff(), f"This kernel is not implemented yet\nCurrent kernel = {available_kernel_rff()}"
        self.kernel = select_kernel_rff(self.kernel_type)

    def fit(self,
            data: LabelledCollection,
            sigma: Union[float, list],
            verbose: bool,
            number_rff: int) -> None:
        """
        data [.data.LabelledCollection] : the Source data
        sigma : Give a list of two elements to use the criteron to find the best sigma inside the interval.
        verbose [bool] : plot the function eigmin(P) w.r.t bandwidth
        number_rff [int] : The number of RFF

        --------

        The fit function stock the mu(P_i) and the kernel.
        """

        self.device_ = data.device
        self.n_points_ = data.n_points
        self.n_classes_ = data.n_classes

        # First we choose the kernel
        if isinstance(sigma, list):
            sigma = bandwidthRFF(P=data,
                                 sigma_min=sigma[0],
                                 sigma_max=sigma[1],
                                 kernel=self.kernel,
                                 verbose=verbose,
                                 device=self.device_,
                                 number_rff=number_rff,
                                 seed=self.seed)

        self.kernel = self.kernel(
            D=number_rff, d=data.dim, sigma=sigma, device=self.device_, seed=self.seed, dtype=data.dtype)

        # Then, compute mu(P_i) with this kernel
        self.mu_ = [self.kernel.fit(data[i]) for i in range(data.n_classes)]
        self.mu_ = torch.stack(self.mu_, axis=1)

    def refit(self, data: LabelledCollection):
        """
        Refit the mu on the data.

        Args:
            data (LabelledCollection): Data
        """
        self.mu_ = [self.kernel.fit(data[i]) for i in range(data.n_classes)]
        self.mu_ = torch.stack(self.mu_, axis=1).float()

    def quantify(self, data: LabelledCollection, target: torch.tensor) -> np.array:
        """

        Args:
            data [LabelledCollection] : Source
            target [torch.tensor] : the target distribution

        -----------------

        return : np.array (self.n_classes_) containing the estimate prop of the target.
        """
        def torch_to_matrix(tensor):
            tensor = tensor.cpu()
            return matrix(np.array(tensor).astype(np.double))

        b = self.kernel.fit(target)
        n = self.n_points_

        P = torch_to_matrix(2*n/(n-1)*self.mu_.T @ self.mu_)
        q = torch_to_matrix(-2 * self.mu_.T @ b)
        # constraint
        G = matrix(-np.eye(self.n_classes_))
        h = matrix(np.zeros(self.n_classes_))
        A = matrix(np.ones(self.n_classes_)).T
        b = matrix([1.0])

        solvers.options['show_progress'] = False
        sol = solvers.qp(P, q, G, h, A, b)

        self.props_ = np.array(sol["x"]).reshape(self.n_classes_)
        return self.props_

    def get_kl(self, pi_target: np.array) -> np.array:
        return super().get_kl(pi_target)

    @property
    def sigma(self):
        return self.kernel.sigma

############################################################################################


class KernelQuantifier(BaseQuantifier):
    """
    KernelQuantifier without RFF.
    """

    def __init__(self, kernel_type: str = "gaussian", seed: float = 123):
        """
        Kernel [str] : "gaussian"
        seed [float] : seed, defaults is 123. 
        """

        self.kernel_ = kernel_type
        self.seed = seed

        assert self.kernel_ in available_kernel(), f"This kernel is not implemented yet\nCurrent kernel = {available_kernel()}"

    def fit(self,
            data: LabelledCollection,
            sigma: Union[float, list],
            verbose: bool = False) -> None:
        """
        data [KMPE.data.LabelledCollection] : the Source data
        sigma :
        verbose [bool] : plot the function eigmin(P) w.r.t bandwidth

        --------

        The fit function stock the kernel and the matrix A.
        """

        self.device_ = data.device
        self.n_points_ = data.n_points
        self.n_classes_ = data.n_classes

        # First we choose the kernel
        if isinstance(sigma, list):
            sigma = bandwidth_kernel(P=data,
                                     sigma_min=sigma[0],
                                     sigma_max=sigma[1],
                                     kernel=GaussianKernel,
                                     verbose=verbose)

        self.kernel = GaussianKernel(sigma=float(sigma))
        self.A = self.kernel.fit_X(data)

    def quantify(self, data: LabelledCollection, target: torch.tensor) -> np.array:
        """
        Generate class prevalence estimates for the sample's target
        :param target: torch.tensor
        :return: `np.ndarray` of shape `(self.n_classes_)` with class prevalence estimates.
        """

        def torch_to_matrix(tensor):
            tensor = tensor.cpu()
            return matrix(np.array(tensor).astype(np.double))

        q = self.kernel.fit_Xy(data, target)

        P = torch_to_matrix(2*self.A)
        q = torch_to_matrix(-2 * q)
        G = matrix(-np.eye(self.n_classes_))
        h = matrix(np.zeros(self.n_classes_))
        A = matrix(np.ones(self.n_classes_)).T
        b = matrix([1.0])

        solvers.options['show_progress'] = False
        sol = solvers.qp(P, q, G, h, A, b)

        self.props_ = np.array(sol["x"]).reshape(self.n_classes_)
        return self.props_

    @property
    def sigma(self):
        return self.kernel.sigma

    def get_kl(self, pi_target: np.array) -> np.array:
        return super().get_kl(pi_target)
