from .measures import (consistency_ratio, minimum_violation, euclidean_distance)
from .utils import rank_data, save_pymc_trace
from .models import priority_vector
import numpy as np
import pandas as pd
from tqdm import tqdm


def compare_statistics(original, estimated, original_theoretical, estimated_theoretical):
    data = {
        "original": [
            consistency_ratio(original),
            minimum_violation(original),
            euclidean_distance(original, original_theoretical),
            str(rank_data(priority_vector(original))),
            str(np.round(priority_vector(original), 3))
        ],
        "estimated": [
            consistency_ratio(estimated),
            minimum_violation(estimated),
            euclidean_distance(estimated, estimated_theoretical),
            str(rank_data(priority_vector(estimated))),
            str(np.round(priority_vector(estimated), 3))
        ],
    }
    index = ["CR", "MV", "ED", "ranking", "priority_vector"]
    return pd.DataFrame(data=data, index=index)


def run_variant(method, data_set, variant_name, mcmc=False):
    if mcmc:
        for key, empirical_matrix in tqdm(data_set.items()):
            estimated_matrix, sigma, model, trace, original_theo, estimated_theo = method(empirical_matrix)
            stats = compare_statistics(empirical_matrix, estimated_matrix, original_theo, estimated_theo)

            with open(f"./results/estimated_matrices/{variant_name}.txt", "a") as fh:
                print(key, file=fh)
                print(estimated_matrix, file=fh)

            with open(f"./results/estimated_matrices/sigma_{variant_name}.txt", "a") as fh:
                print(key, file=fh)
                print(sigma, file=fh)

            with open(f"./results/estimated_matrices/theo_{variant_name}.txt", "a") as fh:
                print(key, file=fh)
                print(original_theo, file=fh)
                print(estimated_theo, file=fh)

            with open(f"./results/stats/{variant_name}.txt", "a") as fh:
                print(key, file=fh)
                print(stats, file=fh)

            save_pymc_trace(f"./results/traces/{key}_{variant_name}.pkl", model, trace)

    else:
        for key, empirical_matrix in tqdm(data_set.items()):
            estimated_matrix, original_theo, estimated_theo = method(empirical_matrix)
            stats = compare_statistics(empirical_matrix, estimated_matrix, original_theo, estimated_theo)

            with open(f"./results/estimated_matrices/{variant_name}.txt", "a") as fh:
                print(key, file=fh)
                print(estimated_matrix, file=fh)

            with open(f"./results/estimated_matrices/theo_{variant_name}.txt", "a") as fh:
                print(key, file=fh)
                print(original_theo, file=fh)
                print(estimated_theo, file=fh)

            with open(f"./results/stats/{variant_name}.txt", "a") as fh:
                print(key, file=fh)
                print(stats, file=fh)