import pickle
import os
from time import time
from collections import defaultdict
from tqdm import tqdm
import numpy as np
from sklearn.linear_model import LogisticRegression

import quapy as qp
from KDEy.kdey_devel import KDEyMLauto, optim_minimize
from method._kdey import KDEBase
from quapy.method.aggregative import PACC, EMQ, KDEyML
from quapy.model_selection import GridSearchQ
from quapy.protocol import UPP
from pathlib import Path
from quapy import functional as F

SEED = 1


def newLR():
    return LogisticRegression(max_iter=1000)#, C=1, class_weight='balanced')

SAMPLE_SIZE=150
qp.environ['SAMPLE_SIZE'] = SAMPLE_SIZE

epsilon = 1e-10
# n_bags_test = 2
DATASETS = [qp.datasets.UCI_MULTICLASS_DATASETS[21]]
# DATASETS = qp.datasets.UCI_MULTICLASS_DATASETS
for i, dataset in enumerate(DATASETS):
    data = qp.datasets.fetch_UCIMulticlassDataset(dataset)
    n_classes = data.n_classes
    print(f'{i=}')
    print(f'{dataset=}')
    print(f'{n_classes=}')
    print(len(data.training))
    print(len(data.test))

    train, test = data.train_test
    train_prev = train.prevalence()
    test_prev = test.prevalence()

    print(f'train-prev = {F.strprev(train_prev)}')
    print(f'test-prev = {F.strprev(test_prev)}')

    # protocol = UPP(test, repeats=n_bags_test)
    #
    # for sample, prev in protocol():
    #     print(f'sample-prev = {F.strprev(prev)}')

    # prev = np.asarray([0.2, 0.3, 0.5])
    # prev = np.asarray([0.33, 0.33, 0.34])
    # prev = train_prev

    # sample = test.sampling(SAMPLE_SIZE, *prev, random_state=1)
    # print(f'sample-prev = {F.strprev(prev)}')

    repeats = 10
    prot = UPP(test, sample_size=SAMPLE_SIZE, repeats=repeats)
    kde = KDEyMLauto(newLR())
    kde.fit(train)
    tr_posteriors, tr_y = kde.classif_predictions.Xy
    for it, (sample, prev) in tqdm(enumerate(prot()), total=repeats):
        te_posteriors = kde.classifier.predict_proba(sample)
        classes = train.classes_

        xaxis = []
        ae_error = []
        rae_error = []
        mse_error = []
        kld_error = []
        likelihood_val = []

        # for bandwidth in np.linspace(0.01, 0.2, 50):
        for bandwidth in np.logspace(-3, 0.5, 50):
            mix_densities = kde.get_mixture_components(tr_posteriors, tr_y, classes, bandwidth)
            test_densities = [kde.pdf(kde_i, te_posteriors) for kde_i in mix_densities]

            def neg_loglikelihood_prev(prev):
                test_mixture_likelihood = sum(prev_i * dens_i for prev_i, dens_i in zip(prev, test_densities))
                test_loglikelihood = np.log(test_mixture_likelihood + epsilon)
                return -np.sum(test_loglikelihood)

            init_prev = np.full(fill_value=1 / n_classes, shape=(n_classes,))
            pred_prev, likelihood = optim_minimize(neg_loglikelihood_prev, init_prev, return_loss=True)

            xaxis.append(bandwidth)
            ae_error.append(qp.error.ae(prev, pred_prev))
            rae_error.append(qp.error.rae(prev, pred_prev))
            mse_error.append(qp.error.mse(prev, pred_prev))
            kld_error.append(qp.error.kld(prev, pred_prev))
            likelihood_val.append(likelihood)

        import matplotlib.pyplot as plt

        # Crear la figura
        fig, ax1 = plt.subplots(figsize=(8, 6))

        # Pintar las series ae_error, rae_error, y kld_error en el primer eje Y
        ax1.plot(xaxis, ae_error, label='AE Error', marker='o', color='b')
        ax1.plot(xaxis, rae_error, label='RAE Error', marker='s', color='g')
        ax1.plot(xaxis, kld_error, label='KLD Error', marker='^', color='r')
        ax1.plot(xaxis, mse_error, label='MSE Error', marker='^', color='c')
        ax1.set_xscale('log')

        # Configurar etiquetas para el primer eje Y
        ax1.set_xlabel('Bandwidth')
        ax1.set_ylabel('Error Value')
        ax1.grid(True)
        ax1.legend(loc='upper left')

        # Crear un segundo eje Y que comparte el eje X
        ax2 = ax1.twinx()

        # Pintar likelihood_val en el segundo eje Y
        ax2.plot(xaxis, likelihood_val, label='(neg)Likelihood', marker='x', color='purple')

        # Configurar etiquetas para el segundo eje Y
        ax2.set_ylabel('Likelihood Value')
        ax2.legend(loc='upper right')

        # Mostrar el gráfico
        plt.title('Error Metrics vs Bandwidth')
        # plt.show()
        os.makedirs('./plots/likelihood/', exist_ok=True)
        plt.savefig(f'./plots/likelihood/fig{it}.png')
        plt.close()