import numpy as np
import h5py
import argparse
import scipy.stats as stats
import unyt

argparser = argparse.ArgumentParser()
argparser.add_argument("input", nargs="+")
argparser.add_argument("output")
args = argparser.parse_args()

nfile = len(args.input)

time = np.zeros(nfile)
radius = np.zeros(nfile)
velocity = np.zeros(nfile)

for ifile, file in enumerate(args.input):
    with h5py.File(file, "r") as handle:
        coords = handle["PartType0/Coordinates"][:]
        rho = handle["PartType0/Densities"][:]
        vs = handle["PartType0/Velocities"][:]
        t = handle["Header"].attrs["Time"][0]
        box = handle["Header"].attrs["BoxSize"][:]

        units = dict(handle["Units"].attrs)
        uM = (units["Unit mass in cgs (U_M)"][0] * unyt.g).in_base("galactic")
        uL = (units["Unit length in cgs (U_L)"][0] * unyt.cm).in_base("galactic")
        ut = (units["Unit time in cgs (U_t)"][0] * unyt.s).in_base("galactic")

        coords = (coords * uL).in_base("galactic")
        rho = (rho * uM / uL ** 3).in_base("galactic")
        vs = (vs * uL / ut).in_base("galactic")
        t = (t * ut).in_base("galactic")
        box = (box * uL).in_base("galactic")

        coords -= 0.5 * box[None, :]

        x = np.sqrt((coords ** 2).sum(axis=1))
        v = (coords * vs).sum(axis=1) / x

        x.convert_to_units("kpc")
        v.convert_to_units("km/s")
        t.convert_to_units("Myr")

        rhohist, _, _ = stats.binned_statistic(x, rho, statistic="median", bins=100)
        vhist, edges, _ = stats.binned_statistic(x, v, statistic="mean", bins=100)
        mids = 0.5 * (edges[1:] + edges[:-1])

        rhohist[np.isnan(rhohist)] = 0.0
        imax = np.argmax(rhohist)

        time[ifile] = t
        radius[ifile] = mids[imax]
        velocity[ifile] = vhist[imax]

isort = np.argsort(time)

time = time[isort]
radius = radius[isort]
velocity = velocity[isort]

with open(args.output, "w") as handle:
    handle.write("# Radial profile of shock wave\n")
    handle.write("# Column 0: time (Myr)\n")
    handle.write("# Column 1: radius (kpc)\n")
    handle.write("# Column 2: velocity (km/s)\n")
    for t, r, v in zip(time, radius, velocity):
        handle.write(f"{t:.5e}\t{r:.5e}\t{v:.5e}\n")