# coding: utf-8
import sys
import numpy as np
import pymses
import os
from shutil import copy
import argparse
import time
import disk_postprocess as dp


storage_in = "/drf/projets/alfven-data/"
storage_out = "/dsm/anais/storageA/"

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("-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(
    "-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,
)


args = parser.parse_args()


user = "nbrucy"
folder = "simus"
project = args.project
runs = args.runs

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


for run in runs:
    path_suffix = user + "/" + folder + "/" + 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)
        copy(path_in + "/disk.nml", path_out)
        copy(path_in + "/output_00001/compilation.txt", path_out)

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

        while not success:
            try:
                dp.make_image_disk(
                    path_in,
                    i,
                    path_out=path_out,
                    tag=run,
                    map_size=1024,
                    force=args.force_redo,
                )
                if args.disk:
                    dp.disk_prop(
                        path_in,
                        i,
                        path_out=path_out,
                        rad_ext=1,
                        nb_bin=50,
                        force=args.force_redo,
                    )
                    dp.plot_disk_prop(path_out, i, tag=run, force=args.force_redo)
                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)
                else:
                    raise