from abc import ABC, abstractmethod
import numpy as np
from sklearn.linear_model import LogisticRegression

from BayesianKDEy._bayeisan_kdey import BayesianKDEy
from BayesianKDEy._bayesian_mapls import BayesianMAPLS
from BayesianKDEy.commons import KDEyCLR, KDEyCLR2
from quapy.method.aggregative import CC, ACC, EMQ, DMy, KDEyML
from quapy.method.base import BaseQuantifier
from quapy.method.confidence import AggregativeBootstrap, BayesianCC


def get_experimental_methods():
    #yield BootsCC()
    #yield BayACC()
    #yield BayEMQ()
    #yield BayKDEyGau()
    #yield BayKDEyAit()
    yield BayKDEyAit2()


# commons
# ------------------------------------------------------------
Cls = LogisticRegression
cls_hyper = {
    'classifier__C': np.logspace(-4,4,9),
    'classifier__class_weight': ['balanced', None]
}
hdy_hyper = {'nbins': [3, 4, 5, 8, 9, 10, 12, 14, 16, 32], **cls_hyper}
kdey_hyper = {'bandwidth': np.logspace(-3, -1, 10), **cls_hyper}
kdey_hyper_clr = {'bandwidth': np.logspace(-2, 2, 10), **cls_hyper}

def hyper2cls(hyperparams):
    return {k.replace('classifier__', ''): v for k, v in hyperparams.items()}


# method descriptor logic
# ------------------------------------------------------------
class MethodDescriptor(ABC):
    def __init__(self,
                 name: str,
                 surrogate_quantifier: BaseQuantifier,
                 hyper_parameters: dict,
                 ):
        self.name = name
        self.surrogate_quantifier_ = surrogate_quantifier
        self.hyper_parameters = hyper_parameters

    @abstractmethod
    def binary_only(self): ...

    def surrogate_quantifier(self):
        return self.surrogate_quantifier_

    def surrogate_quantifier_name(self):
        return self.surrogate_quantifier_.__class__.__name__

    @abstractmethod
    def uncertainty_aware_quantifier(self, hyperparameters): ...


class MulticlassMethodDescriptor(MethodDescriptor):

    def binary_only(self):
        return False


# specific methods definitions
# ------------------------------------------------------------
# ------------------------------------------------------------

# Bootstrap approaches:
# ------------------------------------------------------------
class BootsCC(MulticlassMethodDescriptor):
    def __init__(self):
        super().__init__(name='BoCC', surrogate_quantifier=CC(Cls()), hyper_parameters=cls_hyper)

    def uncertainty_aware_quantifier(self, hyperparameters):
        quantifier = CC(Cls()).set_params(**hyperparameters)
        return AggregativeBootstrap(quantifier, n_test_samples=1000, random_state=0)


# class BootsACC(MulticlassMethodDescriptor):
#     def __init__(self):
#         super().__init__(name='BoACC', surrogate_quantifier=ACC(Cls()), hyper_parameters=cls_hyper)
#
#     def uncertainty_aware_quantifier(self, hyperparameters):
#         quantifier = ACC(Cls()).set_params(**hyperparameters)
#         return AggregativeBootstrap(quantifier, n_test_samples=1000, random_state=0)
#
#
# class BootsEMQ(MulticlassMethodDescriptor):
#     def __init__(self):
#         super().__init__(name='BoEMQ', surrogate_quantifier=EMQ(Cls(), exact_train_prev=False), hyper_parameters=cls_hyper)
#
#     def uncertainty_aware_quantifier(self, hyperparameters):
#         quantifier = EMQ(Cls(), exact_train_prev=False).set_params(**hyperparameters)
#         return AggregativeBootstrap(quantifier, n_test_samples=1000, random_state=0)


# class BootsHDy(MethodDescriptor):
#     def __init__(self):
#         super().__init__(name='BoHDy', surrogate_quantifier=DMy(Cls()), hyper_parameters=hdy_hyper)
#
#     def uncertainty_aware_quantifier(self, hyperparameters):
#         quantifier = DMy(Cls()).set_params(**hyperparameters)
#         return AggregativeBootstrap(quantifier, n_test_samples=1000, random_state=0)
#
#     def binary_only(self):
#         return True

# class BootsKDEy(MulticlassMethodDescriptor):
#     def __init__(self):
#         super().__init__(name='BoKDEy', surrogate_quantifier=KDEyML(Cls()), hyper_parameters=kdey_hyper)
#
#     def uncertainty_aware_quantifier(self, hyperparameters):
#         quantifier = KDEyML(Cls()).set_params(**hyperparameters)
#         return AggregativeBootstrap(quantifier, n_test_samples=1000, random_state=0)


# Bayesian approaches:
# ------------------------------------------------------------
class BayACC(MulticlassMethodDescriptor):
    def __init__(self):
        super().__init__(name='BaACC!', surrogate_quantifier=ACC(Cls()), hyper_parameters=cls_hyper)

    def uncertainty_aware_quantifier(self, hyperparameters):
        classifier = Cls(**hyper2cls(hyperparameters))
        return BayesianCC(classifier, temperature=None, mcmc_seed=0)  # is actually a Bayesian variant of ACC


class BayEMQ(MulticlassMethodDescriptor):
    def __init__(self):
        super().__init__(name='BaEMQ!', surrogate_quantifier=EMQ(Cls(), exact_train_prev=False), hyper_parameters=cls_hyper)

    def uncertainty_aware_quantifier(self, hyperparameters):
        classifier = Cls(**hyper2cls(hyperparameters))
        return BayesianMAPLS(classifier, prior='uniform', temperature=None, exact_train_prev=False)


class BayKDEyGau(MulticlassMethodDescriptor):
    def __init__(self):
        kdey_hyper = {'bandwidth': np.logspace(-3, -1, 10), **cls_hyper}
        super().__init__(name='BaKDE-Gau-T!', surrogate_quantifier=KDEyML(Cls()), hyper_parameters=kdey_hyper)

    def uncertainty_aware_quantifier(self, hyperparameters):
        return BayesianKDEy(Cls(), kernel='gaussian', temperature=None, mcmc_seed=0, **hyperparameters)


class BayKDEyAit(MulticlassMethodDescriptor):
    def __init__(self):
        kdey_hyper = {'bandwidth': np.logspace(-2, 2, 10), **cls_hyper}
        super().__init__(name='BaKDE-Ait-T!', surrogate_quantifier=KDEyCLR(Cls()), hyper_parameters=kdey_hyper)

    def uncertainty_aware_quantifier(self, hyperparameters):
        return BayesianKDEy(Cls(), kernel='aitchison', temperature=None, mcmc_seed=0, **hyperparameters)
    

class BayKDEyAit2(MulticlassMethodDescriptor):
    def __init__(self):
        kdey_hyper = {'bandwidth': np.linspace(0.05, 2., 10), **cls_hyper}
        super().__init__(name='BaKDE-Ait-T!2', surrogate_quantifier=KDEyCLR2(Cls()), hyper_parameters=kdey_hyper)

    def uncertainty_aware_quantifier(self, hyperparameters):
        return BayesianKDEy(Cls(), kernel='aitchison', temperature=None, mcmc_seed=0, **hyperparameters)