import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits

# Load the Cosmicflows-4 grouped velocity field (64^3 grid in supergalactic coordinates)
filename = "cosmic_data/CF4gp_new_64-z008_velocity.fits"
with fits.open(filename) as hdul:
    data = hdul[0].data  # shape expected: (3, N_z, N_y, N_x)

# The data array contains velocity components. Apply scaling to get km/s (CF4 documentation says multiply by ~52):contentReference[oaicite:3]{index=3}.
velocity = data * 52.0  # now velocity components are in km/s
vx, vy, vz = velocity[0], velocity[1], velocity[2]  # define components for clarity

# Define coordinate arrays (in Mpc) for the grid 
N = vx.shape[0]            # number of cells along each axis (e.g., 64)
# Assume the grid spans approximately ±L Mpc from the origin. For example, L ~ 200 Mpc (so total ~400 Mpc across).
L = 200.0  # (Modify L if actual extent is known; here we assume ~200 Mpc half-range)
coords = np.linspace(-L, L, N)  # coordinate values for each index along an axis

# Choose the plane of projection:
plane = "SGX-SGY"  # options: "SGX-SGY" or "SGY-SGZ"

if plane == "SGX-SGY":
    # Take slice at SGZ ≈ 0 (mid-plane index)
    z_index = N // 2  
    # Extract velocity components in this plane (shape: N_y x N_x)
    v_plane_x = vx[z_index, :, :]  # velocity in SGX direction on SGZ=0 plane
    v_plane_y = vy[z_index, :, :]  # velocity in SGY direction on SGZ=0 plane
    X_vals = coords  # SGX coordinates for columns
    Y_vals = coords  # SGY coordinates for rows
    x_label, y_label = "SGX [Mpc]", "SGY [Mpc]"

elif plane == "SGY-SGZ":
    # Take slice at SGX ≈ 0 
    x_index = N // 2
    # Extract velocity components in this plane (here horizontal = SGY, vertical = SGZ)
    v_plane_x = vy[:, :, x_index]  # velocity in SGY direction on SGX=0 plane
    v_plane_y = vz[:, :, x_index]  # velocity in SGZ direction on SGX=0 plane
    X_vals = coords  # SGY coordinates for columns
    Y_vals = coords  # SGZ coordinates for rows
    x_label, y_label = "SGY [Mpc]", "SGZ [Mpc]"

# Create a meshgrid of coordinate values for plotting
X_grid, Y_grid = np.meshgrid(X_vals, Y_vals)

# Compute speed in the plane for color coding (magnitude of v_plane_x and v_plane_y)
speed_plane = np.sqrt(v_plane_x**2 + v_plane_y**2)

# Plot the quiver diagram
plt.style.use('seaborn-v0_8-dark')  # or use Astropy's MPl style for a clean scientific look
fig, ax = plt.subplots(figsize=(8, 8), dpi=300)  # large size and high DPI for publication quality
ax.set_title("Cosmicflows-4 Velocity Field ({} slice)".format(plane), fontsize=14)
ax.set_xlabel(x_label, fontsize=12)
ax.set_ylabel(y_label, fontsize=12)

# Plot quiver (using a colormap to show speed)
q = ax.quiver(X_grid, Y_grid, v_plane_x, v_plane_y, speed_plane, cmap='plasma',
              scale=None, scale_units='xy', pivot='middle', alpha=0.8)
plt.colorbar(q, ax=ax, label='Speed [km/s]')

# Optionally, mark known attractor positions in this plane for reference:
if plane == "SGX-SGY":
    # Great Attractor (Laniakea center) approx at SGX~-60, SGY~0 (SGZ~+40 so just off-plane):contentReference[oaicite:7]{index=7}
    ax.plot(-60, 0, marker='*', color='red', ms=10, label='Great Attractor')
    # Shapley Supercluster approx at SGX~-140, SGY~+60 (SGZ~–16, so somewhat in this plane):contentReference[oaicite:8]{index=8}
    ax.plot(-140, 60, marker='*', color='magenta', ms=10, label='Shapley Supercluster')
    ax.legend(loc='lower left')

# Tight layout and save the figure
plt.tight_layout()
plt.savefig("cf4_velocity_{}_quiver.png".format(plane.replace('-', '')), dpi=300)
plt.show()