from functions import *
from create_directory import *
from Crypto.Cipher import AES
import os
import shutil
import time

home = user_home()

if os.path.exists(home + "DataShareSecure") == False:
    print("\nNous vous prions de lire le fichier \"Readme.txt\" et de suivre ces consignes.\n")
    sys.exit()

print("BIENVENUE DANS CE PROGRAMME DE DECHIFFREMENT DE FICHIERS\n")

print("######### BON À SAVOIR ##########\n")

print("Vous exécutez ce programme stipule que:\n\n"
      "1- Vous avez pris connaissance du fonctionnement de DataShareSecure grâce au \"Readme.txt\" \n"
      "2- Vous avez exécuté le programme \"Public_Key_Manage.py\" au moins une fois et disposer donc d'une "
                                                        "paire de clés\n"
      "3- Vous désirez déchiffrer des fichiers que vous avez reçus d'un correspondant\n")

print("Si vous ne remplissez pas toutes les conditions du \"BON À SAVOIR\", je vous invite à fermer ce programme.\n"
      "Et à prendre le temps de remplir ces conditions.\n")

choix = input("Remplissez-vous les conditions sus-cités ? (O)ui ou (N)on : ")

if choix == 'O' or choix =='o':
    print("\nBien. Nous pouvons donc continuer\n")

    vide_directory(home + "DataShareSecure/Encrypted")
    vide_directory(home + "DataShareSecure/Decrypted")

    os.chdir(home + "DataShareSecure/Received")

    path = home + 'DataShareSecure/Received/key_used'

    with open(path, "r") as file:
        key_encrypted = file.read()

    key = dechiffrer(key_encrypted)
    buffer_size = 65536  # 64kb

    ########## MOVE FILE ############

    print("######## DECHIFFREMENT DES FICHIERS ET VERIFIVATION DES SIGNATURES ######## \n")

    file_dir = []

    file = [f for f in os.listdir(home + "DataShareSecure/Received") if os.path.isfile(f)]
    for f in file:
        if ".dss" in f:
            shutil.copy(f, home + "DataShareSecure/Encrypted")
        elif ".asc" in f:
            shutil.copy(f, home + "DataShareSecure/Decrypted")

    ########## DECRYPT ###############

    print("\n############# DECHIFFREMENT DES FICHIERS REÇUES ############\n")

    os.chdir(home + "DataShareSecure/Encrypted")

    files_dir = []

    files = [f for f in os.listdir(home + "DataShareSecure/Encrypted") if os.path.isfile(f)]
    for f in files:
        files_dir.append(f)

    for x in files_dir:
        with open(home + "DataShareSecure/Encrypted/" + x, "rb") as f:
            f.seek(0)
            path = home + 'DataShareSecure/Decrypted/' + x
            output_file = open(path[:-4], "wb")
            iv = f.read(16)
            cipher_encrypt = AES.new(key, AES.MODE_CFB, iv=iv)
            buffer = f.read(buffer_size)
            while len(buffer) > 0:
                decrypted_bytes = cipher_encrypt.decrypt(buffer)
                output_file.write(decrypted_bytes)
                buffer = f.read(buffer_size)

    print("Vos fichiers déchiffrés sont enregistrés dans le repertoire \"Decrypted\". \n")

    ########## VERIFY SIGNATURE ###############

    print("\n############ VERIFICATION DES FICHERS REÇUES #################\n")

    os.chdir(home + "DataShareSecure/Decrypted/")

    files_dir = []

    files = [f for f in os.listdir(home + "DataShareSecure/Decrypted/") if os.path.isfile(f)]
    for f in files:
        if ".asc" in f:
            files_dir.append(f)

    for x in files_dir:
        with open(home + "DataShareSecure/Decrypted/" + x, "rb") as f:
            file = x[:-4]
            verified = gpg.verify_file(f, file)
            print(file + " : ", verified.status + "")

    print("\nNOUS VOICI À LA FIN\n")


import nltk
from nltk.tokenize import *
import numpy as np

#--------------------------------------------------------
alphabet = ['a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z']
punctuation = ['?',',','!','.',':',';']

char_count= [0] * len(alphabet)
punctuation_count = [0] * len(punctuation)

#--------------------------------------------------------
#             PART OF SPEECH STUFF
#--------------------------------------------------------
#part of speech ratios + lexical variety
# - determiners
# - prepositions
# - pronouns
# - modal auxiliary-verbs -> CAN, COULD, WILL, WOULD
# - adverbs
# - coord-conjuctions
# - nouns
# - proper-nouns
# - adjectives
# - verbs
# - lexical variety = nouns + proper_nouns + adjectives + verbs + adverbs
pronouns_list = ['PRP', 'PRP$', 'WP', 'WP$']
adverbs_list = ['RB' ,'RBR', 'RBS', 'WRB']
adjectives_list = ['JJ', 'JJR', 'JJS']
verbs_list = ['VB', 'VBD', 'VBG', 'VBN', 'VBP', 'VBZ']

pos_ratios = [0] * 11

avg_sentence_length = 0
avg_word_length = 0
total_words = 0
#--------------------------------------------------------

def main():
    np.set_printoptions(suppress=True)

    features = []

    text = open("training_set\Abraham Lincoln\Abraham Lincoln___Lincoln Letters.txt").read()

    #total useful char
    t_u_c = total_useful_char(text)

    total_puctuation = count_punctuation(text)

    total_words = len(word_tokenize(text))

    #FEATURES 1 - 26
    letters_frequency(text, t_u_c)

    #FEATURES 27 - 32
    punctuation_frequency(text, total_puctuation)

    #FEATIRES 33 - 44
    part_of_speech_ratios(text, total_words)

    #FEATURES 44 - 45
    avg_sentence_length = average_sentence_length(text)
    avg_word_length = average_word_length(text)

    features.extend(char_count)
    features.extend(punctuation_count)
    features.extend(pos_ratios)
    features.append(avg_sentence_length)
    features.append(avg_word_length)
    features.append(total_words)

    features = np.array(features).reshape(-1,1)

    print("\n\n FEATURES final array: \n", features)
    print(features.shape)

def average_word_length(text):

    words = word_tokenize(text)
    sum = 0

    for word in words:
        sum += len(word)

    return sum/len(words)

def average_sentence_length(text):

    sentences = sent_tokenize(text)
    sum = 0

    for sentence in sentences:
        sum += len(word_tokenize(sentence))

    return sum/len(sentences)

def count_punctuation(text):
    return text.count('?') + text.count(',') + text.count('!') + text.count('.') + text.count(':') + text.count(';')

def total_useful_char(text):
    return len(text) - text.count(" ") - text.count("\n")

def letters_frequency(text, tChar):

    for char in text.lower():
        if char in alphabet:
            char_count[alphabet.index(char)] += 1

    for letter in char_count:
        char_count[char_count.index(letter)] /= tChar

def punctuation_frequency(text, total_puctuation):

    for char in text:
        if char in punctuation:
            punctuation_count[punctuation.index(char)] += 1

    for element in punctuation_count:
        punctuation_count[punctuation_count.index(element)] /= total_puctuation

def part_of_speech_ratios(text, total_words):

    words = word_tokenize(text)
    tagged_words = nltk.pos_tag(words)

    # lexical variety = nouns + proper_nouns + adjectives + verbs + adverbs

    for tagged_word in tagged_words:

        is_a_pronoun = [pronoun for pronoun in pronouns_list if(pronoun in tagged_word)]
        is_a_adverb = [adverb for adverb in adverbs_list if(adverb in tagged_word)]
        is_a_adjective = [adjective for adjective in adjectives_list if(adjective in tagged_word)]
        is_a_verb = [verb for verb in verbs_list if(verb in tagged_word)]

        if 'DT' in tagged_word:
            pos_ratios[0] += 1
        elif 'IN' in tagged_word:
            pos_ratios[1] += 1
        elif is_a_pronoun:
            pos_ratios[2] += 1
        elif 'MD' in tagged_word:
            pos_ratios[3] += 1
        elif is_a_adverb:
            pos_ratios[4] += 1
            pos_ratios[10] += 1
        elif 'CC' in tagged_word:
            pos_ratios[5] += 1
        elif ('NN' in tagged_word or 'NNS' in tagged_word):
            pos_ratios[6] += 1
            pos_ratios[10] += 1
        elif ('NNP' in tagged_word or 'NNPS' in tagged_word):
            pos_ratios[7] += 1
            pos_ratios[10] += 1
        elif is_a_adjective:
            pos_ratios[8] += 1
            pos_ratios[10] += 1
        elif is_a_verb:
            pos_ratios[9] += 1
            pos_ratios[10] += 1

    for element in pos_ratios:
        pos_ratios[pos_ratios.index(element)] /= total_words

if __name__ == '__main__':
    main()


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable #autograd oops
import torch.optim as optim

# core code of TRADES
# shhhhhh where does the author use l2-norm?????

def squared_l2_norm(x):
    flattened = x.view(x.unsqueeze(0).shape[0], -1)
    return (flattened ** 2).sum(1)

def l2_norm(x):
    return squared_l2_norm(x).sqrt()

# core function for TRADES calculating traded_loss
def trades_loss(model,
                x_natural,
                y,
                optimizer,
                step_size=0.003,
                epsilon=0.031,
                perturb_steps=10,
                beta=1.0,   # the coeff of second term
                distance='l_inf'):
    # define KL-loss for inner maximization https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html?highlight=kldivloss#torch.nn.KLDivLoss
    # If the field size_average is set to False, the losses are instead summed for each minibatch.
    criterion_kl = nn.KLDivLoss(size_average=False)
    # how to use loss :   f_loss(*args)(input) <- two parenthesis
    #eval() for BN and Dropout
    model.eval()

    # feed x_natural here into the loss as a batch
    batch_size = len(x_natural)

    # generate adversarial example
    # initiate an x_adv for skipping the concave
    x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).cpu().detach()
    # detach() tensor won't give it grad calculations anymore.
    if distance == 'l_inf':  # L-infinity ball  # no random start here
        for _ in range(perturb_steps):  # FGSM_k
            x_adv.requires_grad_()         # start from x_adv
            with torch.enable_grad():   # enable_grad vs no_grad
            # For the maximization problem, using torch.nn.KLDivLoss and cross entropy is equivalent because they differ by a constant additive term.

                loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
                                       F.softmax(model(x_natural), dim=1))  # why first term log while second term origin: because in the loss_criteria, there is no "log_target = True"

            grad = torch.autograd.grad(loss_kl, [x_adv])[0] # Computes and returns the sum of gradients of outputs w.r.t. the inputs.
            x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
            #clamp ..
            x_adv = torch.min(torch.max(x_adv, x_natural - epsilon), x_natural + epsilon)

            #clamp original pic
            x_adv = torch.clamp(x_adv, 0.0, 1.0)

    elif distance == 'l_2':# L_2  we will come back later about l_2....not commented yet

        delta = 0.001 * torch.randn(x_natural.shape).cpu().detach()
        delta = Variable(delta.data, requires_grad=True)

        # Setup optimizers
        optimizer_delta = optim.SGD([delta], lr=epsilon / perturb_steps * 2)

        for _ in range(perturb_steps):
            adv = x_natural + delta

            # optimize
            optimizer_delta.zero_grad()
            with torch.enable_grad():
                loss = (-1) * criterion_kl(F.log_softmax(model(adv), dim=1),
                                           F.softmax(model(x_natural), dim=1))
            loss.backward()
            # renorming gradient
            grad_norms = delta.grad.view(batch_size, -1).norm(p=2, dim=1)
            delta.grad.div_(grad_norms.view(-1, 1, 1, 1))
            # avoid nan or inf if gradient is 0
            if (grad_norms == 0).any():
                delta.grad[grad_norms == 0] = torch.randn_like(delta.grad[grad_norms == 0])
            optimizer_delta.step()

            # projection
            delta.data.add_(x_natural)
            delta.data.clamp_(0, 1).sub_(x_natural)
            delta.data.renorm_(p=2, dim=0, maxnorm=epsilon)
        x_adv = Variable(x_natural + delta, requires_grad=False)

    # not implemented for other losses
    else:
        x_adv = torch.clamp(x_adv, 0.0, 1.0)

    # adding two losses: L(fx,y)      ,     L(fx',fx)
    model.train()

    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)

    # not the main part, code related only
    # zero gradient     again, zero_grad -> loss_back -> updae?
    optimizer.zero_grad()

    # calculate robust loss
    logits = model(x_natural)      # pred of fx
    loss_natural = F.cross_entropy(logits, y)   # loss of  fx,y
    # loss of fx' fx
    loss_robust = (1.0 / batch_size) * criterion_kl(F.log_softmax(model(x_adv), dim=1),
                                                    F.softmax(model(x_natural), dim=1))
    loss = loss_natural + beta * loss_robust
    return loss


from django.contrib.auth import mixins
from oauth2_provider.contrib.rest_framework import (
    OAuth2Authentication as BaseOAuth2Authentication,
)

from purplship.server.core.authentication import (
    JWTAuthentication,
    TokenAuthentication,
    get_request_org,
)

class OAuth2Authentication(BaseOAuth2Authentication):
    def authenticate(self, request):
        auth = super().authenticate(request)

        if auth is not None:
            user, _ = auth

            request.org = get_request_org(request, user)

        return auth

class AccessMixin(mixins.AccessMixin):
    """Verify that the current user is authenticated."""

    def dispatch(self, request, *args, **kwargs):
        try:
            auth = (
                OAuth2Authentication().authenticate(request)
                or JWTAuthentication().authenticate(request)
                or TokenAuthentication().authenticate(request)
            )

            if auth is not None:
                user, *_ = auth
                request.user = user

        finally:
            return super().dispatch(request, *args, **kwargs)


import os
from PIL import Image

class ImageUpscaler:
    def __init__(self, image_path, scale_factor):
        self.image_path = image_path
        self.scale_factor = scale_factor

    def upscale_image(self, image_file):
        # Open the image
        image = Image.open(image_file)

        # Calculate the new dimensions
        width, height = image.size
        new_width = int(width * self.scale_factor)
        new_height = int(height * self.scale_factor)

        # Resize the image
        upscaled_image = image.resize((new_width, new_height), Image.BICUBIC)

        # Save the upscaled image
        upscaled_folder = os.path.join(self.image_path, 'upscaled')
        os.makedirs(upscaled_folder, exist_ok=True)

        file_name = os.path.splitext(os.path.basename(image_file))[0]
        save_path = os.path.join(upscaled_folder, f'{file_name}_upscaled.png')
        upscaled_image.save(save_path)

        # print(f"Upscaled image saved: {save_path}")

    def upscale_images_in_directory(self):
        # Get a list of all image files in the directory
        image_files = [
            os.path.join(self.image_path, file_name)
            for file_name in os.listdir(self.image_path)
            if file_name.endswith(('.jpg', '.jpeg', '.png'))
        ]

        for image_file in image_files:
            self.upscale_image(image_file)

if __name__ == '__main__':
    directory_path = '../private_keys'
    scale_factor = 4  # Increase the dimensions by a factor of 4

    upscaler = ImageUpscaler(directory_path, scale_factor)
    upscaler.upscale_images_in_directory()


from calculations import get_standings, get_current_track_data
from utils import get_player_name, get_telegram_name
from plots import timedelta_to_string

def info_about_current_weeks_ladder_changes(old_data, new_data):
    new_data = get_current_track_data(new_data)
    new_data = new_data[new_data["Origin"] == "Player"]
    new_ladder = get_standings(new_data)
    current_track = new_data["track_id"].unique()[0]
    old_data = old_data[old_data["track_id"] == current_track]
    old_data = old_data[old_data["Origin"] == "Player"]
    old_ladder = get_standings(old_data)

    player_overlap = list(set(new_ladder.index) & set(old_ladder.index))
    new_ladder = new_ladder.loc[new_ladder.index.isin(player_overlap)]
    old_ladder = old_ladder.loc[old_ladder.index.isin(player_overlap)]
    changes = new_ladder.index != old_ladder.index
    new_ladder = new_ladder[changes].reset_index().reset_index().set_index("Player")
    old_ladder = old_ladder[changes].reset_index().reset_index().set_index("Player")
    new_ladder["index_change"] = new_ladder["index"] - old_ladder["index"]

    messages = []
    for player in new_ladder.index.values:
        overtakes = new_ladder.loc[player, "index_change"]
        if not overtakes > 0:
            continue
        index = new_ladder.loc[player, "index"]
        overtook = new_ladder[(new_ladder["index"] >= index-overtakes) & (new_ladder["index"] < index)].index.values
        have_scored = old_data.loc[old_data["Origin"] == "Player", "Player"].unique()
        overtook = ", ".join([get_telegram_name(p) for p in overtook if p in have_scored])
        new_record = new_data.groupby(["Player", "track_id"])["Time"].min().loc[player, current_track]
        messages.append(
            f"{get_player_name(player)} scored a {timedelta_to_string(new_record)} and overtook {overtook}."
        )
    return messages

def info_about_new_times(old_data, new_data):
    messages = []
    new_entries_index = new_data[~new_data.isin(old_data)].dropna(how="all").index
    new_entries = new_data.loc[new_entries_index]
    for row_index, entry in new_entries.iterrows():
        player_name = entry["Player"]
        new_record = timedelta_to_string(entry["Time"])
        track = entry["Track"]
        message = f"{track}: {get_player_name(player_name)} scored a new record of {new_record}!"
        messages.append(message)
    return messages


from pyface.qt import QtCore, QtGui, is_qt4
from pyface.image_resource import ImageResource
from pyface.timer.api import do_later
from pyface.ui_traits import Image
from traits.api import (
    Any,
    Bool,
    Button,
    Dict,
    Event,
    List,
    HasTraits,
    Instance,
    Int,
    Property,
    Str,
    cached_property,
    observe,
)

from traitsui.api import (
    EnumEditor,
    InstanceEditor,
    Group,
    Item,
    Label,
    ObjectColumn,
    TableColumn,
    TableFilter,
    UI,
    View,
    default_handler,
    spring,
)
from traitsui.editors.table_editor import (
    BaseTableEditor,
    ReversedList,
    customize_filter,
)
from traitsui.ui_traits import SequenceTypes

from .editor import Editor
from .table_model import TableModel, SortFilterTableModel

if is_qt4:

    def set_qheader_section_resize_mode(header):
        return header.setResizeMode

else:

    def set_qheader_section_resize_mode(header):
        return header.setSectionResizeMode

class TableEditor(Editor, BaseTableEditor):
    """Editor that presents data in a table. Optionally, tables can have
    a set of filters that reduce the set of data displayed, according to
    their criteria.
    """

    # -------------------------------------------------------------------------
    #  Trait definitions:
    # -------------------------------------------------------------------------

    #: The table view control associated with the editor:
    table_view = Any()

    def _table_view_default(self):
        return TableView(editor=self)

    #: A wrapper around the source model which provides filtering and sorting:
    model = Instance(SortFilterTableModel)

    def _model_default(self):
        return SortFilterTableModel(editor=self)

    #: The table model associated with the editor:
    source_model = Instance(TableModel)

    def _source_model_default(self):
        return TableModel(editor=self)

    #: The set of columns currently defined on the editor:
    columns = List(TableColumn)

    #: The currently selected row(s), column(s), or cell(s).
    selected = Any()

    #: The current selected row
    selected_row = Property(Any, observe="selected")

    selected_indices = Property(Any, observe="selected")

    #: Current filter object (should be a TableFilter or callable or None):
    filter = Any()

    #: The indices of the table items currently passing the table filter:
    filtered_indices = List(Int)

    #: Current filter summary message
    filter_summary = Str("All items")

    #: Update the filtered contents.
    update_filter = Event()

    #: The event fired when a cell is clicked on:
    click = Event()

    #: The event fired when a cell is double-clicked on:
    dclick = Event()

    #: The Traits UI associated with the table editor toolbar:
    toolbar_ui = Instance(UI)

    #: The index of the row that was last right clicked on its vertical header
    header_row = Int()

    #: Whether to auto-size the columns or not.
    auto_size = Bool(False)

    #: Dictionary mapping image names to QIcons
    images = Dict()

    #: Dictionary mapping ImageResource objects to QIcons
    image_resources = Dict()

    #: An image being converted:
    image = Image

    def init(self, parent):
        """Finishes initializing the editor by creating the underlying toolkit
        widget."""

        factory = self.factory
        self.filter = factory.filter

        columns = factory.columns[:]
        if (len(columns) == 0) and (len(self.value) > 0):
            columns = [
                ObjectColumn(name=name)
                for name in self.value[0].editable_traits()
            ]
        self.columns = columns

        if factory.table_view_factory is not None:
            self.table_view = factory.table_view_factory(editor=self)
        if factory.source_model_factory is not None:
            self.source_model = factory.source_model_factory(editor=self)
        if factory.model_factory is not None:
            self.model = factory.model_factory(editor=self)

        # Create the table view and model
        self.model.setDynamicSortFilter(True)
        self.model.setSourceModel(self.source_model)
        self.table_view.setModel(self.model)

        # When sorting is enabled, the first column is initially displayed with
        # the triangle indicating it is the sort index, even though no sorting
        # has actually been done. Sort here for UI/model consistency.
        if self.factory.sortable and not self.factory.reorderable:
            self.model.sort(0, QtCore.Qt.SortOrder.AscendingOrder)

        # Connect to the mode specific selection handler and select the first
        # row/column/cell. Do this before creating the edit_view to make sure
        # that it has a valid item to use when constructing its view.
        smodel = self.table_view.selectionModel()
        mode_slot = getattr(self, "_on_%s_selection" % factory.selection_mode)
        smodel.selectionChanged.connect(mode_slot)
        self.table_view.setCurrentIndex(self.model.index(0, 0))

        # Create the toolbar if necessary
        if factory.show_toolbar and len(factory.filters) > 0:
            main_view = QtGui.QWidget()
            layout = QtGui.QVBoxLayout(main_view)
            layout.setContentsMargins(0, 0, 0, 0)
            self.toolbar_ui = self.edit_traits(
                parent=parent,
                kind="subpanel",
                view=View(
                    Group(
                        Item("filter{View}", editor=factory._filter_editor),
                        Item("filter_summary{Results}", style="readonly"),
                        spring,
                        orientation="horizontal",
                    ),
                    resizable=True,
                ),
            )
            self.toolbar_ui.parent = self.ui
            layout.addWidget(self.toolbar_ui.control)
            layout.addWidget(self.table_view)
        else:
            main_view = self.table_view

        # Create auxiliary editor and encompassing splitter if necessary
        mode = factory.selection_mode
        if (factory.edit_view == " ") or mode not in {"row", "rows"}:
            self.control = main_view
        else:
            if factory.orientation == "horizontal":
                self.control = QtGui.QSplitter(QtCore.Qt.Orientation.Horizontal)
            else:
                self.control = QtGui.QSplitter(QtCore.Qt.Orientation.Vertical)
            self.control.setSizePolicy(
                QtGui.QSizePolicy.Policy.Expanding, QtGui.QSizePolicy.Policy.Expanding
            )
            self.control.addWidget(main_view)
            self.control.setStretchFactor(0, 2)

            # Create the row editor below the table view
            editor = InstanceEditor(view=factory.edit_view, kind="subpanel")
            self._ui = self.edit_traits(
                parent=self.control,
                kind="subpanel",
                view=View(
                    Item(
                        "selected_row",
                        style="custom",
                        editor=editor,
                        show_label=False,
                        resizable=True,
                        width=factory.edit_view_width,
                        height=factory.edit_view_height,
                    ),
                    resizable=True,
                    handler=factory.edit_view_handler,
                ),
            )
            self._ui.parent = self.ui
            self.control.addWidget(self._ui.control)
            self.control.setStretchFactor(1, 1)

        # Connect to the click and double click handlers
        self.table_view.clicked.connect(self._on_click)
        self.table_view.doubleClicked.connect(self._on_dclick)

        # Make sure we listen for 'items' changes as well as complete list
        # replacements
        self.context_object.on_trait_change(
            self.update_editor, self.extended_name + "_items", dispatch="ui"
        )

        # Listen for changes to traits on the objects in the list
        self.context_object.on_trait_change(
            self.refresh_editor, self.extended_name + ".-", dispatch="ui"
        )

        # Listen for changes on column definitions
        self.on_trait_change(self._update_columns, "columns", dispatch="ui")
        self.on_trait_change(
            self._update_columns, "columns_items", dispatch="ui"
        )

        # Set up the required externally synchronized traits
        is_list = mode in ("rows", "columns", "cells")
        self.sync_value(factory.click, "click", "to")
        self.sync_value(factory.dclick, "dclick", "to")
        self.sync_value(factory.columns_name, "columns", is_list=True)
        self.sync_value(factory.selected, "selected", is_list=is_list)
        self.sync_value(
            factory.selected_indices, "selected_indices", is_list=is_list
        )
        self.sync_value(factory.filter_name, "filter", "from")
        self.sync_value(factory.filtered_indices, "filtered_indices", "to")
        self.sync_value(factory.update_filter_name, "update_filter", "from")

        self.auto_size = self.factory.auto_size

        # Initialize the ItemDelegates for each column
        self._update_columns()

    def dispose(self):
        """Disposes of the contents of an editor."""
        self.model.beginResetModel()
        self.model.endResetModel()

        # Make sure that the auxiliary UIs are properly disposed
        if self.toolbar_ui is not None:
            self.toolbar_ui.dispose()
        if self._ui is not None:
            self._ui.dispose()

        # Remove listener for 'items' changes on object trait
        self.context_object.on_trait_change(
            self.update_editor, self.extended_name + "_items", remove=True
        )

        # Remove listener for changes to traits on the objects in the list
        self.context_object.on_trait_change(
            self.refresh_editor, self.extended_name + ".-", remove=True
        )

        # Remove listeners for column definition changes
        self.on_trait_change(self._update_columns, "columns", remove=True)
        self.on_trait_change(
            self._update_columns, "columns_items", remove=True
        )

        super().dispose()

    def update_editor(self):
        """Updates the editor when the object trait changes externally to the
        editor."""

        if self._no_notify:
            return

        self.table_view.setUpdatesEnabled(False)
        try:
            filtering = (
                len(self.factory.filters) > 0 or self.filter is not None
            )
            if filtering:
                self._update_filtering()

            # invalidate the model, but do not reset it. Resetting the model
            # may cause problems if the selection sync'ed traits are being used
            # externally to manage the selections
            self.model.invalidate()

            self.table_view.resizeColumnsToContents()
            if self.auto_size:
                self.table_view.resizeRowsToContents()

        finally:
            self.table_view.setUpdatesEnabled(True)

    def restore_prefs(self, prefs):
        """Restores any saved user preference information associated with the
        editor.
        """
        header = self.table_view.horizontalHeader()
        if header is not None and "column_state" in prefs:
            header.restoreState(prefs["column_state"])

    def save_prefs(self):
        """Returns any user preference information associated with the editor."""
        prefs = {}
        header = self.table_view.horizontalHeader()
        if header is not None:
            prefs["column_state"] = header.saveState().data()
        return prefs

    def refresh_editor(self):
        """Requests that the underlying table widget to redraw itself."""

        self.table_view.viewport().update()

    def create_new_row(self):
        """Creates a new row object using the provided factory."""

        factory = self.factory
        kw = factory.row_factory_kw.copy()
        if "__table_editor__" in kw:
            kw["__table_editor__"] = self

        return self.ui.evaluate(
            factory.row_factory, *factory.row_factory_args, **kw
        )

    def items(self):
        """Returns the raw list of model objects."""

        items = self.value
        if not isinstance(items, SequenceTypes):
            items = [items]

        if self.factory and self.factory.reverse:
            items = ReversedList(items)

        return items

    def callx(self, func, *args, **kw):
        """Call a function without notifying the underlying table view or
        model."""

        old = self._no_notify
        self._no_notify = True
        try:
            func(*args, **kw)
        finally:
            self._no_notify = old

    def setx(self, **keywords):
        """Set one or more attributes without notifying the underlying table
        view or model."""

        old = self._no_notify
        self._no_notify = True
        try:
            for name, value in keywords.items():
                setattr(self, name, value)
        finally:
            self._no_notify = old

    def set_selection(self, objects=[], notify=True):
        """Sets the current selection to a set of specified objects."""

        if not isinstance(objects, list):
            objects = [objects]

        mode = self.factory.selection_mode
        indexes = []
        flags = QtGui.QItemSelectionModel.SelectionFlag.ClearAndSelect

        # In the case of row or column selection, we need a dummy value for the
        # other dimension that has not been filtered.
        source_index = self.model.mapToSource(self.model.index(0, 0))
        source_row, source_column = source_index.row(), source_index.column()

        # Selection mode is 'row' or 'rows'
        if mode.startswith("row"):
            flags |= QtGui.QItemSelectionModel.SelectionFlag.Rows
            items = self.items()
            for obj in objects:
                try:
                    row = items.index(obj)
                except ValueError:
                    continue
                indexes.append(self.source_model.index(row, source_column))

        # Selection mode is 'column' or 'columns'
        elif mode.startswith("column"):
            flags |= QtGui.QItemSelectionModel.SelectionFlag.Columns
            for name in objects:
                column = self._column_index_from_name(name)
                if column != -1:
                    indexes.append(self.source_model.index(source_row, column))

        # Selection mode is 'cell' or 'cells'
        else:
            items = self.items()
            for obj, name in objects:
                try:
                    row = items.index(obj)
                except ValueError:
                    continue
                column = self._column_index_from_name(name)
                if column != -1:
                    indexes.append(self.source_model.index(row, column))

        # Perform the selection so that only one signal is emitted
        selection = QtGui.QItemSelection()
        smodel = self.table_view.selectionModel()
        if smodel is None:
            # guard against selection during tear-down
            return
        for index in indexes:
            index = self.model.mapFromSource(index)
            if index.isValid():
                smodel.setCurrentIndex(
                    index, QtGui.QItemSelectionModel.SelectionFlag.NoUpdate
                )
                selection.select(index, index)

        smodel.blockSignals(not notify)
        try:
            if len(selection.indexes()):
                smodel.clear()
                smodel.select(selection, flags)
            else:
                smodel.clear()
        finally:
            smodel.blockSignals(False)

        self.refresh_editor()

    # -------------------------------------------------------------------------
    #  Private methods:
    # -------------------------------------------------------------------------

    def _column_index_from_name(self, name):
        """Returns the index of the column with the given name or -1 if no
        column exists with that name."""

        for i, column in enumerate(self.columns):
            if name == column.name:
                return i
        return -1

    def _customize_filters(self, filter):
        """Allows the user to customize the current set of table filters."""

        filter_editor = TableFilterEditor(editor=self)
        ui = filter_editor.edit_traits(parent=self.control)
        if ui.result:
            self.factory.filters = filter_editor.templates
            self.filter = filter_editor.selected_filter
        else:
            self.setx(filter=filter)

    def _update_filtering(self):
        """Update the filter summary and the filtered indices."""

        items = self.items()
        num_items = len(items)

        f = self.filter
        if f is None:
            self._filtered_cache = None
            self.filtered_indices = list(range(num_items))
            self.filter_summary = "All %i items" % num_items
        else:
            if not callable(f):
                f = f.filter
            self._filtered_cache = fc = [f(item) for item in items]
            self.filtered_indices = fi = [i for i, ok in enumerate(fc) if ok]
            self.filter_summary = "%i of %i items" % (len(fi), num_items)

    def _add_image(self, image_resource):
        """Adds a new image to the image map."""
        image = image_resource.create_icon()

        self.image_resources[image_resource] = image
        self.images[image_resource.name] = image

        return image

    def _get_image(self, image):
        """Converts a user specified image to a QIcon."""
        if isinstance(image, str):
            self.image = image
            image = self.image

        if isinstance(image, ImageResource):
            result = self.image_resources.get(image)
            if result is not None:
                return result
            return self._add_image(image)

        return self.images.get(image)

    def _create_empty_menu(self):
        """Create a QMenu to display in empty space below the rows.

        Returns a QMenu or None if no menu to display.
        """
        if not self.factory.editable or self.factory.row_factory is None:
            return None

        empty_menu = QtGui.QMenu(self.table_view)
        action = empty_menu.addAction("Add new item")
        action.triggered.connect(self._on_context_append)
        return empty_menu

    def _create_header_menu(self):
        """Create a QMenu to display in the vertical header.

        Returns a QMenu or None if no menu to display.
        """
        header_menu = QtGui.QMenu(self.table_view)
        if self.factory.editable:
            if self.factory.row_factory is not None:
                action = header_menu.addAction("Insert new item")
                action.triggered.connect(self._on_context_insert)
            if self.factory.deletable:
                action = header_menu.addAction("Delete item")
                action.triggered.connect(self._on_context_remove)
        if self.factory.reorderable:
            show_up = (self.header_row > 0)
            show_down = (self.header_row < self.model.rowCount() - 1)
            if not header_menu.isEmpty() and (show_up or show_down):
                header_menu.addSeparator()
            if show_up:
                header_menu_up = header_menu.addAction("Move item up")
                header_menu_up.triggered.connect(self._on_context_move_up)
            if show_down:
                header_menu_down = header_menu.addAction("Move item down")
                header_menu_down.triggered.connect(self._on_context_move_down)

        if header_menu.isEmpty():
            return None
        else:
            return header_menu

    # -- Trait Property getters/setters ---------------------------------------

    @cached_property
    def _get_selected_row(self):
        """Gets the selected row, or the first row if multiple rows are
        selected."""

        mode = self.factory.selection_mode

        if mode.startswith("column"):
            return None
        elif mode == "row":
            return self.selected

        try:
            if mode == "rows":
                return self.selected[0]
            elif mode == "cell":
                return self.selected[0]
            elif mode == "cells":
                return self.selected[0][0]
        except IndexError:
            return None

    @cached_property
    def _get_selected_indices(self):
        """Gets the row,column indices which match the selected trait"""
        selection_items = self.table_view.selectionModel().selection()
        indices = self.model.mapSelectionToSource(selection_items).indexes()
        if self.factory.selection_mode.startswith("row"):
            indices = sorted(set(index.row() for index in indices))
        elif self.factory.selection_mode.startswith("column"):
            indices = sorted(set(index.column() for index in indices))
        else:
            indices = [(index.row(), index.column()) for index in indices]

        if self.factory.selection_mode in {"rows", "columns", "cells"}:
            return indices
        elif len(indices) > 0:
            return indices[0]
        else:
            return -1

    def _set_selected_indices(self, indices):
        if not isinstance(indices, list):
            indices = [indices]
        selected = []
        if self.factory.selection_mode.startswith("row"):
            for row in indices:
                selected.append(self.value[row])
        elif self.factory.selection_mode.startswith("column"):
            for col in indices:
                selected.append(self.columns[col].name)
        else:
            for row, col in indices:
                selected.append((self.value[row], self.columns[col].name))

        self.selected = selected
        self.set_selection(self.selected, False)

    # -- Trait Change Handlers ------------------------------------------------

    def _filter_changed(self, old_filter, new_filter):
        """Handles the current filter being changed."""

        if not self._no_notify:
            if new_filter is customize_filter:
                do_later(self._customize_filters, old_filter)
            else:
                self._update_filtering()
                self.model.invalidate()
                self.set_selection(self.selected)

    def _update_columns(self):
        """Handle the column list being changed."""

        self.table_view.setItemDelegate(TableDelegate(self.table_view))
        for i, column in enumerate(self.columns):
            if column.renderer:
                self.table_view.setItemDelegateForColumn(i, column.renderer)

        self.model.invalidate()
        self.table_view.resizeColumnsToContents()
        if self.auto_size:
            self.table_view.resizeRowsToContents()

    def _selected_changed(self, new):
        """Handle the selected row/column/cell being changed externally."""
        if not self._no_notify:
            self.set_selection(self.selected, notify=False)

    def _update_filter_changed(self):
        """The filter has changed internally."""
        self._filter_changed(self.filter, self.filter)

    # -- Event Handlers -------------------------------------------------------

    def _on_row_selection(self, added, removed):
        """Handle the row selection being changed."""

        items = self.items()
        indexes = self.table_view.selectionModel().selectedRows()
        if len(indexes):
            index = self.model.mapToSource(indexes[0])
            selected = items[index.row()]
        else:
            selected = None

        self.setx(selected=selected)
        self.ui.evaluate(self.factory.on_select, self.selected)

    def _on_rows_selection(self, added, removed):
        """Handle the rows selection being changed."""
        items = self.items()
        indexes = self.table_view.selectionModel().selectedRows()
        selected = [
            items[self.model.mapToSource(index).row()] for index in indexes
        ]

        self.setx(selected=selected)
        self.ui.evaluate(self.factory.on_select, self.selected)

    def _on_column_selection(self, added, removed):
        """Handle the column selection being changed."""

        indexes = self.table_view.selectionModel().selectedColumns()
        if len(indexes):
            index = self.model.mapToSource(indexes[0])
            selected = self.columns[index.column()].name
        else:
            selected = ""

        self.setx(selected=selected)
        self.ui.evaluate(self.factory.on_select, self.selected)

    def _on_columns_selection(self, added, removed):
        """Handle the columns selection being changed."""

        indexes = self.table_view.selectionModel().selectedColumns()
        selected = [
            self.columns[self.model.mapToSource(index).column()].name
            for index in indexes
        ]

        self.setx(selected=selected)
        self.ui.evaluate(self.factory.on_select, self.selected)

    def _on_cell_selection(self, added, removed):
        """Handle the cell selection being changed."""

        items = self.items()
        indexes = self.table_view.selectionModel().selectedIndexes()
        if len(indexes):
            index = self.model.mapToSource(indexes[0])
            obj = items[index.row()]
            column_name = self.columns[index.column()].name
        else:
            obj = None
            column_name = ""
        selected = (obj, column_name)

        self.setx(selected=selected)
        self.ui.evaluate(self.factory.on_select, self.selected)

    def _on_cells_selection(self, added, removed):
        """Handle the cells selection being changed."""

        items = self.items()
        indexes = self.table_view.selectionModel().selectedIndexes()
        selected = []
        for index in indexes:
            index = self.model.mapToSource(index)
            obj = items[index.row()]
            column_name = self.columns[index.column()].name
            selected.append((obj, column_name))

        self.setx(selected=selected)
        self.ui.evaluate(self.factory.on_select, self.selected)

    def _on_click(self, index):
        """Handle a cell being clicked."""

        index = self.model.mapToSource(index)
        column = self.columns[index.column()]
        obj = self.items()[index.row()]

        # Fire the same event on the editor after mapping it to a model object
        # and column name:
        self.click = (obj, column)

        # Invoke the column's click handler:
        column.on_click(obj)

    def _on_dclick(self, index):
        """Handle a cell being double clicked."""

        index = self.model.mapToSource(index)
        column = self.columns[index.column()]
        obj = self.items()[index.row()]

        # Fire the same event on the editor after mapping it to a model object
        # and column name:
        self.dclick = (obj, column)

        # Invoke the column's double-click handler:
        column.on_dclick(obj)

    def _on_context_insert(self):
        """Handle 'insert item' being selected from the header context menu."""

        self.model.insertRow(self.header_row)

    def _on_context_append(self):
        """Handle 'add item' being selected from the empty space context
        menu."""

        self.model.insertRow(self.model.rowCount())

    def _on_context_remove(self):
        """Handle 'remove item' being selected from the header context menu."""

        self.model.removeRow(self.header_row)

    def _on_context_move_up(self):
        """Handle 'move up' being selected from the header context menu."""

        self.model.moveRow(self.header_row, self.header_row - 1)

    def _on_context_move_down(self):
        """Handle 'move down' being selected from the header context menu."""

        self.model.moveRow(self.header_row, self.header_row + 1)

# Define the SimpleEditor class.
SimpleEditor = TableEditor

# Define the ReadonlyEditor class.
ReadonlyEditor = TableEditor

# -------------------------------------------------------------------------
#  Qt widgets that have been configured to behave as expected by Traits UI:
# -------------------------------------------------------------------------

class TableDelegate(QtGui.QStyledItemDelegate):
    """A QStyledItemDelegate which fetches Traits UI editors."""

    def createEditor(self, parent, option, index):
        """Reimplemented to return the editor for a given index."""

        model = index.model()
        index = model.mapToSource(index)
        table_editor = model._editor
        column = table_editor.columns[index.column()]
        obj = table_editor.items()[index.row()]

        factory = column.get_editor(obj)
        style = column.get_style(obj)
        if factory is None:
            return None

        target, name = column.target_name(obj)
        handler = default_handler()
        if table_editor.ui.context is None:
            ui = UI(handler=handler)
        else:
            context = table_editor.ui.context.copy()
            context["table_editor_object"] = context["object"]
            context["object"] = target
            ui = UI(handler=handler, context=context)

        # Create and initialize the editor
        factory_method = getattr(factory, style + "_editor")
        editor = factory_method(ui, target, name, "", parent)
        editor.prepare(parent)
        control = editor.control
        control.setParent(parent)

        # Required for QMouseEvents to propagate to the widget
        control.setFocusPolicy(QtCore.Qt.FocusPolicy.StrongFocus)

        # The table view's background will shine through unless the editor
        # paints its own background
        control.setAutoFillBackground(True)

        # Make sure that editors are disposed of correctly
        # will be disposed in closeEditor of the TableView
        control._editor = editor
        return control

    def updateEditorGeometry(self, editor, option, index):
        """Update the editor's geometry."""
        editor.setGeometry(option.rect)

    def paint(self, painter, option, index):
        self.initStyleOption(option, index)

        if (option.state & QtGui.QStyle.StateFlag.State_Selected) and (
            option.state & QtGui.QStyle.StateFlag.State_Active
        ):
            factory = self.parent()._editor.factory
            if factory.selection_bg_color is not None:
                option.palette.setColor(
                    QtGui.QPalette.ColorRole.Highlight, factory.selection_bg_color_
                )
            if factory.selection_color is not None:
                option.palette.setColor(
                    QtGui.QPalette.ColorRole.HighlightedText, factory.selection_color_
                )

        QtGui.QApplication.style().drawControl(
            QtGui.QStyle.ControlElement.CE_ItemViewItem, option, painter, None
        )

class TableView(QtGui.QTableView):
    """A QTableView configured to behave as expected by TraitsUI."""

    _SELECTION_MAP = {
        "row": (
            QtGui.QAbstractItemView.SelectionBehavior.SelectRows,
            QtGui.QAbstractItemView.SelectionMode.SingleSelection,
        ),
        "rows": (
            QtGui.QAbstractItemView.SelectionBehavior.SelectRows,
            QtGui.QAbstractItemView.SelectionMode.ExtendedSelection,
        ),
        "column": (
            QtGui.QAbstractItemView.SelectionBehavior.SelectColumns,
            QtGui.QAbstractItemView.SelectionMode.SingleSelection,
        ),
        "columns": (
            QtGui.QAbstractItemView.SelectionBehavior.SelectColumns,
            QtGui.QAbstractItemView.SelectionMode.ExtendedSelection,
        ),
        "cell": (
            QtGui.QAbstractItemView.SelectionBehavior.SelectItems,
            QtGui.QAbstractItemView.SelectionMode.SingleSelection,
        ),
        "cells": (
            QtGui.QAbstractItemView.SelectionBehavior.SelectItems,
            QtGui.QAbstractItemView.SelectionMode.ExtendedSelection,
        ),
    }

    def __init__(self, editor):
        """Initialise the object."""

        QtGui.QTableView.__init__(self)

        self._initial_size = False
        self._editor = editor
        factory = editor.factory

        # Configure the grid lines.
        self.setShowGrid(factory.show_lines)

        # Configure the selection behaviour.
        self.setCornerButtonEnabled(False)
        behav, mode = self._SELECTION_MAP[factory.selection_mode]
        self.setSelectionBehavior(behav)
        self.setSelectionMode(mode)

        # Configure the editing behavior.
        triggers = (
            QtGui.QAbstractItemView.EditTrigger.DoubleClicked
            | QtGui.QAbstractItemView.EditTrigger.SelectedClicked
        )
        if factory.edit_on_first_click and not factory.reorderable:
            triggers |= QtGui.QAbstractItemView.EditTrigger.CurrentChanged
        self.setEditTriggers(triggers)

        # Configure the reordering and sorting behavior.
        self.setDragEnabled(True)
        self.viewport().setAcceptDrops(True)
        self.setDropIndicatorShown(True)

        if factory.reorderable:
            self.setDragDropMode(QtGui.QAbstractItemView.DragDropMode.InternalMove)
        if factory.sortable:
            self.setSortingEnabled(True)

        if factory._qt_stylesheet is not None:
            self.setStyleSheet(factory._qt_stylesheet)

        self.resizeColumnsToContents()

    def setModel(self, model):
        super().setModel(model)
        self._update_header_sizing()

    def contextMenuEvent(self, event):
        """Reimplemented to create context menus for cells and empty space."""

        # Determine the logical indices of the cell where click occured
        hheader, vheader = self.horizontalHeader(), self.verticalHeader()
        position = event.globalPos()
        row = vheader.logicalIndexAt(vheader.mapFromGlobal(position))
        column = hheader.logicalIndexAt(hheader.mapFromGlobal(position))

        # Map the logical row index to a real index for the source model
        model = self.model()
        row = model.mapToSource(model.index(row, 0)).row()

        # Show a context menu for empty space at bottom of table...
        editor = self._editor
        if row == -1:
            empty_menu = editor._create_empty_menu()
            if empty_menu is not None:
                event.accept()
                empty_menu.exec_(position)

        # ...or show a context menu for a cell.
        elif column != -1:
            obj = editor.items()[row]
            column = editor.columns[column]
            menu_manager = column.get_menu(obj)
            if menu_manager is None:
                menu_manager = editor.factory.menu
            if menu_manager is not None:
                event.accept()
                selected = editor.selected
                if not isinstance(selected, SequenceTypes):
                    selected = [selected]
                if obj not in selected:
                    selected = [obj]
                editor.set_menu_context(selected, obj, column)
                menu = menu_manager.create_menu(self, controller=editor)
                menu.exec_(position)

    def eventFilter(self, obj, event):
        """Reimplemented to create context menu for the vertical header."""

        vheader = self.verticalHeader()
        if obj is vheader and event.type() == QtCore.QEvent.Type.ContextMenu:
            position = event.globalPos()
            editor = self._editor
            row = vheader.logicalIndexAt(event.pos().y())
            if row == -1:
                empty_menu = editor._create_empty_menu()
                if empty_menu is not None:
                    event.accept()
                    empty_menu.exec_(position)
            else:
                editor.header_row = row
                header_menu = editor._create_header_menu()
                if header_menu is not None:
                    event.accept()
                    header_menu.exec_(position)
            return True

        else:
            return QtGui.QTableView.eventFilter(self, obj, event)

    def resizeEvent(self, event):
        """Reimplemented to size the table columns when the size of the table
        changes. Because the layout algorithm requires that the available
        space be known, we have to wait until the UI that contains this table
        gives it its initial size."""

        QtGui.QTableView.resizeEvent(self, event)

        if self._editor.auto_size:
            self.resizeColumnsToContents()
            self.resizeRowsToContents()

        else:
            parent = self.parent()
            if (
                not self._initial_size
                and parent
                and (self.isVisible() or isinstance(parent, QtGui.QMainWindow))
            ):
                self._initial_size = True
                if self._editor.auto_size:
                    self.resizeColumnsToContents()
                    self.resizeRowsToContents()

    def sizeHint(self):
        """Reimplemented to define a better size hint for the width of the
        TableEditor."""

        size_hint = QtGui.QTableView.sizeHint(self)

        # This method is sometimes called by Qt after the editor has been
        # disposed but before this control has been deleted:
        if self._editor.factory is None:
            return size_hint

        width = self.style().pixelMetric(
            QtGui.QStyle.PixelMetric.PM_ScrollBarExtent, QtGui.QStyleOptionHeader(), self
        )
        for column in range(len(self._editor.columns)):
            width += self.sizeHintForColumn(column)
        size_hint.setWidth(width)
        return size_hint

    def sizeHintForColumn(self, column_index):
        """Reimplemented to support absolute width specification via
        TableColumns and to improve the metric for autosizing columns."""

        editor = self._editor
        column = editor.columns[column_index]
        requested_width = column.get_width()

        # Autosize based on column contents and label width. Qt's default
        # implementation of this function does content, we handle the label.
        if requested_width < 1:
            base_width = QtGui.QTableView.sizeHintForColumn(self, column_index)

            # Determine what font to use in the calculation
            font = column.get_text_font(None)
            if font is None:
                font = self.font()
                font.setBold(True)
            else:
                font = QtGui.QFont(font)

            # Determine the width of the column label
            text = column.get_label()
            # QFontMetrics.width() is deprecated and Qt docs suggest using
            # horizontalAdvance() instead, but is only available since Qt 5.11
            if QtCore.__version_info__ >= (5, 11):
                width = QtGui.QFontMetrics(font).horizontalAdvance(text)
            else:
                width = QtGui.QFontMetrics(font).width(text)

            # Add margin to the calculated width as appropriate
            style = self.style()
            option = QtGui.QStyleOptionHeader()
            width += (
                style.pixelMetric(
                    QtGui.QStyle.PixelMetric.PM_HeaderGripMargin, option, self
                )
                * 2
            )
            if editor.factory.sortable and not editor.factory.reorderable:
                # Add size of sort indicator
                width += style.pixelMetric(
                    QtGui.QStyle.PixelMetric.PM_HeaderMarkSize, option, self
                )
                # Add distance between sort indicator and text
                width += style.pixelMetric(
                    QtGui.QStyle.PixelMetric.PM_HeaderMargin, option, self
                )
            return max(base_width, width)

        # Or else set width absolutely
        else:
            return requested_width

    def resizeColumnsToContents(self):
        """Support proportional column width specifications."""

        # TODO: The proportional size specification approach found in the
        # TableColumns is not entirely compatible with the ability to
        # specify the resize_mode.  Namely, there are combinations of
        # specifications that are redundant, and others which are
        # contradictory.  Rework this method so that the various values
        # for **width** have a well-defined, sensible meaning for each
        # of the possible values of resize_mode.

        editor = self._editor
        available_space = self.viewport().width()
        hheader = self.horizontalHeader()

        # Compute sizes for columns with absolute or no size requests
        proportional = []
        for column_index in range(len(editor.columns)):
            column = editor.columns[column_index]
            requested_width = column.get_width()
            if (
                column.resize_mode in ("interactive", "stretch")
                and 0 < requested_width <= 1.0
            ):
                proportional.append((column_index, requested_width))
            elif (
                column.resize_mode == "interactive"
                and requested_width < 0
                and self._initial_size
            ):
                # Keep previous size if initial sizing has been done
                available_space -= hheader.sectionSize(column_index)
            else:
                base_width = hheader.sectionSizeHint(column_index)
                width = max(base_width, self.sizeHintForColumn(column_index))
                hheader.resizeSection(column_index, width)
                available_space -= width

        # Now use the remaining space for columns with proportional width
        # requests
        for column_index, percent in proportional:
            base_width = hheader.sectionSizeHint(column_index)
            width = max(base_width, int(percent * available_space))
            hheader.resizeSection(column_index, width)

    def closeEditor(self, control, hint):
        # dispose traits editor associated with control if any
        editor = getattr(control, "_editor", None)
        if editor is not None:
            editor.dispose()
            delattr(control, "_editor")

        return super().closeEditor(control, hint)

    def _update_header_sizing(self):
        """Header section sizing can be done only after a valid model is set.
        Otherwise results in segfault with Qt5.
        """
        editor = self._editor
        factory = editor.factory
        # Configure the row headings.
        vheader = self.verticalHeader()
        set_resize_mode = set_qheader_section_resize_mode(vheader)
        insertable = factory.row_factory is not None
        if (
            factory.editable and (insertable or factory.deletable)
        ) or factory.reorderable:
            vheader.installEventFilter(self)
            set_resize_mode(QtGui.QHeaderView.ResizeMode.ResizeToContents)
        elif not factory.show_row_labels:
            vheader.hide()
        if factory.row_height > 0:
            vheader.setDefaultSectionSize(factory.row_height)
        self.setAlternatingRowColors(factory.alternate_bg_color)
        self.setHorizontalScrollMode(QtGui.QAbstractItemView.ScrollMode.ScrollPerPixel)
        # Configure the column headings.
        # We detect if there are any stretchy sections at all; if not, then
        # we make the last non-fixed-size column stretchy.
        hheader = self.horizontalHeader()
        set_resize_mode = set_qheader_section_resize_mode(hheader)
        resize_mode_map = dict(
            interactive=QtGui.QHeaderView.ResizeMode.Interactive,
            fixed=QtGui.QHeaderView.ResizeMode.Fixed,
            stretch=QtGui.QHeaderView.ResizeMode.Stretch,
            resize_to_contents=QtGui.QHeaderView.ResizeMode.ResizeToContents,
        )
        stretchable_columns = []
        for i, column in enumerate(editor.columns):
            set_resize_mode(i, resize_mode_map[column.resize_mode])
            if column.resize_mode in ("stretch", "interactive"):
                stretchable_columns.append(i)
        if not stretchable_columns:
            # Use the behavior from before the "resize_mode" trait was added
            # to TableColumn
            hheader.setStretchLastSection(True)
        else:
            # hheader.setSectionResizeMode(
            #     stretchable_columns[-1], QtGui.QHeaderView.ResizeMode.Stretch)
            hheader.setStretchLastSection(False)
        if factory.show_column_labels:
            hheader.setHighlightSections(False)
        else:
            hheader.hide()

# -------------------------------------------------------------------------
#  Editor for configuring the filters available to a TableEditor:
# -------------------------------------------------------------------------

class TableFilterEditor(HasTraits):
    """An editor that manages table filters."""

    # -------------------------------------------------------------------------
    #  Trait definitions:
    # -------------------------------------------------------------------------

    #: TableEditor this editor is associated with
    editor = Instance(TableEditor)

    #: The list of filters
    filters = List(TableFilter)

    #: The list of available templates from which filters can be created
    templates = Property(List(TableFilter), observe="filters")

    #: The currently selected filter template
    selected_template = Instance(TableFilter)

    #: The currently selected filter
    selected_filter = Instance(TableFilter, allow_none=True)

    #: The view to use for the current filter
    selected_filter_view = Property(observe="selected_filter")

    #: Buttons for add/removing filters
    add_button = Button("New")
    remove_button = Button("Delete")

    # The default view for this editor
    view = View(
        Group(
            Group(
                Group(
                    Item("add_button", enabled_when="selected_template"),
                    Item(
                        "remove_button",
                        enabled_when="len(templates) > 1 and "
                        "selected_filter is not None",
                    ),
                    orientation="horizontal",
                    show_labels=False,
                ),
                Label("Base filter for new filters:"),
                Item("selected_template", editor=EnumEditor(name="templates")),
                Item(
                    "selected_filter",
                    style="custom",
                    editor=EnumEditor(name="filters", mode="list"),
                ),
                show_labels=False,
            ),
            Item(
                "selected_filter",
                width=0.75,
                style="custom",
                editor=InstanceEditor(view_name="selected_filter_view"),
            ),
            id="TableFilterEditorSplit",
            show_labels=False,
            layout="split",
            orientation="horizontal",
        ),
        id="traitsui.qt.table_editor.TableFilterEditor",
        buttons=["OK", "Cancel"],
        kind="livemodal",
        resizable=True,
        width=800,
        height=400,
        title="Customize filters",
    )

    # -------------------------------------------------------------------------
    #  Private methods:
    # -------------------------------------------------------------------------

    # -- Trait Property getter/setters ----------------------------------------

    @cached_property
    def _get_selected_filter_view(self):
        view = None
        if self.selected_filter:
            model = self.editor.model
            index = model.mapToSource(model.index(0, 0))
            if index.isValid():
                obj = self.editor.items()[index.row()]
            else:
                obj = None
            view = self.selected_filter.edit_view(obj)
        return view

    @cached_property
    def _get_templates(self):
        templates = [f for f in self.editor.factory.filters if f.template]
        templates.extend(self.filters)
        return templates

    # -- Trait Change Handlers ------------------------------------------------

    def _editor_changed(self):
        self.filters = [
            f.clone_traits()
            for f in self.editor.factory.filters
            if not f.template
        ]
        self.selected_template = self.templates[0]

    @observe('add_button')
    def _create_and_select_new_filter(self, event):
        """Create a new filter based on the selected template and select it."""
        new_filter = self.selected_template.clone_traits()
        new_filter.template = False
        new_filter.name = new_filter._name = "New filter"
        self.filters.append(new_filter)
        self.selected_filter = new_filter

    @observe("remove_button")
    def _delete_selected_filter(self, event):
        """Delete the currently selected filter."""
        if self.selected_template == self.selected_filter:
            self.selected_template = self.templates[0]

        index = self.filters.index(self.selected_filter)
        del self.filters[index]
        if index < len(self.filters):
            self.selected_filter = self.filters[index]
        else:
            self.selected_filter = None

    @observe("selected_filter:name")
    def _update_filter_list(self, event):
        """A hack to make the EnumEditor watching the list of filters refresh
        their text when the name of the selected filter changes.
        """
        filters = self.filters
        self.filters = []
        self.filters = filters


import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_regression
from mpl_toolkits import mplot3d

# Training phase/ Training the LR model/ Find optimal weights
def fit(X, y):
    """
    X: Feature matrix: (n_samples, n_features)
    y: y_true: (n_samples,1)
    Returns: weights
    weights: optimal weights (n_features, 1)
    """
    X = X.copy()
    ones_column = np.ones((len(X),1))
    X = np.concatenate([ones_column, X], axis=1)

    w = np.linalg.inv(X.T.dot(X)).dot(X.T).dot(y)
    return w

# prediction
def predict(X, w):
    """
    X: Feature matrix: (n_samples, n_features)
    w: weight vector: (n_fetures, 1)
    Returns:
    y: y_pred = X.w (n_samples,1)
    """
    X = X.copy()
    ones_column = np.ones((len(X),1))
    X = np.concatenate([ones_column, X], axis=1)

    return X.dot(w)

# r_squared
def r_squared(ytrue, ypred):
    e_method = ((ytrue-ypred)**2).sum() # sum of squares of residuals
    e_baseline = ((ytrue-ytrue.mean())**2).sum() # total sum of squares
    return 1 - e_method/e_baseline

# loss function
def loss(ytrue, ypred):
    return ((ytrue-ypred)**2).sum()

X, y, coeff =  make_regression(n_samples=100, n_features=2, coef=True, noise=0.5, bias=3, random_state=70)
# print(X.shape, y.shape)

# Train the model/ learn the optimal weights
w = fit(X, y)

####################################################

fig = plt.figure(figsize=(8,8))
ax = plt.axes(projection='3d')

ax.scatter(X[:,0], X[:,1], y, c=y, cmap='seismic')

f1 = np.linspace(X[:,0].min(), X[:,0].max(), 50)
f2 = np.linspace(X[:,1].min(), X[:,1].max(), 50)
f1, f2 = np.meshgrid(f1, f2)

# prediction plane
X_ = np.concatenate([f1.reshape(-1,1), f2.reshape(-1,1)], axis=1)
pred = predict(X_, w).reshape(f1.shape)
ax.plot_surface(f1, f2, pred, alpha=0.5, cmap='seismic')

ax.set_xlabel("Feature 1")
ax.set_ylabel("Feature 2")
ax.set_zlabel("Output (y)")

plt.show()


import pytest

from graph_pkg.edit_cost.edit_cost_proteins_tu import EditCostProteinsTU
from graph_pkg.graph.label.label_node_proteins_tu import LabelNodeProteinsTU
from graph_pkg.graph.node import Node

@pytest.mark.parametrize('coord1, e_cost, expected',
                         [
                             ((1,), (1., 1., 1., 1., 'dirac'), 1.),
                             ((0,), (1., 1., 1., 1., 'dirac'), 1.),
                             ((2,), (1., 1., 1., 1., 'dirac'), 1.),
                             ((0,), (11., 1., 1., 1., 'dirac'), 11.),
                             ((0,), (1., 1.9, 1.9, 1.9, 'dirac'), 1.),
                         ])
def test_dirac_proteins_tu_add_node(coord1, e_cost, expected):
    node0 = Node(0, LabelNodeProteinsTU(*coord1))

    edit_cost = EditCostProteinsTU(*e_cost)

    result = edit_cost.cost_insert_node(node0)

    assert result == expected

@pytest.mark.parametrize('coord1, e_cost, expected',
                         [
                             ((1,), (1., 1., 1., 1., 'dirac'), 1.),
                             ((0,), (1., 1., 1., 1., 'dirac'), 1.),
                             ((1,), (16., 12., 18., 17., 'dirac'), 12.),
                         ])
def test_dirac_proteins_tu_delete_node(coord1, e_cost, expected):
    node0 = Node(0, LabelNodeProteinsTU(*coord1))

    edit_cost = EditCostProteinsTU(*e_cost)

    result = edit_cost.cost_delete_node(node0)

    assert result == expected

@pytest.mark.parametrize('coord1, coord2, e_cost, expected',
                         [
                          ((1,), (1,), (1., 1., 1., 1., 'dirac'), 0.),
                          ((0,), (1,), (1., 1., 1., 1., 'dirac'), 2.),
                          ((1,), (0,), (1., 1., 1., 1., 'dirac'), 2.),
                          ((1,), (2,), (3., 2., 2.5, 1., 'dirac'), 5.),

                         ])
def test_dirac_proteins_tu_substitution(coord1, coord2, e_cost, expected):
    node0 = Node(0, LabelNodeProteinsTU(*coord1))
    node1 = Node(1, LabelNodeProteinsTU(*coord2))

    edit_cost = EditCostProteinsTU(*e_cost)

    result = edit_cost.cost_substitute_node(node0, node1)

    assert result == expected


from otree.api import (
    models, widgets, BaseConstants, BaseSubsession, BaseGroup, BasePlayer,
    Currency as c, currency_range
)

from django.utils.translation import ugettext_lazy as _

from GameFeb19_intro.models import add_currency, add_tokens, TRNSL_ERR_MSG, translated_languages

import csv
import random

author = 'Tatiana Mayskaya'

doc = """
Cognitive Reflection Test & IQ Test & GRE-based Test :: whatever counts as cognitive test
"""

class Constants(BaseConstants):
    name_in_url = 'GameFeb19_questions_cognitive'
    players_per_group = None

    # this is done only to count the number of questions in the quiz
    # (assuming Russian and English versions have the same number)
    with open('GameFeb19_questions_cognitive/cognitive_en.csv') as file:
        questions = list(csv.DictReader(file))

    num_rounds = len(questions)

class Subsession(BaseSubsession):
    def creating_session(self):
        assert self.session.config['language'] in translated_languages, TRNSL_ERR_MSG
        if self.round_number == 1:
            if self.session.config['language'] == 'en':
                with open('GameFeb19_questions_cognitive/cognitive_en.csv', encoding='utf-8-sig') as test_file:
                    self.session.vars['test_file_list'] = list(csv.DictReader(test_file))
            else:
                with open('GameFeb19_questions_cognitive/cognitive_ru.csv', encoding='utf-8-sig') as test_file:
                    self.session.vars['test_file_list'] = list(csv.DictReader(test_file))
            for p in self.get_players():
                p.random_questions()
            self.session.vars['num_questions_CT'] = Constants.num_rounds
        for p in self.get_players():
            question_data = p.current_question()
            p.question_id = question_data['id']
            p.question = question_data['question']
            p.solution = int(question_data['solution'])
            if int(question_data['n_choices']) == 0:
                p.solution_text = question_data['solution']
            else:
                p.solution_text = question_data['choice{}'.format(p.solution)]
            p.participant.vars['questions_CT'] = []

    def vars_for_admin_report(self):
        players = []
        for p in self.get_players():
            players.append((p.participant.label, p.question, p.submitted_answer_text, p.solution_text,
                            p.get_is_correct_display()))
        return {'players': players}

class Group(BaseGroup):
    pass

class Player(BasePlayer):
    question_id = models.IntegerField()
    question = models.StringField()
    solution = models.IntegerField()
    solution_text = models.StringField()
    submitted_answer = models.IntegerField()
    submitted_answer_options = models.IntegerField(widget=widgets.RadioSelect)
    submitted_answer_text = models.StringField()
    is_correct = models.BooleanField(initial=False, choices=[[True, _('Yes')], [False, _('No')]])

    def random_questions(self):
        randomized_questions = random.sample(range(1, Constants.num_rounds + 1, 1), Constants.num_rounds)
        self.participant.vars['questions_order_CT'] = randomized_questions

    def current_question(self):
        num = self.participant.vars['questions_order_CT'][self.round_number - 1]
        return self.session.vars['test_file_list'][num - 1]

    def check_correct(self):
        question_data = self.current_question()
        if int(question_data['n_choices']) > 0:
            self.submitted_answer = self.submitted_answer_options
        self.is_correct = (self.submitted_answer == self.solution)
        if int(question_data['n_choices']) == 0:
            self.submitted_answer_text = str(self.submitted_answer)
        else:
            self.submitted_answer_text = question_data['choice{}'.format(self.submitted_answer)]
        self.participant.vars['questions_CT'].append(
            (self.round_number, self.question, self.submitted_answer_text, self.solution_text,
             self.get_is_correct_display()))
        if self.is_correct:
            self.payoff = self.session.vars['rate_CT']

    def set_payoffs(self):
        self.participant.vars['questions_correct_CT'] = sum([int(p.is_correct) for p in self.in_all_rounds()])
        self.participant.vars['payment_formula'] = \
            self.participant.vars['payment_formula'] + \
            ' + ' + str(self.participant.vars['questions_correct_CT']) + '*' + \
            add_currency(self.session.config['currency_used'],
                         self.session.vars['rate_CT'] * self.session.config['real_world_currency_per_point'])


import numpy as np
import scipy.sparse as sp
import torch
import time
import random

from utils.tool import read_data, write_dic, dictionary, normalize, sparse_mx_to_torch_sparse_tensor

def encoding_test(test_graph_path, test_fact_path, train_dataset = "fb237_v1"):

    """load test-graph and test-facts, and do the encoding on the test-graph"""

    t_start = time.time()

    path = "data"

    #these two paths are for loading
    relation_dic_path = "{}/{}/train/relation-dic.txt".format(path, train_dataset)
    type_dic_path = "{}/{}/train/type-dic.txt".format(path, train_dataset)

    test_graph_triples = read_data(test_graph_path)
    test_fact_triples_with_label = read_data(test_fact_path)

    #load relation dic and type dic generated by training
    f_relation_dic = open(relation_dic_path)
    relations = []
    for line in f_relation_dic:
        relation_new = line.strip().split("\t")[1]
        relations.append(relation_new)
    f_type_dic = open(type_dic_path)
    types = []
    for line in f_type_dic:
        type_new = line.strip().split("\t")[1]
        types.append(type_new)

    relation_set = set(relations)

    all_triples_with_label = test_graph_triples + test_fact_triples_with_label

    test_graph_real_triples = []
    test_graph_type_triples = []
    for triple in test_graph_triples:
        if triple[1] != "<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>":
            test_graph_real_triples.append(triple)
        else:
            test_graph_type_triples.append(triple)

    test_fact_real_triples_with_label = []
    test_fact_type_triples_with_label = []
    for triple in test_fact_triples_with_label:
        if triple[1] != "<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>":
            test_fact_real_triples_with_label.append(triple)
        else:
            test_fact_type_triples_with_label.append(triple)

    all_real_triples_with_label = []
    all_type_triples_with_label = []

    constant_set = set()
    for triple in all_triples_with_label:
        if triple[1] != "<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>":
            constant_set.add(triple[0])
            constant_set.add(triple[2])
            all_real_triples_with_label.append(triple)
        else:
            constant_set.add(triple[0])
            all_type_triples_with_label.append(triple)

    constants = list(constant_set)

    constant2index = dictionary(constants)
    relation2index = dictionary(relations)
    type2index = dictionary(types)
    #print("time:",time.time()-t_start)
    #generate list of pairs for encoding

    pairs = []
    pair_set = set()
    for triple in all_real_triples_with_label:
        sub_idx = constant2index[triple[0]]
        obj_idx = constant2index[triple[2]]
        if sub_idx < obj_idx:
            if (sub_idx, obj_idx) not in pair_set:
                pair_set.add((sub_idx, obj_idx))
                pairs.append((sub_idx, obj_idx))
        if sub_idx > obj_idx:
            if (obj_idx, sub_idx) not in pair_set:
                pair_set.add((obj_idx, sub_idx))
                pairs.append((obj_idx, sub_idx))

    for constant_idx in range(len(constants)):
        pairs.append((constant_idx, constant_idx))
        pair_set.add((constant_idx, constant_idx))
    pair2index = dictionary(pairs)

    s_time = time.time()

    #collect related pairs for each constant

    pairs_for_constant = dict([(i,set()) for i in range(len(constants))])
    p_idx = 0
    for pair in pairs:
        p_idx = pair2index[pair]
        c1 = pair[0]
        c2 = pair[1]
        pairs_for_constant[c1].add(p_idx)
        pairs_for_constant[c2].add(p_idx)

    #collect neighbors for each pair node

    pneighbors_for_pair = dict([(i,set()) for i in range(len(pairs))])
    for c_idx in range(len(constants)):
        pairs_c = set(pairs_for_constant[c_idx])
        #pair and n_pair would contain one common constant
        for pair in pairs_c:
            for n_pair in pairs_c:
                if pair != n_pair:
                    pneighbors_for_pair[pair].add(n_pair)

    #generate edge list

    edges = []

    for i in range(len(pairs)):
        pneighbors = pneighbors_for_pair[i]
        for pneighbor in pneighbors:
            edges.append([i, pneighbor])
            edges.append([pneighbor, i])

    #print("Finished generating edges", time.time() - s_time)

    #generate a normalized adjencency matrix (strategy for GCN)
    #print(edges)
    edges = np.array(edges)
    adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), shape=(len(pairs), len(pairs)), dtype=np.float32)
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)

    adj = normalize(adj + sp.eye(adj.shape[0]))
    adj = sparse_mx_to_torch_sparse_tensor(adj)
    del edges
    #print("Total time for adj: {:.4f}s".format(time.time() - s_time))

    #print("Start to generate features, labels, and masks")

    def initialize(test_graph_real_triples, test_graph_type_triples, test_fact_real_triples_with_label, test_fact_type_triples_with_label):

        labels = torch.zeros(len(pairs), len(types) + 2*len(relations))
        masks = torch.zeros(len(pairs),  len(types) + 2*len(relations))
        features = torch.zeros(len(pairs),  len(types) + 2*len(relations))

        #labels and masks are generated for all triples in test-facts (pos&neg)

        for triple in test_fact_type_triples_with_label:
            cons = triple[0]
            typ = triple[2]
            label = triple[3]
            pair_idx= pair2index[(constant2index[cons], constant2index[cons])]
            typ_idx = type2index[typ]
            if label == "1":
                labels[pair_idx][typ_idx] = 1
            elif label == "0":
                labels[pair_idx][typ_idx] = 0
            masks[pair_idx][typ_idx] = 1

        for triple in test_fact_real_triples_with_label:
            sub = triple[0]
            rel = triple[1]
            obj = triple[2]
            label = triple[3]
            sub_idx = constant2index[sub]
            rel_idx = relation2index[rel]
            obj_idx = constant2index[obj]

            try:
                pair_idx = pair2index[(sub_idx, obj_idx)]
            except:
                pair_idx = pair2index[(obj_idx, sub_idx)]
                rel_idx = rel_idx + len(relations)
            if label == "1":
                labels[pair_idx][len(types) + rel_idx] = 1
            elif label == "0":
                labels[pair_idx][len(types) + rel_idx] = 0
            masks[pair_idx][len(types) + rel_idx] = 1

        #features are generated for all triples in test-graph (pos&neg)

        for triple in test_graph_type_triples:
            cons = triple[0]
            typ = triple[2]
            pair_idx= pair2index[(constant2index[cons], constant2index[cons])]
            typ_idx = type2index[typ]
            features[pair_idx][typ_idx] = 1

        for triple in test_graph_real_triples:
            sub = triple[0]
            rel = triple[1]
            obj = triple[2]
            sub_idx = constant2index[sub]
            rel_idx = relation2index[rel]
            obj_idx = constant2index[obj]

            try:
                pair_idx = pair2index[(sub_idx, obj_idx)]
            except:
                pair_idx = pair2index[(obj_idx, sub_idx)]
                rel_idx = rel_idx + len(relations)
            features[pair_idx][len(types) + rel_idx] = 1

        features.requires_grad = True
        labels.requires_grad = False

        return features, labels, masks

    features, labels, masks = initialize(test_graph_real_triples, test_graph_type_triples, test_fact_real_triples_with_label, test_fact_type_triples_with_label)

    num_type = len(types)
    num_relation = len(relations)
    def triple2index(triple_now):
        sub_idx = constant2index[triple_now[0]]
        try:
            relation_idx = relation2index[triple_now[1]]
        except:
            pair_idx = pair2index[(sub_idx, sub_idx)]
            dim_idx = type2index[triple_now[2]]
            return pair_idx, dim_idx

        obj_idx = constant2index[triple_now[2]]
        if (sub_idx, obj_idx) in pair_set:
            pair_idx = pair2index[(sub_idx, obj_idx)]
            dim_idx = len(types) + relation_idx
        elif (obj_idx, sub_idx) in pair_set:
            pair_idx = pair2index[(obj_idx, sub_idx)]
            dim_idx = len(types) + len(relations) + relation_idx
        else:
            print(triple_now, sub_idx, relation_idx, obj_idx)
            print("wrong")
        return pair_idx, dim_idx
    hits_true = []
    for triple in test_fact_triples_with_label:
        if triple[-1] == "1":
            hits_true.append(triple2index(triple))
    #print("Finished generation")

    #print("Total time elapsed for encoding: {:.4f}s".format(time.time() - t_start))

    return adj, features, labels, masks, num_type, num_relation, constants, relations, types, pairs, hits_true