#Copyright Daniel Harding - RomanAILabs
#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # For 3D plotting

# Constants for GR calculations
G = 6.67430e-11  # Gravitational constant (m^3 kg^-1 s^-2)
c = 2.99792458e8  # Speed of light (m/s)
M_sun = 1.989e30  # Mass of the Sun (kg)
R_sun = 6.96e8    # Solar radius (m), used as impact parameter b

# Step 1: GR Deflection Angle Calculation
def gr_deflection_angle(M, b):
    """
    Calculate light deflection angle in GR (radians).
    Args:
    - M: Mass of deflecting object (kg)
    - b: Impact parameter (m)
    Returns: Deflection angle alpha (radians)
    """
    alpha = (4 * G * M) / (c**2 * b)
    return alpha

# Step 2: Simulate straight light path in 2D (for comparison)
def simulate_2d_light_path(start_point, direction, distance, steps=100):
    """
    Simulate a straight light path in 2D.
    Args:
    - start_point: Tuple (x, y)
    - direction: Tuple (dx, dy) - will be normalized
    - distance: Total distance to travel
    - steps: Number of points
    Returns: np.array of points
    """
    dir_norm = np.array(direction) / np.linalg.norm(direction)
    t = np.linspace(0, distance, steps)
    points = np.array(start_point) + t[:, np.newaxis] * dir_norm
    return points

# Step 3: Simulate bent light path in 3D using GR approximation
def simulate_3d_light_bending(start_point, initial_direction, center, M, b_init, distance, steps=1000):
    """
    Simulate light bending in 3D using perturbative GR approximation.
    Args:
    - start_point: Tuple (x, y, z)
    - initial_direction: Tuple (dx, dy, dz) - normalized
    - center: Tuple (cx, cy, cz) - mass center
    - M: Mass (kg)
    - b_init: Initial impact parameter (m)
    - distance: Total path length
    - steps: Integration steps
    Returns: np.array of points along path
    """
    pos = np.array(start_point, dtype=float)
    vel = np.array(initial_direction) / np.linalg.norm(initial_direction)  # Unit velocity
    dt = distance / steps  # Time step (assuming c=1 for simplicity)
    points = [pos.copy()]
    
    for _ in range(steps):
        r_vec = pos - center
        r = np.linalg.norm(r_vec)
        if r < 1e-6:  # Avoid singularity
            break
        
        # GR deflection perturbation (approximate acceleration from geodesic deviation)
        accel_mag = (3 * G * M / (c**2 * r**2)) * (np.dot(vel, r_vec / r)**2 / r)  # Simplified transverse component
        accel = -accel_mag * (r_vec / r)  # Toward center, but adjusted for light
        
        # Update velocity (perturb direction)
        vel += accel * dt
        vel /= np.linalg.norm(vel)  # Keep unit speed (light)
        
        # Update position
        pos += vel * dt
        points.append(pos.copy())
    
    return np.array(points)

# Step 4: 4D Euclidean Distance (for higher-dimensional analogy)
def calculate_4d_distance(point_a, point_b):
    """
    Calculate Euclidean distance in 4D space.
    Args:
    - point_a: Tuple (x1, y1, z1, w1)
    - point_b: Tuple (x2, y2, z2, w2)
    Returns: Distance (float)
    """
    diffs = np.array(point_a) - np.array(point_b)
    return np.sqrt(np.sum(diffs**2))

# Main script to compute, simulate, and visualize
if __name__ == "__main__":
    # Compute GR deflection angle for Sun
    alpha_sun = gr_deflection_angle(M_sun, R_sun)
    alpha_arcsec = np.degrees(alpha_sun) * 3600  # Convert to arcseconds
    print(f"GR Deflection Angle for Sun (radians): {alpha_sun:.6e}")
    print(f"GR Deflection Angle for Sun (arcseconds): {alpha_arcsec:.3f}")
    
    # 2D Straight Path Simulation
    path_2d = simulate_2d_light_path(start_point=(0, 0), direction=(1, 0.5), distance=20)
    plt.figure(figsize=(8, 6))
    plt.plot(path_2d[:, 0], path_2d[:, 1], label="Straight Light Path in 2D")
    plt.title("Simulated Straight Light Path in 2D")
    plt.xlabel("X (arbitrary units)")
    plt.ylabel("Y (arbitrary units)")
    plt.legend()
    plt.grid(True)
    plt.show()
    
    # 3D Bent Path Simulation (scaled down for visualization)
    path_3d = simulate_3d_light_bending(start_point=(0, 0, 0), initial_direction=(1, 0.1, 0),
                                        center=(10, 0, 0), M=M_sun / 1e6, b_init=5, distance=30)
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    ax.plot(path_3d[:, 0], path_3d[:, 1], path_3d[:, 2], label="Bent Light Path in 3D (GR Approx)")
    ax.scatter([10], [0], [0], color='red', s=100, label="Mass Center")
    ax.set_title("Simulated Light Bending in 3D (GR Approximation)")
    ax.set_xlabel("X (arbitrary units)")
    ax.set_ylabel("Y (arbitrary units)")
    ax.set_zlabel("Z (arbitrary units)")
    ax.legend()
    plt.show()
    
    # 4D Distance Example
    point_a = (1, 2, 3, 4)
    point_b = (4, 3, 2, 1)
    distance_4d = calculate_4d_distance(point_a, point_b)
    print(f"Euclidean Distance Between Points in 4D Space: {distance_4d:.3f}")