from sklearn.decomposition import PCA
from dual_spls.get_bin_indices import get_bin_indices
from dual_spls.callist import get_calList
from dual_spls.split import split

def calvalXY(X,
             y=None,
             pcal=None,
             bin_indices=None,
             n_bins=10,
             calList=None,
             center=True,
             method="euclidian",
             dim_reduction='pca',
             pc=9,
             verbose=True):
    """Function to make calibration/validation splits makign sure that both sets
    cover the whole data space and y values

    Args:
        X (_type_): _description_
        y (_type_, optional): _description_. Defaults to None.
        pcal (int, optional): positive integer between 0 and 100.
                              pcal is the percentage of calibration
                              samples to be selected. Defaults to None.
        datatype (_type_, optional): A vector of index specifying each
                                    observation belonging to which group
                                    index. Typically generated by the function
                                    type. If None, will be generated within this
                                    function. Defaults to None.
        n_cells (int, optional): _description_. Defaults to 10.
        calList (_type_, optional): Listecal a numeric vector specifying
                                    how many observations from each group
                                    should be selected as calibration.
                                    Defaults to None.
        center (bool, optional): _description_. Defaults to True.
        method (str, optional): _description_. Defaults to "euclidian".
        pc (int, optional): pc a positive real value indicating
                            the number of component to consider
                            when applying the SVD transformation
                            or the PCA. Defaults to 9.

    Returns:
        (dict) {"indcal": ind_cal, "indval": ind_val}
    """

    # Check inputs
    assert method in ["euclidian", "pca_euclidian", "svd_euclidian"], f"Method {method} invalid."
    assert dim_reduction in ["pca"], f"Dimensionnality reduction methods {dim_reduction} invalid." # Implement SVD as well ? 
    assert pc > 0, f"pc (the number of component for PCA/SVD) must be strictly positive."
    assert (y is None and bin_indices is not None) or (y is not None), "If y is None, then a datatype vector must be given as an argument."

    if pcal is None:
        if verbose:
            print(f"[INFO] pcal not specified. Using default value 90%")
        pcal = 90

    if verbose and pc < 1:
        print(f"[INFO] pc = {pc} < 1 - The number of components will be computed automatically to reach 90% of X variance")

    if (bin_indices is not None) and (n_bins != None) and verbose:
        print(f"[INFO] n_cells is set to {n_bins} but will not be used as a datatype vector was given as an argument.")

    # Center data (if appliable)
    if center:
        X_proc = (X - X.mean(axis=0)).copy()
        if verbose: print("[INFO] Centered data")
    else:
        X_proc = X.copy()

    # Split the observations into groups
    if bin_indices is None:
        bin_indices = get_bin_indices(y, n_bins=n_bins)

    # Decide how much observations to take in each group
    # pcal = pourcentage for calibration
    if calList is None:
        calList = get_calList(bin_indices, pcal, n_bins=n_bins)

    # Apply dimensionnality reduction depending the method:
    if method in ["pca_euclidian", "svd_euclidian"]:
        pca = PCA(n_components=pc)
        scores = pca.fit_transform(X_proc)

        if method == "pca_euclidian":
            X_proc = scores

        elif method == "svd_euclidian":
            """
            SVD: X = U @ Sigma @ V^T

            PCA: T = X @ V 

            ==> T = (U @ Sigma @ V^T) @ V

            V is orthogonal => V^TV = I

            ==> T = U @ Sigma

            ==> U = T @ Sigma^{-1}
            """
            singular_values = pca.singular_values_
            X_proc = scores / singular_values 

        
    result = split(X_proc, X_bin_indices=bin_indices, calList=calList)

    return result

if __name__ == "__main__":
    from dual_spls.simulate import simulate
    ## Configuration:
    # Simulation configuration:
    coefs = [0, 1, 0, 1, 1, 0, 0, 1, 1] 
    n = 300
    p = [1000]
    nondes = [30]
    sigmaondes = [0.40]
    sigma_y = 0.05

    # Split configuration:
    split_mode = "random"

    # Dual-sPLS configuration
    norm = "pseudo-lasso"

    simulation_results = simulate(n=n, p=p, nondes=nondes , sigmaondes=sigmaondes, sigma_y=sigma_y, coefs=coefs)
    X, y = simulation_results["X"], simulation_results["y"]

    # test 3 methods
    calvalXY(X,
             y=y,
             center = True,
             pc=0.9)
    
    calvalXY(X,
             y=y,
             center = True,
             pc=0.9,
             method='pca_euclidian')
    
    calvalXY(X,
             y=y,
             center = True,
             pc=0.9,
             method='svd_euclidian')