from time import time, time_ns
import numpy as np
import torch
import scipy.linalg as slin

def h_norm_func(W):
    d = W.shape[0]
    I = np.eye(d)
    A = W * W
    N = np.linalg.norm(A, ord=1)
    # N = torch.linalg.matrix_norm(A, ord=1)
    h = - np.log(np.linalg.det(1.001 * I - A / N)) + d * np.log(1.001)
    # h = - torch.slogdet(I - A / N)[1]
    delta_h = np.linalg.norm(2 / N * np.linalg.matrix_power(1.001 * I - A / N, -1) * W, np.inf)
    return h, delta_h

def h_log_func(W, s=0.27):
    d = W.shape[0]
    I = np.eye(d)
    A = W * W
    h = - np.log(np.linalg.det(s * I - A)) + d * np.log(s)
    # h = -torch.slogdet(s * I - A)[1] + d * np.log(s)
    delta_h = np.linalg.norm(2 * np.linalg.matrix_power(s * I - A, -1) * W, np.inf)
    return h, delta_h

def h_exp_func(W):
    d = W.shape[0]
    A = W * W
    E = slin.expm(A)
    h = np.trace(E) - d
    delta_h = np.linalg.norm(2 * (E) * W, np.inf)
    return h, delta_h

def h_ploy_func(W):
    d = W.shape[0]
    I = np.eye(d)
    A = W * W
    M = I + A / d
    E = np.linalg.matrix_power(M, d - 1)
    h = (E.transpose() * M).sum() - d
    delta_h = np.linalg.norm(2 * (E) * W, np.inf)
    return h, delta_h

def h_rho_func(W):
    A = W * W
    eigenvalues = np.linalg.eigvals(A)
    h = np.max(np.abs(eigenvalues))
    _, right_eval = np.linalg.eig(A)
    left_eval = np.linalg.inv(right_eval).T
    delta_h_rho = np.linalg.norm(2 * (np.matmul(left_eval, right_eval.T) * np.linalg.matrix_power(np.matmul(left_eval.T, right_eval), -1)) * W, np.inf)
    return h, delta_h_rho

p = 50
np.random.seed(0)
W = np.random.rand(p, p)
h_norms, delta_h_norms, t_norm = [], [], []
h_logs, delta_h_logs, t_logs = [], [], []
h_exps, delta_h_exps, t_exps = [], [], []
h_ploys, delta_h_ploys, t_ploys = [], [], []
h_rhos, delta_h_rhos, t_rhos = [], [], []
for i in range(1, 51):
    sW = W * float(i / 50)
    begin_time = time() * 1000
    h_norm, delta_h_norm = h_norm_func(sW)
    t_norm.append(time() * 1000 - begin_time)
    begin_time = time() * 1000
    h_log, delta_h_log = h_log_func(sW)
    t_logs.append(time() * 1000 - begin_time)
    begin_time = time() * 1000
    h_exp, delta_h_exp = h_exp_func(sW)
    t_exps.append(time() * 1000 - begin_time)
    begin_time = time() * 1000
    h_ploy, delta_h_ploy = h_ploy_func(sW)
    t_ploys.append(time() * 1000 - begin_time)
    begin_time = time() * 1000
    h_rho, delta_h_rho = h_rho_func(sW)
    t_rhos.append(time() * 1000 - begin_time)
    h_norms.append(h_norm)
    delta_h_norms.append(delta_h_norm)
    h_logs.append(h_log)
    delta_h_logs.append(delta_h_log)
    h_exps.append(h_exp)
    delta_h_exps.append(delta_h_exp)
    h_ploys.append(h_ploy)
    delta_h_ploys.append(delta_h_ploy)
    h_rhos.append(h_rho)
    delta_h_rhos.append(delta_h_rho)
print(h_norms)
print(h_logs)
print(h_exps)
print(h_ploys)
print(h_rhos)
print(delta_h_norms)
print(delta_h_logs)
print(delta_h_exps)
print(delta_h_ploys)
print(delta_h_rhos)
print(t_norm)
print(t_logs)
print(t_exps)
print(t_ploys)
print(t_rhos)

# h_norms, delta_h_norms, t_norm = [], [], []
# h_logs, delta_h_logs, t_logs = [], [], []
# h_exps, delta_h_exps, t_exps = [], [], []
# h_ploys, delta_h_ploys, t_ploys = [], [], []
# h_rhos, delta_h_rhos, t_rhos = [], [], []
# for p in [2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000]:
#     np.random.seed(0)
#     W = np.zeros((p, p))
#     for i in range(p - 1):
#         W[i, i + 1] = (-1) ** i
#     W[p-1, 0] = (-1) ** (p-1)
#     W *= 0.5
#     begin_time = time() * 1000
#     h_norm, delta_h_norm = h_norm_func(W)
#     t_norm.append(time() * 1000 - begin_time)
#     begin_time = time() * 1000
#     h_log, delta_h_log = h_log_func(W)
#     t_logs.append(time() * 1000 - begin_time)
#     begin_time = time() * 1000
#     h_exp, delta_h_exp = h_exp_func(W)
#     t_exps.append(time() * 1000 - begin_time)
#     begin_time = time() * 1000
#     h_ploy, delta_h_ploy = h_ploy_func(W)
#     t_ploys.append(time() * 1000 - begin_time)
#     begin_time = time() * 1000
#     h_rho, delta_h_rho = h_rho_func(W)
#     t_rhos.append(time() * 1000 - begin_time)
#     h_norms.append(h_norm)
#     delta_h_norms.append(delta_h_norm)
#     h_logs.append(h_log)
#     delta_h_logs.append(delta_h_log)
#     h_exps.append(h_exp)
#     delta_h_exps.append(delta_h_exp)
#     h_ploys.append(h_ploy)
#     delta_h_ploys.append(delta_h_ploy)
#     h_rhos.append(h_rho)
#     delta_h_rhos.append(delta_h_rho)
# print(h_norms)
# print(h_logs)
# print(h_exps)
# print(h_ploys)
# print(h_rhos)
# print(delta_h_norms)
# print(delta_h_logs)
# print(delta_h_exps)
# print(delta_h_ploys)
# print(delta_h_rhos)
# print(t_norm)
# print(t_logs)
# print(t_exps)
# print(t_ploys)
# print(t_rhos)
