#!/usr/bin/env python3

import argparse
import json
import os
import random
import re
import subprocess
import sys
import glob

from dataclasses import dataclass, field

# Default values, use the CLI to set different values
COMPILERS = [
    "gcomp",
    "grift",
]

GRIFT_BACKENDS = ["gradual", "gradual_mono"]
BENCHMARKS = [
    "array",
    "cps-even-odd",
    "matmult",
    "qsort_mpairs",
    "quicksort",
    "sieve",
    "tak",
    "n_body",
    "fft",
    "ray",
    "blackscholes",
]
INPUTS = {
    "array": "fast",
    "cps-even-odd": "slower",
    "matmult": "400",
    "qsort_mpairs": "descend10000",
    "quicksort": "in_descend_10000",
    "sieve": "fast",
    "tak": "fast",
    "n_body": "slow",
    "fft": "medium1",
    "ray": "fast",
    "blackscholes": "in_4K",
}

# Number of static types in the Grift version of the benchmarks
STATIC_TYPE_COUNT = {
    "array": 23,
    "blackscholes": 128,
    "cps-even-odd": 13,
    "fft": 67,
    "matmult": 33,
    "n_body": 136,
    "qsort_mpairs": 43,
    "quicksort": 44,
    "ray": 282,
    "sieve": 91,
    "tak": 8,
}

SAMPLING_FACTOR = {
    "slow": 10,
    "fast": 1,
}

TYPECHECKING = ["gradual", "gradual_succinct"]
CFLAGS = ["-march=native -O3"]
NUM_BINS = 10
NUM_SAMPLES = 5  # This is now a fallback/default value
NUM_RUNS = 20

DYNAMIZER_CMD = "dune exec dynamizer --"
COMPILER_CMD = "dune exec grev --"
BENCH_RUNNER_CMD = "python3 bench_runner.py"

GRIFT_COMPILER_CMD = "./wrappers/grift_mono"
GRIFT_DYNAMIZER_CMD = "./wrappers/dynamizer"

RACKET_COMPILER_CMD = "raco exe"


@dataclass
class Config:
    compilers: list[str] = field(default_factory=lambda: COMPILERS)
    benchmarks: list[str] = field(default_factory=lambda: BENCHMARKS)
    inputs: dict[str, str] = field(default_factory=lambda: INPUTS)
    typechecking: list[str] = field(default_factory=lambda: TYPECHECKING)
    cflags: list[str] = field(default_factory=lambda: CFLAGS)
    num_bins: int = NUM_BINS
    num_samples: int = NUM_SAMPLES
    num_runs: int = NUM_RUNS
    dynamizer_cmd: str = DYNAMIZER_CMD
    compiler_cmd: str = COMPILER_CMD
    bench_runner_cmd: str = BENCH_RUNNER_CMD
    base_benchmark_dir: str = "benchmarks"
    results_dir: str = "results"
    output_dir: str = "output"
    use_internal_timing: bool = True
    run_baselines: bool = False
    run_experiments: bool = True
    run_extremes: bool = False
    sampling_factor: int | float = SAMPLING_FACTOR["fast"]

    def __post_init__(self):
        if any(bench not in self.inputs for bench in self.benchmarks):
            raise ValueError("at least one benchmark doesn't have a specified input")


def parse_args():
    default = Config()
    """Create configuration from command line arguments"""
    parser = argparse.ArgumentParser(description="Automate benchmarking process")
    parser.add_argument(
        "--compilers",
        nargs="+",
        default=default.compilers,
        help=f"Compilers to include in the experiment, must be one of {COMPILERS}",
    )
    parser.add_argument(
        "--benchmarks",
        nargs="+",
        default=default.benchmarks,
        help="Benchmarks to include in the experiment",
    )
    parser.add_argument(
        "--typechecking",
        nargs="+",
        default=default.typechecking,
        help="Typechecking modes to include in the experiment",
    )
    parser.add_argument(
        "--cflags",
        nargs="+",
        default=default.cflags,
        help="Compiler flag variants to include in the experiment (space-separated)",
    )
    parser.add_argument(
        "--num-bins",
        type=int,
        default=default.num_bins,
        help="Number of bins to be used for the dynamizer",
    )
    parser.add_argument(
        "--num-samples",
        type=int,
        default=default.num_samples,
        help="Number of samples (used as fallback if benchmark not in STATIC_TYPE_COUNT)",
    )
    parser.add_argument(
        "--num-runs",
        type=int,
        default=default.num_runs,
        help="Number of runs per sample",
    )
    parser.add_argument(
        "--sampling-factor",
        default=default.sampling_factor,
        help="Sampling factor: 'fast', 'slow', or a numeric value (multiplied by static type count)",
    )
    parser.add_argument(
        "--internal-timing",
        action=argparse.BooleanOptionalAction,
        default=default.use_internal_timing,
        help="Use internal timing for benchmarks (benchmark explicitly calls `time`)",
    )
    parser.add_argument(
        "--base-benchmark-dir",
        default=default.base_benchmark_dir,
        help="Base directory for benchmarks",
    )
    parser.add_argument(
        "--results-dir", default=default.results_dir, help="Directory to store results"
    )
    parser.add_argument(
        "--run-baselines",
        action=argparse.BooleanOptionalAction,
        default=default.run_baselines,
        help="Whether to run baseline experiments",
    )
    parser.add_argument(
        "--run-extremes",
        action=argparse.BooleanOptionalAction,
        default=default.run_extremes,
        help="Whether to run extreme (fully typed/untyped) experiments",
    )
    parser.add_argument(
        "--run-experiments",
        action=argparse.BooleanOptionalAction,
        default=default.run_experiments,
        help="Whether to run the main experiments",
    )
    args = parser.parse_args()

    sampling_factor_arg = args.sampling_factor
    if sampling_factor_arg in SAMPLING_FACTOR:
        sampling_factor = SAMPLING_FACTOR[sampling_factor_arg]
    else:
        try:
            sampling_factor = float(sampling_factor_arg)
        except ValueError:
            parser.error(
                f"--sampling-factor must be 'fast', 'slow', or a number, got: {sampling_factor_arg}"
            )

    config = Config(
        compilers=args.compilers,
        benchmarks=args.benchmarks,
        typechecking=args.typechecking,
        cflags=args.cflags,
        num_bins=args.num_bins,
        num_samples=args.num_samples,
        num_runs=args.num_runs,
        sampling_factor=sampling_factor,
        use_internal_timing=args.internal_timing,
        base_benchmark_dir=args.base_benchmark_dir,
        results_dir=args.results_dir,
        run_baselines=args.run_baselines,
        run_experiments=args.run_experiments,
        run_extremes=args.run_extremes,
    )
    return config


def get_num_samples(config: Config, benchmark: str) -> int:
    if benchmark not in STATIC_TYPE_COUNT:
        print(
            f"Warning: {benchmark} not in STATIC_TYPE_COUNT, using default num_samples={config.num_samples}"
        )
        return config.num_samples

    static_types = STATIC_TYPE_COUNT[benchmark]
    samples_per_bin = max(
        int(static_types * config.sampling_factor / config.num_bins), 1
    )

    print(
        f"Benchmark {benchmark}: {config.num_bins} bins, {static_types} static types, sampling factor = {config.sampling_factor}, samples per bin = {samples_per_bin}, total samples = {samples_per_bin * config.num_bins}"
    )
    return samples_per_bin


def get_bench_subdir(compiler: str):
    if compiler == "gcomp":
        return "gradual"
    elif compiler == "grift":
        return "grift"
    elif compiler == "racket":
        return "racket"
    else:
        raise ValueError("unexpected compiler")


def get_bench_basename(compiler: str, benchmark: str):
    if compiler == "gcomp":
        return "static"
    elif compiler == "grift":
        return "static"
    elif compiler == "racket":
        return benchmark
    else:
        raise ValueError("unexpected compiler")


def get_bench_suffix(compiler: str):
    if compiler == "gcomp":
        return ".txt"
    elif compiler == "grift":
        return ".grift"
    elif compiler == "racket":
        return ".rkt"
    else:
        raise ValueError("unexpected compiler")


def get_source_path(config: Config, compiler: str, benchmark: str) -> str:
    bench_subdir = get_bench_subdir(compiler)
    basename = get_bench_basename(compiler, benchmark)
    bench_suffix = get_bench_suffix(compiler)
    return f"{config.base_benchmark_dir}/{benchmark}/{bench_subdir}/{basename}{bench_suffix}"


def get_input_path(config: Config, benchmark: str, input: str) -> str:
    return f"{config.base_benchmark_dir}/{benchmark}/inputs/{input}.txt"


def get_result_path(
    config: Config, compiler: str, benchmark: str, typechecking: str, cflags: str
) -> str:
    cflags_safe = cflags.replace(" ", "_").replace("-", "")
    return (
        f"{config.results_dir}/{compiler}_{benchmark}_{typechecking}_{cflags_safe}.json"
    )


def dynamize_cmd(config: Config, compiler: str, benchmark: str) -> str:
    source_file = get_source_path(config, compiler, benchmark)
    num_samples = get_num_samples(config, benchmark)

    if compiler == "gcomp":
        return f"{config.dynamizer_cmd} {source_file} --num-samples={num_samples} --num-bins={config.num_bins}"
    elif compiler == "grift":
        return f"{GRIFT_DYNAMIZER_CMD} {source_file} --fine --bins {config.num_bins} --configurations-count {num_samples}"


def compile_cmd(config: Config, compiler: str, typechecking: str, cflags: str) -> str:
    if compiler == "gcomp":
        flags = " ".join([f"--cflags='{f}'" for f in cflags.split()])
        return (
            f"{config.compiler_cmd} -t {typechecking} {flags} -o {{output}} {{source}}"
        )
    elif compiler == "grift":
        typing_flag = ""
        if typechecking == "static":
            typing_flag = "--static"
        elif typechecking == "gradual":
            typing_flag = "--coercions"
        elif typechecking == "gradual_mono":
            typing_flag = "--monotonic-references"
        elif typechecking == "gradual_casts":
            typing_flag == "--type-based-casts"
        else:
            raise ValueError(
                f"Unsupported typechecking mode '{typechecking}' for grift"
            )
        pattern = r"-O([0-9]*)\s*"
        optimization_flags = re.findall(pattern, cflags)
        cleaned_cflags = re.sub(pattern, "", cflags).strip()
        if len(optimization_flags) > 1:
            raise ValueError(
                "Too many optimization flags for grift: {optimization_flags}"
            )
        opt_level = optimization_flags[0] if len(optimization_flags) == 1 else None
        opt = f" -O {opt_level}" if opt_level is not None else ""
        flags = f" --c-flag '{cleaned_cflags}'" if cleaned_cflags != "" else ""
        cmd = f"{GRIFT_COMPILER_CMD} {typing_flag}{opt}{flags} -o {{output}} --keep-c {{output}}.c {{source}}"
        return cmd
    elif compiler == "racket":
        cmd = f"{RACKET_COMPILER_CMD} -o {{output}} {{source}}"
        return cmd


def exec(cmd: str, description: str) -> bool:
    print(f"{description}:\n{cmd}")
    process = subprocess.run(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
    )
    stdout = process.stdout.decode("utf-8")
    stderr = process.stderr.decode("utf-8")
    print(f"stdout:\n{stdout}")
    print(f"stderr:\n{stderr}")
    if process.returncode != 0:
        print("Failed")
        return False
    return True


def parse_internal_timing_output_gcomp(output: str) -> float:
    timing_pattern = r"Execution time:\s+(\d+\.\d+)\s+seconds"
    match_opt = re.search(timing_pattern, output)
    if match_opt is None:
        print("Failed to parse internal timing output (expected gcomop format)")
        sys.exit(1)
    mean = float(match_opt.group(1))
    return mean


def parse_internal_timing_output_grift(output: str) -> float:
    timing_pattern = r"time \(sec\):\s*(\d+\.\d+)"
    match_opt = re.search(timing_pattern, output)
    if match_opt is None:
        print("Failed to parse internal timing output (expected grift format)")
        sys.exit(1)
    mean = float(match_opt.group(1))
    return mean


def parse_internal_timing_output_racket(output: str) -> float:
    timing_pattern = r"real time:\s*(\d+)"
    match_opt = re.search(timing_pattern, output)
    if match_opt is None:
        print("Failed to parse internal timing output (expected racket format)")
        sys.exit(1)
    # Convert milliseconds to seconds
    mean = float(match_opt.group(1)) / 1000.0
    return mean


def parse_internal_timing_output(compiler, output):
    if compiler == "gcomp":
        return parse_internal_timing_output_gcomp(output)
    elif compiler == "grift":
        return parse_internal_timing_output_grift(output)
    elif compiler == "racket":
        return parse_internal_timing_output_racket(output)
    else:
        raise ValueError("unexpected compiler")


def exec_with_timing(cmd: str, description: str, compiler: str) -> float | None:
    print(f"{description}:\n{cmd}")
    process = subprocess.run(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
    )
    stdout = process.stdout.decode("utf-8")
    stderr = process.stderr.decode("utf-8")

    if compiler == "gcomp":
        timing_output = stderr
    elif compiler == "grift":
        timing_output = stdout
    elif compiler == "racket":
        timing_output = stdout
    else:
        raise ValueError("Unexpected compiler")
    print(f"stdout:\n{stdout}")
    if process.returncode != 0:
        print("Failed")
        print(process.stderr.decode("utf-8"))
        return None
    try:
        timing = parse_internal_timing_output(compiler, timing_output)
        print(f"Execution time: {timing} seconds")
    except:
        print("Failed to parse timing output")
        print(f"output:\n{stdout}")
        return None
    return timing


def run_benchmark(
    config: Config,
    compiler: str,
    benchmark: str,
    typechecking: str,
    cflags: str,
    input: str,
) -> bool:

    dyn_cmd = dynamize_cmd(config, compiler, benchmark)
    if not exec(dyn_cmd, f"Sampling {benchmark} using the {compiler} dynamizer"):
        return False

    source_path = get_source_path(config, compiler, benchmark)
    source_path_without_extension, _ = os.path.splitext(source_path)
    samples_dir = f"{source_path_without_extension}_samples"

    if not os.path.isdir(samples_dir):
        print(
            f"Error, expected dynamizer to have created the directory '{samples_dir}'"
        )
        return False

    compile_command = compile_cmd(config, compiler, typechecking, cflags)
    input_path = get_input_path(config, benchmark, input)
    result_file = get_result_path(config, compiler, benchmark, typechecking, cflags)

    bench_runner_cmd = f"""{config.bench_runner_cmd} \\
    {samples_dir} \\
    --compile-cmd "{compile_command}" \\
    --run-cmd '<{input_path}' \\
    --num-runs {config.num_runs} \\
    --time-format {compiler} \\
    --file-ext {get_bench_suffix(compiler)} \\
    --output-file {result_file}"""

    if not config.use_internal_timing:
        bench_runner_cmd.append(" \\\n    --no-internal-timing")

    return exec(
        bench_runner_cmd,
        f"benchmark {benchmark} with {typechecking} typechecking and {cflags}",
    )


def get_extremes_files(
    config: Config, compiler: str, benchmark: str
) -> tuple[str, str]:
    source_path = get_source_path(config, compiler, benchmark)
    source_path_without_extension, _ = os.path.splitext(source_path)
    samples_dir = f"{source_path_without_extension}_samples"

    if compiler == "grift":
        untyped = os.path.join(samples_dir, "dynamic.grift")
        typed = os.path.join(samples_dir, "static.grift")
        if not os.path.exists(untyped) or not os.path.exists(typed):
            raise FileNotFoundError(
                f"Could not find dynamic.grift or static.grift in {samples_dir}"
            )
        return untyped, typed
    elif compiler == "gcomp":
        # Pattern: sample_{bin}_{num}_{hex}.txt
        pattern = os.path.join(samples_dir, "sample_*.txt")
        files = glob.glob(pattern)
        if not files:
            raise FileNotFoundError(f"No samples found in {samples_dir}")

        def get_hex_val(filepath):
            filename = os.path.basename(filepath)
            match = re.search(r"_([0-9a-fA-F]+)\.txt$", filename)
            if match:
                return int(match.group(1), 16)
            return -1

        sorted_files = sorted(files, key=get_hex_val)
        untyped = sorted_files[0]
        typed = sorted_files[-1]
        if get_hex_val(untyped) != 0:
            print(
                f"Warning: Lowest hex value found for {benchmark} was not 0. It was {get_hex_val(untyped)}"
            )
        return untyped, typed
    else:
        raise ValueError(f"Extremes not supported for compiler: {compiler}")


def run_extremes(config: Config):
    results = []
    failures = []
    combinations = []

    print("Identifying extreme configurations...")

    # generate configs
    for compiler in config.compilers:
        if compiler == "racket":
            continue
        for benchmark in config.benchmarks:
            dyn_cmd = dynamize_cmd(config, compiler, benchmark)
            if not exec(
                dyn_cmd, f"Sampling {benchmark} using the {compiler} dynamizer"
            ):
                print(f"Skipping {benchmark} due to dynamizer failure")
                continue
            try:
                untyped_file, typed_file = get_extremes_files(
                    config, compiler, benchmark
                )
            except Exception as e:
                print(f"Error finding extreme files for {benchmark}: {e}")
                continue
            input_files = {
                "untyped": untyped_file,
                "typed": typed_file,
            }
            backends = config.typechecking if compiler != "grift" else GRIFT_BACKENDS
            for backend in backends:
                for i, cflags in enumerate(config.cflags):
                    for variant in ["untyped", "typed"]:
                        output = os.path.join(
                            config.output_dir,
                            f"extremes_{variant}_{compiler}_{benchmark}_{backend}_{i}",
                        )
                        combinations.append(
                            {
                                "compiler": compiler,
                                "bench": benchmark,
                                "typing": backend,
                                "cflags": cflags,
                                "source": input_files[variant],
                                "output": output,
                                "variant": variant,
                            }
                        )

    # compile
    successful_combinations = []
    for job in combinations:
        cmd = compile_cmd(config, job["compiler"], job["typing"], job["cflags"])
        cmd = cmd.replace("{output}", job["output"])
        cmd = cmd.replace("{source}", job["source"])

        success = exec(
            cmd,
            f"Compiling {job['bench']} (fully-{job['variant']}) with {job['compiler']} ({job['typing']})",
        )
        if not success:
            failures.append((job, "compilation"))
            continue
        successful_combinations.append(job)

    # run
    for i in range(config.num_runs):
        print(f"=== Run {i} ===")
        randomized_combinations = successful_combinations.copy()
        random.shuffle(randomized_combinations)
        for job in randomized_combinations:
            input_file = get_input_path(
                config, job["bench"], config.inputs[job["bench"]]
            )
            run_command = f"./{job['output']} <{input_file}"
            runtime_maybe = exec_with_timing(
                run_command,
                f"Running {job['bench']} (fully-{job['variant']}) compiled with {job['compiler']} ({job['typing']}) run={i+1}",
                job["compiler"],
            )
            if runtime_maybe is None:
                failures.append((job, "execution"))
                continue

            results.append(
                {
                    "compiler": job["compiler"],
                    "bench": job["bench"],
                    "typing": job["typing"],
                    "cflags": job["cflags"],
                    "variant": job["variant"],
                    "runtime": runtime_maybe,
                    "run": i + 1,
                }
            )

    print("\n=== Summary ===")
    print(f"Ran extremes for {len(successful_combinations)} combinations")
    print(f"Attempted {len(combinations)} combinations")
    print(f"Successful compilations: {len(successful_combinations)}")
    print(f"Failed benchmarks: {len(failures)}")
    print(f"Successful runs: {len(results)}")
    return results


def run_baselines(config: Config):
    results = []
    failures = []
    typings = ["static", "dynamic"]
    combinations = []
    for compiler in config.compilers:
        for benchmark in config.benchmarks:
            for typechecking in typings:
                if compiler == "grift" and typechecking not in ["gradual", "static"]:
                    continue
                elif compiler == "racket" and typechecking not in ["dynamic"]:
                    continue
                for i, cflags in enumerate(config.cflags):
                    source_file = get_source_path(config, compiler, benchmark)
                    output_file = os.path.join(
                        config.output_dir,
                        f"baseline_{compiler}_{benchmark}_{typechecking}_{i}",
                    )
                    combinations.append(
                        (compiler, benchmark, typechecking, cflags, output_file)
                    )
    # compile
    successful_combinations = []
    for combination in combinations:
        compiler, benchmark, typechecking, cflags, output_file = combination
        source_file = get_source_path(config, compiler, benchmark)
        cmd = compile_cmd(config, compiler, typechecking, cflags)
        cmd = cmd.replace("{output}", output_file)
        cmd = cmd.replace("{source}", source_file)
        print(cmd)
        success = exec(cmd, f"compiler ({', '.join(combination)})...")
        if not success:
            print(f"Skipping due to compilation failure")
            failures.append(combination + ("compilation",))
            continue
        successful_combinations.append(combination)

    # run
    for i in range(config.num_runs):
        print(f"=== Run {i} ===")
        randomized_combinations = successful_combinations.copy()
        random.shuffle(randomized_combinations)
        for combination in randomized_combinations:
            compiler, benchmark, typechecking, cflags, output_file = combination
            input_file = get_input_path(config, benchmark, config.inputs[benchmark])
            run_command = f"./{output_file} <{input_file}"
            runtime_maybe = exec_with_timing(
                run_command,
                f"executable ({', '.join(combination)}) run={i + 1}...",
                compiler,
            )
            if runtime_maybe is None:
                print(f"Skipping due to execution failure")
                failures.append(combination + ("execution",))
                continue
            runtime = runtime_maybe
            results.append(
                {
                    "compiler": compiler,
                    "bench": benchmark,
                    "typing": typechecking,
                    "cflags": cflags,
                    "runtime": runtime,
                    "run": i + 1,
                }
            )
    print("\n=== Summary ===")
    print(f"Ran baselines for {len(successful_combinations)} combinations")
    for compiler, benchmark, typechecking, cflags, _output_file, reason in failures:
        print(
            f"Failed: {compiler} on benchmark {benchmark} with typing {typechecking} and cflags='{cflags}' due to {reason} failure)"
        )
    print(f"Failed benchmarks: {len(failures)}")
    print(f"Successful runs: {len(results)}")
    return results


def run_experiments(config: Config) -> dict:
    results = {}
    for compiler in config.compilers:
        results[compiler] = {}
        for benchmark in config.benchmarks:
            results[compiler][benchmark] = {}
            for typechecking in config.typechecking:
                results[compiler][benchmark][typechecking] = {}
                for cflags in config.cflags:
                    print(
                        f"\n=== Running {benchmark} with {typechecking} typechecking and {cflags} ==="
                    )
                    success = run_benchmark(
                        config,
                        compiler,
                        benchmark,
                        typechecking,
                        cflags,
                        config.inputs[benchmark],
                    )
                    results[compiler][benchmark][typechecking][cflags] = (
                        "SUCCESS" if success else "FAILED"
                    )
    return results


def main():
    """Main function for command line usage"""
    config = parse_args()
    print("Starting automated benchmarks with configuration:")
    print(f"  Compilers: {config.compilers}")
    print(f"  Benchmarks: {config.benchmarks}")
    print(f"  Typechecking modes: {config.typechecking}")
    print(f"  CFLAGS variants: {config.cflags}")
    print(f"  Num bins: {config.num_bins}")
    print(f"  Num samples (fallback): {config.num_samples}")
    print(f"  Sampling factor: {config.sampling_factor}")
    print(f"  Num runs: {config.num_runs}")
    print(f"  Use internal timing: {config.use_internal_timing}")
    print(f"  Base benchmark dir: {config.base_benchmark_dir}")
    print(f"  Results dir: {config.results_dir}")
    print(f"  Run baselines: {config.run_baselines}")
    print(f"  Run extremes: {config.run_extremes}")
    print(f"  Run experiments: {config.run_experiments}")
    print()

    os.makedirs(config.results_dir, exist_ok=True)
    os.makedirs(config.output_dir, exist_ok=True)

    if config.run_experiments:
        results = run_experiments(config)

        print("\n=== Benchmark Summary ===")
        for benchmark, typechecking_results in results.items():
            print(f"{benchmark}:")
            for typechecking, cflags_results in typechecking_results.items():
                print(f"  {typechecking}:")
                for cflags, status in cflags_results.items():
                    print(f"    {cflags}: {status}")

    if config.run_baselines:
        print("\nRunning baselines...")
        baseline_results = run_baselines(config)
        baseline_path = os.path.join(config.results_dir, "baselines.json")
        with open(baseline_path, "w") as f:
            json.dump(baseline_results, f, indent=4)
        print(f"\nBaseline results written to {baseline_path}")

    if config.run_extremes:
        print("\nRunning extremes...")
        extremes_results = run_extremes(config)
        extremes_path = os.path.join(config.results_dir, "extremes.json")
        with open(extremes_path, "w") as f:
            json.dump(extremes_results, f, indent=4)
        print(f"\nExtreme results written to {extremes_path}")


if __name__ == "__main__":
    main()
