""" This code is used to generate the data in Figure 1 of the paper by:
(i) Generating 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_Lambdavlambda.py

In this code, we have:
N1 = number of samples generated from first manifold
N2 = number of samples generated from second manifold
eta = homogenization parameter
lambd_list = list of values of lambda over which we calculate the quantities involved
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

N1_N2_list = [[120, 40], [150, 40], [160, 40], [200, 40]]
eta = 1
lambd_list = [0.00048+0.00001*i for i in range(200)]
seed_list=[i for i in range(50)]


for N1, N2 in N1_N2_list:
    print(f'N1={N1}, N2={N2}')

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



    for l in range(len(lambd_list)):
        for s in range(len(seed_list)):
            np.random.seed(seed_list[s])
            Y,__,gtruth,__  = manifoldGen('2trefoils-asym',N1,N2)
            n = max(gtruth)
            Y_orig = Y
            Y1 = Y[:,:N1]
            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:N1-1])[:dim]+1
            Y_knn = Y[:,knn]
            w_knn = w[knn-1]

            # Calculate inradius
            r = calc_inradius(y,Y_knn,w_knn,eta)
            
            # Find optimal dual solution
            nu = solve_dual(y,Y_knn,w,lambd_list[l],eta)
            nu_proj = project_on_norm_subspace(nu,Y_knn,eta)
            
            # Calculate gamma
            mu = calc_gamma(Y,nu_proj,N1,N2,w,r,eta)

            # Calculate dist(y,S)
            rhs = np.linalg.norm(np.append(y,eta))*mu/(1+mu)
            dist_x_s = project_on_subspace(y,Y_knn)

            # Calculate upper and lower limits of the interval
            LLi = 1/(r*(np.linalg.norm(np.append(y,eta))-dist_x_s))
            LUi = mu/(r*dist_x_s)

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

    np.savez(f'240519-variables-L-vs-l_N1_{N1}_N2_{N2}.npz', inradius=inradius, gamma_W=gamma_W, ineq=ineq, LL=LL, LU=LU, lambd_list=lambd_list, N1=N1, N2=N2, seed_list=seed_list)