import os
import json
import re

# This correction factor is necessary since the grift dynamizer reports
# incorrect values. The reported precision is off by a constant for each
# benchmark. The constant is correction_factor / static_type_count.
CORRECTION_FACTOR = {
    "array": 3,
    "blackscholes": 6,
    "cps-even-odd": 4,
    "fft": 7,
    "matmult": 3,
    "n_body": 24,
    "qsort_mpairs": 5,
    "quicksort": 4,
    "ray": 20,
    "sieve": 12,
    "tak": 2,
}

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,
}


def extract_precision_from_file(file_path):
    try:
        with open(file_path, "r") as f:
            first_line = f.readline().strip()
            # e.g. ';; 82.61%' or ';; 95.5%'
            match = re.search(r";*\s*([\d.]+)%", first_line)
            if match:
                precision_percent = float(match.group(1))
                return precision_percent / 100.0
            else:
                print(f"Warning: Could not parse precision from {file_path}")
                print(f"First line: '{first_line}'")
                return None
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return None


def correct_precision(precision, bench_name, filename):
    if filename in ["static.grift", "dynamic.grift"]:
        return precision

    if bench_name in CORRECTION_FACTOR and bench_name in STATIC_TYPE_COUNT:
        correction_factor = CORRECTION_FACTOR[bench_name]
        static_type_count = STATIC_TYPE_COUNT[bench_name]
        corrected_percent = precision - (correction_factor / static_type_count)
        return corrected_percent
    else:
        print(
            f"Warning: No offset/baseline found for benchmark {bench_name}, "
            f"using uncorrected value"
        )
        return precision


def find_result_files(results_dir, bench_name):
    result_files = []
    for file in os.listdir(results_dir):
        path = os.path.join(results_dir, file)
        if os.path.isdir(path):
            continue
        basename = os.path.splitext(file)[0]
        parts = basename.split("_")
        compiler = parts[0]
        # Only consider results from grift
        if compiler != "grift":
            continue
        prefix = f"grift_{bench_name}_"
        if basename.startswith(prefix):
            result_files.append(path)
    return result_files


def process_benchmark(bench_name, benchmarks_dir, results_dir):
    samples_dir = os.path.join(benchmarks_dir, bench_name, "grift", "static_samples")
    if not os.path.exists(samples_dir):
        print(f"Warning: Samples directory not found: {samples_dir}")
        return False

    precision_map = {}
    corrections_log = []

    for filename in os.listdir(samples_dir):
        if filename.endswith(".grift"):
            sample_name = os.path.splitext(filename)[0]
            file_path = os.path.join(samples_dir, filename)
            precision = extract_precision_from_file(file_path)
            if precision is None:
                continue
            corrected_precision = correct_precision(precision, bench_name, filename)
            precision_map[sample_name] = corrected_precision
            if filename not in ["static.grift", "dynamic.grift"]:
                corrections_log.append((sample_name, precision, corrected_precision))

    if not precision_map:
        print(f"Warning: No precision data found for benchmark {bench_name}")
        return False

    result_files = find_result_files(results_dir, bench_name)
    print(f"Found {len(result_files)} result files to update:")
    print("\n".join(result_files))
    for result_file in result_files:
        if not os.path.exists(result_file):
            print(f"Warning: Result file not found: {result_file}")
            return False
        try:
            with open(result_file, "r") as f:
                results = json.load(f)

            updated_count = 0
            for run in results:
                basename = run["basename"]
                if basename in precision_map:
                    run["precision"] = precision_map[basename]
                    updated_count += 1
                else:
                    print(
                        f"Warning: No precision found for sample {basename} in {bench_name}"
                    )
                    return False

            with open(result_file, "w") as f:
                json.dump(results, f, indent=4)

            print(f"Updated {updated_count} runs in {result_file}")

        except Exception as e:
            print(f"Error processing result file {result_file}: {e}")
            return False

    if corrections_log:
        print(f"Applied corrections to {len(corrections_log)} samples:")
        for sample_name, raw, corrected in corrections_log[:10]:
            print(f"    {sample_name}: {raw:.2f} -> {corrected:.2f}")
        if len(corrections_log) > 10:
            print(f"    ... and {len(corrections_log) - 10} more")

    return True


def main():
    benchmarks_dir = "benchmarks"
    results_dir = "results"

    if not os.path.exists(benchmarks_dir):
        print(f"Error: Benchmarks directory not found: {benchmarks_dir}")
        return

    if not os.path.exists(results_dir):
        print(f"Error: Results directory not found: {results_dir}")
        return

    benchmarks = []
    for item in os.listdir(benchmarks_dir):
        item_path = os.path.join(benchmarks_dir, item)
        if os.path.isdir(item_path):
            benchmarks.append(item)

    if not benchmarks:
        print("No benchmarks found")
        return

    print(f"Found {len(benchmarks)} benchmarks")

    missing = []
    for bench in benchmarks:
        if bench not in CORRECTION_FACTOR or bench not in STATIC_TYPE_COUNT:
            missing.append(bench)
    if missing:
        print(f"Warning: Missing offset/baseline data for: {', '.join(missing)}")
        print("These benchmarks will use uncorrected precision values.\n")

    success_count = 0
    for bench_name in sorted(benchmarks):
        print(f"Processing benchmark: {bench_name}")
        if process_benchmark(bench_name, benchmarks_dir, results_dir):
            success_count += 1

    print(f"Done, successfully processed {success_count} out of {len(benchmarks)} benchmarks")


if __name__ == "__main__":
    main()
