# Copyright 2018-2019 Maxence Thevenet
#
# This file is part of WarpX.
#
# License: BSD-3-Clause-LBNL

"""
This script loops over 3D plotfiles plt*****, generates a 3D rendering
of the data with fields and particles, and saves one image per plotfile to
img_*****.png. It was written for a beam-driven wakefield acceleration
simulation, and contains a lot of custom values (for transparency,
color intensity etc.), so feel free to modify it to meet your needs.

Execute the file with e.g.
> mpirun -np 12 python yt3d_mpi.py
to generate the images. It can be quite slow for even moderately large
plotfiles.
"""

import glob

import numpy as np
import scipy.constants as scc
import yt
from mpi4py import MPI

yt.funcs.mylog.setLevel(50)

# my_max = 1.e11 # for smooth rendering
my_max = 5.0e10  # for layered rendering
species_to_plot = ["plasma_e", "beam", "driver"]
# For each species, provide [red, green, blue, alpha] between 0. and 1.
species_colors = {
    "plasma_e": [1.0, 1.0, 1.0, 0.15],
    "beam": [1.0, 1.0, 1.0, 0.2],
    "driver": [1.0, 1.0, 1.0, 0.2],
}
# provide these to avoid jitter when using a moving window
use_moving_window = True
plot_mr_patch = False
rendering_type = "layers"  # 'layers' or 'smooth'
maxwell_solver = "ckc"  # 'ckc' or 'yee'
cfl = 0.99
file_list = glob.glob("plotfiles/plt?????")

bounds = (-my_max, my_max)
z_shift = 0.0
w = (0.01 * my_max) ** 2


def jitter_shift(ds, ad, cfl, iteration):
    if maxwell_solver == "yee":
        dt = (
            1.0
            / scc.c
            * 1.0
            / np.sqrt(
                (
                    1.0 / ad["dx"][-1] ** 2
                    + 1.0 / ad["dy"][-1] ** 2
                    + 1.0 / ad["dz"][-1] ** 2
                )
            )
        )
    elif maxwell_solver == "ckc":
        dt = cfl * min([ad["dx"][-1], ad["dy"][-1], ad["dz"][-1]]) / scc.c
    z_front = dt * float(iteration) * scc.c + 7.5e-6 * yt.units.meter
    z_shift = z_front - ds.domain_right_edge[2]
    return z_shift


def get_species_ytpoints(ad, species, color_vec):
    xp = ad[species, "particle_position_x"].v
    yp = ad[species, "particle_position_y"].v
    zp = ad[species, "particle_position_z"].v
    if species == "plasma_e":
        selection = np.abs(xp) < 2.0e-6
        zp = zp[selection]
        yp = yp[selection]
        xp = xp[selection]
    vertices = np.column_stack((xp, yp, zp))
    colors = np.tile(color_vec, (vertices.shape[0], 1))
    points = yt.visualization.volume_rendering.render_source.PointSource(
        vertices, colors=colors, radii=1
    )
    return points


def img_onestep(filename):
    ds = yt.load(filename)
    ad = ds.all_data()
    iteration = int(filename[-5:])
    sc = yt.create_scene(ds, field="Ez")
    if use_moving_window:
        z_shift = jitter_shift(ds, ad, cfl, iteration)
    array_shift = z_shift * np.array([0.0, 0.0, 1.0])
    if plot_mr_patch:
        box_patch = yt.visualization.volume_rendering.render_source.BoxSource(
            left_edge=ds.index.grids[1].LeftEdge + array_shift,
            right_edge=ds.index.grids[1].RightEdge + array_shift,
            color=[1.0, 0.1, 0.1, 0.01],
        )
        sc.add_source(box_patch)
    ########################
    ### volume rendering ###
    ########################
    source = sc[0]
    source.use_ghost_zones = True
    source.grey_opacity = True
    source.set_log(False)
    tf = yt.ColorTransferFunction(bounds)
    if rendering_type == "smooth":
        tf.add_gaussian(-my_max / 4, width=15**2 * w, height=[0.0, 0.0, 1.0, 1])
        tf.add_gaussian(my_max / 4, width=15**2 * w, height=[1.0, 0.0, 0.0, 1])
    if rendering_type == "layers":
        # NEGATIVE
        tf.add_gaussian(-0.04 * my_max, width=8 * w, height=[0.1, 0.1, 1.0, 0.2])
        tf.add_gaussian(-0.2 * my_max, width=5 * w, height=[0.1, 0.1, 1.0, 0.5])
        tf.add_gaussian(-0.6 * my_max, width=w, height=[0.0, 0.0, 1.0, 1.0])
        # POSITIVE
        tf.add_gaussian(0.04 * my_max, width=8 * w, height=[1.0, 1.0, 0.2, 0.2])
        tf.add_gaussian(0.2 * my_max, width=5 * w, height=[1.0, 1.0, 0.2, 0.5])
        tf.add_gaussian(0.6 * my_max, width=w, height=[1.0, 1.0, 0.0, 1.0])

    ######################
    ### plot particles ###
    ######################
    species_points = {}
    for species in species_to_plot:
        species_points[species] = get_species_ytpoints(
            ad, species, species_colors[species]
        )
        sc.add_source(species_points[species])
    source.tfh.tf = tf
    source.tfh.bounds = bounds
    #########################
    ### camera properties ###
    #########################
    cam = sc.camera
    cam.resolution = (2048, 2048)
    cam.width = 0.00018 * yt.units.meter
    cam.focus = (
        ds.domain_center + np.array([0.0, 0.0, 10.0e-6]) * yt.units.meter + array_shift
    )
    cam.position = (
        ds.domain_center
        + np.array([15.0, 15.0, -5.0]) * yt.units.micrometer
        + array_shift
    )
    cam.normal_vector = [-0.3, -0.3, -0.2]
    cam.switch_orientation()
    # save image
    if rendering_type == "smooth":
        sc.save("img_" + str(my_number_list[count]).zfill(5), sigma_clip=5.0)
    if rendering_type == "layers":
        sc.save("img_" + str(my_number_list[count]).zfill(5), sigma_clip=2.0)


file_list.sort()
# Total number of files
nfiles = len(file_list)
# Each file has a unique number
number_list = range(nfiles)
comm_world = MPI.COMM_WORLD
me = comm_world.Get_rank()
nrank = comm_world.Get_size()
# List of files to process for current proc
my_list = file_list[(me * nfiles) / nrank : ((me + 1) * nfiles) / nrank]
# List if file numbers for current proc
my_number_list = number_list[(me * nfiles) / nrank : ((me + 1) * nfiles) / nrank]
for count, filename in enumerate(my_list):
    print("processing " + filename)
    img_onestep(filename)