#!/usr/bin/python3
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy.visualization import astropy_mpl_style
from yt import load_uniform_grid
from yt.visualization.api import Streamlines

plt.style.use(astropy_mpl_style)

# Load Cosmicflows-4 velocity field data
data_path = 'cosmic_data/CF4_new_64-z008_velocity.fits'
hdul = fits.open(data_path)
data = hdul[0].data
hdul.close()

# Create a yt dataset from the FITS data
bbox = np.array([[0, data.shape[0]], [0, data.shape[1]], [0, data.shape[2]]])
ds = load_uniform_grid({"velocity_x": data[0], "velocity_y": data[1], "velocity_z": data[2]}, data.shape[1:], length_unit="Mpc", bbox=bbox)

# Generate streamlines visualization
streamlines = Streamlines(ds, 'velocity_x', 'velocity_y', 'velocity_z', length=100.0)
streamlines.integrate_through_volume()

# Plot streamlines
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlabel('X [Mpc]')
ax.set_ylabel('Y [Mpc]')
ax.set_title('Cosmicflows-4 Galaxy Flow Streamlines')

# Plotting streamlines projected onto the XY plane
for streamline in streamlines.streamlines:
    ax.plot(streamline[:,0], streamline[:,1], lw=0.5, alpha=0.7)

plt.savefig('cosmic_flows_streamlines.png')
plt.show()

print("Visualization complete. Check 'cosmic_flows_streamlines.png'.")