import argparse
import logging
import math
import os
import random

import numpy as np
import torch
import torch.cuda
from scipy.stats import t


def get_stats(
    array, conf_interval=False, name=None, stdout=False, logout=False
):
    """Compute mean and standard deviation from an numerical array

    Args:
        array (array like obj): The numerical array, this array can be
            convert to :obj:`torch.Tensor`.
        conf_interval (bool, optional): If True, compute the confidence interval bound (95%)
            instead of the std value. (default: :obj:`False`)
        name (str, optional): The name of this numerical array, for log usage.
            (default: :obj:`None`)
        stdout (bool, optional): Whether to output result to the terminal.
            (default: :obj:`False`)
        logout (bool, optional): Whether to output result via logging module.
            (default: :obj:`False`)
    """
    eps = 1e-9
    array = torch.Tensor(array)
    std, mean = torch.std_mean(array)
    std = std.item()
    mean = mean.item()
    center = mean

    if conf_interval:
        n = array.size(0)
        se = std / (math.sqrt(n) + eps)
        t_value = t.ppf(0.975, df=n - 1)
        err_bound = t_value * se
    else:
        err_bound = std

    # log and print
    if name is None:
        name = "array {}".format(id(array))
    log = "{}: {:.4f}(+-{:.4f})".format(name, center, err_bound)
    if stdout:
        print(log)
    if logout:
        logging.info(log)

    return center, err_bound


def parse_args():
    parser = argparse.ArgumentParser("Graph Cross Network")
    parser.add_argument(
        "--pool_ratios",
        nargs="+",
        type=float,
        help="The pooling ratios used in graph cross layers",
    )
    parser.add_argument(
        "--hidden_dim",
        type=int,
        default=96,
        help="The number of hidden channels in GXN",
    )
    parser.add_argument(
        "--cross_weight",
        type=float,
        default=1.0,
        help="Weight parameter used in graph cross layer",
    )
    parser.add_argument(
        "--fuse_weight",
        type=float,
        default=1.0,
        help="Weight parameter for feature fusion",
    )
    parser.add_argument(
        "--num_cross_layers",
        type=int,
        default=2,
        help="The number of graph corss layers",
    )
    parser.add_argument(
        "--readout_nodes",
        type=int,
        default=30,
        help="Number of nodes for each graph after final graph pooling",
    )
    parser.add_argument(
        "--conv1d_dims",
        nargs="+",
        type=int,
        help="Number of channels in conv operations in the end of graph cross net",
    )
    parser.add_argument(
        "--conv1d_kws",
        nargs="+",
        type=int,
        help="Kernel sizes of conv1d operations",
    )
    parser.add_argument(
        "--dropout", type=float, default=0.0, help="Dropout rate"
    )
    parser.add_argument(
        "--embed_dim",
        type=int,
        default=1024,
        help="Number of channels of graph embedding",
    )
    parser.add_argument(
        "--final_dense_hidden_dim",
        type=int,
        default=128,
        help="The number of hidden channels in final dense layers",
    )

    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
    parser.add_argument(
        "--weight_decay", type=float, default=0.0, help="Weight decay rate"
    )
    parser.add_argument(
        "--epochs", type=int, default=1000, help="Number of training epochs"
    )
    parser.add_argument(
        "--patience", type=int, default=20, help="Patience for early stopping"
    )
    parser.add_argument(
        "--num_trials", type=int, default=1, help="Number of trials"
    )

    parser.add_argument(
        "--device",
        type=int,
        default=0,
        help="Computation device id, -1 for cpu",
    )
    parser.add_argument(
        "--dataset", type=str, default="DD", help="Dataset used for training"
    )
    parser.add_argument(
        "--seed", type=int, default=-1, help="Random seed, -1 for unset"
    )
    parser.add_argument(
        "--print_every",
        type=int,
        default=10,
        help="Print train log every ? epochs, -1 for silence training",
    )
    parser.add_argument(
        "--dataset_path",
        type=str,
        default="./datasets",
        help="Path holding your dataset",
    )
    parser.add_argument(
        "--output_path",
        type=str,
        default="./output",
        help="Path holding your result files",
    )

    args = parser.parse_args()

    # default value for list hyper-parameters
    if not args.pool_ratios or len(args.pool_ratios) < 2:
        args.pool_ratios = [0.8, 0.7]
        logging.warning(
            "No valid pool_ratios is given, "
            "using default value '{}'".format(args.pool_ratios)
        )
    if not args.conv1d_dims or len(args.conv1d_dims) < 2:
        args.conv1d_dims = [16, 32]
        logging.warning(
            "No valid conv1d_dims is give, "
            "using default value {}".format(args.conv1d_dims)
        )
    if not args.conv1d_kws or len(args.conv1d_kws) < 1:
        args.conv1d_kws = [5]
        logging.warning(
            "No valid conv1d_kws is given, "
            "using default value '{}'".format(args.conv1d_kws)
        )

    # device
    args.device = "cpu" if args.device < 0 else "cuda:{}".format(args.device)
    if not torch.cuda.is_available():
        logging.warning("GPU is not available, using CPU for training")
        args.device = "cpu"
    else:
        logging.warning("Device: {}".format(args.device))

    # random seed
    if args.seed >= 0:
        torch.manual_seed(args.seed)
        random.seed(args.seed)
        np.random.seed(args.seed)
        if args.device != "cpu":
            torch.cuda.manual_seed(args.seed)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False

    # print every
    if args.print_every < 0:
        args.print_every = args.epochs + 1

    # path
    paths = [args.output_path, args.dataset_path]
    for p in paths:
        if not os.path.exists(p):
            os.makedirs(p)

    # datasets ad-hoc
    if args.dataset in ["COLLAB", "IMDB-BINARY", "IMDB-MULTI", "ENZYMES"]:
        args.degree_as_feature = True
    else:
        args.degree_as_feature = False

    return args