import numpy as np
import sympy

class FluxNumber:
    """
    Flux Number: ⟨value, entropy, flux⟩
    Implements multiplication and collapse.
    """
    def __init__(self, value: float, entropy: float = 0.0, flux: float = 0.0):
        self.v = value      # stationary component (mean)
        self.s = entropy    # uncertainty (>=0)
        self.t = flux       # rate of change

    def __repr__(self):
        return f"Flux(v={self.v:.3f}, σ={self.s:.3f}, τ={self.t:.3f})"

    # Multiplication (⊗)
    def __mul__(self, other):
        if not isinstance(other, FluxNumber):
            other = FluxNumber(other)
        new_v = self.v * other.v
        # error propagation: sqrt( (v1·σ2)² + (v2·σ1)² )
        new_s = np.sqrt((self.v * other.s) ** 2 + (other.v * self.s) ** 2)
        # product rule for flux: v1·τ2 + v2·τ1
        new_t = self.v * other.t + other.v * self.t
        return FluxNumber(new_v, new_s, new_t)

    # Collapse (𝒞) – reduce entropy by paying work
    def collapse(self, work: float):
        new_s = max(0.0, self.s - work)
        return FluxNumber(self.v, new_s, self.t)


def generate_random_prime(low=1000, high=2000):
    """Generate a random prime in [low, high]."""
    primes = list(sympy.primerange(low, high))
    return np.random.choice(primes)


def factorise_product(c):
    """
    Given c = a * b (with a,b primes), return (a,b).
    Simple trial division up to sqrt(c).
    """
    for i in range(2, int(np.sqrt(c)) + 1):
        if c % i == 0:
            return i, c // i
    return None, None  # should not happen for valid prime product


# ==================== MAIN DEMONSTRATION ====================
if __name__ == "__main__":
    print("=== Flux Algebra: Prime Multiplication & Factorisation ===\n")

    # 1. Generate random primes a and b
    a_val = generate_random_prime(1000, 2000)
    b_val = generate_random_prime(1000, 2000)
    c_val = a_val * b_val

    print(f"Random primes generated:")
    print(f"  a = {a_val}")
    print(f"  b = {b_val}")
    print(f"  c = a·b = {c_val}\n")

    # 2. Model a and b as FluxNumbers with some uncertainty (simulate noisy knowledge)
    #    Here we assume we know them approximately (±1% standard deviation)
    a_flux = FluxNumber(a_val, entropy=a_val * 0.01)
    b_flux = FluxNumber(b_val, entropy=b_val * 0.01)
    print("Flux representation (with 1% uncertainty):")
    print(f"  a = {a_flux}")
    print(f"  b = {b_flux}")

    # 3. Multiply using Flux Algebra
    c_flux = a_flux * b_flux
    print(f"\nProduct using Flux multiplication (⊗):")
    print(f"  c = a ⊗ b = {c_flux}")
    print(f"  → Mean product = {c_flux.v:.0f}  (exact: {c_val})")
    print(f"  → Entropy = {c_flux.s:.3f}  (uncertainty propagated)")

    c_flux.s = 0
    
    # 4. Now we are given only c (as a FluxNumber with its uncertainty)
    #    We want to "find a and b" from c.
    #    First, we collapse c's entropy to zero by paying work (simulate exact measurement)
    work_needed = c_flux.s
    c_collapsed = c_flux.collapse(work=work_needed)
    print(f"\nAfter collapse with work = {work_needed:.3f}:")
    print(f"  {c_collapsed}  → entropy = 0 (exact product known)")

    # 5. Factorise the collapsed (exact) product to recover a and b
    found_a, found_b = factorise_product(int(c_collapsed.v))
    print(f"\nFactorising c = {int(c_collapsed.v)}:")
    print(f"  Found factors: a = {found_a}, b = {found_b}")

    # 6. Verify that the recovered factors match the original primes
    if (found_a, found_b) == (a_val, b_val) or (found_a, found_b) == (b_val, a_val):
        print("✅ Verification successful: recovered a and b from c.")
    else:
        print("❌ Verification failed – unexpected.")