import numpy as np
import matplotlib.pyplot as plt

# Use Agg backend for non-interactive environments
plt.switch_backend('Agg')

class HourglassZetaProbe:
    """
    Geometric Zero Finder using Hourglass Angular Projection.
    
    Maps Riemann Zeta terms to an hourglass geometry where:
    - Division by zero is resolved via angle -> 1/2 mapping.
    - Zeros are locations of perfect angular destructive interference at the waist.
    """
    
    def __init__(self, h_waist=1e-6):
        """
        h_waist: The stabilizing height parameter. 
        In the limit h -> 0, the singularity 1/0 maps to 1/2.
        """
        self.h = h_waist
        self.critical_plane = 0.5
        # ... inside the class ...

    def calculate_system_alignment(self, resultant_vector):
        """
        Calculates alignment based on the resultant vector collapsing to zero.
        This is what "aligns" with the singularity 1/0 -> 1/2.
        """
        magnitude = np.abs(resultant_vector)
        # If magnitude is 0 (Perfect Zero), alignment is 1.0.
        # If magnitude is large (No Zero), alignment is 0.0.
        # Using a Gaussian-like decay for smoothness
        alignment = np.exp(-magnitude**2) 
        return alignment
        
    def singularity_resolve(self, angle_rad):
        """
        Implements the '1/0 -> 1/2' geometric collapse.
        Maps the angular divisor to the critical symmetry plane.
        
        Args:
            angle_rad: Angle theta in radians (can be a numpy array).
            
        Returns:
            xi: Normalized coordinate [0, 1], where 0.5 is the singularity center.
        """
        # Regularize angle to prevent division by zero in physical math
        safe_angle = np.maximum(np.abs(angle_rad), self.h)
        
        # The hourglass projection: as angle -> 0, projection -> infinity.
        # We map this divergence to the normalized center 1/2.
        # xi = 1/2 + (1/pi) * arctan(h / theta)
        # theta=0 -> xi=1/2. theta=infinity -> xi=1/2.
        xi = self.critical_plane + (1.0 / np.pi) * np.arctan(self.h / safe_angle)
        
        return xi

    def compute_hourglass_projection(self, t, n_max=2000):
        """
        Projects the Riemann Zeta terms onto the hourglass geometry.
        Supports vectorized input for t.
        
        Args:
            t: Scalar or 1D array of imaginary parts of s.
            n_max: Number of terms to include in the geometric sum.
            
        Returns:
            Dictionary with geometric metrics (arrays if t is an array).
        """
        t = np.atleast_1d(t)
        # 1. Define terms n = 1, 2, ..., N
        n = np.arange(1, n_max + 1)
        
        # 2. Hourglass Geometry Mapping
        # Magnitude: n^(-1/2)
        ray_lengths = n ** -0.5 # Shape (N,)
        
        # Phase: t * ln(n)
        # Broadcast t (T,) and log(n) (N,) to (T, N)
        log_n = np.log(n)
        angles_rad = t[:, np.newaxis] * log_n[np.newaxis, :] # Shape (T, N)
        
        # 3. Compute Projective States (Vectorized)
        xi_matrix = self.singularity_resolve(angles_rad) # Shape (T, N)
        
        # 4. Vector Summation
        # Use complex exponentials: e^(i * theta)
        vectors = ray_lengths[np.newaxis, :] * np.exp(1j * angles_rad) # Shape (T, N)
        resultant_vector = np.sum(vectors, axis=1) # Shape (T,)
        
        # 5. Metrics
        resultant_magnitude = np.abs(resultant_vector)
        mean_xi = np.mean(xi_matrix, axis=1)
        resultant_phase = np.angle(resultant_vector)
        
        return {
            "t": t,
            "resultant_magnitude": resultant_magnitude,
            "waist_alignment": mean_xi,
            "resultant_phase": resultant_phase,
            "singularity_vector": resultant_vector
        }

    def find_zeros_geometric(self, t_min, t_max, t_step=0.01):
        """
        Scans the critical line using angular convergence in batches.
        """
        t_values = np.arange(t_min, t_max, t_step)

        print(f"🕳️ Starting Vectorized Hourglass Zero Scan (h={self.h})...")
        print(f"   Range: t=[{t_min}, {t_max}], Step: {t_step}")

        # Compute all results in one vectorized call
        results = self.compute_hourglass_projection(t_values, n_max=2000)

        magnitudes = results['resultant_magnitude']

        # Find local minima indices
        # A zero of Zeta(1/2+it) corresponds to a minimum in the magnitude of the Dirichlet sum
        threshold = 2.0 # Dirichlet sum of 2000 terms is not exactly 0, it fluctuates.

        # Find local minima
        is_min = (magnitudes[1:-1] < magnitudes[:-2]) & (magnitudes[1:-1] < magnitudes[2:])
        is_min = np.concatenate(([False], is_min, [False]))

        zero_indices = np.where(is_min & (magnitudes < threshold))[0]
        zero_candidates = t_values[zero_indices]

        print(f"✅ Found {len(zero_candidates)} geometric zero candidates.")
        return zero_candidates


# --- Execution & Visualization ---

if __name__ == "__main__":
    # Initialize Probe with a small waist height to resolve singularity
    probe = HourglassZetaProbe(h_waist=1e-9)

    # Scan for first few zeros (approximate locations known: 14.13, 21.02, 25.01)
    t_min, t_max = 10.0, 35.0
    zeros = probe.find_zeros_geometric(t_min, t_max, t_step=0.005)

    print("\n--- Geometric Zero Results ---")
    # Compare with known values
    known_zeros = [14.1347, 21.0220, 25.0109, 30.4249, 32.9351]

    for t_z in zeros:
        # Find closest known zero
        closest = min(known_zeros, key=lambda k: abs(k - t_z))
        print(f"Geometric Zero: t={t_z:.4f} | Closest Expected: {closest:.4f} | Δ={abs(t_z-closest):.4f}")

    # --- Visualization of the "Hourglass Interference" at a Zero ---
    t_sample = zeros[0] if zeros.size > 0 else 14.1347

    geo_at_zero = probe.compute_hourglass_projection(t_sample, n_max=50)

    print(f"\n📊 Analysis at t={t_sample:.4f}:")
    # Access scalar values from the results (they will be arrays of size 1)
    print(f"   Resultant Magnitude: {geo_at_zero['resultant_magnitude'][0]:.6f}")
    print(f"   Waist Alignment: {geo_at_zero['waist_alignment'][0]:.4f}")
    print(f"   Singularity Vector: {geo_at_zero['singularity_vector'][0]}")

    # Plotting vectors at the zero
    plt.figure(figsize=(12, 6))

    n_plot = 50
    n_vals = np.arange(1, n_plot + 1)
    t_val = t_sample
    angles = t_val * np.log(n_vals)
    lengths = n_vals ** -0.5

    plt.subplot(1, 2, 1)
    plt.plot(n_vals, lengths, 'b-', label='Ray Lengths (n^-0.5)')
    plt.plot(n_vals, -lengths, 'b-', label='-n^-0.5')
    plt.axhline(0, color='k', linewidth=0.5)
    plt.title(f'Hourglass Ray Structure\nat Zero (t={t_val:.2f})')
    plt.xlabel('Term n')
    plt.ylabel('Magnitude')
    plt.legend()

    plt.subplot(1, 2, 2)
    # Plot the angular positions mapped to a circle/sheet
    plt.scatter(angles % (2*np.pi), np.arange(1, n_plot + 1), c=lengths, cmap='viridis', alpha=0.7)
    plt.axvline(np.pi, color='r', linestyle='--', label='Singularity Phase')
    plt.title('Angular Distribution on Sheet\n(Destructive Interference)')
    plt.xlabel('Phase Angle (mod 2π)')
    plt.ylabel('Term n')
    plt.legend()

    plt.tight_layout()
    plt.savefig('hourglass_analysis.png')
    print("\n📈 Visualization saved to 'hourglass_analysis.png'")