# -*- coding: utf-8 -*-
"""
CODE ADAPTED FROM: https://github.com/sjblim/rmsn_nips_2018

[Treatment Effects with RNNs] core_routines
"""

import logging
import os

import numpy as np
import tensorflow as tf
import treatments.RMSN.configs as configs
import treatments.RMSN.core_routines as core
from treatments.RMSN.core_routines import train

ROOT_FOLDER = configs.ROOT_FOLDER


def rnn_fit(hyperparams_encoder_iptw, dataset_train, dataset_val, networks_to_train, task, MODEL_ROOT):

    specifications = {
        "rnn_propensity_weighted": (
            hyperparams_encoder_iptw["dropout_rate"],
            hyperparams_encoder_iptw["memory_multiplier"],
            hyperparams_encoder_iptw["num_epochs"],
            hyperparams_encoder_iptw["batch_size"],
            hyperparams_encoder_iptw["learning_rate"],
            hyperparams_encoder_iptw["max_norm"],
        ),
        "treatment_rnn_action_inputs_only": (0.1, 3, 100, 128, 0.01, 2.0),
        "treatment_rnn": (0.1, 4, 100, 64, 0.01, 1.0),
    }

    logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)

    # Get the correct networks to train

    if networks_to_train == "propensity_networks":
        logging.info("Training propensity networks")
        net_names = ["treatment_rnn_action_inputs_only", "treatment_rnn"]

    elif networks_to_train == "encoder":
        logging.info("Training R-MSN encoder")
        net_names = ["rnn_propensity_weighted"]

    elif networks_to_train == "user_defined":
        logging.info("Training user defined network")
        raise NotImplementedError("Specify network to use!")

    else:
        raise ValueError("Unrecognised network type")

    logging.info("Running hyperparameter optimisation")

    # Experiment name
    expt_name = "treatment_effects"
    # Possible networks to use along with their activation functions
    if task == "classification":
        activation_map = {
            "rnn_propensity_weighted": ("tanh", "sigmoid"),
            "rnn_propensity_weighted_logistic": ("tanh", "sigmoid"),
            "rnn_model": ("tanh", "sigmoid"),
            "treatment_rnn": ("tanh", "sigmoid"),
            "treatment_rnn_action_inputs_only": ("tanh", "sigmoid"),
        }
    else:
        activation_map = {
            "rnn_propensity_weighted": ("elu", "linear"),
            "rnn_propensity_weighted_logistic": ("elu", "linear"),
            "rnn_model": ("elu", "linear"),
            "treatment_rnn": ("tanh", "sigmoid"),
            "treatment_rnn_action_inputs_only": ("tanh", "sigmoid"),
        }
    # Setup tensorflow
    tf_device = "gpu"
    if tf_device == "cpu":
        config = tf.ConfigProto(log_device_placement=False, device_count={"GPU": 0})
    else:
        config = tf.ConfigProto(log_device_placement=False, device_count={"GPU": 1})
        config.gpu_options.allow_growth = True

    # Unpack simulation outputs
    training_data = dataset_train
    validation_data = dataset_val

    # Start Running hyperparam opt
    opt_params = {}
    for net_name in net_names:

        # Re-run hyperparameter optimisation if parameters are not specified, otherwise train with defined params
        max_hyperparam_runs = 10 if net_name not in specifications else 1

        # Pull datasets
        b_predict_actions = "treatment_rnn" in net_name
        use_truncated_bptt = (
            net_name != "rnn_model_bptt"
        )  # whether to train with truncated backpropagation through time
        b_propensity_weight = "rnn_propensity_weighted" in net_name
        b_use_actions_only = "rnn_action_inputs_only" in net_name

        # Extract only relevant trajs and shift data
        training_processed = core.get_processed_data(training_data, b_predict_actions, b_use_actions_only)
        validation_processed = core.get_processed_data(validation_data, b_predict_actions, b_use_actions_only)

        num_features = training_processed["scaled_inputs"].shape[-1]
        num_outputs = training_processed["scaled_outputs"].shape[-1]

        # Load propensity weights if they exist
        if b_propensity_weight:

            if net_name == "rnn_propensity_weighted_den_only":
                # use un-stabilised IPTWs generated by propensity networks
                propensity_weights = np.load(os.path.join(MODEL_ROOT, "propensity_scores_den_only.npy"))
            elif net_name == "rnn_propensity_weighted_logistic":
                # Use logistic regression weights
                propensity_weights = np.load(os.path.join(MODEL_ROOT, "propensity_scores.npy"))
                tmp = np.load(os.path.join(MODEL_ROOT, "propensity_scores_logistic.npy"))
                propensity_weights = tmp[: propensity_weights.shape[0], :, :]
            else:
                # use stabilised IPTWs generated by propensity networks
                propensity_weights = np.load(os.path.join(MODEL_ROOT, "propensity_scores.npy"))

            logging.info("Net name = {}. Mean-adjusting!".format(net_name))

            propensity_weights /= propensity_weights.mean()

            training_processed["propensity_weights"] = propensity_weights

        # Start hyperparameter optimisation
        hyperparam_count = 0
        while True:

            if net_name not in specifications:

                dropout_rate = np.random.choice([0.1, 0.2, 0.3, 0.4, 0.5])
                memory_multiplier = np.random.choice([0.5, 1, 2, 3, 4])
                num_epochs = 100
                minibatch_size = np.random.choice([64, 128, 256])
                learning_rate = np.random.choice([0.01, 0.005, 0.001])  # ([0.01, 0.001, 0.0001])
                max_norm = np.random.choice([0.5, 1.0, 2.0, 4.0])
                hidden_activation, output_activation = activation_map[net_name]

            else:

                spec = specifications[net_name]
                logging.info("Using specifications for {}: {}".format(net_name, spec))
                dropout_rate = spec[0]
                memory_multiplier = spec[1]
                num_epochs = spec[2]
                minibatch_size = spec[3]
                learning_rate = spec[4]
                max_norm = spec[5]
                hidden_activation, output_activation = activation_map[net_name]

            model_folder = os.path.join(MODEL_ROOT, net_name)

            hyperparam_opt = train(
                net_name,
                expt_name,
                training_processed,
                validation_processed,
                dropout_rate,
                memory_multiplier,
                num_epochs,
                minibatch_size,
                learning_rate,
                max_norm,
                use_truncated_bptt,
                num_features,
                num_outputs,
                model_folder,
                hidden_activation,
                output_activation,
                config,
                "hyperparam opt: {} of {}".format(hyperparam_count, max_hyperparam_runs),
            )

            hyperparam_count = len(hyperparam_opt.columns)

            if hyperparam_count == max_hyperparam_runs:
                opt_params[net_name] = hyperparam_opt.T  # pylint: disable=no-member
                break

        logging.info("Done")
        logging.info(hyperparam_opt.T)  # pylint: disable=no-member

        # Flag optimal params
    logging.info(opt_params)
