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

def run_agcd(input_file):
    cmd = ["./target/release/approximate-gcd", "agcd", input_file]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        output = result.stdout

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

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

def plot_curves(noise_bits, p_bits, test_numbers, success_rates):
    plt.figure(figsize=(10, 6))
    plt.plot(test_numbers, success_rates, marker='o')
    plt.xlabel('Number of Test Values')
    plt.ylabel('Success Rate')
    plt.title(f'Success Rate vs. Number of Test Values\n(noise_bits={noise_bits}, p_bits={p_bits})')
    plt.grid(True)
    plt.ylim(-0.1, 1.1)
    plt.xticks(test_numbers)
    plt.savefig('success_rate_plot.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=0, 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=100, help='Maximum number of test values')
    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 = []
    num_trials = 100

    for num_values in test_numbers:
        successes = 0
        for _ in range(num_trials):
            # Create temporary test file
            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)

                # Run AGCD
                recovered_p = run_agcd(tmp_file.name)

                # Check if recovery was successful
                if recovered_p is not None and abs(recovered_p-true_p) <= 2000:
                    successes += 1

                # Clean up
                os.unlink(tmp_file.name)

        success_rate = successes / num_trials
        success_rates.append(success_rate)
        print(f"Number of values: {num_values}, Success rate: {success_rate:.3f} ({successes}/{num_trials})")

    # Plot the results
    plot_curves(noise_bits, p_bits, test_numbers, success_rates)

if __name__ == "__main__":
    main()