"""
Title: S_smbpls
Date: Created in January 2026

    SMB-PLS main script to show how to use functions.
    
    - Cross validation section
    Cross-validation is performed sequentially.
    Cross-validation is first applied to the first block to select the number of 
    latent variables, then applied to the second block for latent variable selection, 
    and so on for the remaining blocks.
    
    The code is separated into distinct cross-validation sections for each block.
    To visualize the cross-validation results, uncomment the cross-validation 
    loop corresponding to the desired block, as well as all subsequent lines. 
    This reflects the block-wise selection structure on which the code is based.


    - Model calibration section: predictions, limits, VIP and BIP.
    
@author: givernyrobert
"""

#%% Imports

print('\n'*100)

import numpy as np
import matplotlib.pyplot as plt
import RG as rg
import smb_pls as smb
from scipy.stats import f


plt.close('all')


#%% Load data

data = smb.loadmat('JeuData.mat')

# =============================================================================
# # Calibration
# =============================================================================
Y = data['Data']['Ycal']
X_m = data['Data']['XcalB']
Z = data['Data']['Zcal']
X_MV = data['Data']['XcalA']
X_MV = np.column_stack((
    X_MV[:, :2],                    # First two columns
    X_MV[:, 0] ** 2,                # Square of the first column
    X_MV[:, 0] * Z[:, 0]            # Product of first column of X_MV and first column of Z
))

# Save original Y
Y0 = Y.copy()


# =============================================================================
# # Validation
# =============================================================================
X_m_v = data['Data']['XvalB']
Zv = data['Data']['Zval']
Yv = data['Data']['Yval']
X_MVv = data['Data']['XvalA']
X_MVv = np.column_stack((
    X_MVv[:, :2],                    # First two columns
    X_MVv[:, 0] ** 2,                # Square of the first column
    X_MVv[:, 0] * Zv[:, 0]            # Product of first column of X_MV and first column of Z
))

# Save original Yv
Yv0 = Yv.copy()

#%% Data normalisation

bs=1 # Block scaling
X_blocks = [Z, X_MV, X_m] 
Xv_blocks = [Zv, X_MVv, X_m_v] 

# Save copy 
X0_blocks = X_blocks.copy()
Xv0_blocks = Xv_blocks.copy()


Xnorm_blocks, YnormCal, bl, meanBlock, stdBlock = smb.normalisation_blocs_cal(X_blocks, Y=Y)

Xv_norm_blocks, YnormVal = smb.normalisation_blocs_val(Xv_blocks, Y=Yv, bs=1, bl=bl, mean_blocks=meanBlock,std_blocks=stdBlock)

#%% Cross-validation section

# =============================================================================
# Setting parameters for cross-validation
# =============================================================================

nb = 200 # Number of random picks
CVnb = int(0.3*len(Y))  # Number of lines predicted in each CV

n_blocks = len(X_blocks) # number of blocks

# Binary list to modify depending on which block studied in cv
chosen_list = [1] * n_blocks  # By default all blocks at 1

# List containing number of variables per block
A_list = [X.shape[1] for X in X_blocks]


# =============================================================================
# Block 1 - Z
# =============================================================================

# Settings to evaluate block
block_idx = 0    
chosen_list[block_idx] = 0

# ---------- CV loop -------------------------------------------------------

# RMSECV_all = smb.rs_cv(X_blocks, Y, A_list, chosen_list, bs,nb,CVnb)

# count = np.arange(1,len(RMSECV_all)+1)

# fig = plt.figure()
# plt.plot(count, RMSECV_all, '-g', label='RMSECV - RS')
# # Set x-ticks as integers only
# xmin = int(min(count))
# xmax = int(max(count))
# plt.xticks(np.arange(xmin, xmax + 1, 1), fontsize=12)  # step of 1
# plt.yticks(fontsize=12)
# plt.xlabel('Latent variables',fontsize=14)
# plt.ylabel('RMSECV',fontsize=14)
# plt.legend(fontsize=16)
# plt.title('Block '+str(block_idx+1),fontsize=16)
# plt.tight_layout()

# ----------------------------------------------------------------------------

lv1 = 3

print("Number of latent variables selected for block Z: ", lv1)

# Selection of number of latent variables
A_list[block_idx] = lv1
chosen_list = [1] * n_blocks   # all blocks are set back to 1

# =============================================================================
# Block 2 - X_MV
# =============================================================================

# Settings to evaluate block
block_idx = 1
chosen_list[block_idx] = 0

# # ---------- CV loop -------------------------------------------------------

# RMSECV_all = smb.rs_cv(X_blocks, Y, A_list, chosen_list, bs,nb,CVnb)

# count = np.arange(1,len(RMSECV_all)+1)

# fig = plt.figure()
# plt.plot(count, RMSECV_all, '-g', label='RMSECV - RS')
# # Set x-ticks as integers only
# xmin = int(min(count))
# xmax = int(max(count))
# plt.xticks(np.arange(xmin, xmax + 1, 1), fontsize=12)  # step of 1
# plt.yticks(fontsize=12)
# plt.xlabel('Latent variables',fontsize=14)
# plt.ylabel('RMSECV',fontsize=14)
# plt.legend(fontsize=16)
# plt.title('Block '+str(block_idx+1),fontsize=16)
# plt.tight_layout()

# # ----------------------------------------------------------------------------

lv2 = 4

print("Number of latent variables selected for block X_MV: ", lv2)

# Selection of number of latent variables
A_list[block_idx] = lv2
chosen_list = [1] * n_blocks   # all blocks are set back to 1

# =============================================================================
# Block 3 - X_m
# =============================================================================

# Settings to evaluate block
block_idx = 2
chosen_list[block_idx] = 0

# # ---------- CV loop -------------------------------------------------------

# RMSECV_all = smb.rs_cv(X_blocks, Y, A_list, chosen_list, bs,nb,CVnb)

# count = np.arange(1,len(RMSECV_all)+1)

# fig = plt.figure()
# plt.plot(count, RMSECV_all, '-g', label='RMSECV - RS')
# # Set x-ticks as integers only
# xmin = int(min(count))
# xmax = int(max(count))
# plt.xticks(np.arange(xmin, xmax + 1, 1), fontsize=12)  # step of 1
# plt.yticks(fontsize=12)
# plt.xlabel('Latent variables',fontsize=14)
# plt.ylabel('RMSECV',fontsize=14)
# plt.legend(fontsize=16)
# plt.title('Block '+str(block_idx+1),fontsize=16)
# plt.tight_layout()

# # ----------------------------------------------------------------------------

lv3 = 0

print("Number of latent variables selected for block X_MV: ", lv3)

# Selection of number of latent variables
A_list[block_idx] = lv3
chosen_list = [1] * n_blocks   # all blocks are set back to 1

#%% Model calibration section

# X-blocks concatenated for SMB-PLS
Xc = np.hstack(Xnorm_blocks)
Xv = np.hstack(Xv_norm_blocks)

# SMB-PLS model calculation
lv = A_list.copy()
modeleR = smb.smbpls_cal(Xc, YnormCal, bl, lv)

# =============================================================================
# Predictions
# =============================================================================

# Calibration
ym = meanBlock[-1]
ys = stdBlock[-1]
Ypred = (modeleR["ypred"]*ys + ym)

# Validation
resultats_v = smb.smbpls_exploit(Xv, Yv, bl, modeleR)
Ypred_v = (resultats_v["ypred"]*ys + ym)



n_y = Y0.shape[1]
R2c = np.zeros(n_y)
R2v = np.zeros(n_y)

for j in range(n_y):
    R2c[j] = rg.R2(Y0[:, j][:, None], Ypred[:, j][:, None])
    R2v[j] = rg.R2(Yv0[:, j][:, None], Ypred_v[:, j][:, None])


# RMSE calibration
E = Y0 - Ypred              # shape: (n_samples, n_y)
RMSEC = np.sqrt(np.mean(E**2, axis=0))   # shape: (n_y,)

# RMSE validation
Ev = Yv0 - Ypred_v          # shape: (n_samples, n_y)
RMSEV = np.sqrt(np.mean(Ev**2, axis=0))  # shape: (n_y,)



print("\nCalibration results:")
print("-" * 50)
for j in range(n_y):
    print(
        f"Y{j+1:>2} | "
        f"RMSEC = {RMSEC[j]:.4f} | "
        f"R²c = {R2c[j]:.4f}"
    )

print("\nValidation results:")
print("-" * 50)
for j in range(n_y):
    print(
        f"Y{j+1:>2} | "
        f"RMSEV = {RMSEV[j]:.4f} | "
        f"R²v = {R2v[j]:.4f}"
    )

for i in range(len(Y.T)):
    plt.figure()
    plt.plot(Y0[:,i],Ypred[:,i], 'ob', label='Cal')
    plt.plot(Y0[:,i],Y0[:,i], ls='--',c='k',lw=0.04)
    plt.plot(Yv0[:,i],Ypred_v[:,i], 'or', label='Val')
    plt.xlabel(r"Measured $y$")
    plt.ylabel(r"Predicted $\hat{y}$")
    plt.legend()
    plt.grid()
    plt.title(f"SMB-PLS predictions – Y{i+1}", fontweight="bold")
    plt.tight_layout()
    plt.show()


# =============================================================================
# Creation of R parameters
# =============================================================================

# Compute the maximum number of variables in all blocks
k_max = [mat.shape[1] for mat in X_blocks]

#  Compute R
R_keep = smb.CalculR(lv, k_max, modeleR, 3)

R_choisi = R_keep[0]
RXmv = R_keep[1]


# =============================================================================
# Calculation of limits
# =============================================================================

SPEY = (modeleR["ypred"] - YnormCal)**2
SPEYg = np.sum(SPEY, axis=1)

IC_Y_99 = smb.SPELimit(SPEYg,0.99)
IC_Y_95 = smb.SPELimit(SPEYg,0.95)


xlds = modeleR["xlds"]
ylds = modeleR["ylds"]


SPE_blocks = smb.SPECalculation(Xnorm_blocks, xlds, lv, R_keep, bl)

SPEX = SPE_blocks[1]
SPEXg = np.sum(SPEX, axis=1)
IC_X_99 = smb.SPELimit(SPEXg,0.99)
IC_X_95 = smb.SPELimit(SPEXg,0.95)


stdTT = np.std(modeleR['xscrs'], axis=0, ddof=0)  
HT2Cal = np.sum((modeleR['xscrs'] / stdTT)**2, axis=1)  
A = modeleR['xscrs'].shape[1]
N = X_blocks[0].shape[0]

IC_HT2_99 = A * (N**2 - 1) / (N * (N - A)) * f.ppf(0.99, A, N - A)

# =============================================================================
# Yvariance explained
# =============================================================================

# Squared super weights (n_blocs, n_LVs)
Wt = modeleR["Wsuper"]**2

# Block weights
W_blocks = modeleR["Wblock"]

# Yvariance explained per LV
SSY = np.array(modeleR["ssqtot"][2:,-2], dtype=float)/100

# Percentage of Yvariance explained across all blocks per LV
pY_lv = Wt*SSY

# Percentage of Yvariance explained across all blocks
pY_tot = np.sum(pY_lv, axis=1)

# =============================================================================
# Variable and block influence
# =============================================================================

# VIP_squared calculation
VIP = smb.VIP(W_blocks, SSY,lv)

# BIP
BIP = smb.BIP(Wt, SSY)


    
    