# 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.")



#%%