from snapshotprocessor import SnapshotProcessor, U
import pandas as pd
import os

def get_velocity_cubes(pp, unit=None):
    velcubes = [None, None, None]
    for i in range(3):
        velcubes[i] = pp.datacube(getter=lambda dset: dset["vel"][..., i])
        if unit is not None:
            velcubes[i] *= pp.info["unit_velocity"].express(unit)
    return velcubes

def get_density_cube(pp, unit=None):
    dens_cube = pp.datacube(getter=lambda dset: dset["rho"]) 
    if unit is not None:
        dens_cube *= pp.info["unit_density"].express(unit)
    return dens_cube


def write_data(filename, vel, dens):
    # write fields to ramses frig readable ascii file
    f = open(filename, 'w')
    dummy = 1
    size = vel[0].shape[0]
    f.write('{:8}{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n'.format(size, dummy, dummy, dummy, dummy))
    vx, vy, vz = vel
    # This strange order matches the one in the galbox condinit
    for z in range(size):
        for y in range(size):
            for x in range(size):
                f.write('{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n'.format(vx[x,y,z], vy[x,y,z], vz[x,y,z], dens[x, y, z]))

         
        
def extract_from_pp(pp):
    vel = get_velocity_cubes(pp, unit=U.km_s)
    dens = get_density_cube(pp, unit=U.H_cc)
    write_data(f"{pp.path_out}/{pp.run}_{pp.num}_velrho.data", vel, dens)

                
def extract(path, snap_number):
    pp = SnapshotProcessor(path, snap_number, params="../turbox_params.yml")
    extract_from_pp(pp)