""" The function "manifoldGen" is used to generate data on 2-trefoil knots. The code
(i) generates equidistant data points on each 1-dimensional trefoil knot (Note: extrinsic dimension of trefoil knot=3),
(ii) generates a random 1-dimensional orthonormal basis in 100-dimensional space,
(ii) embeds each data sample on trefoil knots into the 100-dimensional subspace representated by the basis """

import numpy as np
from numpy.linalg import svd
from scipy.linalg import orth


def manifoldGen(manifoldType,N=None,N2=None):
    #np.random.seed(10000)
    D = 100  # ambient space dimension
    sigma = 0.00  # noise variance
    if manifoldType == '2trefoils' and (N is None):
        N = 100
    if manifoldType == '2trefoils-asym' and (N is None or N2 is None):
        N = 100
        N2 = 100

    if manifoldType == '2trefoils':
        gtruth = np.concatenate((np.ones(N), 2*np.ones(N)))
        # Generate two trefoil-knots
        d = 3
        Par = 80
        t = np.linspace(0, 2*np.pi, N+1)[:-1]
        Yg = {}
        Yg[1] = np.array([(2+np.cos(3*t))*np.cos(2*t), (2+np.cos(3*t))*np.sin(2*t), np.sin(3*t)])
        Yg[2] = np.array([(2+np.cos(3*t))*np.cos(2*t) + Par, (2+np.cos(3*t))*np.sin(2*t), np.sin(3*t)])
        Y = np.concatenate((Yg[1], Yg[2]), axis=1)
        U = orth(np.random.randn(D, d))
        Yn = U @ Y
        Yn = Yn + sigma * np.random.randn(*Yn.shape)
        x = np.concatenate((np.cos(t), np.cos(t)), axis=0)
        x = np.vstack((x, np.concatenate((np.sin(t), np.sin(t)), axis=0)))
    if manifoldType == '2trefoils-asym':
        gtruth = np.concatenate((np.ones(N), 2*np.ones(N2)))
        # Generate two trefoil-knots
        d = 3
        Par = 8
        t = np.linspace(0, 2*np.pi, N+1)[:-1]
        t2 = np.linspace(0, 2*np.pi, N2+1)[:-1]
        Yg = {}
        Yg[1] = np.array([(2+np.cos(3*t))*np.cos(2*t), (2+np.cos(3*t))*np.sin(2*t), np.sin(3*t)])
        Yg[2] = np.array([(2+np.cos(3*t2))*np.cos(2*t2) +Par, (2+np.cos(3*t2))*np.sin(2*t2), np.sin(3*t2)])
        Y = np.concatenate((Yg[1], Yg[2]), axis=1)
        U = orth(np.random.randn(D, d))
        Yn = U @ Y
        Yn = Yn + sigma * np.random.randn(*Yn.shape)
        x = np.concatenate((np.cos(t), np.cos(t2)), axis=0)
        x = np.vstack((x, np.concatenate((np.sin(t), np.sin(t2)), axis=0)))
    
    return Yn, Y, gtruth, x

