"""Read RAMSES log files and outputs. For each output, print the value of the timestep."""

import argparse
import os
import re
import sys
from collections import defaultdict

from scipy.interpolate import interp1d

parser = argparse.ArgumentParser(description='Extracts the data from the given file')

parser.add_argument("--log-files", nargs="+", type=str, help="Input file(s)", required=True)
parser.add_argument("--outputs", help="Output folder(s)", nargs="+", type=str, required=True)

args = parser.parse_args()

# Extract dt and aexp from lines looking line:
# Fine step= 722821 t=-2.10465E+00 dt= 2.069E-06 a= 3.400E-01 mem=54.0% 66.1%
dt_line_re = re.compile(
    r"Fine step= \d+ t=-?\d+\.\d+E[+-]\d+ dt= (\d+\.\d+E[+-]\d+) a= (\d+\.\d+E[+-]\d+)"
    r" mem=\d+\.\d+% \d+\.\d+%"
)
output_name_re = re.compile(r"output_(\d{5})")

_aexp2dt = []

print("--- READING LOG FILES TO BUILD aexp TO dt MAPPING ---")

data = defaultdict(list)
for log_file in args.log_files:
    print("Extracting data from {}".format(log_file))

    with open(log_file, "r") as f:
        for line in f.readlines():
            m = dt_line_re.match(line.strip())

            if not m:
                continue

            dt, aexp = (float(_) for _ in m.groups())

            _aexp2dt.append((aexp, dt))

aexp2dt = interp1d(
    *zip(*_aexp2dt),
    kind="linear",
)

print("--- READING OUTPUTS TO COMPUTE dt IN EACH OUTPUT ---")
print("#output\tdt (kyr)")
for output in args.outputs:
    parent_folder, output_folder = os.path.split(output)
    m = output_name_re.match(output_folder)
    if not m:
        raise ValueError("Invalid output folder name: {}".format(output_folder))
    
    iout = int(m.group(1))

    info_fname = os.path.join(output, f"info_{iout:05d}.txt")

    with open(info_fname, "r") as f:
        for line in f.readlines():
            if line.startswith("aexp"):
                aexp = float(line.split("=")[1])
            elif line.startswith("unit_t"):
                unit_t = float(line.split("=")[1])
                break
 
        try:
            dt = aexp2dt(aexp) * unit_t
        except ValueError:
            print(
                f"WARNING: aexp={aexp} output={output} not in aexp2dt range. "
                "Are you missing its log?",
                file=sys.stderr,
            )
            continue

        dt_kyr = dt * 3.16880878e-11
        print(f"{output}\t{dt_kyr:.3f}")