from .models import (priority_vector, theoretical_matrix,
                     lin2015_model, normal_inverse_gamma_model,
                     log_mv_normal_model, log_matrix_normal_model,
                     log_mv_normal_double_cholesky_model, mv_normal_double_cholesky_model,
                     mv_normal_model, matrix_normal_model
                     )
from .utils import ensure_reciprocals, inference_step
import pymc3 as pm
import numpy as np


def lin2015_variant(empirical):
    log_pv = np.log(priority_vector(empirical))
    log_empirical = np.log(empirical)
    empirical_theo = theoretical_matrix(log_pv, logs=True)
    estimated_matrix = lin2015_model(log_empirical, empirical_theo)
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=True)
    corrected_estimated_matrix = np.exp(corrected_estimated_matrix)
    estimated_theo = theoretical_matrix(np.log(priority_vector(corrected_estimated_matrix)), logs=True)
    return corrected_estimated_matrix, np.exp(empirical_theo), np.exp(estimated_theo)


def logs_normal_ig_variant(empirical, iters=55_000, burn=5_000):
    log_pv = np.log(priority_vector(empirical))
    log_empirical = np.log(empirical)
    empirical_theo = theoretical_matrix(log_pv, logs=True)
    estimated_matrix = normal_inverse_gamma_model(log_empirical, empirical_theo, iters, burn)
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=True)
    corrected_estimated_matrix = np.exp(corrected_estimated_matrix)
    estimated_theo = theoretical_matrix(np.log(priority_vector(corrected_estimated_matrix)), logs=True)
    return corrected_estimated_matrix, np.exp(empirical_theo), np.exp(estimated_theo)


def normal_ig_variant(empirical, iters=55_000, burn=5_000):
    empirical_theo = theoretical_matrix(priority_vector(empirical), logs=False)
    estimated_matrix = normal_inverse_gamma_model(empirical, empirical_theo, iters, burn)
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=False)
    estimated_theo = theoretical_matrix(priority_vector(corrected_estimated_matrix), logs=False)
    return corrected_estimated_matrix, empirical_theo, estimated_theo


def log_matrix_normal_variant(empirical, samples=22_000, burn=2_000, cores=4):
    model, empirical_theo = log_matrix_normal_model(empirical)
    trace = inference_step(model, samples, burn, cores)
    sigma = pm.summary(trace)[-25:]["mean"].values
    estimated_matrix = pm.summary(trace)["mean"][:25].values
    estimated_matrix = estimated_matrix.reshape((5, 5))
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=True)
    corrected_estimated_matrix = np.exp(corrected_estimated_matrix)
    estimated_theo = theoretical_matrix(np.log(priority_vector(corrected_estimated_matrix)), logs=True)
    return corrected_estimated_matrix, sigma, model, trace, np.exp(empirical_theo), np.exp(estimated_theo)


def log_mv_normal_variant(empirical, samples=22_000, burn=2_000, cores=4):
    model, empirical_theo = log_mv_normal_model(empirical)
    trace = inference_step(model, samples, burn, cores)
    sigma = pm.summary(trace)[-25:]["mean"].values
    estimated_matrix = pm.summary(trace)["mean"][:25].values
    estimated_matrix = estimated_matrix.reshape((5, 5))
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=True)
    corrected_estimated_matrix = np.exp(corrected_estimated_matrix)
    estimated_theo = theoretical_matrix(np.log(priority_vector(corrected_estimated_matrix)), logs=True)
    return corrected_estimated_matrix, sigma, model, trace, np.exp(empirical_theo), np.exp(estimated_theo)


def log_mv_normal_double_chol_variant(empirical, samples=22_000, burn=2_000, cores=4):
    model, empirical_theo = log_mv_normal_double_cholesky_model(empirical)
    trace = inference_step(model, samples, burn, cores)
    sigma = pm.summary(trace)[-25:]["mean"].values
    estimated_matrix = pm.summary(trace)["mean"][:25].values
    estimated_matrix = estimated_matrix.reshape((5, 5))
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=True)
    corrected_estimated_matrix = np.exp(corrected_estimated_matrix)
    estimated_theo = theoretical_matrix(np.log(priority_vector(corrected_estimated_matrix)), logs=True)
    return corrected_estimated_matrix, sigma, model, trace, np.exp(empirical_theo), np.exp(estimated_theo)


def mv_normal_variant(empirical, samples=22_000, burn=2_000, cores=4):
    model, empirical_theo = mv_normal_model(empirical)
    trace = inference_step(model, samples, burn, cores)
    sigma = pm.summary(trace)[-25:]["mean"].values
    estimated_matrix = pm.summary(trace)["mean"][:25].values
    estimated_matrix = estimated_matrix.reshape((5, 5))
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=False)
    estimated_theo = theoretical_matrix(priority_vector(corrected_estimated_matrix), logs=False)
    return corrected_estimated_matrix, sigma, model, trace, empirical_theo, estimated_theo


def mv_normal_double_chol_variant(empirical, samples=22_000, burn=2_000, cores=4):
    model, empirical_theo = mv_normal_double_cholesky_model(empirical)
    trace = inference_step(model, samples, burn, cores)
    sigma = pm.summary(trace)[-25:]["mean"].values
    estimated_matrix = pm.summary(trace)["mean"][:25].values
    estimated_matrix = estimated_matrix.reshape((5, 5))
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=False)
    estimated_theo = theoretical_matrix(priority_vector(corrected_estimated_matrix), logs=False)
    return corrected_estimated_matrix, sigma, model, trace, empirical_theo, estimated_theo


def matrix_normal_variant(empirical, samples=22_000, burn=2_000, cores=4):
    model, empirical_theo = matrix_normal_model(empirical)
    trace = inference_step(model, samples, burn, cores)
    sigma = pm.summary(trace)[-25:]["mean"].values
    estimated_matrix = pm.summary(trace)["mean"][:25].values
    estimated_matrix = estimated_matrix.reshape((5, 5))
    corrected_estimated_matrix = ensure_reciprocals(estimated_matrix, logs=False)
    estimated_theo = theoretical_matrix(priority_vector(corrected_estimated_matrix), logs=False)
    return corrected_estimated_matrix, sigma, model, trace, empirical_theo, estimated_theo