process_data.py

import os
import json
import yaml
import numpy as np
import matplotlib.pyplot as plt
import SignalProcessingTools.time_signal as time_signal
from SignalProcessingTools.time_signal import Windows

from validators import json_validator, yaml_validator, mdpa_validator


COORD_REF = [25, 0.7, 45]
TOL = 1e-6


def main(folder_path: str):
    """
    Main function to process YAML and JSON files in the specified folder.
    It validates the YAML files, checks the corresponding JSON files,
    and generates plots based on the data.

    Parameters:
        folder_path (str): Path to the folder containing YAML files.
    """

    if not os.path.exists(folder_path):
        return

    yaml_files = os.listdir(folder_path)
    yaml_files = [os.path.join(folder_path, file) for file in yaml_files if file.endswith('.yaml')]

    summary = {}

    for yaml_file in yaml_files:
        # validate YAML file
        if not yaml_validator(yaml_file):
            print(f"Validation failed for YAML file: {yaml_file}")
            raise ValueError(f"Invalid YAML file: {yaml_file}")

        with open(yaml_file, 'r') as f:
            meta = yaml.safe_load(f)

        # validate JSON files
        if not json_validator(os.path.join("data", meta["json-file"]), meta["STEM-version"]):
            print(f"Validation failed for JSON file: {meta['json-file']}")
            raise ValueError(f"Invalid JSON file: {meta['json-file']}")

        with open(os.path.join("data", meta["json-file"]), "r") as f:
            data = json.load(f)

        # validate mdpa file
        if not mdpa_validator(os.path.join("data", meta["mdpa-file"])):
            print(f"Validation failed for MDPA file: {meta['mdpa-file']}")
            raise ValueError(f"Invalid MDPA file: {meta['mdpa-file']}")

        with open(os.path.join("data", meta["mdpa-file"]), "r") as f:
            mdpa_content = f.read().splitlines()

        # Plotting the data
        summary[";".join([meta["title"], meta["organisation"]])] = process_plot_data(data, meta, mdpa_content)

    # edit the hugo content files
    edit_content_results(summary)
    edit_content_summary(summary)


def edit_content_results(summary: dict):
    """
    Edits the Hugo content results file to include the summary of processed data.

    Parameters:
        summary (dict): A dictionary containing the summary of processed data.
    """

    # results.md
    with open("STEM-cases/content/results.md", "r") as f:
        content = f.read()

    # Find the markers
    start_marker = "<!-- START AUTOGENERATED -->"
    start_index = content.find(start_marker)

    # Generate the new content
    new_content = []
    for key in sorted(summary.keys(), key=lambda x: int(x.split(" ")[1])):
        new_content.append(f"## {summary[key]['meta']['title']}\n\n")
        new_content.append(f"**Description:** {summary[key]['meta']['test-description']}\n")
        new_content.append(f"**Organization:** {summary[key]['meta']['organisation']}\n\n")
        new_content.append(f"**Date:** {summary[key]['meta']['date']}\n\n")
        new_content.append(f"**STEM Version:** {summary[key]['meta']['STEM-version']}\n\n")
        new_content.append(f"![{summary[key]['meta']['title']}](/TestCases/{summary[key]['plot_location']})\n\n")

    # Replace the content between markers
    before_marker = content[:start_index]

    updated_content = before_marker + "\n" + "".join(new_content)

    # Write back to file
    with open("STEM-cases/content/results.md", "w") as f:
        f.write(updated_content)


def edit_content_summary(summary: dict):
    """
    Edits the Hugo content summary file to include the results of processed data.

    Parameters:
        summary (dict): A dictionary containing the summary of processed data.
    """

    # summary.md
    with open("STEM-cases/content/summary.md", "r") as f:
        content = f.read()
    # Find the markers

    start_marker = "<!-- START AUTOGENERATED -->"
    start_index = content.find(start_marker)

    # Generate the new content
    new_content = ["| Test case | V_y,max | V_eff,max | PSD,max | Freq_PSD,max  |\n"]
    new_content.append("|-----|-----|-----|-----|-----|\n")
    for key in sorted(summary.keys(), key=lambda x: int(x.split(" ")[1])):
        new_content.append(f"| {summary[key]['meta']['title']} | "
                           f"{round(summary[key]['peak_velocity_y'], 3)} | "
                           f"{round(summary[key]['peak_v_eff'], 3)} | "
                           f"{round(summary[key]['peak_psd'], 4)} | "
                           f"{round(summary[key]['freq_peak_psd'], 3)} |\n")

    # Replace the content between markers
    before_marker = content[:start_index]

    updated_content = before_marker + "\n" + "".join(new_content)

    # Write back to file
    with open("STEM-cases/content/summary.md", "w") as f:
        f.write(updated_content)



def process_plot_data(data: dict, meta: dict, mdpa: dict) -> dict:
    """
    Processes and creates a plot from the data and metadata.

    Parameters:
        data (dict): The data dictionary containing the JSON results.
        meta (dict): The metadata dictionary.
        mdpa (dict): The MDPA content as a list of strings.

    Returns:
        dict: A summary dictionary containing peak values, frequencies, and plot location.
    """

    output_folder = "STEM-cases/static"
    name = "_".join(meta['title'].split())
    os.makedirs(output_folder, exist_ok=True)

    # define the node
    node = None

    # read the mdpa file and its nodes
    idx_ini = mdpa.index("Begin SubModelPart json_output")
    idx_end = mdpa[idx_ini:].index("  End SubModelPartNodes")

    mdpa_nodes = mdpa[idx_ini+4:idx_ini + idx_end]

    index_nodes = mdpa.index("Begin Nodes")
    index_end_nodes = mdpa[index_nodes:].index("End Nodes")
    aux = [j for i in mdpa_nodes for j, k in enumerate(mdpa[index_nodes:index_nodes + index_end_nodes]) if k.split()[0] == i.lstrip()]

    mdpa_nodes = [mdpa[index_nodes + i].split() for i in aux]

    for nod in mdpa_nodes:
        if np.linalg.norm(np.array(nod[1:]).astype(float) - np.array(COORD_REF)) < TOL:
            node = f"NODE_{nod[0]}"

    if node is None:
        raise ValueError("The reference node was not found. Please use the reference mesh.")

    if node not in data.keys():
        raise ValueError(f"The node {node} was not found in the data. Please check the JSON file.")

    # process the time signal
    signal = time_signal.TimeSignalProcessing(data["TIME"],
                                              np.array(data[node]["VELOCITY_Y"]),
                                              window=Windows.HAMMING,
                                              window_size=2000)
    signal.psd()
    signal.v_eff_SBR()

    fig, ax = plt.subplots(ncols=3, nrows=1, figsize=(15, 4))
    ax[0].plot(data["TIME"], np.array(data[node]["VELOCITY_Y"])*1000, label=r"v$_{y}$", color="blue")
    ax[1].plot(data["TIME"], signal.v_eff[:len(data["TIME"])], label=r"v$_{eff}$", color="orange")
    ax[2].plot(signal.frequency_Pxx, signal.Pxx*1000**2, label=r"v$_{y}$", color="blue")
    ax[0].set_xlabel("Time (s)")
    ax[0].set_ylabel("Velocity Y (mm/s)")
    ax[1].set_xlabel("Time (s)")
    ax[1].set_ylabel("V$_eff$ (mm/s)")
    ax[2].set_xlabel("Frequency (Hz)")
    ax[2].set_ylabel("PSD ([mm/s]$^2$/Hz)")
    ax[0].set_xlim(left=0)
    ax[1].set_xlim(left=0)
    ax[2].set_xlim(0, 100)
    ax[1].set_ylim(bottom=0)
    ax[2].set_ylim(bottom=0)
    ax[0].grid()
    ax[1].grid()
    ax[2].grid()
    ax[0].legend()
    ax[1].legend()
    ax[2].legend()
    plt.tight_layout()
    plt.savefig(os.path.join(output_folder, f"{name}.png"))
    plt.close()

    # create the summary
    summary = {"peak_velocity_y": np.max(np.abs(np.array(data[node]["VELOCITY_Y"])*1000)),
               "peak_v_eff": np.max(signal.v_eff),
               "peak_psd": np.max(signal.Pxx)*1000**2,
               "freq_peak_psd": signal.frequency_Pxx[np.argmax(signal.Pxx)],
               "plot_location": f"{name}.png",
               "meta": meta}
    return summary


if __name__ == "__main__":
    main("./data")