from build.lib.quapy.data import LabelledCollection
from quapy.method.confidence import WithConfidenceABC
from quapy.protocol import AbstractProtocol
import numpy as np
from tqdm import tqdm
import quapy as qp
from joblib import Parallel, delayed
import copy


def temp_calibration(method:WithConfidenceABC,
                     train:LabelledCollection,
                     val_prot:AbstractProtocol,
                     temp_grid=[.5, 1., 1.5, 2., 5., 10., 100.],
                     nominal_coverage=0.95,
                     amplitude_threshold='auto',
                     n_jobs=1,
                     verbose=True):

    assert (amplitude_threshold == 'auto' or (isinstance(amplitude_threshold, float)) and amplitude_threshold < 1.), \
        f'wrong value for {amplitude_threshold=}, it must either be "auto" or a float < 1.0.'

    if amplitude_threshold=='auto':
        n_classes = train.n_classes
        amplitude_threshold = .1/np.log(n_classes+1)

    if isinstance(amplitude_threshold, float) and amplitude_threshold > 0.1:
        print(f'warning: the {amplitude_threshold=} is too large; this may lead to uninformative regions')

    def evaluate_temperature_job(temp):
        local_method = copy.deepcopy(method)
        local_method.temperature = temp

        coverage = 0
        amplitudes = []
        errs = []

        for i, (sample, prev) in enumerate(val_prot()):
            point_estim, conf_region = local_method.predict_conf(sample)

            if prev in conf_region:
                coverage += 1

            amplitudes.append(conf_region.montecarlo_proportion(n_trials=50_000))
            errs.append(qp.error.mae(prev, point_estim))

        mean_coverage = coverage / val_prot.total()
        mean_amplitude = np.mean(amplitudes)

        if verbose:
            print(f'Temperature={temp} got coverage={mean_coverage*100:.2f}% amplitude={mean_amplitude*100:.2f}%')

        return temp, mean_coverage, mean_amplitude

    temp_grid = sorted(temp_grid)
    method.fit(*train.Xy)

    raw_results = Parallel(n_jobs=n_jobs, backend="loky")(
        delayed(evaluate_temperature_job)(temp)
        for temp in tqdm(temp_grid, disable=not verbose)
    )
    results = [
        (temp, cov, amp)
        for temp, cov, amp in raw_results
        if amp < amplitude_threshold
    ]

    chosen_temperature = 1.
    if len(results) > 0:
        chosen_temperature = min(results, key=lambda x: abs(x[1]-nominal_coverage))[0]

    print(f'chosen_temperature={chosen_temperature:.2f}')

    return chosen_temperature