""" This code is used to generate the data in Figure 2 of the paper by:
(i) Generating N=N1+N2 number of data samples from 2-trefoil knots of intrinsic dimension 1
(ii) Calculating the quantities in Lemma 1 of the paper
(iii) Saving the generated information in an array stored in the same directory as this code.
Once the data is saved, execute the following to generate plots in Figure 1:
python WMC_synthetic/240519_plot_geometric_LambdavN.py

In this code, we have:
eta = homogenization parameter
alpha_list = list of values of alpha
N_list = list of number of data samples (N1=N2) generated from each manifold. Total number of data samples = N1+N2
seed_list = list of seeed values used to generate random embedding of N1, N2 data samples """



import numpy as np
import os
import sys
import matplotlib.pyplot as plt
from cvxopt import matrix, solvers

sys.path.insert(0, './')

from manifoldGen import manifoldGen
from utils import *

# Set input parameters
dim = 2 #=d_l+1
eta = 1
alpha_list = [2,5,50]
N_list=[40+5*i for i in range(100)]
seed_list=[i for i in range(50)]

# Save output metrics
inradius = np.zeros((len(N_list), len(seed_list),len(alpha_list)))
gamma_W = np.ones((len(N_list), len(seed_list),len(alpha_list)))*np.inf
ineq = np.zeros((len(N_list), len(seed_list),len(alpha_list)))
LL = np.zeros((len(N_list), len(seed_list),len(alpha_list)))
LU = np.zeros((len(N_list), len(seed_list),len(alpha_list)))
lambd_list = np.zeros((len(N_list), len(seed_list),len(alpha_list)))

for Ni in range(len(N_list)):
    for s in range(len(seed_list)):
        np.random.seed(seed_list[s])
        N = N_list[Ni]
        N2 = N
        Y,__,gtruth,__  = manifoldGen('2trefoils-asym',N_list[Ni],N2)
        n = max(gtruth)
        y = Y[:,0]

        # Generate W
        dist = np.sqrt(np.sum((Y[:,1:].T-y.T)**2,axis=1)) 
        w = dist/np.sum(dist)

        # Find 2 nearest neighbors in the manifold
        knn = np.argsort(w[0:N-1])[:dim]+1
        Y_knn = Y[:,knn]
        w_knn = w[knn-1]

        for alph in range(len(alpha_list)):

            # Calculate inradius
            r = calc_inradius(y,Y_knn,w_knn,eta)

            # Calculate dist(y,S)
            dist_x_s = project_on_subspace(y,Y_knn)

            # Calculate lower limit of the interval Lambda
            LLi = 1/(r*(np.linalg.norm(np.append(y,eta))-dist_x_s))

            # Set value of lambda
            lambd = LLi*alpha_list[alph] 
            lambd_list[Ni,s,alph] = lambd
            
            # Find optimal dual solution
            nu = solve_dual(y,Y_knn,w,lambd,eta)
            nu_proj = project_on_norm_subspace(nu,Y_knn,eta)
            
            #Calculate gamma
            mu = calc_gamma(Y,nu_proj,N,N2,w,r,eta)

            # Calculate upper limit of the interval Lambda
            rhs = np.linalg.norm(np.append(y,eta))*mu/(1+mu)
            LUi = mu/(r*dist_x_s)

            # Save metrics
            inradius[Ni,s,alph] = r
            gamma_W[Ni,s,alph] = mu
            ineq[Ni,s,alph] = rhs-dist_x_s
            LL[Ni,s,alph] = LLi
            LU[Ni,s,alph] = LUi


np.savez('240519-variables.npz', inradius=inradius, gamma_W=gamma_W, ineq=ineq, LL=LL, LU=LU, lambd_list=lambd_list, N_list=N_list, alpha_list=alpha_list, seed_list=seed_list)

# Calculate mean and standard deviations of output metrics
r_mean = np.zeros((len(N_list),len(alpha_list)))
r_std = np.zeros((len(N_list),len(alpha_list)))
LL_mean = np.zeros((len(N_list),len(alpha_list)))
LL_std = np.zeros((len(N_list),len(alpha_list)))
LU_mean = np.zeros((len(N_list),len(alpha_list)))
LU_std = np.zeros((len(N_list),len(alpha_list)))
gamma_W_mean = np.zeros((len(N_list),len(alpha_list)))
gamma_W_std = np.zeros((len(N_list),len(alpha_list)))
ineq_mean = np.zeros((len(N_list),len(alpha_list)))
ineq_std = np.zeros((len(N_list),len(alpha_list)))
lambd_mean = np.zeros((len(N_list),len(alpha_list)))
lambd_std = np.zeros((len(N_list),len(alpha_list)))
for alph in range(len(alpha_list)):
    for Ni in range(len(N_list)):
        r_mean[Ni,alph] = np.mean(inradius[Ni,:,alph])
        r_std[Ni,alph] = np.std(inradius[Ni,:,alph])
        LL_mean[Ni,alph] = np.mean(LL[Ni,:,alph])
        LL_std[Ni,alph] = np.std(LL[Ni,:,alph])
        LU_mean[Ni,alph] = np.mean(LU[Ni,:,alph])
        LU_std[Ni,alph] = np.std(LU[Ni,:,alph])
        gamma_W_mean[Ni,alph] = np.mean(gamma_W[Ni,:,alph])
        gamma_W_std[Ni,alph] = np.std(gamma_W[Ni,:,alph])
        ineq_mean[Ni,alph] = np.mean(ineq[Ni,:,alph])
        ineq_std[Ni,alph] = np.std(ineq[Ni,:,alph])
        lambd_mean[Ni,alph] = np.mean(lambd_list[Ni,:,alph])
        lambd_std[Ni,alph] = np.std(lambd_list[Ni,:,alph])



# Generate plots
linestyle = [{"linestyle":"solid", "linewidth":2, "markeredgewidth":1, "elinewidth":1, "capsize":4},
             {"linestyle":"dotted", "linewidth":2, "markeredgewidth":1, "elinewidth":1, "capsize":4},
             {"linestyle":(0, (3, 10, 1, 10)),  "markeredgewidth":1, "elinewidth":1, "capsize":4}]
# font = {#'weight' : 'bold',
        # 'size'   : 22}
fig, axes = plt.subplots(nrows=1, ncols=len(alpha_list), figsize=(24, 8))
fig.tight_layout(pad=4)
# plt.rc('font', **font)
N_list = [i*2 for i in N_list]
for alph in range(len(alpha_list)):
    axes[alph].tick_params(axis='both', which='major', labelsize=15)
    axes[alph].errorbar(N_list,lambd_mean[:,alph],lambd_std[:,alph],color="b",**linestyle[0])
    axes[alph].errorbar(N_list,LL_mean[:,alph],LL_std[:,alph],color="g",**linestyle[0])
    axes[alph].errorbar(N_list,LU_mean[:,alph],LU_std[:,alph],color="r",**linestyle[0])
    axes[alph].set_xlabel("N", fontsize = 15)
    axes[alph].legend(["$\lambda$","$\Lambda^l$","$\Lambda^u$"])
    axes[alph].set_title(f"$\lambda={alpha_list[alph]}*\Lambda^l$")


#Save plots
filename = f'./output_WMCsynthetic/Lambda_r{len(seed_list)}.pdf'

if not os.path.exists(os.path.dirname(filename)):
    os.makedirs(os.path.dirname(filename))
plt.savefig(filename, bbox_inches='tight')
plt.close()
