#!/usr/bin/env python3
import numpy as np
import swiftsimio as sw
import unyt
from tqdm import tqdm
import glob

Counter = len(glob.glob1(".", "output_*.hdf5"))
print(Counter)

times = np.zeros(Counter - 2)
accretions = np.zeros(Counter - 2)
mass = np.zeros(Counter - 2)
number_of_gas_neighbours = np.zeros(Counter - 2)
velocity_dispersion = np.zeros(Counter - 2)
N_timesteps = np.zeros(Counter - 2)
h_sm = np.zeros(Counter - 2)


for i in tqdm(range(1, Counter - 1)):
    data = sw.load("output_{:04d}.hdf5".format(i))
    times[i - 1] = data.metadata.time.to("Myr")
    accretions[i - 1] = data.black_holes.accretion_rates.to("Msun/yr")
    mass[i - 1] = data.black_holes.subgrid_masses.to("Msun")
    number_of_gas_neighbours[i - 1] = data.black_holes.number_of_gas_neighbours.value
    velocity_dispersion[i - 1] = data.black_holes.gas_velocity_dispersions.to("km/s")
    N_timesteps[i - 1] = data.black_holes.number_of_time_steps.value
    h_sm[i - 1] = data.black_holes.smoothing_lengths.to("kpc")

np.savetxt(
    "accretion_rates.txt",
    np.transpose(
        [times, accretions, mass, number_of_gas_neighbours, velocity_dispersion, h_sm]
    ),
)

# print(times,accretions)
