import sys
import datashader as ds
import pandas as pd
import datashader.transfer_functions as tf
from datashader.utils import export_image
from datashader.transfer_functions import set_background
import subprocess as SP
import os
import itertools as ITT
from glob import glob

print("Loading stars...")
stars = pd.read_csv("stars.csv", usecols=["id", "x", "z", "mult"], index_col=0)
stars.loc[stars.mult == 1.0, "mult"] = float("nan")

steps = int(sys.argv[1])
size = 1080
mode = "eq_hist"

cvs = ds.Canvas(plot_width=size, plot_height=size)

print("Plotting density")
density_agg = cvs.points(stars, "x", "z")
density = tf.shade(density_agg, cmap=["black", "white"], how=mode)

print("Plotting neutrons")
neutrons_agg = cvs.points(stars, "x", "z", agg=ds.count("mult"))
neutrons = tf.shade(neutrons_agg, cmap=["darkblue", "lightblue"], how=mode)

base = tf.stack(density, neutrons)

# ffplay = SP.Popen([
#     "ffplay","-f","image2pipe","-"
# ],stdin=SP.PIPE,bufsize=0)


for rh_fn in ITT.chain.from_iterable(map(glob, sys.argv[2:])):
    basename = os.path.splitext(os.path.split(rh_fn)[-1])[0]
    filename = "img/{}_{}_{}.mkv".format(basename, size, mode)
    ffmpeg = SP.Popen(
        [
            "ffmpeg",
            "-y",
            "-f",
            "image2pipe",
            "-i",
            "-",
            "-crf",
            "17",
            "-r",
            "25",
            "-pix_fmt",
            "yuv420p",
            filename,
        ],
        stdin=SP.PIPE,
        bufsize=0,
    )

    print("Loading", rh_fn)
    route_hist = pd.read_csv(
        rh_fn, names=["id", "d"], index_col=0, dtype={"d": int}, low_memory=False,
    )
    exp_span = [route_hist.d.min(), route_hist.d.max()]
    stars["d"] = float("nan")
    rng = range(route_hist.d.min(), route_hist.d.max() + 1, steps)
    if steps == 0:
        rng = [route_hist.d.max() + 1]
    for n in rng:
        stars['d'] = route_hist[route_hist.d < n]  # slow
        explored_agg = cvs.points(stars, "x", "z", agg=ds.mean("d"))  # slow
        explored = tf.shade(
            explored_agg, cmap=["darkred", "lightpink"], how="linear", span=exp_span
        )
        img = set_background(tf.stack(base, explored), "black").to_pil()
        img.save(ffmpeg.stdin, "png")
        img.save("img/current.png")
    ffmpeg.stdin.close()
    ffmpeg.wait()