import numpy as np

np.set_printoptions(suppress=True, precision=4)

from tests.unit.fixtures import local_perturbations_model_randomized
from experiments.util import cum_means
from experiments.consistency.model_to_distance_info import (
    main_fitter,
    fitted_model_from_true_model,
    score_sample_with_two_models,
    sample_from_model,
)


def _distance_sequence_from_scores(scores_1, scores_2):
    """
    We want to compute the distance between two -local perturbation- models.

    To do this, we simply take the usual distance between two hmm models; the pretrained predictive model
    stays constant.

    The desired distance is the limiting value of this distance sequence.

    Reference:
        "A Probabilistic Distance Measure for Hidden Markov Models" by Juang and Rabiner.

    Note: for this distance to be correct, we need sample scored by these two models to have
    been generated by model_1!   That's why we keep this a private function and make the (currently awkward)
    get_distance_sequence() function

    :param scores_1: np.array
        Scores from Model 1
    :param scores_2: np.array
        Scores from Model 2
    """
    log_likelihood_ratios = np.log(scores_1) - np.log(scores_2)
    distance_sequence = cum_means(log_likelihood_ratios)
    return distance_sequence


def get_distance_sequence(
    model_1,
    model_2,
    score_seed,
    perturbations_seed_sampler,
    perturbations_seed_scorer,
    perturbations_inv_strength,
    N_samples,
):
    """
    We want to compute the distance between two -local perturbation- models.

    To do this, we simply take the usual distance between two hmm models; the pretrained predictive model
    stays constant.

    The desired distance is the limiting value of this distance sequence.

    Reference:
        "A Probabilistic Distance Measure for Hidden Markov Models" by Juang and Rabiner.
            See Eq. 6 for the distance measure.
    """
    sample_from_model_1 = sample_from_model(
        model_1,
        score_seed,
        perturbations_seed_sampler,
        perturbations_inv_strength,
        N_samples,
    )

    scores_from_model_1, scores_from_model_2 = score_sample_with_two_models(
        sample_from_model_1,
        model_1,
        model_2,
        score_seed,
        perturbations_seed_scorer,
        perturbations_inv_strength,
    )

    distance_sequence = _distance_sequence_from_scores(
        scores_from_model_1, scores_from_model_2
    )

    return (
        distance_sequence,
        sample_from_model_1,
        scores_from_model_1,
        scores_from_model_2,
    )


def distance_from_distance_sequence(
    distance_sequence, convergence_thresh=0.005, convergence_tail_size=20
):
    """

    The convergence criterion for the distance sequence is that the last "convergence_tail_size"
    entries of the distance_sequence must lie within an interval of diameter "convergence_thresh".

        -If it does, we return the last value of the distance_sequence as the distance
        -If it doesn't, we return a ValueError.

    :param convergence_thresh: Float.
        The tail of the distance sequence must lie within a cover whose diameter is less than the
        convergence threshold; otherwise we say the sequence has not converged.
    :param convergence_tail_size: Int.
        The last N entries of the distance_sequence, where N is given by the convergence_tail_size,
        are used to check the convergence criterion.
    """
    distance_sequence_tail = distance_sequence[-convergence_tail_size:]
    tail_cover_diameter = np.max(distance_sequence_tail) - np.min(distance_sequence_tail)
    if tail_cover_diameter < convergence_thresh:
        distance = distance_sequence[-1]
        return distance
    else:
        raise ValueError(
            "Distance sequence does not converge with tail size %d and "
            "convergence thresh %.04f, try again."
            % (convergence_tail_size, convergence_thresh)
        )


def get_distance_ndarray(
    Ns=[100, 1000, 5000, 10000],
    Ws=[10, 20, 40, 80, 160],
    Ks=[3],
    n_reps=5,
    perturbations_inv_strength=1e2,
):

    distance_ndarray = np.zeros((len(Ns), len(Ws), len(Ks)))

    for (w, W) in enumerate(Ws):  # lowercase is index, uppercase is raw value
        for (n, N) in enumerate(Ns):
            for (k, K) in enumerate(Ks):
                distances = []
                for i in range(n_reps):
                    print(
                        "\n Now processing rep %d for n=%d and w=%d and k=%d...."
                        % (i + 1, N, W, K)
                    )
                    true_model = local_perturbations_model_randomized(
                        K=K, W=W, n_Kstars=1, seed=i
                    )
                    fitted_model, sample_train_from_true = fitted_model_from_true_model(
                        true_model,
                        score_seed=i,
                        perturbations_seed_for_training=1111,
                        perturbations_inv_strength=perturbations_inv_strength,
                        N_samples=N,
                    )

                    ### ---- Step 2: calculate distance between perturbed model and the actual model ----
                    distance_sequence, sample_from_true_model, scores_from_true_model, scores_from_fitted_model = get_distance_sequence(
                        true_model,
                        fitted_model,
                        score_seed=i,
                        perturbations_seed_sampler=2222,
                        perturbations_seed_scorer=3333,
                        perturbations_inv_strength=perturbations_inv_strength,
                        N_samples=5000,
                    )
                    distance = distance_from_distance_sequence(
                        distance_sequence,
                        convergence_thresh=0.005,
                        convergence_tail_size=20,
                    )
                    print("Obtained distance: %.04f" % distance)
                    distances.append(distance)
                mean_abs_distance = np.mean([abs(dist) for dist in distances])
                print("-----Mean absolute distance: %.04f-----" % mean_abs_distance)
                distance_ndarray[n, w, k] = mean_abs_distance
    return distance_ndarray
