#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Plots the evolution of the DM particles with the CSDS.
"""

import sys
import os
import argparse
import numpy as np
import matplotlib.pyplot as plt
from velociraptor import load
from velociraptor.particles import load_groups

# Include the CSDS library to the path before importing the CSDS python module
sys.path.append("../../../csds/src/.libs/")

# Now we can load the csds
import libcsds as csds

#%%
def get_redshift(a):
    return 1 / a - 1


def get_scale_factor(z):
    return 1 / (z + 1)


#%% Argparse options
def parse_option():
    description = """"Plots the evolution of the DM particles with the CSDS. """
    epilog = """
Examples:
--------

python3 csds_analysis.py csds_index_0000.dump
python3 csds_analysis.py csds_index_0000.dump --halo 2
python3 csds_analysis.py csds_index_0000.dump --halo 2 -n 250
"""
    parser = argparse.ArgumentParser(description=description, epilog=epilog)

    parser.add_argument("csds_file", type=str, help="CSDS .dump file name.")

    parser.add_argument(
        "--halo",
        action="store",
        type=int,
        dest="halo_index",
        default=1,
        help="Halo id of which we want trace the evolution. The halo id are given by velociraptor.",
    )

    parser.add_argument(
        "-n",
        action="store",
        type=int,
        default=500,
        help="Number of timestamps to fetch the CSDS.",
    )

    parser.add_argument(
        "--stf_file",
        action="store",
        type=str,
        dest="stf_file",
        default="stf_output/snap_0031",
        help="stf file basename.",
    )

    args = parser.parse_args()

    return args


#%% Parse the arguments
args = parse_option()

stf_file = args.stf_file
halo_index = args.halo_index
filename = args.csds_file
n_files = 1  # Number of csds files
n = args.n  # Number of timestamps to fetch the CSDS
dt = 1e-3  # small timestep for csds start (starts at t0+dt, see below)

# These are the number of timestamps we want for z > 2 (n_high_z) and z <=2 (n_low_z). This allows to provide more images outputs for 0 < z < 0 for the movie.
n_high_z = int(n / 10)
n_low_z = int(9 / 10 * n)

print("\n-------------------------------------")
print("Welcome to csds_anlysis.py !\n")

#%% Extract particles from halo
# Load Velociraptor data
print("Extracting halo data...")

catalogue = load(stf_file + ".properties.0")
groups = load_groups(stf_file + ".catalog_groups.0", catalogue=catalogue)

# Get the particles in halo_index
particles, unbound_particles = groups.extract_halo(halo_index=halo_index)

# Get the IDs
IDs = particles.particle_ids
N_particles = len(IDs)

#%% Now, we can open the CSDS, get the evolution of the particles though time with their IDs
cm_positions = np.zeros((n_files, n, 3))
positions = np.empty((n, N_particles, 3))
velocities = np.empty((n, N_particles, 3))
n_particles = np.zeros((n_files, n), dtype=int)
m_tot = np.zeros((n_files, n))

file_index = 0

if filename.endswith(".dump"):
    filename = filename[:-5]

print("Openening the CSDS...")

#%%
with csds.Reader(
    filename, verbose=0, number_index=10, restart_init=False, use_cache=True
) as reader:

    print("The CSDS is opened.")

    # Check the time limits (scale-factors in this example)
    t0, t1 = reader.get_time_limits()

    print("Scale factor limits: [{:e}, {:e}]".format(t0, t1))
    print(np.log10(t0), np.log10(t1))
    print("Redshift limits: [{:e}, {:e}]".format(get_redshift(t0), get_redshift(t1)))

    # Ensure that the fields are present
    fields = ["Coordinates", "Masses", "ParticleIDs", "Velocities"]
    missing = set(fields).difference(
        set(reader.get_list_fields(part_type=csds.particle_types.gas))
    )

    if missing:
        raise Exception("Fields %s not found in the logfile." % missing)

    # Create the list of IDs for *all* particle types
    IDs_list = [None] * csds.particle_types.count
    IDs_list[csds.particle_types.dark_matter] = IDs

    # Read the particles by IDs
    out = reader.get_data(fields=fields, time=t0, filter_by_ids=IDs_list)

    # Print the missing ids
    dm_ids, ids_found = set(IDs), set(out["ParticleIDs"])
    diff_ids = list(dm_ids.difference(ids_found))
    diff_found = list(ids_found.difference(dm_ids))
    print("The following ids were not found: ", np.array(diff_ids))
    print("The following ids are wrongly missing: ", np.array(diff_found))

    # Reverse the time: it is better to start from the end and rewind to the beginning
    high_z = np.linspace(t0 + dt, get_scale_factor(2), num=n_high_z, endpoint=False)
    low_z = np.linspace(get_scale_factor(2), t1, num=n_low_z)
    times = np.concatenate((high_z, low_z))[::-1]

    print("Starting to move through time")
    for i, t in enumerate(times):
        out = reader.get_data(fields=fields, time=t, filter_by_ids=IDs_list)

        # Get positions and velocities
        positions[i] = out["Coordinates"]
        velocities[i] = out["Velocities"]

        # Compute CM of halo
        n_particles[file_index, i] = len(out["ParticleIDs"])
        m_tot[file_index, i] = np.sum(out["Masses"])
        pos_cm = np.sum(out["Masses"][:, np.newaxis] * out["Coordinates"], axis=0)
        pos_cm /= m_tot[file_index, i]
        cm_positions[file_index, i, :] = pos_cm

print("The CSDS is now closed.")
#%% Do a movie showing the evolution of the halo (positions)
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots(num=1, figsize=(5, 5), dpi=300)
fig.tight_layout()


def animate(index):
    # Read the positions in the right order
    i = -index - 1
    x = positions[i, :, 0]
    y = positions[i, :, 1]
    scale_factor = times[i]
    redshift = 1 / scale_factor - 1

    # Clear the previous data
    ax.cla()
    ax.text(
        0.975,
        0.975,
        "z = {:.2f}".format(redshift),
        ha="right",
        va="top",
        transform=ax.transAxes,
        color="k",
    )
    ax.scatter(x, y, c="k", zorder=1, marker=".")
    ax.set_aspect("equal", "box")
    return ax


print("Making movie...")
animation = FuncAnimation(
    fig, animate, range(positions.shape[0]), fargs=[], interval=n / 24
)
animation.save("halo_evolution.mp4")
print("End :)")