# coding: utf-8

import os
import glob
from shutil import copy
import argparse
import time
import numpy as np
from functools import reduce

from pp_params import *
from plotter import *
from postprocessor import *

fake_pp = PostProcessor()

parser = argparse.ArgumentParser()

input_args = parser.add_argument_group("input", "Input selection")
input_args.add_argument("runs", help="name of runs", nargs="*", default=[])

input_args.add_argument(
    "-ip", "--input_path", help="specify input directory", default="/home/nbrucy/simus/"
)
input_args.add_argument(
    "-p",
    "--project",
    help="specify project name (directory within the input directory)",
    default="disk",
)

input_args.add_argument(
    "-c", "--config", help="Path of a default config file", default=None
)

input_args.add_argument(
    "-wo",
    "--which_inputs",
    choices=["all", "id", "time"],
    help="Select inputs by time range, id range or all of them",
    default="all",
)
input_args.add_argument("-b", "--begin", help="id of first input", type=int, default=1)
input_args.add_argument("-e", "--end", help="id of last input", type=int, default=100)
input_args.add_argument(
    "-s", "--step", help="step between two input", type=int, default=1
)
input_args.add_argument(
    "-tb", "--time_begin", help="time of first input", type=float, default=0.0
)
input_args.add_argument(
    "-te", "--time_end", help="time of last input", type=float, default=6.0
)

input_args.add_argument(
    "-w", "--watch", help="wait and watch for missing inputs", action="store_true"
)
input_args.add_argument("--skip", help="skip failed loadings", action="store_true")
input_args.add_argument(
    "-wt",
    "--waiting_time",
    help="time between to successive try when watching new inputs (in second)",
    type=int,
    default=120,
)
input_args.add_argument(
    "-af",
    "--allowed_failures",
    help="number of allowed failures when waiting",
    default=30,
)

output_args = parser.add_argument_group("output", "Output configuration")

output_args.add_argument(
    "-op",
    "--output_path",
    help="specify output directory",
    default="/home/nbrucy/visus/",
)

output_args.add_argument(
    "--tag", help="Add a special tag on output filemanes", default=""
)

output_args.add_argument(
    "--interactive", help="Interactive mode : do not save data", action="store_true"
)

output_args.add_argument(
    "-owr", "--overwrite", help="Overwrite outputs", action="store_true"
)

output_args.add_argument(
    "-owrd",
    "--overwrite_dependencies",
    help="Overwrite outputs for dependencies",
    action="store_true",
    default=False,
)


pp_args = parser.add_argument_group("postproc", "Post Processing configuration")


pp_args.add_argument(
    "--process",
    help="Individual rules to apply",
    choices=fake_pp.rules.keys(),
    default=[],
    nargs="*",
)

pp_args.add_argument(
    "-pargs",
    "--process_args",
    help="Args to give to process rules",
    default=["x", "y", "z"],
    nargs="*",
)

pp_args.add_argument(
    "--compare", help="Time and inter run comparaison", default=[], nargs="*"
)

pp_args.add_argument(
    "-cargs",
    "--compare_args",
    help="Args to give to process rules",
    default=[None],
    nargs="*",
)

pp_args.add_argument("--plot", help="Plot rules", default=[], nargs="*")

pp_args.add_argument(
    "-plargs",
    "--plot_args",
    help="Args to give to plot rules",
    default=[None],
    nargs="*",
)

pp_args.add_argument(
    "-d", "--disk", help="Specify this for disk simulations", action="store_true"
)
pp_args.add_argument(
    "--fft", help="use quick and dirty fft rendering", action="store_true"
)
pp_args.add_argument("--zoom", help="zoom", type=float, default=2.0)
pp_args.add_argument(
    "-ms",
    "--map_size",
    help="size of the maps created in he map mode (in pixel)",
    type=int,
    default=1024,
)

pp_args.add_argument(
    "--nb_bin", help="Number of bins for azimuthal averages", type=int, default=50
)
pp_args.add_argument(
    "--pdf_nb_bin", help="Number of bins for PDF", type=int, default=50
)
pp_args.add_argument(
    "--binning",
    help="Kind of binning (logarithmic or linear)",
    choices=["log", "lin"],
    default="log",
)

plot_args = parser.add_argument_group("plot", "Plot configuration")

plot_args.add_argument(
    "--colormap", help="Colormap used", choices=P.colormaps(), default="plasma"
)
plot_args.add_argument(
    "--format",
    help="Format of the plot images",
    choices=["png", "jpeg", "pdf", "svg", "ps"],
    default="jpeg",
)
plot_args.add_argument(
    "--dpi", help="Resolution of the plot images", type=int, default=400
)
plot_args.add_argument("--beamer", help="Beamer mode", action="store_true")

args = parser.parse_args()

project = args.project
runs = args.runs
storage_in = args.input_path
storage_out = args.output_path

if args.config is None:
    pp_params = default_params()
else:
    pp_params = load_params(args.config)

pp_params.out.zoom = args.zoom
pp_params.out.tag = args.tag
pp_params.out.map_size = args.map_size
pp_params.out.interactive = args.interactive

pp_params.pymses.fft = args.fft


pp_params.disk.on = args.disk
pp_params.disk.binning = args.binning
pp_params.disk.nb_bin = args.nb_bin

pp_params.pdf.nb_bin = args.pdf_nb_bin

# extension for out files
P.style.use("seaborn-deep")
if args.format == "pdf":
    P.style.use("~/.config/matplotlib/pdf.mplstyle")

if args.beamer:
    P.rcParams["font.family"] = "sans-serif"
    P.rcParams["figure.figsize"] = (7, 4.5)

# Plot properties
P.rcParams["image.cmap"] = args.colormap
P.rcParams["savefig.dpi"] = args.dpi
P.rcParams["lines.linewidth"] = 2
P.rcParams["lines.markersize"] = 10
P.rcParams["errorbar.capsize"] = 4

# List of id that were successfully computed
nums_success = {}

# Go through all runs
for run in runs:
    path_suffix = project + "/" + run
    path_in = storage_in + path_suffix
    path_out = storage_out + path_suffix

    if args.tag == "":
        tag_run = run
    else:
        tag_run = run + "_" + args.tag

    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

    nums_success[run] = []

    if args.which_inputs in ["all", "time"]:
        names = glob.glob(path_in + "/output_[0-9][0-9][0-9][0-9][0-9]")
        nums_all = [int(n.split("/")[-1].split("_")[1]) for n in names]
        nums_all = np.sort(nums_all)
        if args.which_inputs == "all":
            nums = nums_all
        else:
            time = [get_time(path_in, n) for n in nums_all]
            nums = [
                n
                for i, n in enumerate(nums_all)
                if time[i] >= args.time_begin and time[i] < args.time_end
            ]
    else:
        nums = range(args.begin, args.end + 1, args.step)

    for num in nums:
        failures = 0
        success = False

        while not success:
            try:
                if len(args.process) > 0:
                    pp = PostProcessor(run, num, pp_params=pp_params)
                    pp.process(
                        args.process,
                        args.process_args,
                        overwrite=args.overwrite,
                        overwrite_dep=args.overwrite_dependencies,
                    )

                # If we are here, success !
                success = True
                nums_success[run].append(num)
            except (ValueError, IOError, KeyError) as e:
                print(e)
                if args.watch and failures < args.allowed_failures:
                    failures = failures + 1
                    print(
                        "ERROR: Unable to proceed for run {} output {}.\
                         Trying again in {} s ({} tries remaining)".format(
                            run,
                            num,
                            args.waiting_time,
                            args.allowed_failures - failures,
                        )
                    )
                    time.sleep(args.waiting_time)
                elif args.skip:
                    break
                else:
                    raise

path_in = storage_in + project
path_out = storage_out + project

if len(args.plot) > 0:
    pl = Plotter(path_in, runs, nums_success, path_out=path_out, pp_params=pp_params)
    pl.process(
        args.plot,
        args.plot_args,
        overwrite=args.overwrite,
        overwrite_dep=args.overwrite_dependencies,
    )

if len(args.compare) > 0:
    cc = Comparator(path_in, runs, nums_success, path_out=path_out, pp_params=pp_params)
    cc.process(
        args.compare,
        args.compare_args,
        overwrite=args.overwrite,
        overwrite_dep=args.overwrite_dependencies,
    )