"""
stellar_timelapse.py
------------------------
Two-pass pipeline for star time-lapse animation over 1 million years.

Pass 1  — Stream CSV in chunks → accumulate pos0.npy + vel.npy (memmap)
Pass 2  — Frame loop: advance positions, subsample, render → MP4

Hardware target: 8-core / 16GB RAM (Baruch lab)
Data:           GAIA DR3, 6-col CSV (ra, dec, parallax, pmra, pmdec, radial_velocity)
"""

import numpy as np
import pandas as pd
import pyvista as pv
import imageio.v2 as imageio
from multiprocessing import Pool, cpu_count
import os
import time

# ── CONFIG ───────────────────────────────────────────────────────────────────
CSV_FILE          = 'radial-vel-6-columns.csv'   # GAIA DR3 source CSV

# I/O
CHUNK_SIZE        = 100_000                      # rows per chunk
POS0_FILE         = 'pos0.npy'                   # memmap output
VEL_FILE          = 'vel.npy'                    # memmap output
COUNTS_FILE       = 'valid_counts.npy'           # saved star count

# Filtering
MIN_PARALLAX      = 0.01                         # mas — avoids near-zero division
MAX_DIST_LY       = 5000                         # filter bad parallax outliers

# Sampling
MAX_STARS         = 10    # reservoir sampling cap
                           # ↑ change to 1_000_000 or 33_000_000 to scale up

# Time
TOTAL_YEARS       = 20_000_000
YEARS_PER_FRAME   = 40_000
N_FRAMES          = TOTAL_YEARS // YEARS_PER_FRAME  # 2000 frames

# Rendering
RENDER_SAMPLE     = 10   # stars per frame (≤ MAX_STARS)
POINT_SIZE        = 16  # size of stars
FPS               = 120
OUTPUT_MP4        = 'stellar_timelapse_10stars.mp4' # filename reflects star count
WINDOW_SIZE       = [640, 360]
CMAP              = 'plasma'                     # colour by initial distance

# Multiprocessing — leave 1 core free for OS
N_WORKERS         = max(1, cpu_count() - 1)      # 7 on Baruch lab
# ── UNIT CONSTANTS ────────────────────────────────────────────────────────────
KM_PER_LY         = 9.461e12
SEC_PER_YR        = 3.156e7
LY_PER_KM_PER_YR  = SEC_PER_YR / KM_PER_LY        # ~3.336e-6
FACTOR            = 4.74047 * LY_PER_KM_PER_YR     # ly/yr per (mas/yr · kpc)

# ═══════════════════════════════════════════════════════════════════════════════
# PASS 1 — ACCUMULATE pos0 + vel from chunked CSV
# ═══════════════════════════════════════════════════════════════════════════════

def process_chunk(chunk_df):
    """
    Given a DataFrame chunk, return (pos0_block, vel_block) arrays.
    Called in worker processes via multiprocessing.Pool.
    Returns None if chunk yields no valid stars.
    """
    df = chunk_df.copy()

    # Filter
    df = df.dropna(subset=['ra', 'dec', 'parallax', 'pmra', 'pmdec'])
    df = df[df['parallax'] >= MIN_PARALLAX]
    df['distance_ly'] = 3262.0 / df['parallax']
    df = df[df['distance_ly'] <= MAX_DIST_LY]
    df = df[np.isfinite(df['distance_ly'])]
    if len(df) == 0:
        return None

    df['radial_velocity'] = df['radial_velocity'].fillna(0.0)

    ra_r  = np.radians(df['ra'].values)
    dec_r = np.radians(df['dec'].values)
    d     = df['distance_ly'].values
    d_kpc = d / 3262.0

    # Cartesian position
    cos_dec = np.cos(dec_r)
    pos0 = np.column_stack([
        d * cos_dec * np.cos(ra_r),
        d * cos_dec * np.sin(ra_r),
        d * np.sin(dec_r),
    ]).astype(np.float32)

    # Velocity vectors
    e_ra = np.column_stack([
        -np.sin(ra_r),
         np.cos(ra_r),
         np.zeros(len(df))
    ])
    e_dec = np.column_stack([
        -np.sin(dec_r) * np.cos(ra_r),
        -np.sin(dec_r) * np.sin(ra_r),
         np.cos(dec_r)
    ])
    e_r = pos0 / d[:, None]

    pm_ra_ly  = FACTOR * df['pmra'].values  * d_kpc
    pm_dec_ly = FACTOR * df['pmdec'].values * d_kpc
    rv_ly     = df['radial_velocity'].values * LY_PER_KM_PER_YR

    # ── LSR correction (Schönrich et al. 2010) ────────────────────────────────
    # Sun's peculiar motion in km/s: U toward GC, V along rotation, W toward NGP
    U_lsr = 11.1   # km/s
    V_lsr = 12.24  # km/s
    W_lsr = 7.25   # km/s

    # Convert to ly/yr
    U_ly = U_lsr * LY_PER_KM_PER_YR
    V_ly = V_lsr * LY_PER_KM_PER_YR
    W_ly = W_lsr * LY_PER_KM_PER_YR

    # Galactic unit vectors in equatorial (ICRS) frame
    # These are standard constants for J2000
    e_U = np.array([-0.0548755604,  0.4941094279, -0.8676661490])
    e_V = np.array([-0.8734370902, -0.4448296300, -0.1980763734])
    e_W = np.array([-0.4838350155,  0.7469822445,  0.4559837762])

    # LSR velocity vector in Cartesian equatorial coords (ly/yr)
    v_lsr = U_ly * e_U + V_ly * e_V + W_ly * e_W   # shape (3,)

    # Subtract from every star's velocity
    vel = (pm_ra_ly[:, None] * e_ra +
           pm_dec_ly[:, None] * e_dec +
           rv_ly[:, None] * e_r
           - v_lsr[None, :]).astype(np.float32)

    return pos0, vel


def run_pass1():
    """Stream CSV via reservoir sampling, write memmap .npy files."""

    print("=" * 60)
    print("PASS 1 — Reservoir sampling positions & velocities")
    print("=" * 60)

    # Pre-scan: count valid stars (used for progress reporting)
    print("Pre-scan: counting valid stars...")
    t0 = time.time()
    total_valid = 0
    for chunk in pd.read_csv(CSV_FILE, chunksize=CHUNK_SIZE,
                             usecols=['ra', 'dec', 'parallax', 'pmra', 'pmdec', 'radial_velocity']):
        sub = chunk.dropna(subset=['ra', 'dec', 'parallax', 'pmra', 'pmdec'])
        sub = sub[sub['parallax'] >= MIN_PARALLAX]
        sub['distance_ly'] = 3262.0 / sub['parallax']
        sub = sub[(sub['distance_ly'] <= MAX_DIST_LY) & np.isfinite(sub['distance_ly'])]
        total_valid += len(sub)
    print(f"Valid stars: {total_valid:,}  (scanned in {time.time()-t0:.1f}s)")

    # This will choose random stars from the 33 million
    reservoir_pos = np.zeros((MAX_STARS, 3), dtype=np.float32)
    reservoir_vel = np.zeros((MAX_STARS, 3), dtype=np.float32)
    n_seen = 0
    rng = np.random.default_rng(42)

    for chunk in pd.read_csv(CSV_FILE, chunksize=CHUNK_SIZE,
                             usecols=['ra', 'dec', 'parallax', 'pmra', 'pmdec', 'radial_velocity']):
        res = process_chunk(chunk)
        if res is None:
            continue
        p, v = res
        for i in range(len(p)):
            n_seen += 1
            if n_seen <= MAX_STARS:
                reservoir_pos[n_seen-1] = p[i]
                reservoir_vel[n_seen-1] = v[i]
            else:
                j = rng.integers(0, n_seen)
                if j < MAX_STARS:
                    reservoir_pos[j] = p[i]
                    reservoir_vel[j] = v[i]

        if n_seen % 500_000 == 0:
            print(f"  Seen {n_seen:,} / {total_valid:,} stars", flush=True)

    # Write reservoir → memmap
    n_final = min(n_seen, MAX_STARS)
    pos0_mm = np.lib.format.open_memmap(POS0_FILE, mode='w+',
                                         dtype=np.float32, shape=(n_final, 3))
    vel_mm  = np.lib.format.open_memmap(VEL_FILE,  mode='w+',
                                         dtype=np.float32, shape=(n_final, 3))
    pos0_mm[:] = reservoir_pos[:n_final]
    vel_mm[:]  = reservoir_vel[:n_final]
    del pos0_mm, vel_mm
    np.save(COUNTS_FILE, np.array([n_final]))
    print(f"Reservoir saved: {n_final:,} stars")
    return n_final

 


# ═══════════════════════════════════════════════════════════════════════════════
# PASS 2 — RENDER FRAMES
# ═══════════════════════════════════════════════════════════════════════════════

def run_pass2(n_stars):
    """Load memmap arrays, render N_FRAMES into MP4."""

    print("\n" + "=" * 60)
    print("PASS 2 — Rendering frames")
    print("=" * 60)

    # Load memmap (read-only — OS will cache in RAM after first access)
    pos0 = np.lib.format.open_memmap(POS0_FILE, mode='r',
                                      dtype=np.float32, shape=(n_stars, 3))
    vel  = np.lib.format.open_memmap(VEL_FILE,  mode='r',
                                      dtype=np.float32, shape=(n_stars, 3))

    print(f"Loaded {n_stars:,} stars from memmap")
    print(f"Rendering {N_FRAMES} frames at {RENDER_SAMPLE:,} stars/frame...")

    # Fixed random subsample indices (same stars across all frames for continuity)
    rng = np.random.default_rng(42)
    if n_stars > RENDER_SAMPLE:
        sample_idx = rng.choice(n_stars, size=RENDER_SAMPLE, replace=False)
        sample_idx.sort()   # sorted for faster memmap access
    else:
        sample_idx = np.arange(n_stars)

    pos0_s = np.array(pos0[sample_idx])   # pull into RAM: (500k, 3)
    vel_s  = np.array(vel [sample_idx])

    # Colour by initial distance from Earth
    dist0  = np.linalg.norm(pos0_s, axis=1)
    norm_d = (dist0 - dist0.min()) / (dist0.max() - dist0.min() + 1e-9)

    # ── PyVista setup (screenshot mode — bypasses PyAV quality bug) ──────────
    
    FRAMES_DIR = 'frames_tmp'
    os.makedirs(FRAMES_DIR, exist_ok=True)

    pl = pv.Plotter(off_screen=True, window_size=WINDOW_SIZE)
    pl.set_background('black')

    # Earth marker
    pl.add_mesh(pv.Sphere(radius=50, center=(0, 0, 0)), color='cyan', opacity=1.0)

    # Initial point cloud
    cloud = pv.PolyData(pos0_s.copy())
    cloud['dist'] = norm_d.astype(np.float32)
    pl.add_points(cloud, scalars='dist', cmap=CMAP,
                  point_size=POINT_SIZE, opacity=0.8,
                  render_points_as_spheres=True,
                  scalar_bar_args={'title': 'Closer to Earth -------> Further from Earth',
                 'color': 'white'})

    # Camera orbit params
    cam_r = np.linalg.norm(pos0_s, axis=1).max() * 2.2
    pl.camera.focal_point = (0, 0, 0)

    # Title — static, added once
    pl.add_text(
        f'{RENDER_SAMPLE:,} stars moving through space for {TOTAL_YEARS:,} years',
        position='upper_edge',
        font_size=10,
        color='white'
    )

    # Time counter — updated every frame
    time_actor = pl.add_text('t = 0 yr', position='lower_left',
                              font_size=10, color='white')

    t0 = time.time()
    frame_paths = []
    
    # Set camera ONCE before the frame loop
    pl.camera.position = (cam_r, 0, cam_r * 0.4)
    pl.camera.focal_point = (0, 0, 0)

    for frame_i in range(N_FRAMES):
        t_yr = frame_i * YEARS_PER_FRAME

        # Advance positions
        new_pos = pos0_s + vel_s * t_yr
        cloud.points = new_pos.astype(np.float32)

        # Update time label
        pl.remove_actor(time_actor)
        time_actor = pl.add_text(
            f't = {t_yr:>10,.0f} yr   ({t_yr/1000:.0f} kyr)',
            position='lower_left', font_size=10, color='white'
        )


        # Screenshot → PNG (bypasses imageio/PyAV entirely)
        png_path = os.path.join(FRAMES_DIR, f'frame_{frame_i:04d}.png')
        pl.screenshot(png_path, window_size=WINDOW_SIZE)
        frame_paths.append(png_path)

        if frame_i % 20 == 0:
            elapsed = time.time() - t0
            eta = (elapsed / (frame_i + 1)) * (N_FRAMES - frame_i - 1)
            print(f"  Frame {frame_i+1:>3}/{N_FRAMES}  "
                  f"{t_yr/1_000_000*100:.1f}%  "
                  f"ETA {eta/60:.1f} min", flush=True)

    pl.close()
    print(f"\nAll frames rendered. Stitching to {OUTPUT_MP4}...")

    # Stitch PNGs → MP4 using imageio directly (no quality kwarg)
    with imageio.get_writer(OUTPUT_MP4, fps=FPS, codec='libx264',
                            output_params=['-crf', '22', '-preset', 'fast',
                                           '-pix_fmt', 'yuv420p']) as writer:
        for i, p in enumerate(frame_paths):
            writer.append_data(imageio.imread(p))
            if i % 20 == 0:
                print(f"  Stitching frame {i+1}/{N_FRAMES}", flush=True)

    # Clean up PNGs
    import shutil
    shutil.rmtree(FRAMES_DIR)
    print(f"\nPass 2 complete → {OUTPUT_MP4}")


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════════════════════

if __name__ == '__main__':

    # Skip Pass 1 if memmap files already exist
    if os.path.exists(POS0_FILE) and os.path.exists(VEL_FILE) and os.path.exists(COUNTS_FILE):
        n_stars = int(np.load(COUNTS_FILE)[0])
        print(f"Found existing memmap files ({n_stars:,} stars). Skipping Pass 1.")
        print("Delete pos0.npy / vel.npy to re-run Pass 1.")
    else:
        n_stars = run_pass1()

    run_pass2(n_stars)

    print("\nAll done.")
    print(f"  Stars in simulation : {n_stars:,}")
    print(f"  Stars rendered/frame: {min(n_stars, RENDER_SAMPLE):,}")
    print(f"  Timespan            : {TOTAL_YEARS:,} years")
    print(f"  Output              : {OUTPUT_MP4}")