import numpy as np
from ssEinterpolator import *
import os
from joblib import Parallel, delayed


ws = np.load('./ws.npy')
sns = np.load('./sns.npy')
f_params = np.stack((ws, sns)).T
s_params = f_params.astype(str)


def load_data(params):
    print(params)
    base_dir = '....' # Replace with the actual base directory where the data files are stored
    sr = np.load(f'{base_dir}w_{params[0]}_sn_{params[1]}_sr.npy')
    slip = np.load(f'{base_dir}w_{params[0]}_sn_{params[1]}_slip.npy')
    state = np.load(f'{base_dir}w_{params[0]}_sn_{params[1]}_state.npy')
    t = np.load(f'{base_dir}w_{params[0]}_sn_{params[1]}_t.npy')
    return sr, slip, state, t


rom = ROM(f_params=f_params, str_params=s_params, load_f=load_data, lf_path='./data/lf.npy', sses_num=19, t_to_u_knot_l=2500, u_to_par_knot_l=250, along_dp_sses_depth_detector = 195, sses_detector_threshold=-4, sses_starts=0, t_start=75, t_end=1e99, base_step_t_interpolate=1e-4, depths_t_interpolate=9, log=True)

n_jobs = 4
output_directory = 'latent_matrix/' # Define the output directory

if not os.path.exists(output_directory):
    os.makedirs(output_directory)

print(f"Starting parallel creation of latent matrices with {n_jobs} processes...")

#
results = Parallel(n_jobs=n_jobs, prefer="processes")(
    delayed(rom.create_one_sse_latent_matrix)(i, output_directory) for i in range(rom.sses_num)
)
print("Finished creating all latent matrices.")

rom.build_pod()
rom.save_pod_to_files(output_directory)