# -*- coding: utf-8 -*-

import argparse
import torch as to
from typing import Tuple


class defaults:
    precision = to.float32
    patch_shift = 1
    shuffle = True
    warmup_Esteps = 0
    interactive_pause = 0.1
    tex_fonts = False
    dpi = 96
    figure_extension = ".png"
    focus_free_energy_last_x = 0.7
    focus_eval_metric_last_x = 0.7
    fontsize_free_energy_eval_metric_plot = 12
    ticksize_free_energy_eval_metric_plot = 8
    linestyle_free_energy_eval_metric_plot = "--"
    marker_free_energy_eval_metric_plot = "+"
    markersize_free_energy_eval_metric_plot = 4
    figsize_single_gf = (0.5, 0.55)
    figsize_lineplots = (4.5, 3)
    no_subplots_yx_datapoints = (2, 25)
    peak_value = 255


def make_parser_and_parse_args(parents: Tuple[argparse.ArgumentParser, ...]):
    p = argparse.ArgumentParser(
        description="Complete image with randomly missing values.\
                     Image will be read from file, degraded by randomly deleting pixels and\
                     segmented into overlapping patches (these correspond to the training data\
                     points). While training the model, data point reconstructions will be computed\
                     and subsequently merged to generate a reconstruction of the input image. Data\
                     output and visualizations will be saved in the directory specified by\
                     by output_directory.",
        parents=parents,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    return p.parse_args()


def get_exp_parser():
    p = argparse.ArgumentParser(add_help=False)

    p.add_argument(
        "image_file",
        type=str,
        help="Full path to input image (standard image formats such as png, jpg, ... and H5\
              (image expected in node called `data`) supported",
    )

    p.add_argument(
        "incomplete_percentage",
        type=int,
        help="Percentage of image pixels that will be set to nan to simulate incomplete data",
    )

    p.add_argument(
        "output_directory",
        type=str,
        help="Full path to directory to save data generated throughout execution of the algorithm\
              (will be created if not exists)",
    )

    p.add_argument(
        "patch_size", type=int, nargs=2, help="Patch size to be used for image segmentation"
    )

    p.add_argument("H", type=int, help="number of generative fields to learn")

    p.add_argument("--batch_size", default=16, type=int, help="batch size")

    p.add_argument("no_epochs", type=int, help="Number of EM-steps to perform")

    p.add_argument(
        "merge_every",
        type=int,
        help="Compute image reconstruction by merging data\
                                                  point reconstructions every X epochs",
    )

    p.add_argument(
        "--interactive",
        action="store_true",
        default=False,
        help="Show interactive visualizations",
    )

    p.add_argument(
        "--interactive_pause",
        type=float,
        default=0.1,
        help="Pause between subsequent interactive visualizations",
    )

    p.add_argument(
        "--keep_training_data_file",
        action="store_true",
        default=False,
        help="Do not remove the H5 file containing the training dataset (i.e. the overlapping\
        image patches)",
    )

    p.add_argument(
        "--keep_training_output_file",
        action="store_true",
        default=False,
        help="Do not remove the H5 file containing the training output",
    )

    p.add_argument(
        "--keep_reco_file",
        action="store_true",
        default=False,
        help="Do not remove the H5 file containing the data reconstructions",
    )

    p.add_argument(
        "--stop_if_eval_metric_decreases_in_x_of_y_epochs",
        default=None,
        nargs=2,
        type=float,
        help="Stop training if evaluation metric (PSNR, RMS, ...) difference between two\
              subsequent reconstructions is negative in x out of y epochs.",
    )

    return p


def get_learning_rate_scheduler_parser():
    p = argparse.ArgumentParser(add_help=False)

    p.add_argument("--min_lr", type=float, default=0.0001, help="mininum learning rate")

    p.add_argument("--max_lr", type=float, default=0.001, help="maximum learning rate")

    p.add_argument(
        "--epochs_per_half_cycle",
        type=int,
        default=20,
        help="epochs per half cycle (TODO: Improve docs)",
    )
    return p


def get_tvae_network_parser():
    p = argparse.ArgumentParser(add_help=False)

    p.add_argument(
        "--inner_net_shape",
        type=int,
        nargs="+",
        default=(512, 512),
        help="Shape of hidden layers; full net shape is \
              (patch_height*patch_width*no_channels,) + inner_net_shape + H",
    )

    return p


def get_eem_parser():
    p = argparse.ArgumentParser(add_help=False)

    p.add_argument(
        "--Ksize",
        type=int,
        default=64,
        help="Size of K sets (hyperparameter of the variational approximation)",
    )

    p.add_argument(
        "--parent_selection",
        default="fitness",
        type=str,
        choices=("fitness", "uniform"),
        help="Parent selection procedure (hyperparameter of the evolutionary optimization\
              in the E-step)",
    )

    p.add_argument(
        "--no_parents",
        type=int,
        default=5,
        help="Number of parents (hyperparameter of the evolutionary optimization in the\
              E-step)",
    )

    p.add_argument(
        "--no_children",
        type=int,
        default=4,
        help="Number of children (hyperparameter of the evolutionary optimization in the\
              E-step)",
    )

    p.add_argument(
        "--no_generations",
        default=1,
        type=int,
        help="Number of generations (hyperparameter of the evolutionary optimization in the\
              E-step)",
    )

    p.add_argument(
        "--crossover",
        action="store_true",
        default=False,
        help="Parent selection procedure (hyperparameter of the evolutionary optimization\
              in the E-step)",
    )

    return p
