clock-curve-math 1.1.3

#!/usr/bin/env python3
"""
SageMath implementations of group order computation algorithms.

This module provides correct implementations of group order algorithms
that can be used to verify and guide the Rust implementations.
"""

import sage.all as sage
from sage.groups.generic import order_from_multiple


class GroupOrderAlgorithms:
    """Group order computation algorithm implementations using SageMath."""

    def __init__(self, p, generator=None):
        """
        Initialize group order computation setup.

        Args:
            p: Prime modulus
            generator: Group generator (optional)
        """
        self.p = sage.Integer(p)
        self.Fp = sage.Integers(self.p)

        if generator is None:
            # Find a generator
            self.generator = self.find_generator()
        else:
            self.generator = self.Fp(generator)

        self.order = self.p - 1  # Order of Fp*

    def find_generator(self):
        """Find a generator of Fp*."""
        for g in range(2, min(100, self.p)):  # Search small values first
            if self.is_generator(g):
                return self.Fp(g)
        raise ValueError("No generator found")

    def is_generator(self, g):
        """Check if g generates Fp*."""
        g = self.Fp(g)
        if g == 1:
            return False

        # Check that g^{(p-1)/q} != 1 for all prime factors q of p-1
        p_minus_1 = self.p - 1
        factors = p_minus_1.factor()

        for q, _ in factors:
            exponent = p_minus_1 // q
            if g**exponent == 1:
                return False

        return True

    def trial_division_order(self, g, max_iterations=10000):
        """
        Compute order using trial division.

        Keep multiplying until we get back to 1.

        Args:
            g: Group element
            max_iterations: Maximum iterations before giving up

        Returns:
            Order of g, or None if not found within limit
        """
        g = self.Fp(g)
        current = self.Fp(1)
        order = 0

        for i in range(1, max_iterations + 1):
            current = current * g
            if current == 1:
                return i

        return None  # Order too large

    def pollard_rho_order(self, g, max_iterations=1000):
        """
        Compute group element order using Pollard's rho algorithm.

        Simplified implementation that may not work for all cases.
        For now, fall back to factorization method for reliability.

        Args:
            g: Group element
            max_iterations: Maximum iterations

        Returns:
            Order of g, or None if not found
        """
        # For now, use factorization method as it's more reliable
        return self.order_from_factorization(g)

    def order_from_factorization(self, g):
        """
        Compute order by factoring p-1 and testing exponents.

        Args:
            g: Group element

        Returns:
            Order of g
        """
        g = self.Fp(g)
        p_minus_1 = self.p - 1
        factors = p_minus_1.factor()

        order = p_minus_1

        # For each prime factor q^e || p-1, find the largest e such that g^{order/q^e} != 1
        for q, e in factors:
            q_to_e = q**e
            exponent = order // q_to_e

            while exponent > 0 and g**exponent == 1:
                order = order // q
                exponent = order // q_to_e

        return order

    def compute_order(self, g, algorithm='factorization'):
        """
        Compute the order of element g.

        Args:
            g: Group element
            algorithm: Algorithm to use ('trial', 'pollard_rho', 'factorization')

        Returns:
            Order of g
        """
        if algorithm == 'trial':
            return self.trial_division_order(g)
        elif algorithm == 'pollard_rho':
            return self.pollard_rho_order(g)
        elif algorithm == 'factorization':
            return self.order_from_factorization(g)
        else:
            raise ValueError(f"Unknown algorithm: {algorithm}")

    def verify_order(self, g, order):
        """Verify that the given order is correct for g."""
        g = self.Fp(g)
        order = sage.Integer(order)

        return g**order == 1 and g**(order // 2) != 1  # Also check it's the minimal order

    def find_element_of_order(self, target_order):
        """
        Find an element of the specified order.

        Args:
            target_order: Desired order

        Returns:
            Element with the specified order, or None
        """
        # The order must divide p-1
        if (self.p - 1) % target_order != 0:
            return None

        # Start with generator and find appropriate power
        g = self.generator
        cofactor = (self.p - 1) // target_order

        # g^cofactor will have order target_order
        return g**cofactor

    def test_order_computation(self):
        """Test order computation algorithms."""
        print("Testing group order computation...")

        # Test with the generator (should have order p-1)
        g = self.generator
        expected_order = self.p - 1

        print(f"Testing order of generator g = {g}")
        print(f"Expected order: {expected_order}")

        # Test trial division
        try:
            order_trial = self.trial_division_order(g, max_iterations=1000)
            print(f"Trial division: {order_trial}")
            print(f"Correct: {order_trial == expected_order}")
        except Exception as e:
            print(f"Trial division failed: {e}")

        # Test factorization method
        try:
            order_factor = self.order_from_factorization(g)
            print(f"Factorization: {order_factor}")
            print(f"Correct: {order_factor == expected_order}")
        except Exception as e:
            print(f"Factorization failed: {e}")

        # Test with an element of smaller order
        if self.p > 100:  # For larger primes
            # Find an element of order dividing a factor of p-1
            p_minus_1_factors = (self.p - 1).factor()
            if len(p_minus_1_factors) > 1:
                small_order = p_minus_1_factors[0][0]  # Smallest prime factor
                h = self.find_element_of_order(small_order)

                if h is not None:
                    print(f"\nTesting element of order {small_order}: h = {h}")
                    order_h = self.order_from_factorization(h)
                    print(f"Computed order: {order_h}")
                    print(f"Correct: {order_h == small_order}")

    def benchmark_order_algorithms(self, test_elements=None):
        """Benchmark different order computation algorithms."""
        if test_elements is None:
            test_elements = [self.generator, self.generator**2, self.generator**5]

        print("\nBenchmarking order computation algorithms...")
        print("Element | Trial (ms) | Factor (ms) | Pollard (ms)")
        print("-" * 50)

        import time

        for g in test_elements:
            # Benchmark trial division
            start = time.time()
            order_trial = self.trial_division_order(g, max_iterations=10000)
            time_trial = (time.time() - start) * 1000

            # Benchmark factorization
            start = time.time()
            order_factor = self.order_from_factorization(g)
            time_factor = (time.time() - start) * 1000

            # Benchmark Pollard rho
            start = time.time()
            order_rho = self.pollard_rho_order(g, max_iterations=10000)
            time_rho = (time.time() - start) * 1000

            print(">8.1f"
                  ">8.1f"
                  ">8.1f")

            # Verify results
            if order_trial != order_factor or (order_rho is not None and order_rho != order_factor):
                print("  WARNING: Inconsistent results!")


if __name__ == "__main__":
    # Test with small prime
    p_small = 101
    go_small = GroupOrderAlgorithms(p_small)
    go_small.test_order_computation()
    go_small.benchmark_order_algorithms()

    # Test with larger prime (but still manageable)
    p_large = 2**31 - 1  # Mersenne prime
    print(f"\nTesting with larger prime: {p_large}")
    go_large = GroupOrderAlgorithms(p_large)
    go_large.test_order_computation()