import pickle
import numpy as np
import os


def _analyze_query_point_assignment(
    query_data_dict: dict,
    init_Rdata_dict: dict,
    init_Edata_dict: dict,
    num_R: int,
    query_point_assignment_array: np.ndarray,
    root: str,
    n_points_to_copy=50,
):
    """
    Analyzes and visualizes qDCA results.
    :param query_data_dict: raw query data.
    :param init_Rdata_dict: raw R data.
    :param init_Edata_dict: raw E data.
    :param num_R: total number of R points.
    :param query_point_assignment_array: query point assignments results.
    :param root: root directory of the experiment.
    :param n_points_to_copy: number of images to save.
    :return: accuracy of qDCA assignments; list of (R, query) points with same label;
    list of (R, query) points with different label
    """
    true_query_data_labels = query_data_dict["labels"]
    assigned_R = query_point_assignment_array[
        query_point_assignment_array[:, 1] < num_R, 1
    ]
    assigned_E = query_point_assignment_array[
        query_point_assignment_array[:, 1] >= num_R, 1
    ]
    assigned_R_labels = init_Rdata_dict["labels"][assigned_R]
    assigned_E_labels = init_Edata_dict["labels"][assigned_E - num_R]

    assigned_query_data_labels = np.empty(
        shape=query_point_assignment_array.shape[0]
    ).astype(np.int32)
    assigned_query_data_labels[
        query_point_assignment_array[:, 1] < num_R
    ] = assigned_R_labels
    assigned_query_data_labels[
        query_point_assignment_array[:, 1] >= num_R
    ] = assigned_E_labels

    accuracy = (
        true_query_data_labels == assigned_query_data_labels
    ).sum() / assigned_query_data_labels.shape[0]

    same_label_idx = np.where(true_query_data_labels == assigned_query_data_labels)[0]
    wrong_label_idx = np.where(true_query_data_labels != assigned_query_data_labels)[0]

    correct_pairs = []
    for i in query_point_assignment_array[same_label_idx]:
        query_idx, init_idx = i
        if init_idx < num_R:
            correct_pairs.append(
                [
                    query_data_dict["paths"].astype(object)[query_idx],
                    init_Rdata_dict["paths"].astype(object)[init_idx],
                    query_data_dict["labels"][query_idx],
                    init_Rdata_dict["labels"][init_idx],
                ]
            )
        else:
            correct_pairs.append(
                [
                    query_data_dict["paths"].astype(object)[query_idx],
                    init_Edata_dict["paths"].astype(object)[init_idx - num_R],
                    query_data_dict["labels"][query_idx],
                    init_Edata_dict["labels"][init_idx - num_R],
                ]
            )

    wrong_pairs = []
    for i in query_point_assignment_array[wrong_label_idx]:
        query_idx, init_idx = i
        if init_idx < num_R:
            wrong_pairs.append(
                [
                    query_data_dict["paths"].astype(object)[query_idx],
                    init_Rdata_dict["paths"].astype(object)[init_idx],
                    query_data_dict["labels"][query_idx],
                    init_Rdata_dict["labels"][init_idx],
                ]
            )
        else:
            wrong_pairs.append(
                [
                    query_data_dict["paths"].astype(object)[query_idx],
                    init_Edata_dict["paths"].astype(object)[init_idx - num_R],
                    query_data_dict["labels"][query_idx],
                    init_Edata_dict["labels"][init_idx - num_R],
                ]
            )

    with open(
        os.path.join(root, "logs", "analyzed_query_point_assignments.pkl"), "wb"
    ) as f:
        pickle.dump(
            {
                "accuracy": accuracy,
                "same_label_idx": same_label_idx,
                "wrong_label_idx": wrong_label_idx,
                "correct_pairs": correct_pairs,
                "wrong_pairs": wrong_pairs,
                "query_point_assignment_array": query_point_assignment_array,
            },
            f,
        )

    same_label_image_path = os.path.join(root, "visualization", "same_label_images")
    wrong_label_image_path = os.path.join(root, "visualization", "wrong_label_images")
    if not os.path.exists(wrong_label_image_path):
        os.mkdir(wrong_label_image_path)

    if not os.path.exists(same_label_image_path):
        os.mkdir(same_label_image_path)

    for i in range(n_points_to_copy):
        query_image_path, init_image_path, query_label, init_label = correct_pairs[i]
        path_to_copy = os.path.join(
            same_label_image_path,
            "i{0}_init_image_querylabel{1}_initlabel{2}.png".format(
                str(i), str(query_label), str(init_label)
            ),
        )

        os.system("cp {0} {1}".format(init_image_path, path_to_copy))

        path_to_copy2 = os.path.join(
            same_label_image_path,
            "i{0}_query_image_querylabel{1}_initlabel{2}.png".format(
                str(i), str(query_label), str(init_label)
            ),
        )
        os.system("cp {0} {1}".format(query_image_path, path_to_copy2))

        (
            w_query_image_path,
            w_init_image_path,
            w_query_label,
            w_init_label,
        ) = wrong_pairs[i]
        path_to_copy_w = os.path.join(
            wrong_label_image_path,
            "i{0}_init_image_querylabel{1}_initlabel{2}.png".format(
                str(i), str(w_query_label), str(w_init_label)
            ),
        )
        os.system("cp {0} {1}".format(w_init_image_path, path_to_copy_w))

        path_to_copy_w2 = os.path.join(
            wrong_label_image_path,
            "i{0}_query_image_querylabel{1}_initlabel{2}.png".format(
                i, w_query_label, w_init_label
            ),
        )
        os.system("cp {0} {1}".format(w_query_image_path, path_to_copy_w2))

    return accuracy, correct_pairs, wrong_pairs


def _generate_query_sets(version: str, N: int = 5000):
    """
    Generates query sets for qDCA experiment in Section 4.3.
    :param version: either version1 (dogs vs kitchen utils) or version2 (random).
    :param N: number of points to sample for R used in DCA.
    """
    with open(f"representations/vgg16/{version}/Rfeatures.pkl", "rb") as f:
        Rdata_v1 = pickle.load(f)

    with open(f"representations/vgg16/{version}/Efeatures.pkl", "rb") as f:
        Edata_v1 = pickle.load(f)

    init_Ridxs = np.random.choice(
        np.arange(len(Rdata_v1["feat_lin1"])), size=N, replace=False
    )
    query_Ridxs = np.setdiff1d(np.arange(len(Rdata_v1["feat_lin1"])), init_Ridxs)

    init_Eidxs = np.random.choice(
        np.arange(len(Edata_v1["feat_lin1"])), size=N, replace=False
    )
    query_Eidxs = np.setdiff1d(np.arange(len(Edata_v1["feat_lin1"])), init_Eidxs)

    with open(f"representations/vgg16/{version}/sampled_Rfeatures.pkl", "wb") as f:
        pickle.dump(
            {
                "feat_lin1": Rdata_v1["feat_lin1"][init_Ridxs],
                "feat_lin2": Rdata_v1["feat_lin2"][init_Ridxs],
                "labels": Rdata_v1["labels"][init_Ridxs],
                "paths": np.array(Rdata_v1["paths"])[init_Ridxs],
                "init_Ridx": init_Ridxs,
                "query_Ridx": query_Ridxs,
            },
            f,
        )

    with open(f"representations/vgg16/{version}/sampled_Efeatures.pkl", "wb") as f:
        pickle.dump(
            {
                "feat_lin1": Edata_v1["feat_lin1"][init_Eidxs],
                "feat_lin2": Edata_v1["feat_lin2"][init_Eidxs],
                "labels": Edata_v1["labels"][init_Eidxs],
                "paths": np.array(Edata_v1["paths"])[init_Eidxs],
                "init_Eidx": init_Eidxs,
                "query_Eidx": query_Eidxs,
            },
            f,
        )

    with open(f"representations/vgg16/{version}/query_features.pkl", "wb") as f:
        pickle.dump(
            {
                "feat_lin1": np.concatenate(
                    [
                        Rdata_v1["feat_lin1"][query_Ridxs],
                        Edata_v1["feat_lin1"][query_Eidxs],
                    ]
                ),
                "feat_lin2": np.concatenate(
                    [
                        Rdata_v1["feat_lin2"][query_Ridxs],
                        Edata_v1["feat_lin2"][query_Eidxs],
                    ]
                ),
                "labels": np.concatenate(
                    [Rdata_v1["labels"][query_Ridxs], Edata_v1["labels"][query_Eidxs]]
                ),
                "paths": np.concatenate(
                    [
                        np.array(Rdata_v1["paths"])[query_Ridxs],
                        np.array(Edata_v1["paths"])[query_Eidxs],
                    ]
                ),
                "init_Eidxs": init_Eidxs,
                "query_Eidxs": query_Eidxs,
                "init_Ridxs": init_Ridxs,
                "query_Ridxs": query_Ridxs,
            },
            f,
        )

[((113, 28, 113, 84), 'os.path.join', 'os.path.join', ({(113, 41, 113, 45): 'root', (113, 47, 113, 62): '"""visualization"""', (113, 64, 113, 83): '"""same_label_images"""'}, {}), "(root, 'visualization', 'same_label_images')", False, 'import os\n'), ((114, 29, 114, 86), 'os.path.join', 'os.path.join', ({(114, 42, 114, 46): 'root', (114, 48, 114, 63): '"""visualization"""', (114, 65, 114, 85): '"""wrong_label_images"""'}, {}), "(root, 'visualization', 'wrong_label_images')", False, 'import os\n'), ((51, 21, 51, 83), 'numpy.where', 'np.where', ({(51, 30, 51, 82): '(true_query_data_labels == assigned_query_data_labels)'}, {}), '(true_query_data_labels == assigned_query_data_labels)', True, 'import numpy as np\n'), ((52, 22, 52, 84), 'numpy.where', 'np.where', ({(52, 31, 52, 83): '(true_query_data_labels != assigned_query_data_labels)'}, {}), '(true_query_data_labels != assigned_query_data_labels)', True, 'import numpy as np\n'), ((101, 8, 111, 9), 'pickle.dump', 'pickle.dump', ({(102, 12, 109, 13): "{'accuracy': accuracy, 'same_label_idx': same_label_idx, 'wrong_label_idx':\n    wrong_label_idx, 'correct_pairs': correct_pairs, 'wrong_pairs':\n    wrong_pairs, 'query_point_assignment_array': query_point_assignment_array}", (110, 12, 110, 13): 'f'}, {}), "({'accuracy': accuracy, 'same_label_idx': same_label_idx,\n    'wrong_label_idx': wrong_label_idx, 'correct_pairs': correct_pairs,\n    'wrong_pairs': wrong_pairs, 'query_point_assignment_array':\n    query_point_assignment_array}, f)", False, 'import pickle\n'), ((115, 11, 115, 49), 'os.path.exists', 'os.path.exists', ({(115, 26, 115, 48): 'wrong_label_image_path'}, {}), '(wrong_label_image_path)', False, 'import os\n'), ((116, 8, 116, 40), 'os.mkdir', 'os.mkdir', ({(116, 17, 116, 39): 'wrong_label_image_path'}, {}), '(wrong_label_image_path)', False, 'import os\n'), ((118, 11, 118, 48), 'os.path.exists', 'os.path.exists', ({(118, 26, 118, 47): 'same_label_image_path'}, {}), '(same_label_image_path)', False, 'import os\n'), ((119, 8, 119, 39), 'os.mkdir', 'os.mkdir', ({(119, 17, 119, 38): 'same_label_image_path'}, {}), '(same_label_image_path)', False, 'import os\n'), ((172, 19, 172, 33), 'pickle.load', 'pickle.load', ({(172, 31, 172, 32): 'f'}, {}), '(f)', False, 'import pickle\n'), ((175, 19, 175, 33), 'pickle.load', 'pickle.load', ({(175, 31, 175, 32): 'f'}, {}), '(f)', False, 'import pickle\n'), ((37, 33, 39, 5), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((99, 8, 99, 74), 'os.path.join', 'os.path.join', ({(99, 21, 99, 25): 'root', (99, 27, 99, 33): '"""logs"""', (99, 35, 99, 73): '"""analyzed_query_point_assignments.pkl"""'}, {}), "(root, 'logs', 'analyzed_query_point_assignments.pkl')", False, 'import os\n'), ((216, 29, 221, 17), 'numpy.concatenate', 'np.concatenate', ({(217, 20, 220, 21): "[Rdata_v1['feat_lin1'][query_Ridxs], Edata_v1['feat_lin1'][query_Eidxs]]"}, {}), "([Rdata_v1['feat_lin1'][query_Ridxs], Edata_v1['feat_lin1'][\n    query_Eidxs]])", True, 'import numpy as np\n'), ((222, 29, 227, 17), 'numpy.concatenate', 'np.concatenate', ({(223, 20, 226, 21): "[Rdata_v1['feat_lin2'][query_Ridxs], Edata_v1['feat_lin2'][query_Eidxs]]"}, {}), "([Rdata_v1['feat_lin2'][query_Ridxs], Edata_v1['feat_lin2'][\n    query_Eidxs]])", True, 'import numpy as np\n'), ((228, 26, 230, 17), 'numpy.concatenate', 'np.concatenate', ({(229, 20, 229, 86): "[Rdata_v1['labels'][query_Ridxs], Edata_v1['labels'][query_Eidxs]]"}, {}), "([Rdata_v1['labels'][query_Ridxs], Edata_v1['labels'][\n    query_Eidxs]])", True, 'import numpy as np\n'), ((193, 25, 193, 52), 'numpy.array', 'np.array', ({(193, 34, 193, 51): "Rdata_v1['paths']"}, {}), "(Rdata_v1['paths'])", True, 'import numpy as np\n'), ((206, 25, 206, 52), 'numpy.array', 'np.array', ({(206, 34, 206, 51): "Edata_v1['paths']"}, {}), "(Edata_v1['paths'])", True, 'import numpy as np\n'), ((233, 24, 233, 51), 'numpy.array', 'np.array', ({(233, 33, 233, 50): "Rdata_v1['paths']"}, {}), "(Rdata_v1['paths'])", True, 'import numpy as np\n'), ((234, 24, 234, 51), 'numpy.array', 'np.array', ({(234, 33, 234, 50): "Edata_v1['paths']"}, {}), "(Edata_v1['paths'])", True, 'import numpy as np\n')]