#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
import sys
import swiftsimio.visualisation.projection as vis
from copy import deepcopy
from shutil import copyfile
import os
from subprocess import call

sys.path.append("../../../csds/.libs/")
import libcsds as csds

boxsize = 80.0
large_width = 15
small_width = 0.4
alpha_particle = 0

width = 0
basename = "index"
resolution = 1080
resolution_mini = resolution * 100 // 512
id_foc = 99839
v_max = 8.0

traj = None


def selectParticles(pos, width):
    ind1 = np.logical_and(pos[:, 0] > 0, pos[:, 0] < width)
    ind2 = np.logical_and(pos[:, 1] > 0, pos[:, 1] < width)
    ind3 = np.logical_and(ind1, ind2)

    # avoid the zoom
    size = resolution_mini * width / resolution
    box = np.logical_and(pos[:, 0] > (width - size), pos[:, 1] < size)
    ind3 = np.logical_and(~box, ind3)

    return ind3


def getImage(parts, width, center, res):
    pos = parts["positions"]
    pos = pos - center + 0.5 * width
    pos /= width
    h = parts["smoothing_lengths"] / width

    m = np.ones(pos.shape[0])
    # Do the projection
    img = vis.scatter(pos[:, 0], pos[:, 1], m, h, res).T
    img /= width ** 3
    ind = img == 0
    img[ind] = img[~ind].min()
    img = np.log10(img)

    return img


def doProjection(parts, t, skip):
    plt.figure(figsize=(8, 8))
    global traj, id_foc

    # evolve in time the particles
    interp = csds.moveForwardInTime(basename, parts, t)

    # Check if some particles where removed
    ind = parts["smoothing_lengths"] == 0
    ind = np.arange(parts.shape[0])[ind]
    if len(ind) != 0:
        parts = np.delete(parts, ind)
        interp = np.delete(interp, ind)
        id_foc -= np.sum(ind < id_foc)

    # Get the position and the image center
    pos = interp["positions"]
    position = pos[id_foc, :]

    # save the trajectory
    if traj is None:
        traj = position[np.newaxis, :]
    else:
        traj = np.append(traj, position[np.newaxis, :], axis=0)

    # Do we want to generate the image?
    if skip:
        return parts

    # Compute the images
    img = getImage(interp, width, position, resolution)
    img[:resolution_mini, -resolution_mini:] = getImage(
        interp, 0.2 * boxsize, position, resolution_mini
    )
    box = width * np.array([0, 1, 0, 1])
    plt.imshow(img, origin="lower", extent=box, cmap="plasma", vmin=0, vmax=v_max)

    # plot the particles
    pos = interp["positions"] - position + 0.5 * width
    ind = selectParticles(pos, width)
    ms = 2
    plt.plot(
        pos[ind, 0],
        pos[ind, 1],
        "o",
        markersize=ms,
        alpha=alpha_particle,
        color="silver",
    )
    plt.plot(
        pos[id_foc, 0], pos[id_foc, 1], "or", markersize=2 * ms, alpha=alpha_particle
    )

    # plot time
    plt.text(
        0.5 * width,
        0.95 * width,
        "Time = %0.2f Hours" % (t / (60 * 60)),
        color="w",
        horizontalalignment="center",
    )

    # plot trajectory
    tr = deepcopy(traj) - position + 0.5 * width
    plt.plot(tr[:, 0], tr[:, 1], "-", color="w", alpha=alpha_particle)

    # plot scale
    plt.plot([0.05 * width, 0.15 * width], [0.05 * width, 0.05 * width], "-w")
    plt.text(
        0.1 * width,
        0.06 * width,
        "%.2f R$_\oplus$" % (0.1 * width),
        horizontalalignment="center",
        color="w",
    )

    # style
    plt.axis("off")
    plt.xlim(box[:2])
    plt.ylim(box[2:])
    plt.tight_layout()
    plt.style.use("dark_background")
    return parts


def skipImage(i):
    if os.path.isfile("output/image_%04i.png" % i):
        print("Skipping image %i" % i)
        return True
    else:
        return False


def doMovie(parts, t0, t1, N, init):
    times = np.linspace(t0, t1, N)
    for i, t in enumerate(times):
        print("Image %i / %i" % (i + 1, N))
        skip = skipImage(i + init)
        parts = doProjection(parts, t, skip)
        if not skip:
            plt.savefig("output/image_%04i.png" % (i + init))
        plt.close()
    return init + N


def doZoom(parts, w_init, w_end, N, init, t0, t1, increase_alpha):
    global width, alpha_particle
    widths = np.linspace(w_init, w_end, N)
    alpha = np.linspace(-8, 0.2, N)
    if not increase_alpha:
        alpha = alpha[::-1]
    k = 5  # parameter for the steepness
    alpha = 1.0 / (1 + np.exp(-2 * k * alpha))
    times = np.linspace(t0, t1, N)
    for i, w in enumerate(widths):
        print("Image %i / %i" % (i + 1, N))
        skip = skipImage(i + init)
        width = w
        alpha_particle = alpha[i]
        parts = doProjection(parts, times[i], skip)
        if not skip:
            plt.savefig("output/image_%04i.png" % (i + init))
        plt.close()
    return init + N


def doStatic(init, number, ref):
    # copy the first picture
    for i in range(number):
        copyfile("output/image_%04i.png" % ref, "output/image_%04i.png" % (i + init))
    return init + number


def doTitle(frames):
    plt.figure()
    plt.axis("off")
    box = np.array([0, 1, 0, 1])
    plt.xlim(box[:2])
    plt.ylim(box[2:])
    plt.tight_layout()
    plt.style.use("dark_background")

    style = {
        "verticalalignment": "center",
        "horizontalalignment": "center",
        "fontweight": "bold",
    }
    plt.text(0.5, 0.6, "Planetary Impact with the CSDS", **style)
    plt.text(0.5, 0.5, "L. Hausammann, J. Kegerreis and P. Gonnet 2020", **style)

    for i in range(frames):
        plt.savefig("output/image_%04i.png" % i)

    plt.close()
    return frames


def main():
    t0, t1 = csds.getTimeLimits(basename)
    parts = csds.loadSnapshotAtTime(basename, t0)

    # Do a few title frames
    init = doTitle(40)

    after_title = init
    frames_after_title = 10
    init += frames_after_title

    # do a zoom while moving forward in time
    init = doZoom(
        parts,
        large_width,
        small_width,
        50,
        init=init,
        t0=t0,
        t1=0.07 * t1,
        increase_alpha=True,
    )

    # Copy a few frames
    doStatic(after_title, frames_after_title, after_title + frames_after_title)
    init = doStatic(init, 10, init - 1)
    # Do the main part of the movie
    init = doMovie(parts, 0.07 * t1, 0.15 * t1, N=250, init=init)
    # copy a few frames
    init = doStatic(init, 10, init - 1)

    # zoom out and finish the movie
    init = doZoom(
        parts,
        small_width,
        large_width,
        607,
        init=init,
        t0=0.15 * t1,
        t1=t1,
        increase_alpha=False,
    )

    # copy a few frames
    init = doStatic(init, 10, init - 1)


if __name__ == "__main__":
    main()

    convert = "ffmpeg -i output/image_%04d.png -y -vcodec libx264 "
    convert += "-profile:v high444 -refs 16 -crf 0 "
    convert += "-preset ultrafast movie.mp4"
    call(convert, shell=True)
