# generate bootstrapped samples for simulated datatask

import os

from pathlib import Path

import argparse

import logging

import torch

import torch.nn as nn

from torch.nn.parameter import Parameter

import pyro

import pyro.distributions as dist

import pandas as pd

import numpy as np

import pickle

import shutil

import utils

import json

from tqdm import tqdm

parser = argparse.ArgumentParser()

parser.add_argument('--setting-dir', default='settings', help="Directory with different settings")

parser.add_argument('--setting', default='collider-prognosticfactor', help="Directory contain setting.json, experimental setting, data-generation, regression model etc")

parser.add_argument('--N', default = '3000', help = "number of units in simulation")

parser.add_argument('--Nvalid', default = '1000', help = "number of units in simulation for validation")

parser.add_argument('--splits', default = 'train.valid', help = "which splits to do, should be separated by .")

parser.add_argument('--counterfactuals', dest='counterfactuals', action='store_true', help="Also generate outcomes for counterfactuals")

parser.add_argument('--no-counterfactuals', dest='counterfactuals', action='store_false', help="Don't generate outcomes for counterfactuals")

parser.add_argument('--sample-imgs', dest='sample_imgs', action = "store_true", help="sample images along with covariate data")

parser.add_argument('--no-imgs', dest='sample_imgs', action="store_false", help="don't get images, matching with units")

parser.add_argument('--seed', default='1234567', help="seed for simluations")

parser.add_argument('--close-range', default=5, type=int, help="when sampling on continuous variables, pick an image from the closest x observations")

parser.add_argument('--replace', action='store_true', help="sample with replacement from images")

parser.add_argument('--debug', action='store_true')

parser.set_defaults(sample_imgs=True, debug=False, counterfactuals=True, replace=False)

class LinearRegressionModel(nn.Module):

    def __init__(self, p, weights = None, bias = None):

        super(LinearRegressionModel, self).__init__()

        self.linear = nn.Linear(p, 1)

        if weights is not None:

            self.linear.weight = Parameter(torch.Tensor([weights]))

        if bias is not None:

            self.linear.bias = Parameter(torch.Tensor([bias]))

    def forward(self, x):

        return self.linear(x)

class LogisticRegressionModel(nn.Module):

    def __init__(self, p, weights = None, bias = None):

        super(LogisticRegressionModel, self).__init__()

        self.linear = nn.Linear(p, 1)

        if weights is not None:

            self.linear.weight = Parameter(torch.Tensor([weights]))

        if bias is not None:

            self.linear.bias = Parameter(torch.Tensor([bias]))

    def forward(self, x):

        return torch.sigmoid(self.linear(x))

class ProductModel(nn.Module):

    def __init__(self, p, weights = None, bias = None):

        super(ProductModel, self).__init__()

        # assert len(list(set([0,1]) - set(list(np.unique(weights, return_counts = False))))) == 0, "only weigths 0 and 1 are implemented for ProductModel"

        if weights is not None:

            self.weights = torch.Tensor([weights]).view(1,-1)

        else:

            self.weights = torch.Tensor([1])

    def forward(self, x):

        # apply weights (1, m)-tensor, broadcast to (n, m) and multiply elementwise

        x = (x*self.weights).clone() # add copy to prevent changing in place

        # select only those with nonzero weights

        x = x[:,self.weights.squeeze().nonzero().squeeze()]

        # multiply everything in column dimension

        return x.prod(1)

distributiondict = {"Bernoulli": dist.Bernoulli,

                    "Normal": dist.Normal}

model_modules = {

    "Linear": LinearRegressionModel,

    "Logistic": LogisticRegressionModel,

    "Product": ProductModel

    }

# create helper for strechting a list to a given size, repeating elements when necessary

def repeat_list(x, N):

    x_len = len(x)

    n_repeats = int(np.ceil(N / x_len))

    x = x * n_repeats

    x = x[:N]

    return x

def repeat_array(x, N):

    """

    Extend the number of rows in an array by repeating elements, to a specified size

    x: np.ndarray

    N: int, out length

    """

    assert isinstance(x, np.ndarray)

    n_repeats = int(np.ceil(N / x.shape[0]))

    # make sure only the first axis gets repeated

    tile_reps = np.ones((x.ndim,), dtype=np.int32)

    tile_reps[0] = n_repeats

    x = np.tile(x, tile_reps)

    return np.take(x, range(N), axis=0)

def grab_closest(x, d, close_range=int(0), replace=False):

    """

    Given a numeric value x, grab an item from dict d that is closest to x.

    For multidimensional x, assumes standard euclidian distance metric

    x: vector

    d: dict with {names: ["name1", "name2", ...], values: np.array([v1, v2, ...])}; values.shape[1] should be x.shape[0]

    close_range: pick an item from the closest n values to x

    replace: don't remove the picked item from d and return updated d

    returns: (name_of_closest_elem, distance_to_x (vector when x is a vector), dict (possibly updated))

    """

    names = d["name"]

    values = d["value"]

    assert type(values) is np.ndarray

    if not isinstance(x, np.ndarray):

        assert x.size==1

    else:

        assert x.shape[0] == values.shape[1]

    dist = (values - x)

    if dist.ndim == 1:

        dist = dist.reshape(-1,1) # reshape to make this work for 1d x, so that dist.shape == (n,1) always

    diff = np.linalg.norm(dist, ord=2, axis=1)

    if close_range > 0:

        closest_idx = np.random.choice(np.argsort(np.abs(diff))[:close_range])

    else:

        closest_idx = np.argmin(np.abs(diff))

    if not replace:

        # print(closest_idx)

        keep_idx = np.array(list(set(np.arange(values.shape[0]).astype(np.int64)) - set([closest_idx])))

        # print(keep_idx[:5])

        # print(keep_idx.shape)

        # print(values.shape)

        assert keep_idx.shape[0] == values.shape[0] - 1

        d = {"name": names[keep_idx],

             "value": np.take(values, keep_idx, axis=0)}

    return names[closest_idx], np.take(dist, closest_idx, axis=0), d

#%% import model specification

def prepare_model(model):

    """

    Prepare a model as defined in a pandas dataframe for sampling

    model: a pandas.DataFrame, see examples

    """

    # TODO add checks on model csv file

    # assert variable has variable_model iff variable_type == dependent

    # assert ordering of structural assignments

    assert isinstance(model, pd.DataFrame)

    model.set_index("variable", drop = False, inplace = True)

    param_cols = [x for x in model.columns if "param" in x]

    model["param_tuple"] = model[param_cols].apply(lambda x: (*x.dropna(),), axis = 1)

    var2label = dict(zip(model.variable.values, model.label.values))

    label2var = dict(zip(model.label.values, model.variable.values))

    return  model, var2label, label2var

def prepare_image_sets(model, img_path = "data", split = "train", N = 1000):

    """

    Prepare a dict of img names which are matched on variables present in 

    the generative model. Presently only works for binary variables

    Generate vectors of length N-samples, of which items can be picked one 

    by one to reduce reduncancy

    """

    gen_labels = model.label.tolist()

    # create list of variable roots, since there are labels with different names, e.g.:

    # malignancy_binary, malignancy_isborderline, malignancy_mean etc

    gen_variable_roots = [x.split("_")[0] for x in gen_labels]

    img_df = pd.read_csv(os.path.join(img_path, "labels.csv"))

    img_df = img_df[img_df.split == split]

    if split in img_df.name.values[0]: # some imgs can contain the split in the name: train/img_01.png

        # img_df["name"] = img_df.name.apply(lambda x: x.split("/")[1])

        img_df["name"] = img_df.name.apply(lambda x: os.path.basename(x))

    # print(img_df["name"].values[:10])

    # some generative variables should correspond to image features as recorded in data/labels.csv

    # keep only the columns that appear in the generative model, and name

    img_df = img_df[[x for x in img_df.columns if x.split("_")[0] in gen_variable_roots] + ["name"]]

    # define image vars that are in gen model and image labels

    img_vars = [x for x in img_df.columns if x in gen_labels]

    img_gen_model = model[model.label.isin(img_vars)]

    # distinguish continous generative variables

    img_cont_vars = img_gen_model[img_gen_model.distribution=="Normal"].label.tolist()

    img_disc_vars = [x for x in img_vars if x not in img_cont_vars]

    # assert len(img_cont_vars) < 3, "Currently only implemented for max 2 continuous variables"

    print("img_vars: {}".format(img_vars))

    print("img_cont_vars: {}".format(img_cont_vars))

    print("img_disc_vars: {}".format(img_disc_vars))

    # for the discrete image variables, ensure that for every group, there

    # are enough rows to accomodate the required simulation size

    img_disc_dict = {}

    if len(img_disc_vars) > 0:

        # remove possible 'borderline' images for removing noise in labels

        img_df = img_df[img_df[[x.split("_")[0] + "_isborderline" for x in img_disc_vars]].max(axis=1)==0]

        df_grp = img_df.groupby(img_disc_vars, sort=False)

        img_disc_dict = {}

        img_cont_dict = {}

        for name, group in df_grp:

            print("{} original items for key {}".format(group.shape[0], name))

            # names.append(name)

            img_disc_dict[name]  = repeat_list(group["name"].tolist(), 2*N)

            img_cont_dict[name] = {

                "name":  np.array(repeat_list(group["name"].tolist(), 2*N)),

                "value": np.array(repeat_array(group[img_cont_vars].values, 2*N))

            }

        # TODO remake pretty df with proper variable names

    else:

        img_cont_dict = {

            "name": np.array(repeat_list(img_df["name"].tolist(), 2*N)),

            "value":np.array(repeat_array(img_df[img_cont_vars].values, 2*N))

        }

    # print(img_cont_dict)

    return img_df, img_cont_vars, img_disc_vars, img_disc_dict, img_cont_dict    

def build_dataset(model, args, setting, N = 100):

    model_vars = model.variable.tolist()

    dep_vars = model[model.type == "dependent"].variable.tolist()

    # create dicts for going from variable name to column index and back

    if args.counterfactuals:

        model_vars = model_vars + ["y0", "y1"]

        dep_vars = dep_vars + ["y0", "y1"]

    n_vars = len(model_vars)

    var2idx = dict(zip(model_vars, range(n_vars)))

    idx2var = dict(zip(range(n_vars), model_vars))

    # TODO inject noise, base on how well the cnn-model can predict a feature 

    # which we are sampling on, to model the expected loss

    # initialize tensor

    X = torch.zeros([N, n_vars], requires_grad = False)

    for var, row in model.iterrows():

        column_idx = var2idx[var]

        # for noise variables, sample from distribution

        if row["type"] == "noise":

            distribution = distributiondict[row["distribution"]]

            params = row["param_tuple"]

            fn = distribution(*params)

            X[:, column_idx] = fn.sample(torch.Size([N])).requires_grad_(False)

        # for dependent variables, sample according to distribution parameterized via noise variables

        else:

            betas = model["b_"+var].values

            if args.counterfactuals:

                betas = np.append(betas, [0.,0.])

            bias = row["param_1"]

            model_type = row["variable_model"]

            variable_model = model_modules[model_type](len(betas), betas, bias)

            distribution = row["distribution"]

            MU = variable_model.forward(X.detach()).squeeze()

            if distribution == "Normal":

                X[:, column_idx] = MU

            #  NB possibility to use Bernoulli(logits = ...) here

            elif distribution == "Bernoulli":

                fn = distributiondict[distribution](MU)

                X[:, column_idx] = fn.sample().squeeze().requires_grad_(False)

        # df = pd.DataFrame(X.detach().numpy(), columns = model_vars)

        # print(df)

    # TODO update counterfactuals to include interactions

    if args.counterfactuals:

        # fill column with 0s and 1s

        X_0 = X.scatter(1, var2idx["t"]*torch.ones((N, 1)).long(), 0.)

        X_1 = X.scatter(1, var2idx["t"]*torch.ones((N, 1)).long(), 1.)

        if "interaction" in model.label.tolist():

            X_0[:,var2idx["zt"]] = X_0[:,var2idx["t"]] # all zeros

            X_1[:,var2idx["zt"]] = X_1[:,var2idx["z"]] # all equal to z 

        # get outcome model

        betas          = np.append(model["b_y"].values, [0.,0.])

        bias           = model.loc["y", "param_1"]

        model_type     = model.loc["y", "variable_model"]

        variable_model = model_modules[model_type](len(betas), betas, bias)

        distribution   = model.loc["y", "distribution"]

        MU_0 = variable_model.forward(X_0).squeeze()

        MU_1 = variable_model.forward(X_1).squeeze()

        if distribution == "Normal":

            X[:, var2idx["y0"]] = MU_0

            X[:, var2idx["y1"]] = MU_1

        #  NB possibility to use Bernoulli(logits = ...) here

        elif distribution == "Bernoulli":

            fn_0 = distributiondict[distribution](MU_0)

            fn_1 = distributiondict[distribution](MU_1)

            X[:, var2idx["y0"]] = fn_0.sample().squeeze()

            X[:, var2idx["y1"]] = fn_1.sample().squeeze()

    for var in dep_vars:

        print("mean (sd) {}: {:.3f} ({:.3f})".format(var, X[:, var2idx[var]].mean(), X[:, var2idx[var]].std()))

    return X, var2idx, idx2var

if __name__ == '__main__':

    # Load the parameters from json file

    args         = parser.parse_args()

    # Load information from last setting if none provided:

    if args.setting == "" and Path('last-defaults.json').exists():

        print("using last default setting")

        last_defaults = utils.Params("last-defaults.json")

        args.setting = last_defaults.dict["setting"]

        for param, value in last_defaults.dict.items():

            print("{}: {}".format(param, value))

    else:

        with open("last-defaults.json", "r+") as jsonFile:

            defaults = json.load(jsonFile)

            tmp = defaults["setting"]

            defaults["setting"] = args.setting

            jsonFile.seek(0)  # rewind

            json.dump(defaults, jsonFile)

            jsonFile.truncate()

    setting_home = os.path.join(args.setting_dir, args.setting)

    setting      = utils.Params(os.path.join(setting_home, "setting.json"))

    data_dir     = os.path.join(setting_home, "data")

    mode3d       = setting.mode3d

    GEN_MODEL    = setting.gen_model

    N_SAMPLES    = {"train": int(args.N), "valid": int(args.Nvalid), "test": int(args.Nvalid)}

    SPLITS       = str(args.splits).split(".")

    SAMPLE_IMGS  = args.sample_imgs

    MANUAL_SEED  = int(args.seed)

    if mode3d: 

        IMG_DIR = "data" # source location of all images

    else:

        IMG_DIR = Path("data","slices")

    # load and prepare generative model dataframe

    # model_df = pd.read_csv(os.path.join(HOME_PATH, "experiments", "sims", GEN_MODEL + ".csv"))

    model_df = pd.read_csv(os.path.join("experiments", "sims", GEN_MODEL + ".csv"))

    model_df, var2label, label2var = prepare_model(model_df)

    shutil.copy(os.path.join("experiments", "sims", GEN_MODEL + ".csv"),

                os.path.join(setting_home, "generating_model.csv"))

    dfs = {}

    dfs_oracle = {}

    # associate an image with each unit

    for i, split in enumerate(SPLITS):

        # remove earlier possible images

        if os.path.isdir(os.path.join(data_dir, split)) and SAMPLE_IMGS:

            shutil.rmtree(os.path.join(data_dir, split))

        # simulate data

        # logging.info("generating data for %s split" % (split))

        print("generating data for %s split" % (split))

        torch.manual_seed(MANUAL_SEED + i)

        X, var2idx, idx2var = build_dataset(model_df, args, setting, N_SAMPLES[split])

        df_oracle = pd.DataFrame(X.detach().numpy(), columns = list(var2idx.keys()))

        # extract Y and treatment

        y = X[:, var2idx["y"]]

        y = y.detach().numpy()

        t = X[:, var2idx["t"]]

        t = t.detach().numpy()

        if args.counterfactuals:

            y0 = X[:, var2idx["y0"]]

            y0 = y0.detach().numpy()

            y1 = X[:, var2idx["y1"]]

            y1 = y1.detach().numpy()

        # export

        if not os.path.isdir(os.path.join(data_dir, split)):

            logging.info("making dirs")

            os.makedirs(os.path.join(data_dir, split))

        torch.save(X, os.path.join(data_dir, split, "X.pt"))

        np.save(os.path.join(data_dir, split, "X.npy"), X.detach().numpy())

        if SAMPLE_IMGS:

            img_df, img_cont_vars, img_disc_vars, img_disc_dict, img_cont_dict = prepare_image_sets(model_df, IMG_DIR, split, N_SAMPLES[split])

            # when no discrete generative image variables provided,

            # no grouping is necessary

            if len(img_disc_vars) == 0:

                # extract columns from simulated data, corresponding to image vars

                img_cont_var_col_ids = [var2idx[label2var[x]] for x in img_cont_vars]

                x_cont = X[:, img_cont_var_col_ids]

                x_cont = x_cont.detach().squeeze().numpy()

                print(f"number of continuous variables: {len(img_cont_vars)}")

                diffs = np.zeros_like(x_cont)

                if diffs.ndim == 1:

                    diffs = diffs.reshape(-1,1)

                    x_cont = x_cont.reshape(-1,1)

                # sample images for each simulated unit

                img_names_out = []

                for i in tqdm(range(x_cont.shape[0])):

                    img_name, diff, img_cont_dict = grab_closest(x_cont[i,:], img_cont_dict, args.close_range, args.replace)

                    diffs[i,:] = diff

                    # print("image name: {}, x_value: {:.3f}, difference: {:.3f}".format(img_name, x[i], diff))

                    img_name_out = os.path.join(str(i) + "_" + img_name)

                    if "imgs/" in img_name:

                        img_name_out = os.path.basename(img_name_out)

                    # print(img_name_out)

                    # print(img_name)

                    img_names_out.append(img_name_out)

                    shutil.copy(os.path.join(IMG_DIR, split, img_name), 

                                os.path.join(data_dir, split, img_name_out))

            else:

                # sample based on discrete variables

                print(f"number of continuous variables: {len(x_img_cont_vars)}")

                n_img_disc_vars = len(img_disc_vars)

                img_disc_var_col_ids = [var2idx[label2var[x]] for x in img_disc_vars]

                x_disc = X[:, img_disc_var_col_ids].reshape(-1, n_img_disc_vars)

                x_disc = x_disc.detach().numpy().astype(int)

                if len(img_cont_vars) > 0:

                    img_cont_var_col_ids = [var2idx[label2var[x]] for x in img_cont_vars]

                    x_cont = X[:, img_cont_var_col_ids]

                    x_cont = x_cont.detach().squeeze().numpy()

                    diffs = np.zeros_like(x_cont)

                    if diffs.ndim == 1:

                        diffs = diffs.reshape(-1,1)

                        x_cont = x_cont.reshape(-1,1)

                img_names_out = []

                for i in tqdm(range(N_SAMPLES[split])):

                    key = tuple(x_disc[i, :])

                    if n_img_disc_vars == 1:

                        key = key[0]

                    if len(img_cont_vars) == 0:

                        img_names = img_disc_dict[key]

                        # pick first in list, then split this one off

                        img_name = img_names[0]

                        img_dict[key] = img_names[1:]

                    else:

                        # grab the continuous variable dict corresponding to discrete setting

                        cont_var_dict = img_cont_dict[key]

                        img_name, diff, cont_var_dict = grab_closest(x_cont[i,:], cont_var_dict, args.close_range, args.replace)

                        diffs[i,:] = diff

                        img_cont_dict[key] = cont_var_dict

                    img_name_out = os.path.join(str(i) + "_" + img_name)

                    img_names_out.append(img_name_out)

                    shutil.copy(os.path.join(IMG_DIR, split, img_name), 

                                os.path.join(data_dir, split, img_name_out))

            df_oracle["name"] = img_names_out

            if len(img_cont_vars) > 0:

                for i, cont_var in enumerate(img_cont_vars):

                    df_oracle["diff_"+label2var[cont_var]] = diffs[:,i]

                    df_oracle[label2var[cont_var]+"_actual"] = df_oracle[label2var[cont_var]].values + diffs[:,i]

            dict_out = {

                'name': img_names_out,

                't': t,

                'y': y}

            if args.counterfactuals:

                dict_out["y0"] = y0

                dict_out["y1"] = y1

            if "x" in var2idx.keys():

                dict_out["x"] = X[:, var2idx["x"]].detach().numpy()

            if "z" in var2idx.keys():

                dict_out["z"] = X[:, var2idx["z"]].detach().numpy()

            print("unique number of images sampled for split {}: {}".format(split, len(set([x.split("_")[-1] for x in img_names_out]))))

            print("sampling difference sd: {:.3f}".format(diffs.std()))

            df_out = pd.DataFrame(dict_out)

            dfs[split] = df_out

            df_out.to_csv(os.path.join(data_dir, split, "labels.csv"), index = False)

            df_oracle.to_csv(os.path.join(data_dir, split, "oracle.csv"), index = False)

        # add oracle data frame to dict

        dfs_oracle[split] = df_oracle

    # save data frame with all splits, and vardicts

    with open(os.path.join(data_dir, "vardicts.pt"), 'wb') as f:

        pickle.dump((var2idx, idx2var, var2label, label2var), f)

    oracle = pd.concat(dfs_oracle, axis = 0)

    oracle.reset_index(inplace=True)

    oracle.rename(index = str, columns = {"level_0": "split"}, inplace=True)

    if SAMPLE_IMGS:

        oracle["name"] = oracle[["split", "name"]].apply(lambda x: os.path.join(x[0], x[1]), axis = 1)

    oracle.to_csv(os.path.join(data_dir, "oracle.csv"), index = False)

    if SAMPLE_IMGS:

        df = pd.concat(dfs, axis = 0)

        df = df.reset_index()

        df["split"] = df.level_0

        df["name"] = df[["split", "name"]].apply(lambda x: os.path.join(x[0], x[1]), axis = 1)

        df = df.drop(["level_0", "level_1"], axis=1)

        df.to_csv(os.path.join(data_dir, "labels.csv"), index = False)

    if args.debug:

        x_train = torch.load(os.path.join(data_dir, "train", "X.pt"))

        x_train = x_train.detach().numpy()

        np.savetxt("scratch/X.csv", x_train, delimiter=',')

    logging.info("- done.")

#%%