from prsctrl.utility.prsdata import PrsData
import os
import re
import numpy as np
import matplotlib.pyplot as plt

def process_results(
        data_dir,
        dir_regex=r".*",
        out_dir=None,
        get_varied_param=lambda data: data.metadata["lock-in_settings"]["frequency_Hz"],
        fig_suptitle="Frequency scan: <qty>",
        xlabel="$f$ [Hz]",
        what=["theta", "stheta", "dR_R", "sdR_R"],
):
    """
    Process the results of several spectra recorded with one varying measurement or lock-in parameter.
    
    :param data_dir: Directory containing all data directories of the individual files
    :param dir_regex: Regex for filtering data directories
    :param out_dir: Directory into which the plots are saved
    :param get_varied_param: lambda function getting the varied parameter value from a loaded PrsData object
    :param fig_suptitle: Figure suptitle, <qty> will be replaced by the name of the quantity
    :param xlabel: xlabel for the varied parameter
    :param what: For which quantities to create plots
    :return: 
    """
    data_dir = os.path.expanduser(data_dir)
    if out_dir is None:
        out_dir = data_dir
    paths = os.listdir(data_dir)
    data_dirs = []
    for p in paths:
        full_path = os.path.join(data_dir, p)
        if not os.path.isdir(full_path): continue
        m = re.fullmatch(dir_regex, p)
        if m:
            data_dirs.append(full_path)
        else:
            print(f"Unmatched directory {p}")
    assert len(data_dirs) > 0
    data_dirs.sort()

    varied_params = []
    data = {}
    shape = None
    wls = None
    for d in data_dirs:
        print(f"Getting data from {d}")
        pd = PrsData(load_data_path=d)
        # f = pd.metadata["lock-in_settings"]["frequency_Hz"]
        varied_param = get_varied_param(pd)
        # print(d, f)
        sdata = pd.get_spectrum_data()
        print(pd.wavelengths)
        print(pd.data.keys())
        if wls is None: wls = sdata[:,0]
        if shape is None: shape = sdata.shape
        else:
            if shape != sdata.shape:
                print(f"ERROR Shape mismatch for param={varied_param}: {shape} != {sdata.shape}")
                continue
                # raise ValueError(f"Shape mismatch for {d}: {shape} != {sdata.shape}")
        varied_params.append(varied_param)
        data[varied_param] = sdata
    data_per_wl_and_vp = np.empty((shape[0], len(varied_params), shape[1]))
    varied_params.sort()
    for i in range(shape[0]):
        for j, f in enumerate(varied_params):
            data_per_wl_and_vp[i, j, :] = data[f][i,:]
    print(f"Found wavelengths: {wls}")
    n_cols = 2
    n_rows = wls.shape[0] // n_cols + wls.shape[0] % n_cols
    for qty in what:
        fig, axs = plt.subplots(n_rows, n_cols, figsize=(8, 8))
        axs = axs.flatten()
        qty_idx = PrsData.default_spectrum_columns.index(qty)
        fig.suptitle(fig_suptitle.replace("<qty>", PrsData.key_names[qty]))
        for i, wl in enumerate(wls):
            ax = axs[i]
            ax.set_xlabel(xlabel)
            ax.set_ylabel(PrsData.labels[qty])
            ax.plot(varied_params, data_per_wl_and_vp[i, :, qty_idx])
            ax.set_title(f"$\\lambda = {wl}$ nm")
        fig.tight_layout()
        fig.savefig(os.path.join(out_dir, f"result_{qty}.pdf"))
    print(varied_params)