import json
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.colors import LogNorm

# Constants from our Rust implementation
G = 6.67430e-11  # Gravitational constant
R_0 = 149597870700     # Earth radius (normalization length)
M_0 = 5.972e24   # Earth mass (normalization mass)
T_0 = 5060.0     # Characteristic time

def load_bodies(config_file):
    """Load bodies from config file."""
    with open(config_file, 'r') as f:
        data = json.load(f)
    return data['bodies']

def calculate_potential(x, y, bodies):
    """Calculate gravitational potential at point (x,y)."""
    potential = 0.0
    for body in bodies:
        dx = x - body['position'][0]
        dy = y - body['position'][1]
        r = np.sqrt(dx*dx + dy*dy)
        if r > 0:  # Avoid division by zero
            potential -= G * body['mass'] / r
    return potential

def calculate_lagrange_points(bodies):
    """Calculate approximate positions of Lagrange points."""
    if len(bodies) != 2:
        return []
    
    # Find primary and secondary bodies
    if bodies[0]['mass'] > bodies[1]['mass']:
        primary, secondary = bodies[0], bodies[1]
    else:
        primary, secondary = bodies[1], bodies[0]
    
    # Calculate distance between bodies
    dx = secondary['position'][0] - primary['position'][0]
    dy = secondary['position'][1] - primary['position'][1]
    r = np.sqrt(dx*dx + dy*dy)
    
    # Calculate mass ratio
    mu = secondary['mass'] / (primary['mass'] + secondary['mass'])
    
    # Calculate Lagrange points
    L1 = np.array([primary['position'][0] + (1 - mu**(1/3)) * dx,
                   primary['position'][1] + (1 - mu**(1/3)) * dy])
    
    L2 = np.array([primary['position'][0] + (1 + mu**(1/3)) * dx,
                   primary['position'][1] + (1 + mu**(1/3)) * dy])
    
    L3 = np.array([primary['position'][0] - (1 + 5*mu/12) * dx,
                   primary['position'][1] - (1 + 5*mu/12) * dy])
    
    L4 = np.array([primary['position'][0] + 0.5 * dx - 0.866 * dy,
                   primary['position'][1] + 0.5 * dy + 0.866 * dx])
    
    L5 = np.array([primary['position'][0] + 0.5 * dx + 0.866 * dy,
                   primary['position'][1] + 0.5 * dy - 0.866 * dx])
    
    return [L1, L2, L3, L4, L5]

def plot_potential(config_file, grid_size=200, extent=1.5):
    """Plot gravitational potential field."""
    # Load bodies
    bodies = load_bodies(config_file)
    
    # Find Earth and Sun
    earth = next((body for body in bodies if body['name'] == 'Earth'), None)
    sun = next((body for body in bodies if body['name'] == 'Sun'), None)
    
    if not (earth and sun):
        print("Error: Earth and Sun must be present in the config file")
        return
    
    # Calculate Lagrange points for Earth-Sun system
    lagrange_points = calculate_lagrange_points([sun, earth])
    
    # Create coordinate grid centered on Earth
    x = np.linspace(earth['position'][0] - extent*R_0, earth['position'][0] + extent*R_0, grid_size)
    y = np.linspace(earth['position'][1] - extent*R_0, earth['position'][1] + extent*R_0, grid_size)
    X, Y = np.meshgrid(x, y)
    
    # Calculate potential at each grid point
    Z = np.zeros_like(X)
    for i in range(grid_size):
        for j in range(grid_size):
            Z[i,j] = calculate_potential(X[i,j], Y[i,j], bodies)
    
    # Take absolute value for logarithmic scale
    Z_abs = np.abs(Z)
    
    # Calculate reference potential at Earth's position
    earth_potential = abs(calculate_potential(earth['position'][0], earth['position'][1], [sun]))
    
    # Set scale to focus on Earth-Sun system
    vmin = earth_potential / 10  # Show potential variations within an order of magnitude
    vmax = earth_potential * 10
    
    # Create plot
    plt.figure(figsize=(12, 10))
    ax = plt.gca()
    
    # Plot potential field with logarithmic scale
    im = ax.imshow(Z_abs, extent=[x[0], x[-1], y[0], y[-1]],
                   origin='lower', cmap='viridis', norm=LogNorm(vmin=vmin, vmax=vmax))
    
    # Add colorbar
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax, label='|Gravitational Potential| (J/kg)')
    
    # Plot bodies with different colors
    body_colors = {
        'Sun': 'yellow',
        'Earth': 'blue',
        'Moon': 'gray',
        'Mars': 'red',
        'Venus': 'orange',
        'Mercury': 'brown',
        'Jupiter': 'orange',
        'Saturn': 'gold',
        'Uranus': 'lightblue',
        'Neptune': 'blue'
    }
    
    for body in bodies:
        color = body_colors.get(body['name'], 'white')  # Default to white if name not found
        ax.plot(body['position'][0], body['position'][1], 'o', 
                color=color, markersize=15, label=body['name'])
    
    # Plot Lagrange points
    for i, point in enumerate(lagrange_points, 1):
        ax.plot(point[0], point[1], 'r+', markersize=12, label=f'L{i}')
    
    # Customize plot
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_title('Gravitational Potential Field with Lagrange Points (Log Scale)')
    ax.legend()
    
    # Save plot
    plt.savefig('potential_field.png', dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Plot gravitational potential field from config file')
    parser.add_argument('config_file', help='Path to config file')
    parser.add_argument('--grid-size', type=int, default=200, help='Grid size for potential calculation')
    parser.add_argument('--extent', type=float, default=1.5, help='Plot extent in Earth radii')
    args = parser.parse_args()
    
    plot_potential(args.config_file, args.grid_size, args.extent)