import numpy as np
from itertools import combinations_with_replacement
from sklearn.feature_selection import f_regression, r_regression, mutual_info_regression
from scipy.cluster.vq import kmeans, vq
from scipy.spatial.distance import cdist
from sklearn.model_selection import train_test_split
import numpy as np
from sklearn.neighbors import KernelDensity
from sklearn.model_selection import GridSearchCV
import numpy as np
from itertools import combinations_with_replacement
from scipy.special import legendre, eval_chebyt
from sklearn.preprocessing import PolynomialFeatures
from sklearn.gaussian_process.kernels import RBF
import matplotlib.pyplot as plt


# Helper functions
def create_fourier_basis_functions(x, n_harmonics):
    """Create Fourier basis functions up to the nth harmonic."""
    fourier_basis = []
    for n in range(1, n_harmonics + 1):
        fourier_basis.append(np.sin(n * x))
        fourier_basis.append(np.cos(n * x))
    return np.hstack(fourier_basis)

def create_legendre_polynomials(x, order):
    """Create Legendre polynomials up to the specified order."""
    return np.hstack([legendre(n)(x) for n in range(order + 1)])

def create_chebyshev_polynomials(x, order):
    """Create Chebyshev polynomials of the first kind up to the specified order."""
    return np.hstack([eval_chebyt(n, x) for n in range(order + 1)])

def rbf_transform(x, centers, width):
    """Compute Gaussian RBF values given the input x, RBF centers, and width."""
    kernel = RBF(length_scale=width)
    return np.hstack([kernel(np.atleast_2d(center), np.atleast_2d(x))[0] for center in centers])


# Extended create_feats function
def create_feats(trajs, use_fourier, use_rbf, num_feats, n_harmonics=2, n_rbf_centers=5, rbf_width=1.0, legendre_order=3, chebyshev_order=3):
    feats = []
    
    # RBF centers (if using RBF)
    if use_rbf:
        centers = np.linspace(-1, 1, n_rbf_centers)

    for traj in trajs:
        feat = np.zeros((traj.shape[0], num_feats))
        for idx in range(traj.shape[0]):
            curr_point = traj[idx]

            # First, second, and third degree polynomials
            second_degree = [x * y for x, y in combinations_with_replacement(curr_point, 2)]
            third_order = [x * y * z for x, y, z in combinations_with_replacement(curr_point, 3)]
            
            arr_values = curr_point.tolist() + second_degree + third_order

            # Fourier basis
            if use_fourier:
                fourier_series = create_fourier_basis_functions(curr_point, n_harmonics).tolist()
                arr_values += fourier_series
            
            # Gaussian RBFs
            if use_rbf:
                rbf_values = rbf_transform(curr_point, centers, rbf_width).tolist()
                arr_values += rbf_values

            # Legendre polynomials
            legendre_basis = create_legendre_polynomials(curr_point, legendre_order).tolist()
            arr_values += legendre_basis

            # Chebyshev polynomials
            chebyshev_basis = create_chebyshev_polynomials(curr_point, chebyshev_order).tolist()
            arr_values += chebyshev_basis

            feat[idx, :] = np.array(arr_values)
        feats.append(feat)

    return feats


############################
##   FEATS NAME FUNCTION   ##
############################

from itertools import combinations_with_replacement

def create_feats_name(use_fourier=False, use_rbf=False, use_chebyshev=False, use_legendre=False, dims=3, n_harmonics=2, n_rbf_centers=5, legendre_order=3, chebyshev_order=3):
    """
    Generate feature names for first-degree, second-degree, third-order polynomials,
    Fourier series, Gaussian RBFs, Legendre polynomials, and Chebyshev polynomials.
    """
    # Variable names for dimensions
    variables = [f'f_{i}' for i in range(dims)]

    # First-degree polynomial terms
    first_degree = variables
    second_degree = [f"{x} * {y}" for x, y in combinations_with_replacement(variables, 2)]
    third_order = [f"{x} * {y} * {z}" for x, y, z in combinations_with_replacement(variables, 3)]
    # forth_order = [f"{x} * {y} * {z} * {w}" for x, y, z, w in combinations_with_replacement(variables, 4)]  

    # Fourier series names
    fourier_series_names = [f"fourier_sin_{i}_{axis}" for axis in variables for i in range(1, n_harmonics + 1)] + \
                           [f"fourier_cos_{i}_{axis}" for axis in variables for i in range(1, n_harmonics + 1)]

    # Gaussian RBF names
    rbf_names = [f"rbf_center_{i}_{axis}" for axis in variables for i in range(n_rbf_centers)]

    # Legendre polynomial names
    legendre_names = [f"legendre_{i}_{axis}" for axis in variables for i in range(legendre_order + 1)]

    # Chebyshev polynomial names
    chebyshev_names = [f"chebyshev_{i}_{axis}" for axis in variables for i in range(chebyshev_order + 1)]

    # Combine all features
    all_polynomials = first_degree + second_degree + third_order

    if use_fourier:
        all_polynomials += fourier_series_names

    if use_rbf:
        all_polynomials += rbf_names
        
    if use_chebyshev:
        all_polynomials += chebyshev_names
    
    if use_legendre:
        all_polynomials += legendre_names

    num_feats = len(all_polynomials)
    
    return all_polynomials, num_feats
