import numpy as np
import pickle
from tqdm import tqdm

avg_num_zeros = 0.0
from model_h0 import ModelCifar
from functools import partial
from hyper_params import load_hyper_params
import argparse
import torch
import os
import torch.nn as nn
from torchvision.datasets import FashionMNIST
import torch.nn.functional as F
from utils import dataset_mapper, image2flatten
from data import load_data_fast

np.random.seed(2023)
torch.manual_seed(2023)


def one_hot(arr, num_classes):
    new = []
    for i in tqdm(range(len(arr))):
        temp = np.zeros(num_classes)
        temp[arr[i]] = 1.0
        new.append(temp)
    return np.array(new)


def save_obj(obj, name):
    with open(name + ".pkl", "wb") as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)


def unpickle(file):
    import pickle

    with open(file, "rb") as fo:
        dict = pickle.load(fo, encoding="bytes")
    return dict


# def get_probs(label):
#     probs = np.power(np.add(list(range(10)), 1), 0.5)
#     probs[label] = probs[label] * 10

#     probs = probs / probs.sum()

#     return probs


# def get_biased_probs(label):
#     probs = np.power(np.add(list(range(10)), 1), 0.5)
#     probs[label] = probs[label] * 1.5

#     probs = probs / probs.sum()

#     return probs


# def get_wrong_probs(label):
#     probs = np.ones(10)
#     probs[(label + 1) % 10] *= 10

#     probs = probs / probs.sum()

#     return probs


# def get_uniform_probs(label):
#     probs = np.ones(10)

#     probs = probs / probs.sum()

#     return probs


# def get_model_probs(model, image):
#     with torch.no_grad():
#         image = torch.tensor(
#             np.repeat(image.reshape(1, 1, 28, 28), repeats=3, axis=1)
#         ).float()
#         probs = (
#             torch.softmax(model(image.to(device)), dim=-1)
#             .cpu()
#             .numpy()
#             .squeeze(0)
#             .astype(np.float32)
#         )
#     return probs


def get_model_probs_we(image, model, eps, device):
    with torch.no_grad():
        # if dataset == "fmnist":
        #     image = torch.tensor(
        #         image.reshape(-1, 3, 28, 28).repeat(1, 3, 1, 1)
        #     ).float()
        # elif dataset == "cifar":
        #     image = torch.tensor(image.reshape(-1, 3, 32, 32)).float()
        # else:
        #     raise ValueError(f"Dataset {dataset} not valid.")
        probs = (
            torch.softmax(model(image.to(device)), dim=-1)
            .cpu()
            .numpy()
            .astype(np.float32)
        )
    probs[probs < eps] = eps
    probs /= np.sum(probs, axis=-1, keepdims=True)
    return probs


parser = argparse.ArgumentParser()
parser.add_argument(
    "-c", "--config", required=True, help="Path to experiment config file."
)
parser.add_argument(
    "-l",
    "--linear",
    action="store_true",
    help="If used, the logging policy is a linear model",
)
parser.add_argument(
    "--tau",
    type=float,
    required=True,
    help="Softmax inverse temperature for training the logging policy.",
)
parser.add_argument(
    "--ul",
    type=float,
    required=True,
    help="The ratio of missing-reward to known-reward samples.",
)
parser.add_argument(
    "--dataset",
    type=str,
    required=True,
    help="Dataset to generate the bandit dataset from. Can be either fmnist or cifar.",
)
parser.add_argument(
    "--raw_image",
    action="store_true",
    help="If used, raw flatten image is given to the model instead of pretrained features.",
)
parser.add_argument(
    "--feature_size",
    required=False,
    type=int,
    help="If used, given feature size is supposed for the context.",
)

args = parser.parse_args()
linear = args.linear
hyper_params = load_hyper_params(args.config)
hyper_params["raw_image"] = args.raw_image
ul_ratio = None
labeled_proportion = 1.0
tau = args.tau
eps = 0
dataset = args.dataset
ul_ratio = args.ul
full_dataset = hyper_params["dataset"]
hyper_params["dataset"] = dataset_mapper[dataset]
hyper_params["dataset"]["name"] = full_dataset
feature_size = args.feature_size
# dataset = 'biased'
if feature_size is not None:
    hyper_params["feature_size"] = feature_size
else:
    hyper_params["feature_size"] = np.prod(hyper_params["dataset"]["data_shape"])

if ul_ratio is not None:
    labeled_proportion = 1 / (ul_ratio + 1)
print("Hyperparams: ", hyper_params)
print("Labeled proportion =", labeled_proportion)

if ul_ratio is None:
    exit()
if eps > 0:
    exit()

store_folder = dataset
if hyper_params["raw_image"]:
    store_folder += "_raw"
if linear:
    store_folder += "_linear/"
else:
    store_folder += "/"
store_folder += f"{tau}"
store_folder += f"_{int(ul_ratio)}"

print(store_folder)
device = "cuda:1"

print("Is Linear?", linear)

if linear:
    model = nn.Linear(
        hyper_params["feature_size"], hyper_params["dataset"]["num_classes"]
    )
    model_path = f"models/{dataset}/log_policy_{'deep' if not linear else 'linear'}{'_raw' if args.raw_image else ''}_1.0_tau{tau}.pth"
else:
    model = ModelCifar(hyper_params)
    model_path = f"models/{dataset}/log_policy_{'deep' if not linear else 'linear'}{'_raw' if args.raw_image else ''}_1.0_tau{tau}.pth"
print("Model Path: ", model_path)
model.load_state_dict(torch.load(model_path))
model.to(device)
model.eval()
probs_fn = partial(get_model_probs_we, model=model, eps=eps, device=device)
stored_feature = None

os.makedirs(f"../data/{store_folder}", exist_ok=True)

x_train, x_val, x_test = [], [], []
y_train, y_val, y_test = [], [], []

# train_dataset = FashionMNIST(root="data/fmnist/", train=True)
# test_dataset = FashionMNIST(root="data/fmnist/", train=False)
# N, M = len(train_dataset), len(test_dataset)
# print("Len Train =", N)
# print("Len Test =", M)

# # Train
# for i in range(N):
#     image, label = train_dataset[i]
#     image = np.array(image).reshape(28 * 28)
#     x_train.append(image)
#     y_train.append(label)
# x_train = np.stack(x_train)
# y_train = np.stack(y_train)

# # Test
# for i in range(M):
#     image, label = test_dataset[i]
#     image = np.array(image).reshape(28 * 28)
#     x_test.append(image)
#     y_test.append(label)
# x_test = np.stack(x_test)
# y_test = np.stack(y_test)

# # Normalize X data
# x_train = x_train.astype(float) / 255.0
# x_test = x_test.astype(float) / 255.0

# # One hot the rewards
# y_train = one_hot(y_train)
# y_test = one_hot(y_test)

# # Shuffle the dataset once
# indices = np.arange(len(x_train))
# np.random.shuffle(indices)
# assert len(x_train) == len(y_train)
# x_train = x_train[indices]
# y_train = y_train[indices]
# print(x_train.shape)
# N = len(x_train)
# n = int(N * labeled_proportion)

train_reader, test_reader, val_reader = load_data_fast(
    hyper_params, device=device, labeled=False
)
loaders = {"train": train_reader, "val": val_reader, "test": test_reader}
data = dict()
# Start creating bandit-dataset
for num_sample in [hyper_params["num_sample"]]:  # [1, 2, 3, 4, 5]:
    print("Pre-processing for num sample = " + str(num_sample))

    final_x, final_y, final_actions, final_prop, final_labeled = [], [], [], [], []

    avg_num_zeros = 0.0
    expected_reward = 0.0
    total = 0.0
    neg_cost_count = 0
    data = {}
    for mode in ["train", "val"]:
        print(f"####### MODE {mode}:")
        final_x, final_y, final_actions, final_prop, final_labeled = (
            [],
            [],
            [],
            [],
            [],
        )
        avg_num_zeros = 0.0
        expected_reward = 0.0
        total = 0.0
        neg_cost_count = 0
        for epoch in range(num_sample):
            s = 0
            N = len(loaders[mode].dataset)
            n = int(N * labeled_proportion)
            print(N, n)
            for x, y, action, delta, prop, _ in loaders[mode]:
                x, y, action, delta, prop = (
                    x.to(device),
                    y.to(device),
                    action.to(device),
                    delta.to(device),
                    prop.to(device),
                )
                image = x
                label = y.cpu().numpy()

                probs = probs_fn(image)
                u = probs.astype(np.float64)
                actions = []
                labeled = []
                for i in range(len(u)):
                    actionvec = np.random.multinomial(1, u[i] / np.sum(u[i]))
                    # print(actionvec)
                    # print(actionvec)
                    act = np.argmax(actionvec)
                    actions.append([act])
                    if s < n:
                        labeled.append([1.0])
                    else:
                        if label[i] == act:
                            neg_cost_count += 1
                        labeled.append([0.0])
                    s += 1
                # print(np.argmax(probs, axis=-1))
                # print(np.argmax(probs, axis=-1))
                actions = np.array(actions)
                labeled = np.array(labeled)
                u = image2flatten(
                    image.cpu().numpy(),
                    hyper_params["dataset"],
                    hyper_params["raw_image"],
                )
                # print("################## X SHAPE:", u.shape)
                final_x.append(u)
                final_actions.append(actions)
                final_prop.append(probs)
                final_labeled.append(labeled)
                # print(final_prop[0].shape)

                final_y.append(
                    F.one_hot(y, num_classes=hyper_params["dataset"]["num_classes"])
                    .cpu()
                    .numpy()
                )
                # print(actions, label)
                expected_reward += (actions[:, 0] == label).sum()
                total += len(label)

        avg_num_zeros /= float(N)
        avg_num_zeros = round(avg_num_zeros, 4)
        print(
            "Num sample = "
            + str(num_sample)
            + "; Acc = "
            + str(100.0 * expected_reward / total)
        )

        print(
            "Neg reward proportion = "
            + str(
                neg_cost_count / total / (1 - labeled_proportion)
                if labeled_proportion < 1.0
                else 0
            )
        )
        print()

        # Save as CSV
        # if labeled_proportion < 1.0:
        final_x = np.concatenate(final_x, axis=0)
        final_y = np.concatenate(final_y, axis=0)
        final_prop = np.concatenate(final_prop, axis=0)
        final_actions = np.concatenate(final_actions, axis=0)
        final_labeled = np.concatenate(final_labeled, axis=0)
        print(
            final_x.shape,
            final_y.shape,
            final_prop.shape,
            final_actions.shape,
            final_labeled.shape,
        )
        data[mode] = np.concatenate(
            (final_x, final_y, final_prop, final_actions, final_labeled), axis=1
        )
        print(f"################ Number of labeled = {final_labeled.sum()}")

    for mode in ["test"]:
        print(f"####### MODE {mode}:")
        final_x, final_y, final_actions, final_prop = (
            [],
            [],
            [],
            [],
        )
        avg_num_zeros = 0.0
        expected_reward = 0.0
        total = 0.0
        neg_cost_count = 0
        for epoch in range(num_sample):
            s = 0
            N = len(loaders[mode].dataset)
            n = int(N * labeled_proportion)
            for x, y, action, delta, prop in loaders[mode]:
                x, y, action, delta, prop = (
                    x.to(device),
                    y.to(device),
                    action.to(device),
                    delta.to(device),
                    prop.to(device),
                )
                image = x
                label = y.cpu().numpy()

                probs = probs_fn(image)
                u = probs.astype(np.float64)
                actions = []
                for i in range(len(u)):
                    actionvec = np.random.multinomial(1, u[i] / np.sum(u[i]))
                    act = np.argmax(actionvec)
                    actions.append([act])
                    s += 1
                actions = np.array(actions)
                u = image2flatten(
                    image.cpu().numpy(),
                    hyper_params["dataset"],
                    hyper_params["raw_image"],
                )
                # print("################## X SHAPE:", u.shape)
                final_x.append(u)
                final_actions.append(actions)
                final_prop.append(probs)
                # print(final_prop[0].shape)

                final_y.append(
                    F.one_hot(y, num_classes=hyper_params["dataset"]["num_classes"])
                    .cpu()
                    .numpy()
                )
                # print(actions, label)
                expected_reward += (actions[:, 0] == label).sum()
                total += len(label)

        avg_num_zeros /= float(N)
        avg_num_zeros = round(avg_num_zeros, 4)
        print(
            "Num sample = "
            + str(num_sample)
            + "; Acc = "
            + str(100.0 * expected_reward / total)
        )
        print(
            "Neg reward proportion = "
            + str(
                neg_cost_count / total / (1 - labeled_proportion)
                if labeled_proportion < 1.0
                else 0
            )
        )
        print()
        final_x = np.concatenate(final_x, axis=0)
        final_y = np.concatenate(final_y, axis=0)
        final_prop = np.concatenate(final_prop, axis=0)
        final_actions = np.concatenate(final_actions, axis=0)

        # Save as CSV
        # if labeled_proportion < 1.0:
        data[mode] = np.concatenate(
            (final_x, final_y, final_prop, final_actions), axis=1
        )
    filename = f"../data/{store_folder}/bandit_data_"
    filename += "sampled_" + str(num_sample) + "_"
    print(f"Saving dataset in {filename}...")
    for mode in ["train", "val", "test"]:
        save_obj(data[mode], filename + mode)
