# Import useful libraries
import numpy as np
import scipy as sc
import networkx as nx
import pandas as pd


n = 10
d = 3
T = 250
p = 3
beta = np.array([2, -1, 3])
sigma2 = 1

# Set a seed if desired
seed = 2025
rng = np.random.default_rng(seed = seed)

def build_graph(d, n, seed):

    g = nx.random_regular_graph(d, n, seed)
    W = nx.to_numpy_array(g)

    evals, evecs = np.linalg.eigh(W)

    hW = sc.linalg.expm(W)
    Ws = np.array([
        np.outer(evecs[:, i], np.linalg.inv(evecs)[i, :]) for i in range(n)
    ])

    print("Evec version is Idempotent:", np.allclose(Ws[0], Ws[0] @ Ws[0]))


    return evals, Ws, hW

def generate_linear_component(T, n, p, beta, sigma2, rng):
    X = np.concatenate([
    np.ones((T, n, 1)),
    rng.normal(size = (T, n, p - 1))
        ], axis = -1)
    beta = np.array([2, -1, 3])

    sigma2 = 1
    e = np.sqrt(sigma2) * rng.normal(size = (T, n))

    # Linear component
    eta = X @ beta + e

    return X, eta

def transform_linear_component(eta, hW):
    Y = np.einsum("ij,hj", np.linalg.inv(hW), eta)

    return Y

def computation(X, n, T):
    
    X_stack = X.reshape(-1, 3)
    A_stack = np.linalg.inv(X_stack.T @ X_stack) @ X_stack.T
    M_stack = np.eye(T * n) - X_stack @ np.linalg.inv(X_stack.T @ X_stack) @ X_stack.T

    return A_stack, M_stack

def estimate_kappa(Y, Ws, T, n, A_stack, M_stack):

    def loss(kappa_1p_hat):
        kappa_hat = np.concatenate([[1], kappa_1p_hat], axis = 0)

        # Slow method:
        # hY_hat = np.zeros((T, n))
        # for t in range(Y.shape[0]):
        #     for i in range(kappa_hat.shape[0]):
        #         hY_hat[t, :] += kappa_hat[i] * (Ws[i] @ Y[t])

        # Fast method
        hY_hat = np.einsum("i,ijk,hk", kappa_hat, Ws, Y)
        hY_hat_stack = hY_hat.reshape(-1)

        # Compute log determinant term
        S_kappa_hat = np.sum([
            kappa_hat[i] * Ws[i] for i in range(n)
        ], axis = 0)
        sign, abslogdet = np.linalg.slogdet(S_kappa_hat)
        logdet = sign * abslogdet

        # Compute some garbage
        beta_hat = A_stack @ hY_hat_stack
        sigma2_hat = np.mean(np.square(M_stack @ hY_hat_stack))

        # Concentrated log likelihood
        L = -0.5 * n * T * np.log(sigma2_hat) + T * logdet

        return -L

    kappa_1p_init = np.ones(n - 1)
    res = sc.optimize.minimize(loss, x0 = kappa_1p_init)

    return np.concatenate([[1], res.x])

#choose number of simulations

num_sims = 250

evals, Ws, hW = build_graph(d, n, seed)
true_kappa = np.exp(evals) / np.exp(evals[0])
estimates = []

for sim in range(num_sims):
    X, eta = generate_linear_component(T, n, p, beta, sigma2, rng)
    Y = transform_linear_component(eta, hW)
    A_stack, M_stack = computation(X, n, T)
    estimated_kappa = estimate_kappa(Y, Ws, T, n, A_stack, M_stack)
    estimates.append(estimated_kappa)

# average the estimated kappa
estimates = np.array(estimates)
mean_kappa = estimates.mean(axis=0)

print("True kappa:", np.exp(evals) / np.exp(evals[0]))
print("estimated Kappa average: ", mean_kappa)

Evec version is Idempotent: True
True kappa: [  1.           1.57805456   2.3417808    4.28959703   8.06781785
  11.66033367  31.69607313  42.06966776  99.76747962 234.20406245]
estimated Kappa average:  [  1.           1.57456493   2.34271076   4.29437674   8.0650336
  11.67029994  31.68573327  42.07039518  99.7791601  233.47800453]