#!/bin/python3
import argparse
import subprocess
import re
import tempfile
import os
import numpy as np
import matplotlib.pyplot as plt
import time
import math
from gen_values import generate_test_file

def run_agcd(input_file):
    cmd = ["./target/release/approximate-gcd", "agcd", input_file]
    try:
        start = time.perf_counter()
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        duration = time.perf_counter() - start
        output = result.stdout

        match = re.search(r"Recovered p: (\d+)", output)
        return int(match.group(1)) if match else None, duration

    except subprocess.CalledProcessError as e:
        print(f"Error running command for input_file={input_file}: {e}")
        return None, 0.0
    except (AttributeError, ValueError) as e:
        print(f"Error parsing output for input_file={input_file}: {e}")
        return None, 0.0

def plot_all(noise_bits, p_bits, test_numbers, success_rates, mean_distances, mean_times):
    num_points = len(test_numbers)
    max_ticks = 20
    step = max(1, num_points // max_ticks)
    xticks = list(test_numbers)[::step]
    if test_numbers[-1] not in xticks:
        xticks.append(test_numbers[-1])

    fig, axs = plt.subplots(3, 1, figsize=(10, 15), sharex=True)

    # Success rate plot
    axs[0].plot(test_numbers, success_rates, linestyle='-')
    axs[0].set_ylabel('Success Rate')
    axs[0].set_ylim(-0.1, 1.1)
    axs[0].grid(True)
    axs[0].set_title(f'Success Rate (noise_bits={noise_bits}, p_bits={p_bits})')

    # Mean distance plot
    axs[1].plot(test_numbers, mean_distances, linestyle='-')
    axs[1].set_ylabel('Mean Distance to True p')
    axs[1].grid(True)
    axs[1].set_title(f'Mean Distance (noise_bits={noise_bits}, p_bits={p_bits})')

    # Mean runtime plot
    axs[2].plot(test_numbers, mean_times, linestyle='-')
    axs[2].set_ylabel('Mean Runtime (s)')
    axs[2].set_xlabel('Number of Test Values')
    axs[2].grid(True)
    axs[2].set_title(f'Mean Runtime (noise_bits={noise_bits}, p_bits={p_bits})')

    axs[2].set_xticks(xticks)

    fig.tight_layout()
    plt.savefig('agcd_plots.png')
    plt.show()

def main():
    parser = argparse.ArgumentParser(description='Test AGCD with varying number of test values.')
    parser.add_argument('--noise-bits', type=int, default=1, help='Number of noise bits')
    parser.add_argument('--p-bits', type=int, default=10000, help='Number of bits for p')
    parser.add_argument('--min-values', type=int, default=2, help='Minimum number of test values')
    parser.add_argument('--max-values', type=int, default=50, help='Maximum number of test values')
    parser.add_argument('--trials', type=int, default=50, help='Number of trials per setting')
    args = parser.parse_args()

    noise_bits = args.noise_bits
    p_bits = args.p_bits
    test_numbers = range(args.min_values, args.max_values + 1)

    success_rates = []
    mean_distances = []
    mean_times = []

    for num_values in test_numbers:
        successes = 0
        distances = []
        runtimes = []

        for _ in range(args.trials):
            with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as tmp_file:
                true_p = generate_test_file(noise_bits, num_values, p_bits, tmp_file.name)
            recovered_p, duration = run_agcd(tmp_file.name)
            os.unlink(tmp_file.name)

            if recovered_p is not None:
                diff = abs(recovered_p - true_p)
                threshold = math.isqrt(true_p.bit_length())
                if diff <= threshold:
                    successes += 1
                    distances.append(diff)
                    runtimes.append(duration)

        success_rates.append(successes / args.trials)
        mean_distances.append(np.mean(distances) if distances else float('nan'))
        mean_times.append(np.mean(runtimes))

        print(f"Values: {num_values}, Success rate: {success_rates[-1]:.3f} ({successes}/{args.trials}), "
              f"Mean distance: {mean_distances[-1]:.2f}, Mean time: {mean_times[-1]:.4f}s")

    plot_all(noise_bits, p_bits, list(test_numbers), success_rates, mean_distances, mean_times)

if __name__ == "__main__":
    main()