import os
import sys
from sklearn.preprocessing import StandardScaler
from sklearn.feature_extraction import DictVectorizer
from d3_feature_extractor import extract
from sklearn import tree
from sklearn.cross_validation import StratifiedKFold
from sklearn.ensemble import RandomForestClassifier
from sklearn.externals import joblib
import numpy as np
import random
import time
import json
import traceback
import argparse
from PIL import Image
from utility import reverse_dict

# mapping of integers to the chart type they represent
symbol_dict = {
    1: "line", 2: "scatter", 4: "bar", 19: "geographic_map", 35: "graph",
    14: "chord", 10: "bubble", 37: "parallel_coordinates", 13: "sankey",
    9: "box", 16: "area", 31: "stream_graph", 7: "heat_map",
    15: "radial", 33: "hexabin", 38: "sunburst", 22: "treemap",
    40: "voronoi", 18: "donut", 39: "waffle", 41: "word_cloud", 29: "pie"
}
num_dict = reverse_dict(symbol_dict)

text_features = [
    "text_word_count", "text_max_font_size", "text_min_font_size", "text_var_font_size",
    "text_unique_font_size_count", "text_unique_x_count","text_unique_y_count"
]

class BeagleClassifier(object):
    """
    Object that contains the classifier for Beagle.

    NOTE: When initializing args requires at least the collections name
    """

    def __init__(self, args, training=None):
        # set presets
        self.num_charts = 10 # number of charts to sample per vis type
        self.max_charts = 400 # maximum number of charts to include per vis type
        self.runs = 1 # number of times to repeat the experiment
        self.use_text = True

        # set paths for files
        self.collection = args[0].split(',')
        self.training = training.split(',') if training else None
        self.features = {}
        self.train_features = {}

        # assumes all data in this directory, TODO make more general
        self.base_path = "/data/scidb/000/2/learnvis_data/"

        # override defaults based on args
        if len(args) > 1:
            self.runs = int(args[1])  # number of times to repeat the experiment

        if len(args) > 2:
            self.num_charts = int(args[2]) # number of charts to sample per vis type

        if len(args) > 3:
            self.use_text = args[3].lower() == "true"
            print("using text? ", self.use_text)

    def load_features(self):
        """
        Loads the features from the features.txt file if it exists. Otherwise
        it extracts the features from the SVGs and creates a features.txt file
        that holds all of the features.
        """
        # handle base set (test set if training exists)
        for collection in self.collection:
            output_file = os.path.join(self.base_path, collection, "charts", "features.txt")

            if not os.path.isfile(output_file):
                print("creating features from scratch")
                self.features.update(self.create_features(collection))
            else:
                print("loading features file")
                with open(output_file) as f:
                    self.features.update(json.load(f))

        # handle training set if specified
        if self.training is not None:
            for collection in self.training:
                output_file = os.path.join(self.base_path, collection, "charts", "features.txt")

                if not os.path.isfile(output_file):
                    print("creating features from scratch for training")
                    self.train_features.update(self.create_features(collection))
                else:
                    print("loading features file for training set")
                    with open(output_file) as f:
                        self.train_features.update(json.load(f))

    def create_features(self, collection):
        """
        Creates the features for all of the SVGs in the collection and
        stores the file in self.output_file which is a json of features
        """
        features = {}
        badCharts = {}

        # files for reading/storing info
        data_path = os.path.join(self.base_path, collection, "charts")
        urls_file_path = os.path.join(self.base_path, collection, "urls.txt")
        images_path = os.path.join(self.base_path, collection, "images")
        output_file = os.path.join(data_path,"features.txt")
        badfile = os.path.join(data_path,"bad.txt")

        if os.path.isfile(badfile):
            with open(badfile) as f:
                badCharts = json.load(f)
        urls_file = open(urls_file_path, 'r')
        for line in urls_file:
            flags = line.strip().rstrip('\n').rstrip('\r').split(" ")
            chart = collection + flags[0]
            chartsubname = flags[0]
            if chart in badCharts:
                continue
            if not os.path.exists(os.path.join(data_path,chartsubname,"svg.txt")):
                badCharts[chart] = True
                continue
            multiple_labels = flags[2].split(",")

            for i, label in enumerate(multiple_labels):
                if label in num_dict:
                    multiple_labels[i] = num_dict[label]
            try:
                label = int(multiple_labels[0])
            except:
                print("bad label", multiple_labels[0])
                continue
            # ignore charts with bad images
            if label not in symbol_dict: # ignore unsupported chart types
                badCharts[chart] = True
                continue
            if "i" in flags:
                badCharts[chart] = True
                continue
            if not self.test_image(os.path.join(images_path,chartsubname+".png")):
              badCharts[chart] = True
              continue
            feature_dict = extract(os.path.join(data_path,chartsubname,"svg.txt"))
            if isinstance(feature_dict, str):
                print("wrong output")
                badCharts[chart] = True
                continue
            features[chart] = feature_dict
        output_json = open(output_file, 'w')
        json.dump(features, output_json)
        with open(badfile,'w') as f:
            json.dump(badCharts, f)
        output_json.close()
        return features

    def test_image(self, image_path):
        """
        Tests if the file at image_path is an image file
        """
        try:
            Image.open(image_path)
            return True
        except:
            pass
        return False # didn't work

    def organize_samples(self, collection, types_lists, secondary_labels):
        """
        Goes over all samples and filters for only valid vis types
        each vis type has a list of (chart,url) pairs.
        Modifies types_lists with additions of charts of that type.
        Only gets called once!

        collection has to either be self.collection or self.training
        """
        print("organizing samples...")
        urls_file_path = os.path.join(self.base_path, collection, "urls.txt")
        urls_file = open(urls_file_path, 'r')
        try:
            features = self.features
        except:
            # is probably from training, so just pass for now
            pass
        if self.training and collection in self.training:
            features = self.train_features

        for line in urls_file:
            line_list = line.rstrip("\n").rstrip("\r").split(" ")
            chart = collection + line_list[0]
            if chart not in features:
                continue
            secondary_label = None
            url = line_list[1]
            multiple_labels = line_list[2].split(",")
            try:
                label = int(multiple_labels[0])
            except:
                try:
                    multiple_labels = [num_dict[x] for x in multiple_labels]
                    label = int(multiple_labels[0])
                except:
                    # label not a label we do
                    continue
            if label not in symbol_dict:
                continue
            if len(multiple_labels) > 1:
                secondary_label = multiple_labels[1:]
            types_lists[label].append((chart, url))
            if secondary_label is not None:
                secondary_labels[chart] = secondary_label

    def select_subsets(self, feature_dicts, labels, urls, types_lists, repeat=True, training=False):
        """
        Selects subsets from each vis type and extracts the features for this type.
        Samples a vis type if has >num_charts amount
        Gets called on each new run.

        If repeat is false, then it will not repeat when there are <num_charts
        """
        features = self.features if not training else self.train_features

        print("\tselecting sample subsets ("+str(self.num_charts)+" samples per chart type) and extracting features...")
        for label in types_lists: # for each valid vis type
            chosenSamples = types_lists[label]

            if repeat:
                if len(types_lists[label]) == 0:
                    continue
                # not enough samples, repeat samples evenly
                if len(types_lists[label]) < self.num_charts:
                    diff = self.num_charts / len(types_lists[label])
                    remainder = self.num_charts - diff * len(types_lists[label])
                    chosenSamples = list(types_lists[label]) * diff
                    for i in range(remainder):
                        chosenSamples.append(types_lists[label][i])
                # enough samples
                else:
                    chosenSamples = random.sample(types_lists[label], self.num_charts)
            else:
                if len(chosenSamples) > self.max_charts:
                    chosenSamples = random.sample(types_lists[label], self.max_charts)

            print("\t\textracting features for label "+str(label)+" ("+symbol_dict[label]+")...")
            for chart, url in chosenSamples:
                try:
                    if chart in features:
                        feature_dict = features[chart]
                    else:
                        continue

                    if isinstance(feature_dict, str):
                        continue
                    if not self.use_text:
                        for k in text_features:
                            if k in feature_dict:
                                feature_dict.pop(k)
                    # for k in features_to_ignore:
                    #   if k in feature_dict:
                    #     print k
                    #     feature_dict.pop(k)
                    feature_dicts.append(feature_dict)
                    labels.append(label)
                    urls.append((chart, url))
                except Exception as e:
                    print(e)
                    traceback.print_exc()
                    pass

    def do_cross_validation(self, feature_dicts, labels, secondary_labels, urls):
        """
        Performs the cross validation for the chosen samples
        """
        vec = DictVectorizer(sparse=False)
        scaler = StandardScaler()
        print("\tcreating features array...")

        print("\tperforming stratified k-fold...")
        skf = StratifiedKFold(labels, n_folds=5)
        features_array = scaler.fit_transform(vec.fit_transform(feature_dicts))
        # clf = RandomForestClassifier(n_estimators=14)
        correct_count = {}
        total_count = {}
        wrong = {}
        result = {}
        for train_index, test_index in skf:
            training_dicts = [features_array[t] for t in train_index]
            training_labels = [labels[t] for t in train_index]
            testing_points = [features_array[t] for t in test_index]
            testing_labels = [labels[t] for t in test_index]

            clf = RandomForestClassifier(n_estimators=14)
            clf.fit(training_dicts, training_labels)

            for i in range(len(testing_points)):
                label = testing_labels[i]
                for d in (correct_count, total_count):
                    if label not in list(d.keys()):
                        d[label] = 0
                total_count[label] += 1
                prediction = clf.predict(np.array(testing_points[i]).reshape(1,-1))[0]
                if prediction == label:
                    correct_count[label] += 1
                else:
                    failed = False
                    folder = urls[test_index[i]][0]
                    if folder in list(secondary_labels.keys()):
                        if str(prediction) in secondary_labels[folder]:
                            correct_count[label] += 1
                        else:
                            failed = True
                    else:
                        failed = True
                    if failed:
                        real_label = symbol_dict[label]
                        if real_label not in wrong:
                            wrong[real_label] = []
                        wrong[real_label].append({
                            'feature_dict': feature_dicts[test_index[i]],
                            'url': urls[test_index[i]],
                            "predicted": symbol_dict[np.asscalar(prediction)],
                            "real": symbol_dict[label]
                        })

        for k in total_count: # for each label
            result[symbol_dict[k]] = 1.0*correct_count[k]/total_count[k]
        return {"accuracy":result,"wrong":wrong,"correct":correct_count,"total":total_count}

    def do_cross_validation_trained(self, feature_dicts, labels, secondary_labels, urls, clft):
        """
        Performs the cross validation for the chosen samples
        """
        vec = DictVectorizer(sparse=False)
        scaler = StandardScaler()
        print("\tcreating features array...")

        print("\tperforming stratified k-fold...")
        skf = StratifiedKFold(labels, n_folds=5)
        features_array = scaler.fit_transform(vec.fit_transform(feature_dicts))
        clf = RandomForestClassifier(n_estimators=14)

        correct_count = {}
        total_count = {}
        wrong = {}
        result = {}

        for train_index, test_index in skf:
            training_dicts = [features_array[t] for t in train_index]
            training_labels = [labels[t] for t in train_index]
            testing_points = [features_array[t] for t in test_index]
            testing_labels = [labels[t] for t in test_index]
            t = training_dicts + testing_points
            l = training_labels + testing_labels
            ind = list(train_index) + list(test_index)
            # clf.fit(training_dicts, training_labels)
            # clf.fit(testing_points, testing_labels)
            # clf.fit(t, l)
            for i in range(len(t)):
                label = l[i]
                for d in (correct_count, total_count):
                    if label not in list(d.keys()):
                        d[label] = 0
                total_count[label] += 1
                prediction = clft.predict(np.array(t[i]).reshape(1,-1))[0]
                if prediction == label:
                    correct_count[label] += 1
                else:
                    failed = False
                    folder = urls[ind[i]][0]
                    if folder in list(secondary_labels.keys()):
                        if str(prediction) in secondary_labels[folder]:
                            correct_count[label] += 1
                        else:
                            failed = True
                    else:
                        failed = True
                    if failed:
                        real_label = symbol_dict[label]
                        if real_label not in wrong:
                            wrong[real_label] = []
                        wrong[real_label].append({
                            'feature_dict': feature_dicts[ind[i]],
                            'url': urls[ind[i]],
                            "predicted": symbol_dict[np.asscalar(prediction)],
                            "real": symbol_dict[label]
                        })

        for k in total_count: # for each label
            result[symbol_dict[k]] = 1.0*correct_count[k]/total_count[k]
        return {"accuracy":result,"wrong":wrong,"correct":correct_count,"total":total_count}

    def consolidate_accuracy_results_across_runs(self, run_results):
        """
        Gets all of the accuracy results from all of the runs
        """
        final_correct = 0
        final_total = 0
        final_accuracy_per_label = {}
        final_wrong = {}
        for run_result in run_results:
            accuracy = run_result['accuracy'] # TODO clean up using tuple probably
            wrong = run_result['wrong']
            correct_count = run_result['correct']
            total_count = run_result['total']
            for label in accuracy:
                if label not in final_accuracy_per_label:
                    final_accuracy_per_label[label] = []
                final_accuracy_per_label[label].append(accuracy[label])
            print(correct_count,total_count)
            for label in correct_count:
                final_correct += correct_count[label]
            for label in total_count:
                final_total += total_count[label]
            for label in wrong:
                if label not in final_wrong:
                    final_wrong[label] = []
                final_wrong[label].extend(wrong[label])
        print("average accuracy across all runs:")
        avg_accuracy = []
        with open("accuracy.txt", "w") as myfile:
            myfile.write("\t".join(['label','accuracy'])+"\n")
            for label in sorted(final_accuracy_per_label):
                final_average = np.mean(final_accuracy_per_label[label])
                print(label+"\t"+str(final_average)+"\t",final_accuracy_per_label[label])
                avg_accuracy.append(final_average)
                myfile.write("\t".join([label,str(final_average)])+"\n")
        print("average:")
        print(np.mean(np.array(avg_accuracy)))

        print("overall:")
        print(1.0 * final_correct / final_total)

        with open("wrong.txt", "w") as myfile:
            myfile.write("\t".join(['visid','real','predicted','url','feature_dict'])+"\n")
            for label in sorted(final_wrong):
                for w in final_wrong[label]:
                    myfile.write("\t".join([w['url'][0],w['real'],w['predicted'],w['url'][1],json.dumps(w['feature_dict'])])+"\n")
        wrong_counts = {}
        for label in sorted(final_wrong):
            for w in final_wrong[label]:
                real = w['real']
                predicted = w['predicted']
                if wrong_counts.get(real):
                    if wrong_counts[real].get(predicted):
                        wrong_counts[real][predicted] += 1
                    else:
                        wrong_counts[real][predicted] = 1
                else:
                    wrong_counts[real] = { predicted:1 }
        wrong_counts_file = open("wrong_counts.txt", 'w')
        for chart in list(wrong_counts.keys()):
            wrong_counts_file.write(chart + ":\n")
            sorted_items = sorted(list(wrong_counts[chart].items()), key=lambda x: x[1])
            for chart_type, count in sorted_items:
                wrong_counts_file.write(chart_type + ": " + str(count) + '\n')
            wrong_counts_file.write("\n")
        wrong_counts_file.close()

    def print_chart_stats(self, types_lists):
        """
        Tells how many valid charts there are for each chart type
        """
        print("chart stats:")
        outp = {}
        for l in types_lists:
            outp[symbol_dict[l]] = len(types_lists[l])
        print(outp)

    def run_verification(self):
        """
        Runs the verification of the classifer given all of the presets
        """
        # holds all available samples collected from urls.txt, separated by vis type
        # valid vis types only
        types_lists = {}
        for k in symbol_dict:
            types_lists[k] = []

        # if vis sample has more than one label, put the second label in here
        secondary_labels = {}

        total_time = []
        accuracy = []
        s = time.time()
        self.load_features() # build the feature dicts, if they don't exist yet
        for collection in self.collection:
            self.organize_samples(collection, types_lists, secondary_labels) # only call once
        # print secondary_labels

        if self.training:
            train_types_lists = {}
            for k in symbol_dict:
                train_types_lists[k] = []
            train_secondary_labels = {}
            for collection in self.training:
                self.organize_samples(collection, train_types_lists, train_secondary_labels)

            # make it so can only use samples in both
            for t in types_lists:
                if len(types_lists[t]) == 0:
                    train_types_lists[t] = []

            for t in train_types_lists:
                if len(train_types_lists[t]) == 0:
                    types_lists[t] = []

            # get features to train on all of training set
            t_feature_dicts, t_labels, t_urls = [], [], []
            self.select_subsets(t_feature_dicts, t_labels, t_urls, train_types_lists, repeat=False, training=True)

            t_feature_dicts, t_labels, t_urls = [], [], []
            self.select_subsets(t_feature_dicts, t_labels, t_urls, train_types_lists, repeat=False, training=True)

            vec = DictVectorizer(sparse=False)
            scaler = StandardScaler()
            skft = StratifiedKFold(t_labels, n_folds=5)
            features_arrayt = scaler.fit_transform(vec.fit_transform(t_feature_dicts))
            clft = RandomForestClassifier(n_estimators=14)
            clft.fit(features_arrayt, t_labels)

            total_time.append(time.time() - s)
            print("running experiment "+str(self.runs)+" time(s)...")

            # train on training set + test on test set
            for run_index in range(self.runs):
                print("executing run "+str(run_index)+"...")
                feature_dicts, labels, urls = [], [], []
                s = time.time()

                self.select_subsets(feature_dicts, labels, urls, types_lists, repeat=False)

                total_time.append(time.time() - s)

                result_dict = self.do_cross_validation_trained(feature_dicts, labels, secondary_labels, urls, clft)
                # result_dict = self.do_cross_validation(feature_dicts,labels,secondary_labels,urls)
                accuracy.append(result_dict)
        else:
            total_time.append(time.time() - s)
            print("running experiment "+str(self.runs)+" time(s)...")

            for run_index in range(self.runs):
                print("executing run "+str(run_index)+"...")
                feature_dicts, labels, urls = [], [], []

                s = time.time()
                if self.num_charts == -1:
                    self.select_subsets(feature_dicts, labels, urls, types_lists, repeat=False)
                else:
                    self.select_subsets(feature_dicts,labels,urls,types_lists)
                total_time.append(time.time() - s)
                result_dict = self.do_cross_validation(feature_dicts,labels,secondary_labels,urls)
                accuracy.append(result_dict)

        print("average extraction time:",np.sum(total_time)/self.runs)
        self.consolidate_accuracy_results_across_runs(accuracy)
        self.print_chart_stats(types_lists)

def main():
    """
    Runs the verification for some experiment on the classifier
    """
    parser = argparse.ArgumentParser(prog='Beagle Classifier')
    parser.add_argument('--train', help="Tells you the set that it is trained on")
    parser.add_argument('args', nargs="+", help="List of args for the classifier")

    sysargs = parser.parse_args()
    args = sysargs.args
    training = sysargs.train

    classifier = BeagleClassifier(args, training)
    classifier.run_verification()

if __name__ == '__main__':
    main()


from __future__ import annotations
import collections
import enum

import quokka
from quokka.types import (
    AddressT,
    DefaultDict,
    Dict,
    Index,
    Iterator,
    List,
    LocationValueType,
    Mapping,
    MutableMapping,
    Optional,
    ReferenceTarget,
    ReferenceType,
    Tuple,
    Union,
)

class Reference:
    """Reference: a link between a source and a destination

    Arguments:
        source: Source of the reference
        destination: Destination (target) of the reference
        type_: Type of the reference

    Attributes:
        source: Source of the reference
        destination: Destination (target) of the reference
        type: Type of the reference
    """

    def __init__(
        self,
        source: ReferenceTarget,
        destination: ReferenceTarget,
        type_: ReferenceType,
    ) -> None:
        """Constructor"""
        self.source: ReferenceTarget = source
        self.destination: ReferenceTarget = destination
        self.type: ReferenceType = type_

class ReferencesLocation(enum.Enum):
    """Reference location

    A reference may be attached to one of the following :
        * an instruction index (that means every instance of the instruction)
        * a data index
        * a structure position (e.g. a structure or a member inside a structure)
        * an instruction : a tuple with (chunk, block, inst_index) identifying precisely
            one instruction
        * a function
        * a chunk
    """

    INSTRUCTION = "inst_idx"
    DATA = "data_idx"
    STRUCTURE = "struct_position"
    INST_INSTANCE = "instruction_position"
    FUNCTION = "function_idx"
    CHUNK = "chunk_idx"

    @staticmethod
    def from_proto(location_type: str) -> ReferencesLocation:
        """Convert reference location from proto"""
        # These are the name of the fields in the protobuf
        mapping = {
            "inst_idx": ReferencesLocation.INSTRUCTION,
            "data_idx": ReferencesLocation.DATA,
            "struct_position": ReferencesLocation.STRUCTURE,
            "instruction_position": ReferencesLocation.INST_INSTANCE,
            "function_idx": ReferencesLocation.FUNCTION,
            "chunk_idx": ReferencesLocation.CHUNK,
        }

        try:
            return mapping[location_type]
        except IndexError as exc:
            raise ValueError("Unknown location type") from exc

class References(Mapping):
    """References bucket : maintain the list of all references inside the program

    The instantiation of the class will create a mapping for every reference but not
    resolve them yet.

    This class is probably the most messy one in Quokka and needs a global refactoring.
    However, since it's somehow working, let's try to not break it yet.

    Arguments:
        program: Reference to the program

    Attributes:
        program: Reference to the program
        proto_ref: Protobuf data
        references_category: A mapping that contains every reference for each direction.
            For instance, a ref X from inst a to struct b will be stored in both the
            structure key and the instruction key

    """

    SOURCE = "source"
    DESTINATION = "destination"

    def __init__(self, program: quokka.Program) -> None:
        """Init method

        Arguments:
            program: A backref to program
        """

        # Define the update method
        def update(
            mapping: Mapping,  # type: ignore
            key: Union[Tuple[int], LocationValueType],
            value: Index,
        ) -> None:
            """Recursive method to update the `self.references_category` mapping"""
            if isinstance(key, int):
                mapping[key].append(value)
                return

            current_key: int = key[0]
            sub_key: Union[Tuple[()], Tuple[int], Tuple[int, int]] = key[1:]
            if sub_key:
                update(mapping[current_key], sub_key, value)
            elif not isinstance(mapping.get(current_key), MutableMapping):
                mapping[current_key].append(value)

        self.program: quokka.Program = program
        self.proto_ref = program.proto.references

        # Init the references mapping
        self.references_category: Dict[ReferencesLocation, DefaultDict] = {  # type: ignore
            ReferencesLocation.INSTRUCTION: collections.defaultdict(list),
            ReferencesLocation.DATA: collections.defaultdict(list),
            ReferencesLocation.STRUCTURE: collections.defaultdict(
                lambda: collections.defaultdict(list)
            ),
            ReferencesLocation.INST_INSTANCE: collections.defaultdict(
                lambda: collections.defaultdict(lambda: collections.defaultdict(list))
            ),
            ReferencesLocation.FUNCTION: collections.defaultdict(list),
            ReferencesLocation.CHUNK: collections.defaultdict(list),
        }

        reference: quokka.pb.Quokka.Reference
        for index, reference in enumerate(self.proto_ref):
            source_ref = self.references_category[self.location_type(reference.source)]
            update(source_ref, self.get_location_value(reference.source), index)

            destination_ref = self.references_category[
                self.location_type(reference.destination)
            ]
            update(
                destination_ref, self.get_location_value(reference.destination), index
            )

    @staticmethod
    def location_type(
        location: quokka.pb.Quokka.Location,
    ) -> ReferencesLocation:
        """Convert the proto location type"""
        return ReferencesLocation.from_proto(location.WhichOneof("LocationType"))

    def get_location_value(
        self, location: quokka.pb.Quokka.Location
    ) -> LocationValueType:
        """Resolve a location

        This method resolves a location for quokka.

        Arguments:
            location: Reference location

        Returns:
            A LocationTypeValue
        """
        location_type: ReferencesLocation = self.location_type(location)
        location_value: LocationValueType = getattr(location, location_type.value)

        if location_type in (
            ReferencesLocation.INSTRUCTION,
            ReferencesLocation.DATA,
            ReferencesLocation.FUNCTION,
            ReferencesLocation.CHUNK,
        ):
            return location_value

        elif location_type == ReferencesLocation.STRUCTURE:
            structure_idx: int = getattr(location_value, "structure_idx")

            if getattr(location_value, "no_member") is False:
                return structure_idx, getattr(location_value, "member_idx")

            return structure_idx, -1

        elif location_type == ReferencesLocation.INST_INSTANCE:
            return (
                getattr(location_value, "func_chunk_idx"),
                getattr(location_value, "block_idx"),
                getattr(location_value, "instruction_idx"),
            )

        raise ValueError

    def __getitem__(self, k: Index) -> Reference:
        """Lazy loading for references"""
        return self.create_reference(self.proto_ref[k])

    def __len__(self) -> int:
        """References count"""
        return len(self.proto_ref)

    def __iter__(self) -> Iterator[Dict[AddressT, Reference]]:
        """Iterator over references"""
        raise NotImplementedError

    @property
    def structures(self) -> List[quokka.Structure]:
        """Accessor to the program structures"""
        return self.program.structures

    @property
    def data(self):
        """Accessor to the program PROTO data"""
        return self.program.proto.data

    @property
    def instructions(self):
        """Accessor to the program proto instruction"""
        return self.program.proto.instructions

    def create_reference(self, reference: quokka.pb.Quokka.Reference) -> "Reference":
        """Create a reference

        Start with resolving both the source and the destination.

        Arguments:
            reference: A protobuf reference object

        Returns:
            A python reference with both source and destination resolved
        """
        source: ReferenceTarget = self.resolve_location(reference.source)
        destination: ReferenceTarget = self.resolve_location(reference.destination)

        return Reference(
            source, destination, ReferenceType.from_proto(reference.reference_type)
        )

    def find_instruction(
        self,
        instruction_identifier: "quokka.pb.Quokka.Location.InstructionIdentifier",
    ) -> Tuple[quokka.Chunk, quokka.Block, Index]:
        """Search an instruction from an instruction identifier.

        The parameters are a chunk index, a block index, and an inst index inside the
        block.

        Arguments:
            instruction_identifier: Protobuf instruction identifier

        Returns:
            A instruction tuple (Chunk, Block, Instruction Index)
        """
        chunk = self.program.get_chunk(
            instruction_identifier.func_chunk_idx, instruction_identifier.block_idx
        )
        block = chunk[chunk.index_to_address[instruction_identifier.block_idx]]
        return chunk, block, instruction_identifier.instruction_idx

    def resolve_location(self, location: quokka.pb.Quokka.Location) -> ReferenceTarget:
        """Resolve a location

        This convert a proto location to the actual instance of the object.

        Arguments:
            location: A protobuf Location (sigh) object

        Returns:
            A ReferenceTarget

        Raises:
            ValueError: When no appropriate location has been found.
        """
        location_type: ReferencesLocation = self.location_type(location)
        if location_type == ReferencesLocation.INSTRUCTION:
            raise DeprecationWarning("Not used anymore")

        elif location_type == ReferencesLocation.DATA:
            return self.program.data_holder[location.data_idx]

        elif location_type == ReferencesLocation.STRUCTURE:
            structure: quokka.Structure = self.structures[
                location.struct_position.structure_idx
            ]
            if location.struct_position.no_member is False:
                return structure[
                    structure.index_to_offset[location.struct_position.member_idx]
                ]
            return structure

        elif location_type == ReferencesLocation.INST_INSTANCE:
            return self.find_instruction(location.instruction_position)

        elif location_type == ReferencesLocation.CHUNK:
            return self.program.chunks[location.chunk_idx]

        raise ValueError("No location found")

    @staticmethod
    def get_direction(towards: bool = True) -> Tuple[str, str]:
        """Helper method to analyze if we take the reference or unwind it"""
        target = "destination" if towards else "source"
        wanted = "destination" if not towards else "source"

        return target, wanted

    def resolve_block_references(
        self,
        chunk_index: Index,
        block_index: Index,
        reference_type: ReferenceType,
        towards: bool = True,
    ) -> List[Reference]:
        """Return a list of references from/towards a block.

        Arguments:
            chunk_index: Index of the chunk
            block_index: Index of the block
            reference_type: TYpe of reference wanted (e.g. Data)
            towards: True if we want the reference *to* this block.
                False if we want the references *from* the block.

        Returns:
            A List of references matching the criteria
        """
        target, _ = self.get_direction(towards)

        return_list = []
        for reference_ids in self.references_category[ReferencesLocation.INST_INSTANCE][
            chunk_index
        ][block_index].values():
            for reference_id in reference_ids:
                reference: Reference = self[reference_id]
                if reference.type == reference_type and isinstance(
                    getattr(reference, target), tuple
                ):
                    target_loc = getattr(reference, target)
                    if (target_loc[0].proto_index, target_loc[1].proto_index) == (
                        chunk_index,
                        block_index,
                    ):
                        return_list.append(reference)

        return return_list

    def resolve_inst_instance(
        self,
        inst_tuple: Tuple[int, int, int],
        reference_type: Optional[ReferenceType] = None,
        towards: bool = True,
    ) -> List[Reference]:
        """Resolve instruction references

        Return a list of references from/towards an instruction.
        If a reference_type is specified the references will be filtered by their type

        Arguments:
            inst_tuple: A tuple (Chunk Index, Block Index, Instruction Index)
            reference_type: Reference type to consider (filter result)
            towards: In which sense to search the reference?

        Returns:
            A list of reference matching the criteria
        """
        return_list = []
        target, _ = self.get_direction(towards)

        for reference_idx in self.references_category[ReferencesLocation.INST_INSTANCE][
            inst_tuple[0]
        ][inst_tuple[1]][inst_tuple[2]]:
            reference: "Reference" = self[reference_idx]
            if (not reference_type or reference.type == reference_type) and isinstance(
                getattr(reference, target), tuple
            ):
                target_loc = getattr(reference, target)
                if (
                    target_loc[0].proto_index,
                    target_loc[1].proto_index,
                    target_loc[2],
                ) == inst_tuple:
                    return_list.append(reference)

        return return_list

    def resolve_calls(
        self, chunk: quokka.Chunk, towards: bool = True
    ) -> List[Union[Tuple[quokka.Chunk, quokka.Block, int], quokka.Chunk]]:
        """Resolve calls to a chunk

        This method resolves calls initiated or received from every instruction in a
        chunk. Of note, a call must be from an Instruction Tuple to another.

        Arguments:
            chunk: Target chunk
            towards: Do we look at calls towards this chunk {X,Y,Z} -> Chunk or at calls
                from this chunk Chunk -> {X,Y,Z} ?

        Returns:
            A list of objects that are either tuple (Chunk, Block, Instruction Index) or
            Chunk (only for fake chunks that are created from an imported function)
        """
        return_list = []
        target, wanted = self.get_direction(towards)

        for block_index in self.references_category[ReferencesLocation.INST_INSTANCE][
            chunk.proto_index
        ]:
            # For super-chunks, filter out block that do not belong to this chunk
            # anymore
            if block_index not in chunk.index_to_address:
                continue

            block_references = self.references_category[
                ReferencesLocation.INST_INSTANCE
            ][chunk.proto_index][block_index]
            for references in block_references.values():
                for reference_idx in references:
                    reference: "Reference" = self[reference_idx]
                    if reference.type == ReferenceType.CALL and isinstance(
                        getattr(reference, target), tuple
                    ):
                        target_loc = getattr(reference, target)
                        if target_loc[0].proto_index == chunk.proto_index:
                            location = getattr(reference, wanted)
                            # Fix: Only returns valid tuple
                            # TODO(dm): check why sometimes a Data might be returned
                            if isinstance(location, (tuple, quokka.Chunk)):
                                return_list.append(location)

        return return_list

    def resolve_data(
        self,
        data_index: Index,
        reference_type: Union[ReferenceType, None] = None,
    ) -> List[Reference]:
        """Resolve data references

        Returns a list of reference towards a data
        If a reference_type is specified the references will be filtered by their type

        Arguments:
            data_index: Index of the data in the protobuf
            reference_type: Type of reference

        Returns:
            A list of reference matching the criteria
        """

        references: List[Reference] = []
        for reference_idx in self.references_category[ReferencesLocation.DATA][
            data_index
        ]:
            reference = self[reference_idx]

            if reference_type is not None and reference_type != reference.type:
                continue

            references.append(reference)

        return references


from __future__ import absolute_import, print_function

import json
from copy import deepcopy
from functools import partial

from flask import Blueprint, abort, current_app, make_response, request, \
    url_for
from invenio_db import db
from invenio_oauth2server import require_api_auth, require_oauth_scopes
from invenio_pidstore.errors import PIDInvalidAction
from invenio_records_rest.utils import obj_or_import_string
from invenio_records_rest.views import \
    create_error_handlers as records_rest_error_handlers
from invenio_records_rest.views import \
    create_url_rules as records_rest_url_rules
from invenio_records_rest.views import need_record_permission, pass_record
from invenio_rest import ContentNegotiatedMethodView
from invenio_rest.views import create_api_errorhandler
from webargs import fields
from webargs.flaskparser import use_kwargs
from werkzeug.utils import secure_filename

from ..api import Deposit
from ..errors import FileAlreadyExists, WrongFile
from ..scopes import write_scope
from ..search import DepositSearch
from ..signals import post_action
from ..utils import extract_actions_from_class

def create_error_handlers(blueprint):
    """Create error handlers on blueprint."""
    blueprint.errorhandler(PIDInvalidAction)(create_api_errorhandler(
        status=403, message='Invalid action'
    ))
    records_rest_error_handlers(blueprint)

def create_blueprint(endpoints):
    """Create Invenio-Deposit-REST blueprint.

    See: :data:`invenio_deposit.config.DEPOSIT_REST_ENDPOINTS`.

    :param endpoints: List of endpoints configuration.
    :returns: The configured blueprint.
    """
    blueprint = Blueprint(
        'invenio_deposit_rest',
        __name__,
        url_prefix='',
    )
    create_error_handlers(blueprint)

    for endpoint, options in (endpoints or {}).items():
        options = deepcopy(options)

        if 'files_serializers' in options:
            files_serializers = options.get('files_serializers')
            files_serializers = {mime: obj_or_import_string(func)
                                 for mime, func in files_serializers.items()}
            del options['files_serializers']
        else:
            files_serializers = {}

        if 'record_serializers' in options:
            serializers = options.get('record_serializers')
            serializers = {mime: obj_or_import_string(func)
                           for mime, func in serializers.items()}
        else:
            serializers = {}

        file_list_route = options.pop(
            'file_list_route',
            '{0}/files'.format(options['item_route'])
        )
        file_item_route = options.pop(
            'file_item_route',
            '{0}/files/<path:key>'.format(options['item_route'])
        )

        options.setdefault('search_class', DepositSearch)
        search_class = obj_or_import_string(options['search_class'])

        # records rest endpoints will use the deposit class as record class
        options.setdefault('record_class', Deposit)
        record_class = obj_or_import_string(options['record_class'])

        # backward compatibility for indexer class
        options.setdefault('indexer_class', None)

        for rule in records_rest_url_rules(endpoint, **options):
            blueprint.add_url_rule(**rule)

        search_class_kwargs = {}
        if options.get('search_index'):
            search_class_kwargs['index'] = options['search_index']

        if options.get('search_type'):
            search_class_kwargs['doc_type'] = options['search_type']

        ctx = dict(
            read_permission_factory=obj_or_import_string(
                options.get('read_permission_factory_imp')
            ),
            create_permission_factory=obj_or_import_string(
                options.get('create_permission_factory_imp')
            ),
            update_permission_factory=obj_or_import_string(
                options.get('update_permission_factory_imp')
            ),
            delete_permission_factory=obj_or_import_string(
                options.get('delete_permission_factory_imp')
            ),
            record_class=record_class,
            search_class=partial(search_class, **search_class_kwargs),
            default_media_type=options.get('default_media_type'),
        )

        deposit_actions = DepositActionResource.as_view(
            DepositActionResource.view_name.format(endpoint),
            serializers=serializers,
            pid_type=options['pid_type'],
            ctx=ctx,
        )

        blueprint.add_url_rule(
            '{0}/actions/<any({1}):action>'.format(
                options['item_route'],
                ','.join(extract_actions_from_class(record_class)),
            ),
            view_func=deposit_actions,
            methods=['POST'],
        )

        deposit_files = DepositFilesResource.as_view(
            DepositFilesResource.view_name.format(endpoint),
            serializers=files_serializers,
            pid_type=options['pid_type'],
            ctx=ctx,
        )

        blueprint.add_url_rule(
            file_list_route,
            view_func=deposit_files,
            methods=['GET', 'POST', 'PUT'],
        )

        deposit_file = DepositFileResource.as_view(
            DepositFileResource.view_name.format(endpoint),
            serializers=files_serializers,
            pid_type=options['pid_type'],
            ctx=ctx,
        )

        blueprint.add_url_rule(
            file_item_route,
            view_func=deposit_file,
            methods=['GET', 'PUT', 'DELETE'],
        )
    return blueprint

class DepositActionResource(ContentNegotiatedMethodView):
    """Deposit action resource."""

    view_name = '{0}_actions'

    def __init__(self, serializers, pid_type, ctx, *args, **kwargs):
        """Constructor."""
        super(DepositActionResource, self).__init__(
            serializers,
            default_media_type=ctx.get('default_media_type'),
            *args,
            **kwargs
        )
        for key, value in ctx.items():
            setattr(self, key, value)

    @pass_record
    @need_record_permission('update_permission_factory')
    def post(self, pid, record, action):
        """Handle deposit action.

        After the action is executed, a
        :class:`invenio_deposit.signals.post_action` signal is sent.

        Permission required: `update_permission_factory`.

        :param pid: Pid object (from url).
        :param record: Record object resolved from the pid.
        :param action: The action to execute.
        """
        record = getattr(record, action)(pid=pid)

        db.session.commit()
        # Refresh the PID and record metadata
        db.session.refresh(pid)
        db.session.refresh(record.model)
        post_action.send(current_app._get_current_object(), action=action,
                         pid=pid, deposit=record)
        response = self.make_response(pid, record,
                                      202 if action == 'publish' else 201)
        endpoint = '.{0}_item'.format(pid.pid_type)
        location = url_for(endpoint, pid_value=pid.pid_value, _external=True)
        response.headers.extend(dict(Location=location))
        return response

class DepositFilesResource(ContentNegotiatedMethodView):
    """Deposit files resource."""

    view_name = '{0}_files'

    def __init__(self, serializers, pid_type, ctx, *args, **kwargs):
        """Constructor."""
        super(DepositFilesResource, self).__init__(
            serializers,
            *args,
            **kwargs
        )
        for key, value in ctx.items():
            setattr(self, key, value)

    @pass_record
    @need_record_permission('read_permission_factory')
    def get(self, pid, record):
        """Get files.

        Permission required: `read_permission_factory`.

        :param pid: Pid object (from url).
        :param record: Record object resolved from the pid.
        :returns: The files.
        """
        return self.make_response(obj=record.files, pid=pid, record=record)

    @require_api_auth()
    @require_oauth_scopes(write_scope.id)
    @pass_record
    @need_record_permission('update_permission_factory')
    def post(self, pid, record):
        """Handle POST deposit files.

        Permission required: `update_permission_factory`.

        :param pid: Pid object (from url).
        :param record: Record object resolved from the pid.
        """
        # load the file
        uploaded_file = request.files['file']
        # file name
        key = secure_filename(
            request.form.get('name') or uploaded_file.filename
        )
        # check if already exists a file with this name
        if key in record.files:
            raise FileAlreadyExists()
        # add it to the deposit
        record.files[key] = uploaded_file.stream
        record.commit()
        db.session.commit()
        return self.make_response(
            obj=record.files[key].obj, pid=pid, record=record, status=201)

    @require_api_auth()
    @require_oauth_scopes(write_scope.id)
    @pass_record
    @need_record_permission('update_permission_factory')
    def put(self, pid, record):
        """Handle the sort of the files through the PUT deposit files.

        Expected input in body PUT:

        .. code-block:: javascript

            [
                {
                    "id": 1
                },
                {
                    "id": 2
                },
                ...
            }

        Permission required: `update_permission_factory`.

        :param pid: Pid object (from url).
        :param record: Record object resolved from the pid.
        :returns: The files.
        """
        try:
            ids = [data['id'] for data in json.loads(
                request.data.decode('utf-8'))]
        except KeyError:
            raise WrongFile()

        record.files.sort_by(*ids)
        record.commit()
        db.session.commit()
        return self.make_response(obj=record.files, pid=pid, record=record)

class DepositFileResource(ContentNegotiatedMethodView):
    """Deposit files resource."""

    view_name = '{0}_file'

    get_args = dict(
        version_id=fields.UUID(
            location='headers',
            load_from='version_id',
        ),
    )
    """GET query arguments."""

    def __init__(self, serializers, pid_type, ctx, *args, **kwargs):
        """Constructor."""
        super(DepositFileResource, self).__init__(
            serializers,
            *args,
            **kwargs
        )
        for key, value in ctx.items():
            setattr(self, key, value)

    @use_kwargs(get_args)
    @pass_record
    @need_record_permission('read_permission_factory')
    def get(self, pid, record, key, version_id=None, **kwargs):
        """Get file.

        Permission required: `read_permission_factory`.

        :param pid: Pid object (from url).
        :param record: Record object resolved from the pid.
        :param key: Unique identifier for the file in the deposit.
        :param version_id: File version. Optional. If no version is provided,
            the last version is retrieved.
        :returns: the file content.
        """
        try:
            obj = record.files[str(key)].get_version(version_id=version_id)
            return self.make_response(
                obj=obj or abort(404), pid=pid, record=record)
        except KeyError:
            abort(404)

    @require_api_auth()
    @require_oauth_scopes(write_scope.id)
    @pass_record
    @need_record_permission('update_permission_factory')
    def put(self, pid, record, key):
        """Handle the file rename through the PUT deposit file.

        Permission required: `update_permission_factory`.

        :param pid: Pid object (from url).
        :param record: Record object resolved from the pid.
        :param key: Unique identifier for the file in the deposit.
        """
        try:
            data = json.loads(request.data.decode('utf-8'))
            new_key = data['filename']
        except KeyError:
            raise WrongFile()
        new_key_secure = secure_filename(new_key)
        if not new_key_secure or new_key != new_key_secure:
            raise WrongFile()
        try:
            obj = record.files.rename(str(key), new_key_secure)
        except KeyError:
            abort(404)
        record.commit()
        db.session.commit()
        return self.make_response(obj=obj, pid=pid, record=record)

    @require_api_auth()
    @require_oauth_scopes(write_scope.id)
    @pass_record
    @need_record_permission('update_permission_factory')
    def delete(self, pid, record, key):
        """Handle DELETE deposit file.

        Permission required: `update_permission_factory`.

        :param pid: Pid object (from url).
        :param record: Record object resolved from the pid.
        :param key: Unique identifier for the file in the deposit.
        """
        try:
            del record.files[str(key)]
            record.commit()
            db.session.commit()
            return make_response('', 204)
        except KeyError:
            abort(404, 'The specified object does not exist or has already '
                  'been deleted.')


from typing import Callable, Dict, Optional, List

from graph.graph import Graph

def minus(l, r):
    if l:
        return False if r else True
    return False

class Operation:
    operation: Dict[str, Callable[[int, int], int]] = {
        'and': lambda left, right: left and right,
        'or': lambda left, right: left or right,
        'xor': lambda left, right: left != right,
        'minus': minus,
        'eq': lambda left, right: left == right,
        'impl': lambda left, right: (not left) or right,
        'coimpl': lambda left, right: not ((not left) or right),
        'sheph': lambda left, right: not (left and right),
        'pirs': lambda left, right: not (left or right)
    }

    def __init__(self, op: str):
        self.op = op

    def __call__(self, left: int, right: int) -> int:
        return self.operation[self.op](left, right)

def binary_operation(a: Graph, b: Graph, op: str) -> Optional[Graph]:
    if a.size() != b.size():
        return None
    if a.size() < b.size():
        a, b = b, a
    size = a.size()
    g = [[0 for i in range(size)] for i in range(size)]
    operation = Operation(op)
    a = a.to_matrix()
    b = b.to_matrix()

    for i in range(size):
        for j in range(size):
            if i == j:
                continue
            g[i][j] = int(operation(a[i][j], b[i][j]))

    return Graph.from_matrix(g)

def binary_operation_with_matrix(a: List[List[int]], b: List[List[int]], op: str) -> Optional[List[List[int]]]:
    if len(a) != len(b):
        return None
    size = len(a)
    g = [[0 for i in range(size)] for i in range(size)]
    operation = Operation(op)
    for i in range(size):
        for j in range(size):
            if i == j:
                continue
            g[i][j] = int(operation(a[i][j], b[i][j]))
    return g


import jax.numpy as jnp
import numpy as np

num_classes = 3

vocab = {'very0': jnp.array([2,0,0]),
         'very1': jnp.array([0,2,0]),
         'very2': jnp.array([0,0,2]),
         'some0': jnp.array([1,0,0]),
         'some1': jnp.array([0,1,0]),
         'some2': jnp.array([0,0,1]),
         'neutral0': jnp.array([0,0,0]),
         'neutral1': jnp.array([1,1,1]),
         'not0': jnp.array([-1,0,0]),
         'not1': jnp.array([0,-1,0]),
         'not2': jnp.array([0,0,-1])}

numerical_vocab = {0: jnp.array([2,0,0]),
                   1: jnp.array([0,2,0]),
                   2: jnp.array([0,0,2]),
                   3: jnp.array([1,0,0]),
                   4: jnp.array([0,1,0]),
                   5: jnp.array([0,0,1]),
                   6: jnp.array([0,0,0]),
                   7: jnp.array([1,1,1]),
                   8: jnp.array([-1,0,0]),
                   9: jnp.array([0,-1,0]),
                   10: jnp.array([0,0,-1])}

def uniform_length_sampler(length):

  def sample(num):
    return jnp.array([length]*num)

  return sample

def variable_length_sampler(mean, std):

  MIN_LENGTH = 20
  def sample(num):
    return np.round(np.clip(np.random.randn(num)*std + mean, MIN_LENGTH, np.inf)).astype(int)

  return sample

def ground_truth_label(sentences, lengths):

  scores = batch_score(sentences, lengths)
  return jnp.argmax(scores, axis=1)

def score(sentence, length):

  return sum([numerical_vocab[si] for si in sentence[:length]])

def batch_score(sentences, lengths):

  return jnp.array([score(s, l) for s,l in zip(sentences, lengths)])

def sample_batch(batch_size=64, length_sampler = None):

  if length_sampler == None:
    DEFAULT_LENGTH = 20
    length_sampler = uniform_length_sampler(DEFAULT_LENGTH)

  lengths = length_sampler(num=batch_size)
  max_length = max(lengths)

  num_tokens = len(numerical_vocab)
  batch = {'inputs': np.random.randint(num_tokens, size=(batch_size, max_length))}

  batch['labels'] = ground_truth_label(batch['inputs'], lengths)
  batch['index'] = lengths

  return batch


import json
from utility import *
import logging as logger

logger.basicConfig(level=logger.INFO, format='> %(message)s')

input_files = ['gopatriots', 'gohawks', 'nfl', 'patriots', 'sb49', 'superbowl']  # tweet-data file names

logger.info("EXECUTING: QUESTION 4 - CROSS VALIDATION LINEAR REGRESSION")

for file_name in input_files:
    logger.info("Working on #{}".format(file_name))
    tweets = open('../Dataset/tweet_data/tweets_#{0}.txt'.format(file_name), 'rb')
    first_tweet = json.loads(tweets.readline())
    start_time = first_tweet.get('firstpost_date')

    # features for model construction
    features = extra_feature_dict()
    flag_between = True  # flag to set new window start point
    flag_after = True  # flag to set new window start point
    current_window = 1
    end_time_of_window = start_time + current_window * 3600
    tweets.seek(0, 0)

    # store data
    tweet_features_before = []
    tweet_features_between = []
    tweet_features_after = []
    tweet_class_before = []
    tweet_class_between = []
    tweet_class_after = []

    logger.info("Extracting features from tweets")

    for tweet in tweets:
        tweet_data = json.loads(tweet)
        end_time = tweet_data.get('firstpost_date')

        '''
        features : retweets, followers, max_followers, impressions, favorite_count,
                   ranking_score, hour_of_day, number_of_users, number_of_long tweets
        '''

        if end_time < 1422777600:  # tweeted Before Feb. 1, 8:00 a.m.
            if end_time < end_time_of_window:
                features = calculate_statistic(features, tweet_data)
            else:
                extracted_features = get_features(features)
                tweet_class_before.append(extracted_features[0])
                tweet_features_before.append(extracted_features[1:])
                features = reset_features_dict()
                features = calculate_statistic(features, tweet_data)
                current_window += 1
                end_time_of_window = start_time + current_window * 3600

        elif (end_time > 1422777600) and (end_time < 1422820800):  # tweeted Between Feb. 1, 8:00 a.m to 8 p.m
            if flag_between:  # reset start_time
                flag_between = False
                current_window = 1
                start_time = end_time
                end_time_of_window = start_time + (current_window * 3600)

            if end_time < end_time_of_window:
                features = calculate_statistic(features, tweet_data)
            else:
                extracted_features = get_features(features)
                tweet_class_between.append(extracted_features[0])
                tweet_features_between.append(extracted_features[1:])
                features = reset_features_dict()
                features = calculate_statistic(features, tweet_data)
                current_window += 1
                end_time_of_window = start_time + current_window * 3600

        else:  # tweeted After Feb. 1, 8 p.m
            if flag_after:  # reset start_time
                flag_after = False
                current_window = 1
                start_time = end_time
                end_time_of_window = start_time + current_window * 3600

            if end_time < end_time_of_window:
                features = calculate_statistic(features, tweet_data)
            else:
                extracted_features = get_features(features)
                tweet_class_after.append(extracted_features[0])
                tweet_features_after.append(extracted_features[1:])
                features = reset_features_dict()
                features = calculate_statistic(features, tweet_data)
                current_window += 1
                end_time_of_window = start_time + current_window * 3600

    tweet_class_before = np.roll(np.array(tweet_class_before), -1)
    tweet_class_between = np.roll(np.array(tweet_class_between), -1)
    tweet_class_after = np.roll(np.array(tweet_class_after), -1)

    # perform 10 fold cross validation for each time-period
    logger.info("BEFORE MODEL - Classification")
    perform_classification(np.array(tweet_features_before), np.array(tweet_class_before))
    logger.info("BETWEEN MODEL - Classification")
    perform_classification(np.array(tweet_features_between), np.array(tweet_class_between))
    logger.info("AFTER MODEL - Classification")
    perform_classification(np.array(tweet_features_after), np.array(tweet_class_after))

    logger.info("*************************************")

Project 4 - Popularity Prediction on Twitter/Scripts/q4_superbowl.py


# @Time : 2020/10/12 11:21
# @Author : 30037
# @Email : 960364395@qq.com
# @File : search_order_case.py
# @Project : test_ecshop
from testcase.base_case import BaseCase
from page.home_page import HomePage
from time import sleep
from page.base_page import BasePage
from page.orderlist_page import OrderListPage
from page.order_query_page import OrderQueryPage
from selenium.webdriver.common.by import By
import unittest
from page.orderlist_page import OrderListPage

"""搜索订单"""
class SearchOrder(BaseCase):

    consignee_name = "dtt"
    phonenumber = "15342601420"
    adress_assert = '海德'
    lookover_locator = (By.XPATH,'//*[@id="listDiv"]/table[1]/tbody/tr[3]/td[7]/a')
    assert_adress_locator = (By.XPATH,'/html/body/form/div[1]/table/tbody/tr[18]/td[2]')
    assert_adress = '东城区'

    #根据收货人搜索订单，用例编号（ECshop_ST_ddgl_003）
    def test_search_order_by_consignee(self):
        #进入到订单页面
        hp = HomePage(self.driver)
        sleep(1)
        hp.click(hp.orderlist_locator)
        sleep(1)
        # 切换到main-frame
        bp = BasePage(self.driver)
        bp.switch_main_frame()
        sleep(1)
        #输入收货人
        op = OrderListPage(self.driver)
        op.send_keys(op.input_consignee_locator,self.consignee_name)
        #订单状态选择最上面一项
        select = op.select(op.select_status_locator)
        select.select_by_value('-1')
        sleep(1)
        #点击搜索
        op.click(op.search_locator)
        sleep(1)
        #断言
        get_consignee_name = op.text(op.consignee_locator)
        self.assertEqual(self.consignee_name,get_consignee_name)

    #根据手机号码搜索订单，用例编号“ECshop_ST_ddgl_016”
    def test_search_order_by_phonenumber(self):
        # 进入到订单页面
        hp = HomePage(self.driver)
        sleep(1)
        hp.click(hp.order_qurey_locator)
        sleep(1)
        # 切换到main-frame
        bp = BasePage(self.driver)
        bp.switch_main_frame()
        sleep(1)
        # 输入收手机号码
        oqp = OrderQueryPage(self.driver)
        hp.send_keys(oqp.input_phonenumber_locator,self.phonenumber)
        # 点击搜索
        hp.click(oqp.click_search_locator)
        sleep(1)
        #断言
        op = OrderListPage(self.driver)
        text = hp.text(op.phonenumber_locator)
        self.assertIn(self.phonenumber,text)

    #通过地址查询订单，用例编号，ECshop_ST_ddgl_014
    def test_search_order_by_adress(self):
        hp = HomePage(self.driver)
        sleep(1)
        hp.click(hp.order_qurey_locator)
        sleep(1)
        # 切换到main-frame
        bp = BasePage(self.driver)
        bp.switch_main_frame()
        sleep(1)
        # 输入地址
        oqp = OrderQueryPage(self.driver)
        hp.send_keys(oqp.input_adress_locator,self.adress_assert)
        # 点击搜索
        hp.click(oqp.click_search_locator)
        sleep(1)
        #断言
        op = OrderListPage(self.driver)
        text = hp.text(op.adress_locator)
        self.assertIn(self.adress_assert,text)

    #根据收货人查询订单，用例编号：ECshop_ST_ddgl_012
    def test_search_order_by_consignee_query(self):
        hp = HomePage(self.driver)
        sleep(1)
        hp.click(hp.order_qurey_locator)
        sleep(1)
        # 切换到main-frame
        bp = BasePage(self.driver)
        bp.switch_main_frame()
        sleep(1)
        # 输入收货人
        oqp = OrderQueryPage(self.driver)
        hp.send_keys(oqp.input_consignee_locator,self.consignee_name)
        # 点击搜索
        hp.click(oqp.click_search_locator)
        sleep(1)
        #断言
        op = OrderListPage(self.driver)
        text = hp.text(op.consignee_locator)
        self.assertIn(self.consignee_name,text)

    #根据所在地区查找订单
    def test_search_order_by_area(self):
        hp = HomePage(self.driver)
        sleep(1)
        hp.click(hp.order_qurey_locator)
        sleep(1)
        # 切换到main-frame
        bp = BasePage(self.driver)
        bp.switch_main_frame()
        sleep(1)
        # 选择地区
        oqp = OrderQueryPage(self.driver)
        hp.select(oqp.coutry_select_locator).select_by_index(1)
        sleep(1)
        hp.select(oqp.province_select_locator).select_by_index(1)
        sleep(1)
        hp.select(oqp.city_select_locator).select_by_index(1)
        sleep(1)
        hp.select(oqp.district_select_locator).select_by_index(1)
        sleep(1)
        # 点击搜索
        hp.click(oqp.click_search_locator)
        sleep(1)
        #断言
        hp.click(self.lookover_locator)
        sleep(1)
        text = hp.text(self.assert_adress_locator)
        self.assertIn(self.assert_adress,text)

if __name__ == '__main__':
    unittest.main()


import random
import os
import re
import math
import xml.etree.ElementTree as ET
import nltk
from nltk import sent_tokenize, word_tokenize
import torch

# Dataloader for CoNLL 2003 and i2b2 NER datasets
class CoNLL2003Reader:

    def __init__(self, train_file, dev_file, test_file, task_string):
        train_examples = self.read_conll_format_file(train_file)
        dev_examples = self.read_conll_format_file(dev_file) if dev_file != '' else []
        test_examples = self.read_conll_format_file(test_file)
        if not dev_examples:
            dev_size = int(math.ceil(0.1 * len(train_examples)))
            dev_examples = train_examples[:dev_size]
            train_examples = train_examples[dev_size:]

        # Dataset construction
        self.dataset = SeqDataset(train_examples, dev_examples, test_examples, task_string)

    def read_conll_format_file(self, filepath):
        reader = open(filepath)
        ex_id = 0
        tokens = []
        labels = []
        examples = []   # Each example is a dict containing ID, tokens and NER labels

        for line in reader:
            if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                if tokens:
                    examples.append({'id': ex_id, 'tokens': tokens, 'labels': labels})
                    ex_id += 1
                    tokens = []
                    labels = []
            else:
                ex_data = line.split()
                tokens.append(ex_data[0])
                label = ex_data[-1].rstrip()
                labels.append(label if label != 'O' else 'OUT')
        # Last example may be left out
        if tokens:
            examples.append({'id': ex_id, 'tokens': tokens, 'labels': labels})
        return examples

# Dataloader for i2b2 2006 DEID dataset
class i2b22006Reader:

    def __init__(self, train_file, test_file, task_string):
        train_text = self.embed_tags_in_text(train_file)
        test_text = self.embed_tags_in_text(test_file)

        # Example construction
        train_examples = self.construct_examples(train_text)
        test_examples = self.construct_examples(test_text)
        random.shuffle(train_examples)

        dev_size = int(math.ceil(0.1 * len(train_examples)))
        dev_examples = train_examples[:dev_size]
        train_examples = train_examples[dev_size:]

        # Label conversion into common format
        train_examples = self.correct_labels(train_examples)
        dev_examples = self.correct_labels(dev_examples)
        test_examples = self.correct_labels(test_examples)

        # Dataset construction
        self.dataset = SeqDataset(train_examples, dev_examples, test_examples, task_string)

    def embed_tags_in_text(self, file):
        root = ET.parse(open(file))
        doc_text = str(ET.tostring(root.getroot())).replace('\\n', '\n')
        phi_tags = re.findall(re.compile("<PHI [^/]*>[^<]*</PHI>[^\\s]*"), str(doc_text))
        for tag in phi_tags:
            phi_type = tag.split("TYPE=\"")[1].split('"')[0]
            pre_tag_text = tag.split("<")[0]
            post_tag_text = tag.split(">")[-1]
            tag_text = tag.split('</')[0].split('>')[-1].split()
            replacement_tag = pre_tag_text + ' ' + tag_text[0]+'||B-'+phi_type + ' ' + ' '.join([x+'||I-'+phi_type for x in tag_text[1:]]) + ' ' + post_tag_text
            replacement_tag = replacement_tag.strip()
            doc_text = str(doc_text).replace(tag, replacement_tag+' ')
        return doc_text

    def construct_examples(self, tagged_text):
        examples = []
        ex_id = 0
        for line in tagged_text.split('\n'):
            if line.startswith('<ROOT>') or line.startswith('<RECORD') or line.startswith('<TEXT') \
                    or line.startswith('</ROOT>') or line.startswith('</TEXT>') or line.startswith('</RECORD>'):
                continue
            words = line.split()
            tokens = []
            labels = []
            for word in words:
                if '||' not in word:
                    tokens.append(word)
                    labels.append('OUT')
                else:
                    token, label = word.split('||')
                    tokens.append(token)
                    labels.append(label)
            examples.append({'id': ex_id, 'tokens': tokens, 'labels': labels})
            ex_id += 1
        return examples

    def correct_labels(self, examples):
        label_dict = {
            'AGE': 'MISC',
            'ID': 'MISC',
            'PHONE': 'MISC',
            'DOCTOR': 'PER',
            'PATIENT': 'PER',
            'LOCATION': 'LOC',
            'HOSPITAL': 'ORG',
            'DATE': 'MISC'
        }
        for example in examples:
            corrected_labels = []
            for label in example['labels']:
                if '-' not in label:
                    corrected_labels.append(label)
                else:
                    start, tag = label.split('-')
                    corrected_labels.append('{}-{}'.format(start, label_dict[tag]))
            example['labels'] = corrected_labels
        return examples

# Dataloader for i2b2 2014 DEID dataset
class i2b22014Reader:

    def __init__(self, train_folder, test_folder, task_string):
        train_examples = self.read_xml_files(train_folder, 'train')
        test_examples = self.read_xml_files(test_folder, 'test')

        # Example construction
        train_examples = self.construct_examples(train_examples)
        test_examples = self.construct_examples(test_examples)
        random.shuffle(train_examples)

        dev_size = int(math.ceil(0.1 * len(train_examples)))
        dev_examples = train_examples[:dev_size]
        train_examples = train_examples[dev_size:]

        # Label conversion into common format
        train_examples = self.correct_labels(train_examples)
        dev_examples = self.correct_labels(dev_examples)
        test_examples = self.correct_labels(test_examples)

        # Dataset construction
        self.dataset = SeqDataset(train_examples, dev_examples, test_examples, task_string)

    def read_xml_files(self, folder, split):
        texts = []
        for file in os.listdir(folder):
            root = ET.parse(open(os.path.join(folder, file)))
            text, tags = '', []
            for child in root.getroot():
                if child.tag == 'TEXT':
                    text = child.text
                if child.tag == 'TAGS':
                    for node in child:
                        tag_type = node.tag
                        tag_start = int(node.attrib['start'])
                        tag_end = int(node.attrib['end'])
                        tag_text = node.attrib['text']
                        compare_text = text[tag_start : tag_end]
                        tags.append([tag_type, tag_start, tag_end, tag_text])
                        if ' '.join(compare_text.split()) != ' '.join(tag_text.split()):
                            if '&' in compare_text:
                                continue
                            print('ERROR: Annotation offsets seem to be incorrect!')
            for node in reversed(tags):
                tokens = word_tokenize(node[-1])
                tokens = [x + '||' + 'I-' + node[0] for x in tokens]
                tokens[0] = tokens[0].replace('||I-', '||B-')
                text = text[:node[1]] + ' ' + ' '.join(tokens) + ' ' + text[node[2]:]
            texts.append(text)
        return texts

    def construct_examples(self, texts):
        examples = []
        ex_id = 0
        for text in texts:
            sents = sent_tokenize(text)
            for sent in sents:
                words = word_tokenize(sent)
                tokens = []
                labels = []
                for word in words:
                    if '||' not in word:
                        tokens.append(word)
                        labels.append('OUT')
                    else:
                        token, label = word.split('||')
                        tokens.append(token)
                        labels.append(label)
                examples.append({'id': ex_id, 'tokens': tokens, 'labels': labels})
                ex_id += 1
        return examples

    def correct_labels(self, examples):
        label_dict = {
            'DATE': 'MISC',
            'PROFESSION': 'MISC',
            'CONTACT': 'MISC',
            'LOCATION': 'LOC',
            'NAME': 'PER',
            'PHI': 'MISC',
            'ID': 'MISC',
            'AGE': 'MISC'
        }
        for example in examples:
            corrected_labels = []
            for label in example['labels']:
                if '-' not in label:
                    corrected_labels.append(label)
                else:
                    start, tag = label.split('-')
                    corrected_labels.append('{}-{}'.format(start, label_dict[tag]))
            example['labels'] = corrected_labels
        return examples

# Dataloader for Timebank and MTSamples event extraction datasets
class EventTSVReader:

    def __init__(self, folder, task_string):
        train_ids = open(os.path.join(folder, 'train_ids.txt')).read().splitlines()
        dev_ids = open(os.path.join(folder, 'dev_ids.txt')).read().splitlines()
        test_ids = open(os.path.join(folder, 'test_ids.txt')).read().splitlines()

        examples = self.read_examples_from_files(folder)
        train_examples, dev_examples, test_examples = [], [], []
        for file_id in examples:
            if file_id in train_ids:
                train_examples += examples[file_id]
            if file_id in dev_ids:
                dev_examples += examples[file_id]
            if file_id in test_ids:
                test_examples += examples[file_id]

        # Dataset construction
        self.dataset = SeqDataset(train_examples, dev_examples, test_examples, task_string)

    def read_examples_from_files(self, folder):
        examples = {}
        ex_id = 0
        for file in os.listdir(folder):
            examples[file.split('.tsv')[0]] = []
            tokens, labels = [], []
            reader = open(os.path.join(folder, file))
            for line in reader:
                if line == '\n':
                    if tokens:
                        examples[file.split('.tsv')[0]].append({'id': ex_id, 'tokens': tokens, 'labels': labels})
                        ex_id += 1
                        tokens, labels = [], []
                else:
                    if '\t' not in line:
                        continue
                    if len(line.strip().split('\t')) == 1:  # Some misaligned entity annotations in MTSamples cause this
                        continue
                    token, label = line.strip().split('\t')
                    tokens.append(token)
                    labels.append(label if label != 'ENT' else 'O')
            # Last example may be left out
            if tokens:
                examples[file.split('.tsv')[0]].append({'id': ex_id, 'tokens': tokens, 'labels': labels})
        return examples

# Dataloader for i2b2 2012 event extraction dataset
class i2b22012Reader:

    def __init__(self, train_folder, test_folder, task_string):
        train_examples = self.read_examples_from_files(train_folder)
        test_examples = self.read_examples_from_files(test_folder)
        random.shuffle(train_examples)

        dev_size = int(math.ceil(0.1 * len(train_examples)))
        dev_examples = train_examples[:dev_size]
        train_examples = train_examples[dev_size:]

        # Dataset construction
        self.dataset = SeqDataset(train_examples, dev_examples, test_examples, task_string)

    def read_examples_from_files(self, folder):
        examples = []
        ex_id = 0
        for file in os.listdir(folder):
            if not file.endswith('.extent'):
                continue
            file_id = file.split('.xml')[0]
            text_reader = open(os.path.join(folder, file_id + '.xml.txt'))
            lines = []
            for line in text_reader:
                lines.append(line.strip().split())
            anno_reader = open(os.path.join(folder, file_id + '.xml.extent'), 'r')
            events = [['O']*len(i) for i in lines]
            for line in anno_reader:
                if not line.startswith('EVENT'):
                    continue
                event, _, _, _ = line.split('||')
                _, event_text, event_offsets = event.split('"')
                start_offset, end_offset = event_offsets.strip().split()
                start_sent, start_token = start_offset.split(':')
                end_sent, end_token = end_offset.split(':')
                start_sent, start_token, end_sent, end_token = int(start_sent), int(start_token), int(end_sent), int(end_token)
                if ' '.join(lines[start_sent-1][start_token : end_token+1]) != event_text:
                    print('ERROR: Annotation offsets seem to be incorrect!')
                events[start_sent-1][start_token : end_token+1] = ['EVENT'] * (end_token+1 - start_token)
            for token_list, event_list in zip(lines, events):
                examples.append({'id': ex_id, 'tokens': token_list, 'labels': event_list})
                ex_id += 1
        return examples

# Dataloader for file containing unlabeled data
class UnlabeledDataset:

    def __init__(self, text_file):

        reader = open(text_file)
        self.text = []
        for line in reader:
            self.text.append(line.strip())
        print('Loaded {} lines'.format(len(self.text)))

    def batch_and_tokenize_data(self, examples, tokenizer, batch_size):
        random.shuffle(examples)
        batches = []
        for i in range(0, len(examples), batch_size):
            start = i
            end = min(start+batch_size, len(examples))
            batch = self.tokenize_examples(examples[start:end], tokenizer)
            batches.append(batch)
        return batches

    def tokenize_examples(self, examples, tokenizer):
        tokenized_inputs = tokenizer(
            examples,
            padding='max_length',
            truncation=True,
            max_length=128,
            return_tensors='pt'
        )
        return tokenized_inputs

# Class to mix labeled and unlabeled data for domain classification in Instance Weighting and Adversarial Training
class CombinedDataset:

    def __init__(self, source_sents, target_sents):
        min_len = min(len(source_sents), len(target_sents))
        random.shuffle(source_sents)
        random.shuffle(target_sents)
        source_data = source_sents[:min_len]
        target_data = target_sents[:min_len]
        self.domain_labels = [0] * len(source_data) + [1] * len(target_data)
        self.domain_data = source_data + target_data

    def batch_and_tokenize_data(self, tokenizer, batch_size):
        combined_data = list(zip(self.domain_data, self.domain_labels))
        random.shuffle(combined_data)
        domain_data, domain_labels = zip(*combined_data)
        domain_data, domain_labels = list(domain_data), list(domain_labels)
        batches = []
        for i in range(0, len(domain_data), batch_size):
            start = i
            end = min(start+batch_size, len(domain_data))
            batch = self.tokenize_examples(domain_data[start:end], domain_labels[start:end], tokenizer)
            batches.append(batch)
        return batches

    def tokenize_examples(self, examples, labels, tokenizer):
        tokenized_inputs = tokenizer(
            examples,
            padding='max_length',
            truncation=True,
            max_length=128,
            return_tensors='pt'
        )
        tokenized_inputs["labels"] = torch.LongTensor(labels)
        return tokenized_inputs

# Class to store sequence labeling dataset (either for NER or event extraction)
class SeqDataset:

    def __init__(self, train_data, dev_data, test_data, task):

        self.train_data = train_data
        self.dev_data = dev_data
        self.test_data = test_data
        self.task = task

        if task == 'ner-fine':
            self.label_vocab = {
                'B-PER': 0,
                'I-PER': 1,
                'B-LOC': 2,
                'I-LOC': 3,
                'B-ORG': 4,
                'I-ORG': 5,
                'B-MISC': 6,
                'I-MISC': 7,
                'OUT': 8
            }
        elif task == 'ee':
            self.label_vocab = {'EVENT': 1, 'O': 0}
        elif task == 'ner-coarse':
            self.label_vocab = {
                'B-ENT': 0,
                'I-ENT': 1,
                'OUT': 2
            }
        else:
            print('Task must be either NER (fine or coarse) or event extraction!!')
            exit(1)

        self.construct_label_sequences(self.train_data)
        self.construct_label_sequences(self.dev_data)
        self.construct_label_sequences(self.test_data)

    def construct_label_sequences(self, examples):
        if self.task != 'ner-coarse':
            for example in examples:
                label_seq = [self.label_vocab[x] for x in example['labels']]
                example['gold_seq'] = label_seq
        else:
            for example in examples:
                label_seq = []
                labels = []
                for label in example['labels']:
                    if label.startswith('O'):
                        label_seq.append(self.label_vocab['OUT'])
                        labels.append('OUT')
                    elif label.startswith('B'):
                        label_seq.append(self.label_vocab['B-ENT'])
                        labels.append('B-ENT')
                    elif label.startswith('I'):
                        label_seq.append(self.label_vocab['I-ENT'])
                        labels.append('I-ENT')
                example['gold_seq'] = label_seq
                example['labels'] = labels

    def batch_and_tokenize_data(self, tokenizer, batch_size):
        final_batches = []
        self.tokenized_examples = []
        for examples in [self.train_data, self.dev_data, self.test_data]:
            random.shuffle(examples)
            self.tokenized_examples.append(examples)
            batches = []
            for i in range(0, len(examples), batch_size):
                start = i
                end = min(start+batch_size, len(examples))
                batch = self.tokenize_and_align_labels(examples[start:end], tokenizer)
                batches.append(batch)
            final_batches.append(batches)
        return final_batches

    def tokenize_and_align_labels(self, examples, tokenizer):
        example_texts = [x['tokens'] for x in examples]
        tokenized_inputs = tokenizer(
            example_texts,
            padding='max_length',
            truncation=True,
            # We use this argument because the texts in our dataset are lists of words (with a label for each word).
            is_split_into_words=True,
            max_length=128,
            return_tensors='pt'
        )
        labels = []
        for i, example in enumerate(examples):
            label_seq = example['gold_seq']
            word_ids = tokenized_inputs.word_ids(batch_index=i)
            previous_word_idx = None
            label_ids = []
            for word_idx in word_ids:
                # Special tokens have a word id that is None. We set the label to -100 so they are automatically
                # ignored in the loss function.
                if word_idx is None:
                    label_ids.append(-100)
                # We set the label for the first token of each word.
                elif word_idx != previous_word_idx:
                    label_ids.append(label_seq[word_idx])
                # For the other tokens in a word, we set the label to -100, but we might want to change that?
                else:
                    label_ids.append(-100)
                previous_word_idx = word_idx
            labels.append(label_ids)
        tokenized_inputs["labels"] = torch.LongTensor(labels)
        tokenized_inputs["example_ids"] = [x['id'] for x in examples]
        return tokenized_inputs


import ctypes
from ctypes import c_void_p, c_uint64, c_int32, c_ulonglong, c_size_t, c_float, c_longlong, c_bool
import numpy as np

import os
import time

src_dir = os.path.join(os.path.dirname(__file__))

start_time = time.time()
print("LOADING....")
if os.environ.get("COMPILE_MPM") or not os.path.isfile(f"{src_dir}/libmaniskill_mpm.so"):
    command = f"nvcc -O3 -Xptxas -O3,-v -ccbin=g++ --compiler-options -fPIC -shared {src_dir}/csrc/integrator.cu -o {src_dir}/libmaniskill_mpm.so"
    print(command)
    if os.system(command) != 0:
        print("Failed to compile mpm library")
        exit()

lib = ctypes.cdll.LoadLibrary(f"{src_dir}/libmaniskill_mpm.so")
print(f"Loaded with {time.time() - start_time} secs.")

cuda_stream_t = c_void_p
texture_t = c_ulonglong

class vec3(ctypes.Array):
    _length_ = 3
    _type_ = ctypes.c_float

    length = 3
    ctype = ctypes.c_float
    size = 12
    np_type = np.float32

    def __repr__(self):
        return f"ivec3({self[0]}, {self[1]}, {self[2]})"

class quat(ctypes.Array):
    _length_ = 4
    _type_ = ctypes.c_float

    length = 4
    ctype = ctypes.c_float
    size = 16
    np_type = np.float32

    def __repr__(self):
        return f"quat({self[0]}, {self[1]}, {self[2]}, {self[3]})"

class ivec3(ctypes.Array):
    _length_ = 3
    _type_ = ctypes.c_int

    length = 3
    ctype = ctypes.c_int
    size = 12
    np_type = np.int32

    def __repr__(self):
        return f"ivec3({self[0]}, {self[1]}, {self[2]})"

class mat3(ctypes.Array):
    _length_ = 9
    _type_ = ctypes.c_float

    length = 9
    ctype = ctypes.c_float
    size = 36
    np_type = np.float32

    def __repr__(self):
        return f"ivec3({self[0]}, {self[1]}, {self[2]}, {self[3]}, {self[4]}, {self[5]}, {self[6]}, {self[7]}, {self[8]})"

class texture_resources(ctypes.Structure):
    _fields_ = [("array", c_void_p), ("texture", texture_t)]

class float32:
    length = 1
    ctype = ctypes.c_float
    size = 4
    np_type = np.float32

class int32:
    length = 1
    ctype = ctypes.c_int
    size = 4
    np_type = np.int32

class int64:
    length = 1
    ctype = ctypes.c_int64
    size = 8
    np_type = np.int64

lib.cuda_alloc.restype = c_void_p
lib.cuda_alloc.argtypes = [c_size_t]

lib.cuda_free.argtypes = [c_void_p]

lib.cuda_upload.argtypes = [c_void_p, c_void_p, c_size_t]
lib.cuda_upload_async.argtypes = [c_void_p, c_void_p, c_size_t, cuda_stream_t]

lib.cuda_download.argtypes = [c_void_p, c_void_p, c_size_t]
lib.cuda_download_async.argtypes = [c_void_p, c_void_p, c_size_t, cuda_stream_t]

lib.cuda_copy.argtypes = [c_void_p, c_void_p, c_size_t]

lib.cuda_copy2d.argtypes = [c_void_p, c_size_t, c_void_p, c_size_t, c_size_t, c_size_t]

lib.cuda_zero.argtypes = [c_void_p, c_size_t]
lib.cuda_zero_async.argtypes = [c_void_p, c_size_t, cuda_stream_t]

lib.create_volume.argtypes = [c_void_p, c_int32, c_int32]
lib.create_volume.restype = texture_resources

lib.destroy_volume.argtypes = [texture_resources]

lib.cuda_stream_create.restype = cuda_stream_t
lib.cuda_stream_destroy.argtypes = [cuda_stream_t]
lib.cuda_stream_sync.argtypes = [cuda_stream_t]

lib.compute_grid_lower.argtypes = [
    c_void_p,
    c_float,
    c_float,
    c_void_p,
    c_int32,
    cuda_stream_t,
]

lib.compute_svd.argtypes = [
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_float,
    c_int32,
    cuda_stream_t,
]

lib.compute_svd_grad.argtypes = [
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,
    c_float,
    c_int32,
    cuda_stream_t,
]

lib.p2g.argtypes = [
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,

    ctypes.POINTER(ivec3),
    # ivec3,

    c_float,
    c_float,
    c_float,
    c_void_p,
    c_void_p,
    c_void_p,
    c_int32,
    cuda_stream_t,
]

lib.p2g_grad.argtypes = [
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,

    ctypes.POINTER(ivec3),
    # ivec3,

    c_float,
    c_float,
    c_float,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    c_int32,
    cuda_stream_t,
]

lib.grid_op_v2.argtypes = [c_void_p] * 11 + [c_float] * 4 + [c_void_p, ctypes.POINTER(ivec3), c_int32, cuda_stream_t]

lib.grid_op_v2_grad.argtypes = [c_void_p] * 11 + [c_void_p] * 6 + [c_float] * 4 + [c_void_p, c_void_p, ctypes.POINTER(ivec3), c_int32, cuda_stream_t]

lib.g2p.argtypes = [
    c_void_p,
    c_void_p,
    c_void_p,
    c_float,
    c_float,
    c_float,
    ctypes.POINTER(ivec3),
    c_void_p,
    c_float,
    c_void_p,
    c_void_p,
    c_int32,
    cuda_stream_t,
]

lib.g2p_grad.argtypes = [
    c_void_p,
    c_void_p,
    c_void_p,
    c_float,
    c_float,
    c_float,
    ctypes.POINTER(ivec3),
    c_void_p,
    c_float,
    c_void_p,
    c_void_p,
    c_int32,

    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,
    c_void_p,

    cuda_stream_t,
]

lib.render.argtypes = [c_void_p] * 13 + [c_float, ctypes.POINTER(ivec3), c_int32, c_bool, ctypes.POINTER(ivec3), c_int32, c_int32, c_float, c_int32, c_void_p, cuda_stream_t]

lib.particle_sdf.argtypes = [c_void_p] * 5 + [c_int32, ctypes.POINTER(ivec3), c_float] + [c_void_p] * 3 + [c_int32, cuda_stream_t]

lib.compute_dist.argtypes = [c_void_p] * 6 + [c_int32] + [c_void_p] * 4 + [c_int32, c_int32, cuda_stream_t]

lib.particle2mass.argtypes = [c_void_p] * 3 + [ctypes.POINTER(ivec3), c_float, c_float] + [c_void_p] * 4 + [c_int32, c_int32, c_int32, cuda_stream_t]

class array:
    def __init__(self, dtype=float32, length=0):
        assert length > 0
        if dtype == float:
            dtype = float32
        if dtype == int:
            dtype = int32
        self.dtype = dtype

        self.bytes = dtype.size
        self.nbytes = dtype.size * length
        self.data_ptr = lib.cuda_alloc(self.nbytes)

        shape = (length,)
        if dtype.length != 1:
            shape = shape + (dtype.length,)
        self.shape = shape

        self.zero()

    def upload_async(self, arr: np.ndarray, stream):
        # only for pos
        n = arr.shape[0]
        assert arr.shape[0] == self.shape[0]
        data_size = self.bytes * n

        assert self.dtype.np_type == arr.dtype
        assert self.shape == arr.shape
        assert self.nbytes == arr.nbytes
        assert arr.data.contiguous
        ptr = arr.__array_interface__["data"][0]
        lib.cuda_upload_async(self.data_ptr, ptr, data_size, stream)

    def upload(self, arr: np.ndarray, strict=False):
        arr = np.ascontiguousarray(arr, dtype=self.dtype.np_type)
        n = arr.shape[0]
        if strict:
            assert n == self.shape[0]
        assert n <= self.shape[0], f"{n} {self.shape[0]}"
        data_size = self.bytes * n

        assert self.dtype.np_type == arr.dtype, f"{self.dtype.np_type}, {arr.dtype}"
        assert (arr.shape[0],) + self.shape[1:] == arr.shape, f"{self.shape}, {arr.shape}"

        ptr = arr.__array_interface__["data"][0]
        lib.cuda_upload(self.data_ptr, ptr, data_size)

    def download(self, n=None, device='numpy', stream=None):
        if device != 'numpy':
            import torch
            x = self.download(n, 'numpy')
            return torch.tensor(x, device=device)

        if n is not None:
            shape = (n,) + self.shape[1:]
            data_size = self.bytes * n
        else:
            shape = self.shape
            data_size = self.nbytes
        #shape = (n,) + self.shape[1:] if n is not None else self.shape
        arr = np.empty(shape, self.dtype.np_type)
        ptr = arr.__array_interface__["data"][0]

        if stream is not None:
            lib.cuda_download(ptr, self.data_ptr, data_size)
        else:
            lib.cuda_download_async(ptr, self.data_ptr, data_size, stream)
        return arr

    """
    def download_async(self, stream):
        arr = np.empty(self.shape, self.dtype.np_type)
        ptr = arr.__array_interface__["data"][0]
        lib.cuda_download_async(ptr, self.data_ptr, self.nbytes, stream)
        return arr
    """

    def copy_to_torch(self, tensor):
        assert tensor.device.type == "cuda"
        assert tensor.is_contiguous()
        width = spitch = self.dtype.size
        dpitch = tensor.stride()[0] * tensor.element_size()
        height = self.nbytes // width
        lib.cuda_copy2d(tensor.data_ptr(), dpitch, self.data_ptr, spitch, width, height)

    def cuda_add(self, values, stream=None):
        # TODO: make it faster
        x = self.download(n=len(values), stream=stream)
        x += values
        self.upload(x, stream)

    def zero(self, stream=None):
        if stream is None:
            lib.cuda_zero(self.data_ptr, self.nbytes)
        else:
            lib.cuda_zero_async(self.data_ptr, self.nbytes, stream)

    def zero_async(self, stream):
        lib.cuda_zero_async(self.data_ptr, self.nbytes, stream)

    def __del__(self):
        lib.cuda_free(self.data_ptr)

    def __repr__(self):
        return self.download().__repr__()


#STEP I
#Import the JSON Library and print out definition

#STEP II
#Wir erweitern die Funktion retrieve_definition um die Abfrage, ob das Wort
#überhaupt in dem Dictionary existiert und geben es nur dann aus, wenn
#es auch da ist

#STEP III
#Wir erweitern die Funktion so, dass Nutzer 'Rain' oder 'rain' eingeben und
#trotzdem den gleichen output bekommen, using lower()
#Nutzer können jetzt also schreiben, wie sie wollen, also z.B
#'football', 'Football' oder 'FOOTBALL' eingeben und kriegen trd die passende
#definition ausgegeben, sofern sie denn existiert ('else') und der Input
#keinen Tippfehler enthält

import json

#JSON or JavaScript Object Notationm is a minimal, readable (for both humans
# and computers) format for structuring data. It mainly has two things, a
#key and a value associated with the key.

#The important thing here: We'll load the data in json Format, but it will be
#stored in the 'data' variable as python dictionary data format.

#Loading the json data as python dictionary
#Try typing 'type(data)' in terminal after executing the first two lines
#of this snippet
data=json.load(open('dictionary.json'))

#print(type(data))
#out: <class 'dict'>

#Function for retrieving definition
def retrieve_definition(word):
    #Alle Wörter erstmal in LowerCase, weil unsere Daten in dem Format sind
    word = word.lower()

    #Abfrage, ob das Wort existiert
    if word in [data]:
        return data[word]

    #Aus dem kleingeschriebenen fragen wir jetzt ab, ob es im Datensatz das
    #Wort mit großem Anfangsbuchstaben gibt. Also z.B. Delhi oder Texas.
    #Wenn ja, geben wir die Definition des Worts aus.
    elif word.title() in data:
        return data[word.title()]

    #Jetzt transformieren wir das LowerCase Wort in UPPERCASE (komplett) und
    #prüfen, ob es dieses WORT im Datensatz gibt, um auch Akronyme wie USA, NATO
    #mit einzubeziehen. Wenn ja, geben wir die Definition des WORTS aus.
    elif word.upper() in data:
        return data[word.upper()]

    else:
        return("Word doesn't exist, please double check it.")

#Input from user:
word_user = input("Enter a word to get definition: ")
print(retrieve_definition(word_user))

LPBRWApps | Part I | Dictionary/step3.py