'''
SMB-PLS utility functions.

    Includes:
        - data normalization: normalisation_data, normalisation_data_val
        - block data normalization: normalisation_blocs_cal, normalisation_blocs_val
        - SMB-PLS: smbplsnipals, smbpls_cal, smbpls_exploit  
        - Random-subset (RS) cross-validation: rs_cv
        - Predictive performances: CalculR, SPELimit, SPECalculation

    
    Version History:
    ----------------
    - Original smbpls codes in matlab: Julien Lauzon-Gauthier
    - Future versions of smbpls codes in matlab: Adeline Paris
    - Version 0: january 2026
        SMB-PLS matlab codes translated to Python by Giverny Robert
            
           
            
   '''
   
#%% Import modules   
   
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as st
import RG as rg
from scipy.stats import chi2


#%% Data scaling -  based on rg.autoscale

# Calibration
def normalisation_data(cal):
    """
    Normalize a matrix using column-wise z-score normalization.
    
    Parameters
    ----------
    cal : ndarray
        Calibration matrix (n_samples × n_variables).

    Returns
    -------
    calnorm : ndarray
        Normalized calibration matrix.
    mean_d : ndarray
        Column-wise mean of cal.
    std_d : ndarray
        Column-wise standard deviation of cal.

    """

    # Copy matrix to notoverwrite original cal
    cal_a = cal.copy()
    
    # Autoscale performed with Ryan's function for other methods
    calnorm, mean_d, std_d = rg.autoscale(cal_a)
    
    return calnorm, mean_d, std_d

# Validation
def normalisation_data_val(val, mean_d, std_d):
    """
    Normalize a matrix using predefined mean and standard deviation.
    
    Parameters
    ----------
    val : ndarray
        Validation matrix (n_samples × n_variables).
    mean_d : ndarray
        Column-wise mean of cal (n_variables,).
    std_d : ndarray
        Column-wise standard deviation of cal (n_variables,).

    Returns
    -------
    valnorm : ndarray
        Normalized validation matrix.

    """
    
    # Change to fit autoscale performed with Ryan's function for other methods
    valnorm  = (val - mean_d) / std_d
         
    return valnorm

#%% Block data scaling

# Calibration
def normalisation_blocs_cal(X_blocks, Y, bs=1):
    """
    Normalize X blocks and Y for SMB-PLS calibration.

    Each X block is mean-centered and scaled. Optionally, block scaling
    is applied by dividing each normalized block by sqrt(number of variables).
    Y is normalized using the same procedure (without block scaling).

    Parameters
    ----------
    X_blocks : list of ndarray
        List of predictor blocks, each of shape (n_samples, n_variables).
    Y : ndarray
        Response matrix or vector of shape (n_samples, n_y_variables).
    bs : int, optional
        Block scaling flag (1 = apply block scaling, 0 = no block scaling).
        Default is 1.

    Returns
    -------
    X_blocks_norm : list of ndarray
        List of normalized (and optionally block-scaled) X blocks.
    Y_norm : ndarray
        Normalized response matrix.
    bl : list of lists
        List of variable indices for each X block. Each sublist corresponds to a block 
        and contains the indices of its variables in the global X matrix.
    mean_blocks : list of ndarray
        Column-wise means for each X block, followed by Y mean.
    std_blocks : list of ndarray
        Column-wise standard deviations for each X block, followed by Y std.
    """
    
    # Initialization
    X_blocks_norm = []
    bl = []
    mean_blocks = []
    std_blocks = []

    # X normalization
    start_idx = 0
    for X in X_blocks:
        X_norm, X_mean, X_std = normalisation_data(X)
        n_vars = X.shape[1]
        # Block saling
        if bs == 1:
            X_norm /= np.sqrt(n_vars)
        X_blocks_norm.append(X_norm)
        mean_blocks.append(X_mean)
        std_blocks.append(X_std)
        bl.append(list(range(start_idx, start_idx + n_vars)))
        start_idx += n_vars

    # Y normalization
    Y_norm, Y_mean, Y_std = normalisation_data(Y)
    mean_blocks.append(Y_mean)
    std_blocks.append(Y_std)

    return X_blocks_norm, Y_norm, bl, mean_blocks, std_blocks

# Validation
def normalisation_blocs_val(X_blocks, Y, bs, bl, mean_blocks, std_blocks):
    """
    Normalize X blocks and Y for SMB-PLS validation using calibration parameters.

    Each X block is normalized using the mean and standard deviation obtained
    during calibration. Optional block scaling is applied using the number of
    variables in each block. Y is normalized using calibration parameters only.

    Parameters
    ----------
    X_blocks : list of ndarray
        List of predictor blocks to be validated.
    Y : ndarray
        Response matrix or vector to be validated.
    bs : int
        Block scaling flag (1 = apply block scaling, 0 = no block scaling).
    bl : list of lists
        List of variable indices for each X block. Each sublist corresponds to a block 
        and contains the indices of its variables in the global X matrix.
    mean_blocks : list of ndarray
        Column-wise means from calibration (X blocks followed by Y).
    std_blocks : list of ndarray
        Column-wise standard deviations from calibration (X blocks followed by Y).

    Returns
    -------
    X_blocks_norm : list of ndarray
        List of normalized (and optionally block-scaled) X blocks.
    Y_norm : ndarray
        Normalized response matrix.
    """

    # X normalization
    X_blocks_norm = []
    for i, X in enumerate(X_blocks):
        X_norm = normalisation_data_val(X, mean_blocks[i], std_blocks[i])
        if bs == 1:
            X_norm /= np.sqrt(len(bl[i]))
        X_blocks_norm.append(X_norm)

    # Y normalization
    Y_norm = normalisation_data_val(Y, mean_blocks[-1], std_blocks[-1])

    return X_blocks_norm, Y_norm


#%% SMB-PLS

# Correlation matrix for computation in SMB-PLS at l.355
def safe_corrcoef(tb, u):
    """
    Compute correlation matrix between tb and u safely, ignoring zero-variance columns.

    Parameters
    ----------
    tb : ndarray
        Matrix of block scores (observations × components).
    u : ndarray
        Vector of current Y-score (observations,).

    Returns
    -------
    tucorr : ndarray
        Correlation matrix between tb and u, with zero-variance columns ignored.
    """
    
    # Ensure u is 2D column vector
    u_col = u.reshape(-1, 1)

    # Compute variance of each column
    tb_var = np.var(tb, axis=0)
    u_var = np.var(u_col, axis=0)

    # Keep only columns with non-zero variance
    tb_nonzero = tb[:, tb_var > 0]
    u_nonzero = u_col[:, u_var > 0]

    if tb_nonzero.shape[1] == 0 or u_nonzero.shape[1] == 0:
        print("[WARNING] All columns have zero variance. Returning zeros.")
        return np.zeros((tb.shape[1] + 1, tb.shape[1] + 1))  # fallback

    # Compute correlation safely
    tucorr = np.corrcoef(np.hstack([tb_nonzero, u_nonzero]).T)

    return tucorr


# SMB-NIPALS
def smbnipals(x, y, bl, blnow):
    """
    Based on basic MBPLS code for 1 component and no missing data.
    
    References:
    -----------
    - Wold S, Sjöström M, Eriksson L. "PLS-regression: a basic tool of chemometrics." 
      Chemom Intell Lab Syst. 2001; 58(2):109–130.
    - Modified for multiblock PLS:
      Westerhuis JA, Kourti T, MacGregor JF. "Analysis of multiblock and hierarchical PCA and PLS models." 
      J Chemom. 1998; 12(5):301–321.
    - Modified for sequential orthogonalization of the block data.
    
    
    Version History:
    ----------------
    Version 3  - 12 September 2022: Saves Beta coefficients for correlation blocks
    Version 2  - March 2014: Includes orthogonalization of block scores
    Version 1  - November 2013
    
    © Julien Lauzon-Gauthier  
    
    Parameters
    ----------
    x : array-like, shape (M, Nx)
        X-blocks concatenated for SMB-PLS.      
    y : array-like, shape (M, Ny)
        Y-block.
    bl : list of lists
        List of variable indices for each X block. Each sublist corresponds to a block 
        and contains the indices of its variables in the global X matrix.
    blnow : int
        Current block index being processed (used in sequential calculations).

    Returns
    -------
    p : list of ndarray
        Loadings for each block (variables × latent variables) in the current decomposition.
    q : ndarray
        Y-block loadings (responses × latent variables).
    w : list of ndarray
        Weight vectors for each block (variables × latent variables).
    tT : ndarray
        Scores for the concatenated X-blocks (observations × latent variables).
    u : ndarray
        Scores for the Y-block (observations × latent variables).
    tb : list of ndarray
        Scores for individual X-blocks (observations × latent variables per block).
    wb : list of ndarray
        Orthogonal weights for each X-block (variables × latent variables per block).
    pb : list of ndarray
        Block-specific loadings (variables × latent variables per block).
    wt : ndarray
        Overall weight vector for concatenated X (variables × latent variables).
    tucorr : ndarray
        Correlation between X-block scores and Y scores (latent variables × 1).
    blcorvar : list
        Indices of X variables that are correlated with the current block Y.
    count : int
        Number of iterations or latent variables extracted (depending on implementation).
    Beta : ndarray
        Regression coefficients for predicting Y from concatenated X blocks.

    """
    # y_copy = y.copy()
    # u = y_copy[:, 0]
    u = y[:, 0]
    tT = x[:, 0]
    nbl = len(bl)
    m, nx = x.shape

    tb = np.zeros((m, nbl))
    wb = np.zeros((nx,))
    pb = np.zeros((nx,))
    Beta = [None] * nbl
    xcor = x.copy()
    blvar = np.array([np.sum(x[:, blk] ** 2) for blk in bl])

    count = 0
    convr = 1e8

    while convr >= 1e-8 and count <= 300:
        told = tT.copy()
        count += 1

        blk = bl[blnow]
        wb[blk] = x[:, blk].T @ u / (u.T @ u)
        wb[blk] /= np.linalg.norm(wb[blk])
        tb[:, blnow] = x[:, blk] @ wb[blk]

        for i in range(blnow + 1, nbl):
            xcor[:, bl[i]] = np.outer(tb[:, blnow], (tb[:, blnow].T @ x[:, bl[i]]) / (tb[:, blnow].T @ tb[:, blnow]))
            wb[bl[i]] = xcor[:, bl[i]].T @ u / (u.T @ u)
            wb[bl[i]] /= np.linalg.norm(wb[bl[i]])
            tb[:, i] = xcor[:, bl[i]] @ wb[bl[i]]

        # tucorr = np.corrcoef(np.hstack([tb, u.reshape(-1, 1)]).T)
        tucorr = safe_corrcoef(tb, u) # replaced 30 dec. 2025
        
        wt = tb.T @ u / (u.T @ u)
        wt /= np.linalg.norm(wt)
        tT = tb @ wt / (wt.T @ wt)
        
        c = y.T @ tT / (tT.T @ tT)
        u = y @ c / (c.T @ c)
        
        convr = np.linalg.norm(told - tT) / np.linalg.norm(tT)

    for i in range(blnow + 1, nbl):
        Beta[i] = (tb[:, blnow].T @ x[:, bl[i]]) / (tb[:, blnow].T @ tb[:, blnow])

    blcorvar = np.zeros(nbl)
    for i in range(nbl):
        blcorvar[i] = np.sum(xcor[:, bl[i]] ** 2) * 100.0 / blvar[i]

    for i in range(blnow, nbl):
        pb[bl[i]] = x[:, bl[i]].T @ tb[:, i] / (tb[:, i].T @ tb[:, i])

    p = np.zeros((nx,))
    for i in range(blnow, nbl):
        p[bl[i]] = x[:, bl[i]].T @ tT / (tT.T @ tT)

    w = x.T @ u / (u.T @ u)
    w /= np.linalg.norm(w)
    q = c.squeeze()

    return p, q, w, tT, u, tb, wb, pb, wt, tucorr, blcorvar, count, Beta

# SMB-PLS calibration
def smbpls_cal(x, y, bl, lv, display=0):
    """
    Multi-block SEQUENTIAL PLS - mbseqplsnipals 
    Performs SMB-PLS regression using NIPALS-like algorithm.
    
    Version History:
    ----------------
    Version 7 (Adéline Paris) - 20/09/2022
        - Calculation of the beta coefficient for the extraction of correlated information
    
    Version 6 (New Algorithm)
        - Orthogonalization of block data
        - Superscore computation: tT = [tx1 tx2 ... txb]
        - Convergence of tT (super-score)
        - Super score deflation
    
    Modified from Eigenvector Research, Inc. 2002-2013
    
    Additional Versions:
    --------------------
    Version 3  - 12 September 2022: Saves Beta coefficients for correlation blocks
    Version 5  - March 2014
    Version 1  - November 2013
    
    © Julien Lauzon-Gauthier

    Parameters
    ----------
    x : array-like, shape (M, Nx)
        X-blocks concatenated for SMB-PLS.      
    y : array-like, shape (M, Ny)
        Y-block.
    bl : list of lists
        List of variable indices for each X block. Each sublist corresponds to a block 
        and contains the indices of its variables in the global X matrix.
    lv : list or array
        Number of components per block.
    display : int, optional
        If 1, prints the SSQ table. Default is 0.

    Returns
    -------
    modelstruct : dict
        Dictionary containing SMB-PLS model outputs, including:
            ssq : the sum of squares captured (ssq),
            xlds : X-block loadings,
            ylds : Y-block loadings,
            wts : X-block weights, (now in non-deflated X format)
            xsupscrs : X-block scores,
            yscrs : Y-block scores,
            bin : the inner relation coefficients,
            nipwts : X-block weights, (original deflated X format)
            Wblock : X-block block weights,
            Tblock : X-block block scores,
            Wsuper : X-block super weights,
            T : super-scores,
            lv : number of principal components pers block (b,),
            Tcorr : Correlated part of block scores across blocks,
            lvtbvarcapt : Percentage of variance captured per block,
            lvconvcount : Number of iterations for convergence of each latent variable,
            xres : x residual,
            yres : y residual,
            ypred : y prediction,
            beta : beta coefficients for SMB-PLS model.
            
            


    """
    y_copy = y.copy()
    m, nx = x.shape
    ny = y.shape[1]
    nbl = len(bl)
    assert len(lv) == nbl, "X blocks in bl and lv must match"

    olv = sum(lv)
    lvcount = 0

    loads_xscrs = np.zeros((m, olv))
    loads_yscrs = np.zeros((m, olv))
    loads_xlds = np.zeros((nx, olv))
    loads_ylds = np.zeros((ny, olv))
    wts = np.zeros((nx, olv))
    bin_vals = np.zeros((olv,))
    ypred = np.zeros_like(y)

    Wblock = [np.zeros((len(bl[i]), olv)) for i in range(nbl)]
    Pblock = [np.zeros((len(bl[i]), olv)) for i in range(nbl)]
    Tblock = [np.zeros((m, olv)) for _ in range(nbl)]
    Wsuper = np.zeros((nbl, olv))
    T_list = [None] * olv
    Tcorr = [None] * olv
    lvconvcount = np.zeros(olv, dtype=int)
    lvtbvarcapt = np.zeros((olv, nbl))
    Beta_Block = [{} for _ in range(nbl)]

    ssqx = np.sum(x ** 2)
    ssqy = np.sum(y ** 2)
    ssq = np.zeros((olv, 2))
    blssq = np.zeros((olv, nbl))
    blssqx = [np.sum(x[:, bl[i]] ** 2) for i in range(nbl)]

    for j in range(nbl):
        for i in range(lv[j]):
            lvcount += 1
            p, q, w, tT, u, tb, wb, pb, wt, tbucorr, tbvarcapt, count, Beta = smbnipals(x, y_copy, bl, j)

            bin_vals[lvcount - 1] = np.dot(u.T, tT) / np.dot(tT.T, tT)
            x -= np.outer(tT, p)
            y_copy -= bin_vals[lvcount - 1] * np.outer(tT, q)
            ypred += bin_vals[lvcount - 1] * np.outer(tT, q)

            ssq[lvcount - 1, 0] = np.sum(x ** 2) * 100 / ssqx
            ssq[lvcount - 1, 1] = np.sum(y_copy ** 2) * 100 / ssqy
            for k in range(nbl):
                blssq[lvcount - 1, k] = np.sum(x[:, bl[k]] ** 2) * 100 / blssqx[k]

            loads_xscrs[:, lvcount - 1] = tT
            loads_yscrs[:, lvcount - 1] = u
            loads_xlds[:, lvcount - 1] = p
            loads_ylds[:, lvcount - 1] = q
            wts[:, lvcount - 1] = w
            T_list[lvcount - 1] = tb
            Wsuper[:, lvcount - 1] = wt
            Tcorr[lvcount - 1] = tbucorr
            lvtbvarcapt[lvcount - 1, :] = tbvarcapt
            lvconvcount[lvcount - 1] = count

            for l in range(nbl):
                Wblock[l][:, lvcount - 1] = wb[bl[l]]
                Pblock[l][:, lvcount - 1] = pb[bl[l]]
                Tblock[l][:, lvcount - 1] = tb[:, l]
                if Beta[l] is not None:
                    Beta_Block[l][lvcount - 1] = Beta[l]


    nipwts = wts.copy()
    wts = wts @ np.linalg.pinv(loads_xlds.T @ wts)

    ssqdif = np.zeros((olv, 2))
    ssqdif[0, :] = 100 - ssq[0, :]
    for i in range(1, olv):
        ssqdif[i, :] = ssq[i - 1, :] - ssq[i, :]

    ssq_combined = np.column_stack([
        np.arange(1, olv + 1),
        ssqdif[:, 0],
        np.cumsum(ssqdif[:, 0]),
        ssqdif[:, 1],
        np.cumsum(ssqdif[:, 1])
    ])

    blssqdif = np.zeros((olv, nbl))
    blssqdif[0, :] = 100 - blssq[0, :]
    for i in range(1, olv):
        blssqdif[i, :] = blssq[i - 1, :] - blssq[i, :]

    blssq_combined = np.hstack([
        np.arange(1, olv + 1).reshape(-1, 1),
        np.column_stack([
            np.column_stack([blssqdif[:, k], np.cumsum(blssqdif[:, k])])
            for k in range(nbl)
        ])
    ])

    for j in range(1,nbl):
        Beta_Block[j] = np.vstack([Beta_Block[j][k] for k in sorted(Beta_Block[j])])

    # ssqtot: list of lists to simulate a cell array
    n_comp = olv
    n_col = 2 * nbl + 5  # 2 per block + 2 (X), 2 (Y), and 1 for 'Comp'
    ssqtot = [['' for _ in range(n_col)] for _ in range(n_comp + 2)]
    
    # Header
    for k in range(nbl):
        ssqtot[0][2*k + 1] = f'X{k+1}'
        ssqtot[0][2*k + 2] = f'X{k+1}'
    ssqtot[0][2*nbl + 1] = 'X'
    ssqtot[0][2*nbl + 2] = 'X'
    ssqtot[0][2*nbl + 3] = 'Y'
    ssqtot[0][2*nbl + 4] = 'Y'
    
    # Sub-header
    ssqtot[1][0] = 'Comp'
    for l in range(1, n_col):
        ssqtot[1][l] = 'This' if l % 2 == 1 else 'Total'
    
    # Data
    for i in range(n_comp):
        row = [i+1]
        for k in range(nbl):
            row += [blssqdif[i, k], np.cumsum(blssqdif[:, k])[i]]
        row += [ssqdif[i, 0], np.cumsum(ssqdif[:, 0])[i]]
        row += [ssqdif[i, 1], np.cumsum(ssqdif[:, 1])[i]]
        ssqtot[i+2] = row

    ssqtot = np.array(ssqtot)

    modelstruct = {
        'ssq': ssq_combined,
        'ssqtot': ssqtot,
        'blssq': blssq_combined,
        'xlds': loads_xlds,
        'ylds': loads_ylds,
        'wts': wts,
        'xscrs': loads_xscrs,
        'yscrs': loads_yscrs,
        'bin': bin_vals,
        'nipwts': nipwts,
        'Wblock': Wblock,
        'Tblock': Tblock,
        'Pblock': Pblock,
        'Wsuper': Wsuper,
        'T': T_list,
        'lv': lv,
        'Tcorr': Tcorr,
        'lvtbvarcapt': lvtbvarcapt,
        'lvconvcount': lvconvcount,
        'xres': x,
        'yres': y_copy,
        'ypred': ypred,
        'beta': Beta_Block
    }

    if display:
        print("SSQ Table:\n", ssq_combined)

    return modelstruct

# SMB-PLS validation
def smbpls_exploit(x, y, bl, modelstruct):
        
    """
    Multi-block SEQUENTIAL PLS - mbseqplsnipals 
    Performs SMB-PLS regression using NIPALS-like algorithm.
    
    Version History:
    ----------------
    Version 7 (Adéline Paris) - 20/09/2022
        - Calculation of the beta coefficient for the extraction of correlated information
    
    Version 6 (New Algorithm)
        - Orthogonalization of block data
        - Superscore computation: tT = [tx1 tx2 ... txb]
        - Convergence of tT (super-score)
        - Super score deflation
    
    Modified from Eigenvector Research, Inc. 2002-2013
    
    Additional Versions:
    --------------------
    Version 3  - 12 September 2022: Saves Beta coefficients for correlation blocks
    Version 5  - March 2014
    Version 1  - November 2013
    
    © Julien Lauzon-Gauthier
    
    Parameters
    ----------
    x : array-like, shape (M, Nx)
        X-blocks concatenated for SMB-PLS.      
    y : array-like, shape (M, Ny)
        Y-block.
    bl : list of lists
        List of variable indices for each X block. Each sublist corresponds to a block 
        and contains the indices of its variables in the global X matrix.
    modelstruct : dict
        Dictionary containing SMB-PLS model outputs.
        
    Returns
    -------
    xscrs : X-block loadings,
    Tblock : X-block block scores
    T : T_list,
    lv : list or array
          Number of components per block.
    xres : x residual
    yres : y residual
    ypred : y prediction
    xpred : x prediction

    """
    
    lv = modelstruct['lv']
    # x = np.atleast_2d(x)
    # if x.shape[0] == 1:
    #     x = x.T    
    
    m, _ = x.shape
    assert y.shape[0] == m, "Number of y-block rows must match x"
    nbl = len(bl)
    assert nbl == len(lv), "X blocks and number of components must match"

    p = modelstruct['xlds']
    q = modelstruct['ylds']
    bin_vals = modelstruct['bin']
    Wblock = modelstruct['Wblock']
    Wsuper = modelstruct['Wsuper']
    Beta = modelstruct['beta']

    olv = sum(lv)
    loads_xscrs = np.zeros((m, olv))
    Tblock = [np.zeros((m, olv)) for _ in range(nbl)]
    T_list = [None] * olv
    ypred = np.zeros_like(y)
    xpred = np.zeros_like(x)

    lvcount = 0
    for j in range(nbl):
        for i in range(lv[j]):
            lvcount += 1
            tb = np.zeros((m, nbl))
            tb[:, j] = x[:, bl[j]] @ Wblock[j][:, lvcount - 1]
            xcor = x.copy()
            for k in range(j + 1, nbl):
                xcor[:, bl[k]] = tb[:, [j]] @ Beta[k][[lvcount - 1],:]
                tb[:, k] = xcor[:, bl[k]] @ Wblock[k][:, lvcount - 1]

            tT = tb @ Wsuper[:, lvcount - 1]
            ypred += bin_vals[lvcount - 1] * np.outer(tT, q[:, lvcount - 1])
            xpred += np.outer(tT, p[:, lvcount - 1])
            x -= np.outer(tT, p[:, lvcount - 1])
            y -= bin_vals[lvcount - 1] * np.outer(tT, q[:, lvcount - 1])

            loads_xscrs[:, lvcount - 1] = tT
            T_list[lvcount - 1] = tb
            for l in range(nbl):
                Tblock[l][:, lvcount - 1] = tb[:, l]

    return {
        'xscrs': loads_xscrs,
        'Tblock': Tblock,
        'T': T_list,
        'lv': lv,
        'xres': x,
        'yres': y,
        'ypred': ypred,
        'xpred': xpred
    }


#%% Random-subset (RS) cross-validation

def rs_cv(X_blocs, Y, A_list, chosen_list, bs,nb,CVnb):
    """
    Perform repeated SMB-PLS cross-validation and compute RMSECV.

    Parameters
    ----------
    X_blocs : list of ndarray
        List of X-block matrices. Each element is a block (observations × variables).
    Y : ndarray
        Response matrix (observations × responses).
    A_list : list of int
        Number of components to extract per block.
    chosen_list : list of int
        Indices of blocks chosen for modeling or validation.
    bs : int
        Number of bootstrap or cross-validation repetitions.
    
    Returns
    -------
    RMSECV_all : ndarray
        Matrix of root mean squared errors of cross-validation (RMSECV) for each block and repetition.
    """
    
    RMSECV_all = []
    n_samples = Y.shape[0]
    n_blocks = len(X_blocs)
    
    for block_idx in range(n_blocks):
        A = A_list[block_idx]
        chosen = chosen_list[block_idx]
            
        # Set nbPC to 0 for all posterior blocks
        if chosen==0:    
            lv_base = A_list.copy()
            for j in range(block_idx + 1, n_blocks):
                lv_base[j] = 0
        
            for k in range(1, A + 1):  
                E = []
                lv = lv_base.copy()
                lv[block_idx] = k
                
                for i in range(nb):
                    r = np.random.choice(n_samples, CVnb, replace=False)
                
                    X_blocks_v = [Xb[r, :] for Xb in X_blocs]
                    Y_v = Y[r,:]
                    
                    X_blocks_c = [np.delete(Xb, r, axis=0) for Xb in X_blocs]
                    Y_c = np.delete(Y, r, axis=0)
                    
                    norm_blocks_c, Ynorm_c, bl, means, stds = normalisation_blocs_cal(X_blocks_c, Y_c, bs=bs)
                    norm_blocks_v, Ynorm_v = normalisation_blocs_val(X_blocks_v, Y_v, bs, bl, means, stds)
                    
                    Xc = np.hstack(norm_blocks_c)
                    Xv = np.hstack(norm_blocks_v)
                    model = smbpls_cal(Xc, Ynorm_c, bl, lv, display=0)
                    prediction = smbpls_exploit(Xv, Ynorm_v, bl, model)
                    
                    Y_pred = prediction["ypred"]*stds[-1] + means[-1]
                    e = Y_pred - Y_v
                    E.append(np.squeeze(e))
                    
                E = np.array(E).ravel()    
                RMSECV = np.sqrt(np.mean(E**2))
                RMSECV_all.append(RMSECV)
    
    RMSECV_all = np.array(RMSECV_all)
    
    return RMSECV_all      


#%% Predictive performances

def CalculR(nc, nx, Modele, B):
    """
    Compute block-wise regression matrices R in a multiblock PLS model.

    Parameters
    ----------
    nc : list of int
        Number of components per X block.
    nx : list of int
        Number of variables per X block.
    Modele : dict
        SMB-PLS model dictionary containing:
        - beta: list of beta matrices for each block
        - Wblock: list of block weight matrices
        - Wsuper: super weights matrix
        - xlds: X-block loadings matrix
        - lv: list of latent variables per block
    B : int
        Total number of blocks.

    Returns
    -------
    R_keep : list of ndarray
        List of regression matrices R for each block. If a block has no latent variables,
        the corresponding entry is None.

    """
    R_keep = [None] * B
    for i in range(B):
        bi = i
        alpha_c_i = sum(nc[:bi])
        alpha_c_f = sum(nc[:bi + 1])
        alpha_x_i = sum(nx[:bi])
        alpha_x_f = sum(nx[:bi + 1])
        A_int = []
        for j in range(bi + 1, B):
            part = np.diag(Modele['beta'][j][alpha_c_i:alpha_c_f, :] @ Modele['Wblock'][j][:, alpha_c_i:alpha_c_f])
            A_int.append(part)
        A_int = np.column_stack(A_int) if A_int else np.empty((nc[i], 0))
        if bi == 0:
            A = np.column_stack((np.ones((nc[i], 1)), A_int))
        else:
            A = np.column_stack((np.zeros((nc[i], bi)), np.ones((nc[i], 1)), A_int))
        if Modele['lv'][i] != 0:
            Gamma1 = Modele['Wblock'][bi][:, alpha_c_i:alpha_c_f] * (A @ Modele['Wsuper'][:, alpha_c_i:alpha_c_f]).T
            inv_part = np.linalg.inv(Modele['xlds'][alpha_x_i:alpha_x_f, alpha_c_i:alpha_c_f].T @ Gamma1)
            R = Gamma1 @ inv_part
            R_keep[i] = R
    return R_keep


def SPELimit(SPEY, LC):
    """
    Compute the statistical limit for squared prediction errors (SPE) based on a chi-squared distribution.

    Parameters
    ----------
    SPEY : ndarray
        Squared prediction errors for the calibration set.
    LC : float
        Confidence level for the limit (e.g., 0.95 for 95% limit).

    Returns
    -------
    IC : float
        Calculated SPE control limit at the given confidence level.

    """
    vSPEY = np.var(SPEY)
    mSPEY = np.mean(SPEY)
    d1 = 2 * mSPEY ** 2 / vSPEY
    d2 = vSPEY / (2 * mSPEY)
    IC = d2 * chi2.ppf(LC, d1)
    return IC


#%%

def SPECalculation(Xnorm_blocks, xlds, lv, R_keep, bl):
    """
    Compute SPE for n-block SMB-PLS with safety for blocks with lv=0
    
    Parameters
    ----------
    Xnorm_blocks : list of np.array
        Normalized data matrices for each block [X1, X2, ..., Xn]
    xlds : np.array
        Full loading matrix (all blocks stacked)
    lv : list of int
        Number of latent variables per block
    R_keep : list of np.array or None
        R matrices for each block, None if lv=0
    bl : list of slice or indices
        Row indices of xlds corresponding to each block
    
    Returns
    -------
    SPE_blocks : list of np.array
        Squared prediction errors for each block
    """
    
    n_blocks = len(Xnorm_blocks)
    SPE_blocks = []
    t_prev = None  # will hold all previous scores
    
    lv_cumsum = np.cumsum([0] + lv)  # cumulative LV sum for slicing
    
    for i in range(n_blocks):
        Xblock = Xnorm_blocks[i]
        
        # If lv[i] == 0 or R_keep[i] is None, skip calculations
        if lv[i] == 0 or R_keep[i] is None:
            # SPE_blocks.append((Xblock - 0) ** 2)  
            # SPE_blocks.append(np.full_like(Xblock, np.nan)) # or fill nan matrix
            continue
        
        # First block: only correlated part
        if i == 0:
            t_t = Xblock @ R_keep[i]
            z_pred = t_t[:, :lv[i]] @ xlds[bl[i], :lv[i]].T
            SPE_blocks.append((z_pred - Xblock) ** 2)
            t_prev = t_t[:, :lv[i]]
        else:
            # Compute residual after removing contributions of previous blocks
            Xresidual = Xblock.copy()
            
            # correlated part from previous blocks
            xcorr_chap = np.zeros_like(Xblock)
            
            for j in range(i):
                if lv[j] == 0 or R_keep[j] is None:
                    continue  # skip previous block if no LV
                start = lv_cumsum[j]
                end = lv_cumsum[j+1]
                xcorr_chap += t_prev[:, start:end] @ xlds[bl[i], start:end].T
                Xresidual -= t_prev[:, start:end] @ xlds[bl[i], start:end].T
            
            # scores for orthogonal part of current block
            t_t_x = Xresidual @ R_keep[i]
            
            # predicted orthogonal contribution
            start_lv = lv_cumsum[i]
            end_lv = lv_cumsum[i+1]
            x_orth_chapeau = t_t_x[:, :lv[i]] @ xlds[bl[i], start_lv:end_lv].T
            
            # full predicted block
            x_pred = x_orth_chapeau + xcorr_chap
            SPE_blocks.append((x_pred - Xblock) ** 2)
            
            # append new scores to t_prev for next block
            t_prev = np.hstack([t_prev, t_t_x[:, :lv[i]]])
    
    return SPE_blocks


#%%

def VIP(W_blocks, SSY, lv):
    """ 
    Parameters
    ----------
    W : array-like, shape (n_variables, n_components)
        Weight matrix.
    SSY : array-like, shape (n_components,&)
        Explained Y variance (per component).
    lv : number of principal components pers block (b,)
    
    Returns
    -------
    VIP : array, shape (n_variables,1)
        VIP scores for each predictor variable.
    """
    
    cum_lv = np.cumsum(lv)
    SSY_cum = np.array([
        SSY[:end].sum()
        for end in cum_lv])
      
    VIP_block = []
    
    # Loop over each block
    for b, (Wb, A_b) in enumerate(zip(W_blocks, cum_lv)):
        # print('\nBlock: ', b)
        
        p_b = Wb.shape[0]
        VIP_b = np.zeros(p_b)
    
        for j in range(A_b):
            # global component index
            SSY_j = SSY[j]
    
            norm_wj = np.sum(Wb[:, j] ** 2)
    
            VIP_b += p_b * SSY_j * (Wb[:, j] ** 2 / norm_wj)
    
        VIP_b = VIP_b/SSY_cum[b]
        VIP_block.append(VIP_b)
    
    
    VIP2 = np.concatenate(VIP_block)
    
    VIP = np.sqrt(VIP2)

    return VIP


#%%

def BIP(Wt, SSY):
    """ 
    Parameters
    ----------
    Wt : array-like, shape (n_blocks, n_components)
        Squared super weights.
    SSY : array-like, shape (n_components,&)
        Explained Y variance (per component).
    
    Returns
    -------
    BIP : array, shape (n_variables,1)
        BIP scores for each block.
    """
    # BIP calculation
    n_blocks = len(Wt)
    BIP = np.zeros((n_blocks))
    for b in range(n_blocks):
        numerateur = np.sum(SSY * Wt[b, :])
        denominateur = np.sum(SSY)
        BIP[b] = np.sqrt(n_blocks * numerateur / denominateur)


    return BIP


#%% Loadmat

import scipy.io as spio

def loadmat(filename):
    '''
    this function should be called instead of direct spio.loadmat
    as it cures the problem of not properly recovering python dictionaries
    from mat files. It calls the function check keys to cure all entries
    which are still mat-objects
    '''
    def _check_keys(d):
        '''
        checks if entries in dictionary are mat-objects. If yes
        todict is called to change them to nested dictionaries
        '''
        for key in d:
            if isinstance(d[key], spio.matlab.mio5_params.mat_struct):
                d[key] = _todict(d[key])
        return d

    def _has_struct(elem):
        """Determine if elem is an array and if any array item is a struct"""
        return isinstance(elem, np.ndarray) and any(isinstance(e, spio.matlab.mio5_params.mat_struct) for e in elem)

    def _todict(matobj):
        '''
        A recursive function which constructs from matobjects nested dictionaries
        '''
        d = {}
        for strg in matobj._fieldnames:
            elem = matobj.__dict__[strg]
            if isinstance(elem, spio.matlab.mio5_params.mat_struct):
                d[strg] = _todict(elem)
            elif _has_struct(elem):
                d[strg] = _tolist(elem)
            else:
                d[strg] = elem
        return d

    def _tolist(ndarray):
        '''
        A recursive function which constructs lists from cellarrays
        (which are loaded as numpy ndarrays), recursing into the elements
        if they contain matobjects.
        '''
        elem_list = []
        for sub_elem in ndarray:
            if isinstance(sub_elem, spio.matlab.mio5_params.mat_struct):
                elem_list.append(_todict(sub_elem))
            elif _has_struct(sub_elem):
                elem_list.append(_tolist(sub_elem))
            else:
                elem_list.append(sub_elem)
        return elem_list
    data = spio.loadmat(filename, struct_as_record=False, squeeze_me=True)
    return _check_keys(data)


