import glob
import re

# Directories where the profile files are stored
data_dir = "Outputs/md/"

# Get all profile files
energy_files = glob.glob(data_dir + "energyKspace_*.txt")
dcoeff_files = glob.glob(data_dir + "dcoeff_*.txt")

# Regular expressions to extract pressBar from filenames
patterns = {
    "energy": re.compile(r'energyKspace_(\d+)\.txt'),
    "dcoeff": re.compile(r'dcoeff_(\d+)\.txt')
}

# Function to read data from files
def read_data(files, pattern, n_cols):
    results = []
    for file in files:
        match = pattern.search(file)
        if match:
            pressBar = float(match.group(1))
            try:
                with open(file, 'r') as f:
                    lines = f.readlines()[1:]  # Skip header line
                    for line in lines:
                        data = line.split()
                        if len(data) == n_cols:
                            results.append((pressBar, *map(float, data)))
            except (ValueError, IndexError):
                continue
    return results

# Read and sort energy and dcoeff data
energy_results = read_data(energy_files, patterns["energy"], 5)
dcoeff_results = read_data(dcoeff_files, patterns["dcoeff"], 5)

energy_results.sort(key=lambda x: (x[0], x[1]))
dcoeff_results.sort(key=lambda x: (x[0], x[1]))

# Write results to single files
def write_results(filename, header, results, fmt):
    with open(filename, 'w') as outfile:
        outfile.write(header)
        for result in results:
            outfile.write(fmt.format(*result))

write_results("Outputs/md/combined_energy_data.txt", "#PressBar Timestep Energy1 Energy2 Energy3 Energy4\n",
              energy_results, "{:.2f} {:.2f} {:.4f} {:.4f} {:.4f} {:.4f}\n")

write_results("Outputs/md/combined_dcoeff_data.txt", "#PressBar Timestep Dcoeff1 Dcoeff2 Dcoeff3 Dcoeff4\n",
              dcoeff_results, "{:.2f} {:.2f} {:.8f} {:.8f} {:.8f} {:.8f}\n")