import numpy as np
import numpy.random as npr
import pymc3 as pm
from math import sqrt
from .utils import triangle_indices, priority_vector, theoretical_matrix


def lin2015_model(empirical, theoretical):
    tau = 1 / 0.5  # 2, tau ^ -2
    sigma = 1 / 0.25  # 4,  sigma ^ -2
    denominator = sigma + tau
    n_rows, n_cols = empirical.shape
    mu_hat = np.empty((n_rows, n_cols))
    for i in range(n_rows):
        for j in range(n_cols):
            mu_hat[i, j] = (empirical[i, j] * sigma + theoretical[i, j] * tau) / denominator
    return mu_hat


def normal_inverse_gamma_model(empirical, theoretical, iters, burn):
    def model(empirical_val, theoretical_val):
        """ Sample from the joint posterior (mu, sigma | data) """
        sigma_0 = 4  # 2 ** 2
        mu = np.zeros(iters)
        sigma = np.zeros(iters)
        mu[0] = 0
        sigma[0] = 1

        for i in range(1, iters):
            weight = sigma_0 / (sigma[i - 1] ** 2 + sigma_0)

            mu_mean = weight * empirical_val + (1 - weight) * theoretical_val
            mu_sd = sqrt(weight * sigma[i - 1] ** 2)
            mu[i] = npr.normal(mu_mean, mu_sd)

            a = 0.001 + 0.5
            b = 0.001 + 0.5 * (theoretical_val - mu[i]) ** 2
            sigma[i] = sqrt(1 / npr.gamma(a, b))

        # discard burn-in
        mu = mu[burn:]
        return np.mean(mu)

    n, m = theoretical.shape
    indices = triangle_indices(empirical, triangle="upper")
    estimated_vector = np.array([model(empirical[r][c],
                                       theoretical[r][c])
                                 for r, c in indices])

    estimated_matrix = np.zeros((n, m))
    estimated_matrix[np.triu_indices(n, 1)] = estimated_vector
    return estimated_matrix


# noinspection PyPep8Naming
def log_matrix_normal_model(empirical):
    """
    https://stackoverflow.com/questions/53222664/pymc3-passing-stochastic-covariance-matrix-to-pm-mvnormal

    https://docs.pymc.io/notebooks/LKJ.html
    """
    with pm.Model() as model:
        m, n = empirical.shape

        log_pv = np.log(priority_vector(empirical))
        log_empirical = np.log(empirical)
        theoretical = theoretical_matrix(log_pv, logs=True)

        # mu = pm.Uniform('mu', lower=1/9, upper=9, shape=5)
        sd = pm.HalfCauchy.dist(1)
        packed_L = pm.LKJCholeskyCov('packed_L', n=m, eta=2., sd_dist=sd)
        L = pm.expand_packed_triangular(m, packed_L)
        sigma = pm.Deterministic('sigma', L.dot(L.T))
        # before
        mu_1 = pm.MvNormal('mu_1', mu=theoretical, chol=L, shape=(m, n))
        # after
        # mu_1 = pm.MvNormal('mu_1', mu=theoretical, cov=np.eye(m), shape=(m, n))

        rowcov = np.eye(m)
        colcov = np.eye(n)
        # before
        obs = pm.MatrixNormal('obs', mu=mu_1, colcov=colcov, rowcov=rowcov, shape=(m, n), observed=log_empirical)
        # after
        # obs = pm.MatrixNormal('obs', mu=mu_1, colchol=L, rowcov=rowcov, shape=(m, n), observed=log_empirical)

    return model, theoretical


# noinspection PyPep8Naming
def log_mv_normal_model(empirical):
    with pm.Model() as model:
        m, n = empirical.shape
        log_pv = np.log(priority_vector(empirical))
        log_empirical = np.log(empirical)
        theoretical = theoretical_matrix(log_pv, logs=True)

        packed_L = pm.LKJCholeskyCov('packed_L', n=m, eta=2., sd_dist=pm.HalfCauchy.dist(1))
        L = pm.expand_packed_triangular(m, packed_L)
        sigma = pm.Deterministic('sigma', L.dot(L.T))
        mu_1 = pm.MvNormal('mu_1', mu=theoretical, chol=L, shape=(m, n))

        cov = np.eye(5)
        obs = pm.MvNormal('obs', mu=mu_1, cov=cov, observed=log_empirical)
    return model, theoretical


def log_mv_normal_double_cholesky_model(empirical):
    with pm.Model() as model:
        m, n = empirical.shape
        log_pv = np.log(priority_vector(empirical))
        log_empirical = np.log(empirical)
        theoretical = theoretical_matrix(log_pv, logs=True)

        packed_L = pm.LKJCholeskyCov('packed_L', n=m, eta=2., sd_dist=pm.HalfCauchy.dist(1))
        L = pm.expand_packed_triangular(m, packed_L)
        sigma = pm.Deterministic('sigma', L.dot(L.T))
        mu_1 = pm.MvNormal('mu_1', mu=theoretical, chol=L, shape=(m, n))
        obs = pm.MvNormal('obs', mu=mu_1, chol=L, observed=log_empirical)
    return model, theoretical


# noinspection PyPep8Naming
def mv_normal_model(empirical):
    with pm.Model() as model:
        m, n = empirical.shape
        pv = priority_vector(empirical)
        theoretical = theoretical_matrix(pv, logs=False)

        packed_L = pm.LKJCholeskyCov('packed_L', n=m, eta=2., sd_dist=pm.HalfCauchy.dist(1))
        L = pm.expand_packed_triangular(m, packed_L)
        sigma = pm.Deterministic('sigma', L.dot(L.T))
        cov = np.eye(5)
        mu_1 = pm.MvNormal('mu_1', mu=theoretical, cov=cov, shape=(m, n))
        obs = pm.MvNormal('obs', mu=mu_1, chol=L, observed=empirical)
    return model, theoretical


def mv_normal_double_cholesky_model(empirical):
    with pm.Model() as model:
        m, n = empirical.shape
        pv = priority_vector(empirical)
        theoretical = theoretical_matrix(pv, logs=False)

        packed_L = pm.LKJCholeskyCov('packed_L', n=m, eta=2., sd_dist=pm.HalfCauchy.dist(1))
        L = pm.expand_packed_triangular(m, packed_L)
        sigma = pm.Deterministic('sigma', L.dot(L.T))
        mu_1 = pm.MvNormal('mu_1', mu=theoretical, chol=L, shape=(m, n))
        obs = pm.MvNormal('obs', mu=mu_1, chol=L, observed=empirical)
    return model, theoretical


# noinspection PyPep8Naming
def matrix_normal_model(empirical):
    """
    https://stackoverflow.com/questions/53222664/pymc3-passing-stochastic-covariance-matrix-to-pm-mvnormal

    https://docs.pymc.io/notebooks/LKJ.html
    """
    with pm.Model() as model:
        m, n = empirical.shape

        pv = priority_vector(empirical)
        theoretical = theoretical_matrix(pv, logs=False)

        # mu = pm.Uniform('mu', lower=1/9, upper=9, shape=5)
        packed_L = pm.LKJCholeskyCov('packed_L', n=m, eta=2., sd_dist=pm.HalfCauchy.dist(1))
        L = pm.expand_packed_triangular(m, packed_L)
        sigma = pm.Deterministic('sigma', L.dot(L.T))
        mu_1 = pm.MvNormal('mu_1', mu=theoretical, chol=L, shape=(m, n))

        rowcov = np.eye(m)
        colcov = np.eye(n)
        obs = pm.MatrixNormal('obs', mu=mu_1, colcov=colcov, rowcov=rowcov, shape=(m, n), observed=empirical)

    return model, theoretical