# coding: utf-8

import os
from shutil import copy
import argparse
import time
import numpy as np
import disk_postprocess as dp


storage_in = "/home/nbrucy/simus/"
storage_out = "/home/nbrucy/visus/"

parser = argparse.ArgumentParser()
parser.add_argument("runs", help="name of runs", nargs="*", default=["015_iso"])
parser.add_argument("-b", "--begin", help="id of first output", type=int, default=1)
parser.add_argument("-e", "--end", help="id of last output", type=int, default=100)
parser.add_argument("-s", "--step", help="step between two output", type=int, default=1)
parser.add_argument(
    "-d", "--disk", help="do specific disk radial analysis", action="store_true"
)
parser.add_argument(
    "--nb_bin", help="Number of bins for azimuthal averages", type=int, default=50
)


parser.add_argument("-m", "--maps", help="do generic maps", action="store_true")

parser.add_argument("-p", "--project", help="specify project name", default="disk")

parser.add_argument(
    "-fr",
    "--force_redo",
    help="redo plots even if the files already exist",
    action="store_true",
)

parser.add_argument(
    "-w", "--watch", help="wait and watch for missing outputs", action="store_true"
)

parser.add_argument("--skip", help="skip failed loadings", action="store_true")

parser.add_argument(
    "-wt",
    "--waiting_time",
    help="time between to successive try when watching new outputs (in second)",
    type=int,
    default=120,
)
parser.add_argument(
    "-af",
    "--allowed_failures",
    help="number of allowed failures when waiting",
    default=30,
)
parser.add_argument("--level", help="plot levels", action="store_true")
parser.add_argument("--cpu", help="plot cpu", action="store_true")

parser.add_argument(
    "-x", help="x position of the central point", type=float, default=1.0
)
parser.add_argument(
    "-y", help="y position of the central point", type=float, default=1.0
)
parser.add_argument(
    "-z", help="z position of the central point", type=float, default=1.0
)
parser.add_argument(
    "--fft", help="use quick and dirty fft rendering", action="store_true"
)
parser.add_argument("-i", "--interactive", help="Interactive mode", action="store_true")

parser.add_argument(
    "--colormap", help="Colormap used", choices=dp.P.colormaps(), default="plasma"
)
parser.add_argument(
    "--format",
    help="Format of the output",
    choices=["png", "jpeg", "pdf", "ps"],
    default="jpeg",
)
parser.add_argument("--dpi", help="Resolution of the output", type=int, default=400)


args = parser.parse_args()

project = args.project
runs = args.runs

first = args.begin
last = args.end
step = args.step

# extension for out files
dp.out_ext = "." + args.format
# Plot properties
dp.P.rcParams["image.cmap"] = args.colormap
dp.P.rcParams["savefig.dpi"] = args.dpi

for run in runs:
    path_suffix = project + "/" + run
    path_in = storage_in + path_suffix
    path_out = storage_out + path_suffix

    if not os.path.exists(path_out):
        os.makedirs(path_out)
        try:
            copy(path_in + "/disk.nml", path_out)
            copy(path_in + "/output_00001/compilation.txt", path_out)
        except:
            pass

    for i in range(first, last + 1, step):
        failures = 0
        success = False

        while not success:
            try:
                if args.maps:
                    dp.make_image_disk(
                        path_in,
                        i,
                        path_out=path_out,
                        tag=run,
                        map_size=1024,
                        force=args.force_redo,
                        level=args.level,
                        cpu=args.cpu,
                        pos_star=np.array([args.x, args.y, args.z]),
                        fft=args.fft,
                        interactive=args.interactive,
                    )
                if args.disk:
                    dp.disk_prop(
                        path_in,
                        i,
                        path_out=path_out,
                        rad_ext=1,
                        nb_bin=args.nb_bin,
                        force=args.force_redo,
                        pos_star=np.array([args.x, args.y, args.z]),
                    )
                    dp.plot_disk_prop(
                        path_out,
                        i,
                        tag=run,
                        force=args.force_redo,
                        interactive=args.interactive,
                    )
                success = True
            except ValueError:
                if args.watch and failures < args.allowed_failures:
                    failures = failures + 1
                    print(
                        "Unable to proceed for run {} output {}. Trying again in {} s ({} tries remaining)".format(
                            run, i, args.waiting_time, args.allowed_failures - failures
                        )
                    )
                    time.sleep(args.waiting_time)
                elif args.skip:
                    break
                else:
                    raise