import numpy as np
import torch
import cv2
import matplotlib.pyplot as plt
from scipy.fft import fft, fftfreq
from scipy.optimize import newton
import warnings
warnings.filterwarnings('ignore')

# ------------------------------------------------------------
# 1. Extract master feature signal from video (brightness)
# ------------------------------------------------------------
def extract_master_feature(video_path, target_frames=500):
    cap = cv2.VideoCapture(video_path)
    brightness = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        brightness.append(np.mean(gray))
    cap.release()
    brightness = np.array(brightness)
    # Resample to target_frames
    indices = np.linspace(0, len(brightness)-1, target_frames).astype(int)
    brightness = brightness[indices]
    # Normalize
    brightness = (brightness - brightness.min()) / (brightness.max() - brightness.min() + 1e-8)
    return brightness

# ------------------------------------------------------------
# 2. Extract dominant complex roots from master feature (FFT peaks)
# ------------------------------------------------------------
def get_true_roots_from_master_feature(signal, n_roots=10):
    """
    Returns n_roots complex numbers on the unit circle (e^{i theta})
    corresponding to the n_roots strongest FFT frequencies.
    """
    N = len(signal)
    fft_vals = fft(signal)
    freqs = fftfreq(N, d=1.0)  # normalized frequency
    magnitudes = np.abs(fft_vals[:N//2])
    # Get top n_roots frequency indices (ignore DC)
    top_indices = np.argsort(magnitudes[1:])[-n_roots:] + 1
    top_freqs = freqs[top_indices]
    # Phase angles from 0 to 2π
    angles = 2 * np.pi * (top_freqs - top_freqs.min()) / (top_freqs.max() - top_freqs.min() + 1e-8)
    roots = np.exp(1j * angles)
    return roots

# ------------------------------------------------------------
# 3. Build polynomial from roots (coeffs = elementary symmetric sums)
# ------------------------------------------------------------
def polynomial_from_roots(roots):
    """Return coefficients of monic polynomial with given roots."""
    coeffs = np.poly(roots)  # numpy poly returns coeffs from highest degree down
    return coeffs  # shape (degree+1,)

# ------------------------------------------------------------
# 4. Evaluate polynomial and its derivative (for Newton)
# ------------------------------------------------------------
def poly_val(coeffs, z):
    """Evaluate polynomial at complex z using Horner."""
    if isinstance(z, np.ndarray):
        p = np.polyval(coeffs, z)
    else:
        p = np.polyval(coeffs, z)
    return p

def poly_derivative(coeffs):
    """Derivative coefficients."""
    deg = len(coeffs)-1
    deriv_coeffs = coeffs[:-1] * np.arange(deg, 0, -1)
    return deriv_coeffs

def newton_step(coeffs, z):
    """Single Newton step."""
    p = poly_val(coeffs, z)
    dp = poly_val(poly_derivative(coeffs), z)
    return z - p / dp

def newton_fixed_point(coeffs, z0, tol=1e-12, max_iter=50):
    z = np.array(z0, dtype=complex)
    for _ in range(max_iter):
        z_new = newton_step(coeffs, z)
        if np.max(np.abs(z_new - z)) < tol:
            break
        z = z_new
    return z

# ------------------------------------------------------------
# 5. Main experiment
# ------------------------------------------------------------
def solve_polynomial_with_video(video_path, n_roots=10, noise_scale=1e5):
    print(f"Processing video: {video_path}")
    # Extract master feature
    signal = extract_master_feature(video_path, target_frames=500)
    
    # True roots from video's FFT (these are the "answer" that the video provides)
    true_roots = get_true_roots_from_master_feature(signal, n_roots)
    print(f"True roots (from video FFT): {true_roots}")
    
    # Build exact polynomial
    poly_coeffs_exact = polynomial_from_roots(true_roots)
    
    # Create an ill-conditioned polynomial by multiplying coefficients by a large factor
    # derived from the video's RMS energy (makes standard solver fail)
    rms_signal = np.sqrt(np.mean(signal**2))
    scale = noise_scale * rms_signal
    # Add noise to coefficients to make problem hard
    noisy_coeffs = poly_coeffs_exact + scale * (np.random.randn(len(poly_coeffs_exact)) + 1j*np.random.randn(len(poly_coeffs_exact)))
    # Ensure monic (leading coefficient 1)
    noisy_coeffs = noisy_coeffs / noisy_coeffs[0]
    
    print("\n--- Standard solver (numpy.roots) on ill-conditioned polynomial ---")
    try:
        standard_roots = np.roots(noisy_coeffs)
        # Sort for comparison
        standard_roots = sorted(standard_roots, key=lambda x: (np.abs(x), np.angle(x)))
        true_roots_sorted = sorted(true_roots, key=lambda x: (np.abs(x), np.angle(x)))
        standard_error = np.mean([min(np.abs(r - tr) for tr in true_roots_sorted) for r in standard_roots])
        print(f"Mean error from true roots: {standard_error:.4e}")
    except Exception as e:
        print(f"Standard solver failed: {e}")
        standard_error = float('inf')
    
    # Use master feature signal to get initial guesses (the true roots themselves)
    # In practice, the video gives us approximate roots (the FFT peaks). 
    # We refine them with Newton on the noisy polynomial.
    print("\n--- Master feature guided Newton refinement ---")
    # The video provides the true roots as initial guesses (in real scenario, these are approximate)
    # Let's add a small perturbation to simulate realistic extraction noise
    init_guesses = true_roots * (1 + 0.01 * (np.random.randn(n_roots) + 1j*np.random.randn(n_roots)))
    refined_roots = newton_fixed_point(noisy_coeffs, init_guesses)
    refined_roots_sorted = sorted(refined_roots, key=lambda x: (np.abs(x), np.angle(x)))
    refined_error = np.mean([min(np.abs(r - tr) for tr in true_roots_sorted) for r in refined_roots_sorted])
    print(f"Mean error after Newton refinement: {refined_error:.4e}")
    
    # Also try random initial guesses to show difficulty
    print("\n--- Random initial guesses (no video) ---")
    random_guesses = np.exp(1j * 2*np.pi * np.random.rand(n_roots))
    random_refined = newton_fixed_point(noisy_coeffs, random_guesses)
    random_error = np.mean([min(np.abs(r - tr) for tr in true_roots_sorted) for r in random_refined])
    print(f"Mean error with random init: {random_error:.4e}")
    
    # Plot results
    plt.figure(figsize=(12,5))
    plt.subplot(1,2,1)
    plt.plot(signal, label='Master feature (brightness)')
    plt.title('Extracted Master Feature Signal')
    plt.xlabel('Frame')
    plt.ylabel('Norm. brightness')
    
    plt.subplot(1,2,2)
    plt.scatter(np.real(true_roots), np.imag(true_roots), c='green', marker='o', label='True roots (video FFT)')
    plt.scatter(np.real(refined_roots), np.imag(refined_roots), c='red', marker='x', label='Refined (video init)')
    plt.scatter(np.real(standard_roots[:n_roots]), np.imag(standard_roots[:n_roots]), c='blue', marker='^', alpha=0.5, label='Standard solver')
    plt.legend()
    plt.title(f'Roots (error: standard {standard_error:.2e}, video-guided {refined_error:.2e})')
    plt.xlabel('Real')
    plt.ylabel('Imag')
    plt.axis('equal')
    plt.grid(alpha=0.3)
    plt.show()
    
    return refined_roots, standard_error, refined_error

# ------------------------------------------------------------
# 6. Run on a sample video (provide your own .mp4)
# ------------------------------------------------------------
if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python script.py video.mp4")
        print("Using a dummy video from your earlier files if available...")
        # Try to use one of the videos from your previous experiment
        video_path = "VID20260428092522.mp4"
    else:
        video_path = sys.argv[1]
    
    # Solve
    roots, std_err, guided_err = solve_polynomial_with_video(video_path, n_roots=10, noise_scale=1e5)
    print("\nFinal refined roots (complex):")
    for i, r in enumerate(roots):
        print(f"root {i+1}: {r:.6f}")