from astropy.io import fits
import os

# Input FITS cube filename
input_cube = "cube.fits"

# Output directory
output_dir = "frames"
os.makedirs(output_dir, exist_ok=True)

# Open the FITS cube
with fits.open(input_cube) as hdul:
    data = hdul[0].data  # assuming cube is in primary HDU
    header = hdul[0].header

# Check dimensions
print("Cube shape:", data.shape)

# Loop through frames
n_frames = data.shape[0]  # first axis = frames
for i in range(n_frames):
    frame_data = data[i, :, :]
    frame_header = header.copy()

    # Update header to be 2D
    frame_header['NAXIS'] = 2
    frame_header['NAXIS1'] = frame_data.shape[1]
    frame_header['NAXIS2'] = frame_data.shape[0]
    for k in list(frame_header.keys()):
        if k.startswith("NAXIS") and k not in ("NAXIS", "NAXIS1", "NAXIS2"):
            del frame_header[k]

    # Add frame index
    frame_header['FRAME'] = i

    # Output filename
    output_file = os.path.join(output_dir, f"frame_{i:04d}.fits")
    fits.writeto(output_file, frame_data, frame_header, overwrite=True)
    print(f"Saved {output_file}")

print(f"✅ Done! Extracted {n_frames} frames into '{output_dir}/'")