#!/usr/bin/env python

import sys
import argparse
import logging
logging.basicConfig()

import numpy as np

import ROOT
ROOT.gROOT.SetBatch()
ROOT.PyConfig.IgnoreCommandLineOptions = True

from KerasROOTClassification import load_from_dir
from KerasROOTClassification.plotting import (
    get_mean_event,
    plot_NN_vs_var_2D,
    plot_profile_2D,
    plot_hist_2D_events,
    plot_cond_avg_actmax_2D,
    plot_NN_vs_var_2D_all,
)
from KerasROOTClassification.utils import get_single_neuron_function, get_max_activation_events, weighted_quantile

parser = argparse.ArgumentParser(description='Create various 2D plots for a single neuron')
parser.add_argument("project_dir")
parser.add_argument("output_filename")
parser.add_argument("varx")
parser.add_argument("vary")
parser.add_argument("-m", "--mode",
                    choices=["mean_sig", "mean_bkg", "profile_sig", "profile_bkg", "hist_sig", "hist_bkg", "hist_actmax", "cond_actmax"],
                    default="mean_sig")
parser.add_argument("-l", "--layer", type=int, help="Layer index (takes last layer by default)")
parser.add_argument("-n", "--neuron", type=int, default=0, help="Neuron index (takes first neuron by default)")
parser.add_argument("-a", "--all-neurons", action="store_true", help="Create a summary plot for all neurons in all hidden layers")
parser.add_argument("--log", action="store_true", help="Plot in color in log scale")
parser.add_argument("--contour", action="store_true", help="Interpolate with contours")
parser.add_argument("-b", "--nbins", default=20, type=int, help="Number of bins in x and y direction")
parser.add_argument("-x", "--xrange", type=float, nargs="+", help="xrange (low, high)")
parser.add_argument("-y", "--yrange", type=float, nargs="+", help="yrange (low, high)")
parser.add_argument("-p", "--profile-metric", help="metric for profile modes", default="average", choices=["mean", "average", "max"])
parser.add_argument("--ntries-cond-actmax", help="number of random events to be maximised and averaged per bin", default=20, type=int)
parser.add_argument("--nit-cond-actmax", help="number of iterations for maximisation per bin", default=1, type=int)
parser.add_argument("--ntries-actmax", help="number of random events to be maximised for hist_actmax", default=10000, type=int)
parser.add_argument("-t", "--threshold", help="minimum activation threshold", default=0.2, type=float)
parser.add_argument("-v", "--verbose", action="store_true")
parser.add_argument("-s", "--step-size", help="step size for activation maximisation", default=1., type=float)

args = parser.parse_args()

if args.all_neurons and (not args.mode.startswith("mean")):
    parser.error("--all-neurons currently only supported for mean_sig and mean_bkg")

if args.verbose:
    logging.getLogger().setLevel(logging.DEBUG)

c = load_from_dir(args.project_dir)

plot_vs_activation = (args.vary == "activation")

layer = args.layer
neuron = args.neuron

if layer is None:
    layer = len(c.model.layers)-1

varx_index = c.fields.index(args.varx)
if not plot_vs_activation:
    vary_index = c.fields.index(args.vary)
else:
    vary_index = 0 # dummy value in this case

varx_label = args.varx
vary_label = args.vary

total_weights = c.w_test*np.array(c.class_weight)[c.y_test.astype(int)]

try:
    mask_value = c.mask_value
except NameError:
    mask_value = None

# varx_test = c.x_test[:,varx_index]
# vary_test = c.x_test[:,vary_index]

# x_not_masked = np.where(varx_test != mask_value)[0]
# y_not_masked = np.where(vary_test != mask_value)[0]

# percentilesx = weighted_quantile(varx_test[x_not_masked], [0.01, 0.99], sample_weight=total_weights[x_not_masked])
# percentilesy = weighted_quantile(vary_test[y_not_masked], [0.01, 0.99], sample_weight=total_weights[y_not_masked])

def get_ranges(x, quantiles, weights, mask_value=None, filter_index=None):
    ranges = []
    for var_index in range(x.shape[1]):
        x_var = x[:,var_index]
        not_masked = np.where(x_var != mask_value)[0]
        ranges.append(weighted_quantile(x_var[not_masked], quantiles, sample_weight=weights[not_masked]))
    return ranges

percentilesx = get_ranges(c.x_test, [0.01, 0.99], total_weights, mask_value=mask_value, filter_index=varx_index)[0]
percentilesy = get_ranges(c.x_test, [0.01, 0.99], total_weights, mask_value=mask_value, filter_index=vary_index)[0]


if args.xrange is not None:
    if len(args.xrange) < 3:
        args.xrange.append(args.nbins)
    varx_range = args.xrange
else:
    varx_range = (percentilesx[0], percentilesx[1], args.nbins)

if args.yrange is not None:
    if len(args.yrange) < 3:
        args.yrange.append(args.nbins)
    vary_range = args.yrange
else:
    vary_range = (percentilesy[0], percentilesy[1], args.nbins)
    if plot_vs_activation:
        vary_range = (0, 1, args.nbins)

if args.mode.startswith("mean"):

    if args.mode == "mean_sig":
        means = get_mean_event(c.x_test, c.y_test, 1, mask_value=mask_value)
    elif args.mode == "mean_bkg":
        means = get_mean_event(c.x_test, c.y_test, 0, mask_value=mask_value)

    print(means)

    if hasattr(c, "get_input_list"):
        input_transform = lambda x : c.get_input_list(c.transform(x))
    else:
        input_transform = c.transform

    if not args.all_neurons:
        plot_NN_vs_var_2D(
            args.output_filename,
            means=means,
            varx_index=varx_index,
            vary_index=vary_index,
            scorefun=get_single_neuron_function(
                c.model, layer, neuron,
                input_transform=input_transform
            ),
            xmin=varx_range[0], xmax=varx_range[1], nbinsx=varx_range[2],
            ymin=vary_range[0], ymax=vary_range[1], nbinsy=vary_range[2],
            varx_label=varx_label, vary_label=vary_label,
            logscale=args.log, only_pixels=(not args.contour)
        )
    else:
        if hasattr(c, "get_input_list"):
            transform_function = lambda inp : c.get_input_list(c.scaler.transform(inp))
        else:
            transform_function = c.scaler.transform
        plot_NN_vs_var_2D_all(
            args.output_filename,
            means=means,
            model=c.model,
            transform_function=transform_function,
            varx_index=varx_index,
            vary_index=vary_index,
            xmin=varx_range[0], xmax=varx_range[1], nbinsx=varx_range[2],
            ymin=vary_range[0], ymax=vary_range[1], nbinsy=vary_range[2],
            logz=args.log,
            plot_last_layer=False,
        )

elif args.mode.startswith("profile"):

    def my_average(x, weights):
        if weights.sum() <= 0:
            return np.nan
        else:
            return np.average(x, weights=weights)

    metric_dict = {
        "mean" : np.mean,
        "max" : np.max,
        "average" : my_average,
    }

    if args.mode == "profile_sig":
        class_index = 1
    else:
        class_index = 0
    valsx = c.x_test[c.y_test==class_index][:,varx_index]
    valsy = c.x_test[c.y_test==class_index][:,vary_index]
    scores = c.scores_test[c.y_test==class_index].reshape(-1)

    opt_kwargs = dict()

    if args.profile_metric == "average":
        opt_kwargs["weights"] = c.w_test[c.y_test==class_index]

    plot_profile_2D(
        args.output_filename,
        valsx, valsy, scores,
        xmin=varx_range[0], xmax=varx_range[1], nbinsx=varx_range[2],
        ymin=vary_range[0], ymax=vary_range[1], nbinsy=vary_range[2],
        metric=metric_dict[args.profile_metric],
        varx_label=varx_label, vary_label=vary_label,
        log=args.log,
        **opt_kwargs
    )

elif args.mode.startswith("hist"):

    if not args.mode == "hist_actmax":
        if args.mode == "hist_sig":
            class_index = 1
        else:
            class_index = 0

        valsx = c.x_test[c.y_test==class_index][:,varx_index]
        if not plot_vs_activation:
            valsy = c.x_test[c.y_test==class_index][:,vary_index]
        else:
            valsy = c.scores_test[c.y_test==class_index].reshape(-1)
        weights = c.w_test[c.y_test==class_index]
    else:
        # ranges in which to sample the random events
        x_test_scaled = c.transform(c.x_test)
        ranges = get_ranges(x_test_scaled, [0.01, 0.99], total_weights, mask_value=mask_value)
        losses, events = get_max_activation_events(c.model, ranges, ntries=args.ntries_actmax, step=args.step_size, layer=layer, neuron=neuron, threshold=args.threshold)
        events = c.inverse_transform(events)
        valsx = events[:,varx_index]
        if not plot_vs_activation:
            valsy = events[:,vary_index]
        else:
            valsy = losses
        weights = None

    plot_hist_2D_events(
        args.output_filename,
        valsx, valsy,
        xmin=varx_range[0], xmax=varx_range[1], nbinsx=varx_range[2],
        ymin=vary_range[0], ymax=vary_range[1], nbinsy=vary_range[2],
        weights=weights,
        varx_label=varx_label, vary_label=vary_label,
        log=args.log,
    )

elif args.mode.startswith("cond_actmax"):

    x_test_scaled = c.transform(c.x_test)

    # ranges in which to sample the random events
    ranges = get_ranges(x_test_scaled, [0.01, 0.99], total_weights, mask_value=mask_value)

    plot_cond_avg_actmax_2D(
        args.output_filename,
        c.model, layer, neuron, ranges,
        varx_index,
        vary_index,
        xmin=varx_range[0], xmax=varx_range[1], nbinsx=varx_range[2],
        ymin=vary_range[0], ymax=vary_range[1], nbinsy=vary_range[2],
        scaler=c.scaler,
        ntries=args.ntries_cond_actmax,
        maxit=args.nit_cond_actmax,
        step=args.step_size,
        varx_label=varx_label, vary_label=vary_label,
        log=args.log,
    )