## test_attack.py -- sample code to test attack procedure
##
## Copyright (C) 2016, Nicholas Carlini <nicholas@carlini.com>.
##
## This program is licenced under the BSD 2-Clause licence,
## contained in the LICENCE file in this directory.

import tensorflow as tf
import numpy as np
import time
import random
import argparse
from cv2 import imwrite

from setup_cifar import CIFAR, CIFARModel
from setup_mnist import MNIST, MNISTModel
from setup_mnets import MNETSModel, ImagenetTF, NodeLookup

from l2_attack import CarliniL2
from l0_attack import CarliniL0
from li_attack import CarliniLi

slim = tf.contrib.slim

def show(img_input, name=None):
    """
    Show MNSIT digits in the console.
    """
    remap = "  .*#"+"#"*100
    img = (img_input.flatten()+.5)*3
    if len(img) != 784 and name is not None:
        scaled = (0.5+img_input)*255
        imwrite(name, scaled)
        return
    print("START")
    for i in range(28):
        print("".join([remap[int(round(x))] for x in img[i*28:i*28+28]]))

def generate_data(data, samples, num_targets=1, targeted=True, target_classes=None, start=0, imagenet=True):
    """
    Generate the input data to the attack algorithm.

    data: the images to attack
    samples: number of samples to use
    targeted: if true, construct targeted attacks, otherwise untargeted attacks
    start: offset into data to use
    inception: if targeted and inception, randomly sample 100 targets intead of 1000
    """
    targets = []
    inputs_scaled = []

    print(data.test_data.shape)
    print(data.test_labels.shape)
    if not imagenet:
        seq = range(data.test_labels.shape[1])
    else:
        seq = target_classes if target_classes else random.sample(range(1,1001), num_targets)
    print("Target classes", seq)

    for i in range(samples):
        if targeted:
            for j in seq:
                if (j == np.argmax(data.test_labels[start+i])) and (imagenet == False):
                    continue
                inputs_scaled.append(data.test_data[start+i]/2.0)
                targets.append(np.eye(data.test_labels.shape[1])[j])
        else:
            inputs_scaled.append(data.test_data[start+i]/2.0)
            targets.append(data.test_labels[start+i])

    inputs_scaled = np.array(inputs_scaled)
    #i = 1
    #for img in inputs_scaled:
    #    show(img, 'orig' + str(i) + '.png')
    #    i += 1
    targets = np.array(targets)
    return inputs_scaled, targets

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Checkpoint file')
    parser.add_argument('--ckpt', dest='ckpt', type=str,
                        help='Checkpoint file for model')
    parser.add_argument('--model', dest='model_name', type=str,
                        help='model_name')
    parser.add_argument('--norm', dest='norm', type=str,
                        help='distance metric')
    parser.add_argument('--conf', dest='conf', type=int,
                       help='confidence')
    #parser.add_argument('--targets', dest='num_targets', type=int,
    #                    help='number of targets')
    args = parser.parse_args()
    batch_size = 5

    with tf.Session() as sess:
        print("Running model {}".format(args.model_name))
        model = MNETSModel(args.ckpt, args.model_name, batch_size, sess)
        data =  ImagenetTF(args.model_name, model.image_size)

        tf.train.start_queue_runners(sess)
        data.get_batch(sess)

        if args.norm == "0":
            norm = CarliniL0
        elif args.norm == "i":
            norm = CarliniLi
        else:
            norm = CarliniL2

        target_classes = [893, 858, 350, 71, 948, 715, 558, 408, 349, 215]
        target_classes = target_classes
        attack = norm(sess, model, max_iterations=1000, confidence=args.conf)
        inputs, targets = generate_data(data, samples=len(target_classes),
                                        targeted=True, target_classes=target_classes,
                                        start=0, imagenet=True)

        print("Attack constructed")

        #print(tf.global_variables())
        if args.model_name == "inception_v3":
            variables_to_restore = slim.get_variables(scope="InceptionV3")
        elif args.model_name == "resnet_v2_152":
            variables_to_restore = slim.get_variables(scope="ResnetV2152")
        elif args.model_name.startswith("mobilenet"):
            variables_to_restore = slim.get_variables(scope="MobilenetV1")

        saver = tf.train.Saver(variables_to_restore)
        saver.restore(sess, args.ckpt)
        print("Checkpoint restored")

        print("Running attack")
        timestart = time.time()
        adv = attack.attack(inputs, targets)
        timeend = time.time()

        print("Took",timeend-timestart,"seconds to run",len(inputs),"samples.")
        adv = adv.astype(np.float32)

    	# Creates node ID --> English string lookup.
        node_lookup = NodeLookup()

        def topk(predictions):
            top_k = predictions.argsort()#[-FLAGS.num_top_predictions:][::-1]
            top_k = top_k[::-1]
            count = 1
            for node_id in top_k[:10]:
                #print('ID {}, score {}'.format(node_id, predictions[node_id]))
                human_string = node_lookup.id_to_string(node_id)
                score = predictions[node_id]
                print("{}. {} (score = {})".format(count, human_string, score))
                count += 1

        for i in range(len(adv)):
            print("Types:", inputs[i].dtype, adv[i].dtype)
            print("Valid:")
            show(inputs[i], name="input" + str(i) + ".png")
            print("Adversarial:")
            show(adv[i], name="adv" + str(i) + ".png")

            pred = model.predict(tf.convert_to_tensor(inputs[i:i+1]), reuse=True)
            pred_adv = model.predict(tf.convert_to_tensor(adv[i:i+1]), reuse=True)

            pred = sess.run(pred)
            pred_adv = sess.run(pred_adv)

            pred = np.squeeze(pred)
            pred_adv = np.squeeze(pred_adv)

            print("Original classification:")
            topk(pred)
            print("Adversarial classification:")
            topk(pred_adv)

            print("Total distortion:", np.sum((adv[i]-inputs[i])**2)**.5)


#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
dedupe provides the main user interface for the library the
Dedupe class
"""

import itertools
import logging
import pickle
import multiprocessing
import warnings
import os
import sqlite3
import tempfile

import numpy
import json
import rlr

import dedupe.core as core
import dedupe.serializer as serializer
import dedupe.blocking as blocking
import dedupe.clustering as clustering
import dedupe.datamodel as datamodel
import dedupe.labeler as labeler
import dedupe.predicates

from typing import (Mapping,
                    Optional,
                    List,
                    Tuple,
                    Set,
                    Dict,
                    Union,
                    Generator,
                    Iterable,
                    Sequence,
                    BinaryIO,
                    cast,
                    TextIO)
from typing_extensions import Literal
from dedupe._typing import (Data,
                            Clusters,
                            RecordPairs,
                            RecordID,
                            RecordDict,
                            Blocks,
                            TrainingExample,
                            LookupResults,
                            Links,
                            TrainingData,
                            Classifier,
                            JoinConstraint)

logger = logging.getLogger(__name__)

class Matching(object):
    """
    Base Class for Record Matching Classes
    """

    def __init__(self,
                 num_cores: Optional[int],
                 in_memory: bool = False,
                 **kwargs) -> None:

        if num_cores is None:
            self.num_cores = multiprocessing.cpu_count()
        else:
            self.num_cores = num_cores

        self.in_memory = in_memory
        self._fingerprinter: Optional[blocking.Fingerprinter] = None
        self.data_model: datamodel.DataModel
        self.classifier: Classifier
        self.predicates: List[dedupe.predicates.Predicate]

    @property
    def fingerprinter(self) -> blocking.Fingerprinter:

        if self._fingerprinter is None:
            raise ValueError('the record fingerprinter is not intialized, '
                             'please run the train method')

        return self._fingerprinter

class IntegralMatching(Matching):
    """
    This class is for linking class where we need to score all possible
    pairs before deciding on any matches
    """

    def score(self,
              pairs: RecordPairs) -> Union[numpy.memmap, numpy.ndarray]:
        """
        Scores pairs of records. Returns pairs of tuples of records id and
        associated probabilities that the pair of records are match

        Args:
            pairs: Iterator of pairs of records

        """
        try:
            matches = core.scoreDuplicates(pairs,
                                           self.data_model,
                                           self.classifier,
                                           self.num_cores)
        except RuntimeError:
            raise RuntimeError('''
                You need to either turn off multiprocessing or protect
                the calls to the Dedupe methods with a
                `if __name__ == '__main__'` in your main module, see
                https://docs.python.org/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods''')

        return matches

class DedupeMatching(IntegralMatching):
    """
    Class for Deduplication, extends Matching.

    Use DedupeMatching when you have a dataset that can contain
    multiple references to the same entity.

    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

    def partition(self,
                  data: Data,
                  threshold: float = 0.5) -> Clusters:  # pragma: no cover
        """
        Identifies records that all refer to the same entity, returns
        tuples containing a sequence of record ids and corresponding
        sequence of confidence score as a float between 0 and 1. The
        record_ids within each set should refer to the same entity and the
        confidence score is a measure of our confidence a particular entity
        belongs in the cluster.

        For details on the confidence score, see :func:`dedupe.Dedupe.cluster`.

        This method should only used for small to moderately sized
        datasets for larger data, you need may need to generate your
        own pairs of records and feed them to :func:`~score`.

        Args:
            data: Dictionary of records, where the keys are record_ids
                  and the values are dictionaries with the keys being
                  field names

            threshold: Number between 0 and 1.  We
                       will only consider put together records into
                       clusters if the `cophenetic similarity
                       <https://en.wikipedia.org/wiki/Cophenetic>`_ of
                       the cluster is greater than the threshold.

                       Lowering the number will increase recall,
                       raising it will increase precision

        .. code:: python

           > clusters = matcher.partition(data, threshold=0.5)
           > print(duplicates)
           [((1, 2, 3), (0.790, 0.860, 0.790)),
            ((4, 5), (0.720, 0.720)),
            ((10, 11), (0.899, 0.899))]

        """
        pairs = self.pairs(data)
        pair_scores = self.score(pairs)
        clusters = self.cluster(pair_scores, threshold)

        clusters = self._add_singletons(data, clusters)

        clusters = list(clusters)

        try:
            mmap_file = pair_scores.filename  # type: ignore
        except AttributeError:
            pass
        else:
            del pair_scores
            if mmap_file:
                os.remove(mmap_file)

        return clusters

    def _add_singletons(self, data, clusters):

        singletons = set(data.keys())

        for record_ids, score in clusters:
            singletons.difference_update(record_ids)
            yield (record_ids, score)

        for singleton in singletons:
            yield (singleton, ), (1.0, )

    def pairs(self, data):
        '''
        Yield pairs of records that share common fingerprints.

        Each pair will occur at most once. If you override this
        method, you need to take care to ensure that this remains
        true, as downstream methods, particularly :func:`cluster`, assumes
        that every pair of records is compared no more than once.

        Args:
            data: Dictionary of records, where the keys are record_ids
                  and the values are dictionaries with the keys being
                  field names

        .. code:: python

            > pairs = matcher.pairs(data)
            > print(list(pairs))
            [((1, {'name' : 'Pat', 'address' : '123 Main'}),
              (2, {'name' : 'Pat', 'address' : '123 Main'})),
             ((1, {'name' : 'Pat', 'address' : '123 Main'}),
              (3, {'name' : 'Sam', 'address' : '123 Main'}))
             ]

        '''

        self.fingerprinter.index_all(data)

        id_type = core.sqlite_id_type(data)

        # Blocking and pair generation are typically the first memory
        # bottlenecks, so we'll use sqlite3 to avoid doing them in memory
        with tempfile.TemporaryDirectory() as temp_dir:
            if self.in_memory:
                con = sqlite3.connect(':memory:')
            else:
                con = sqlite3.connect(temp_dir + '/blocks.db')

            # Set journal mode to WAL.
            con.execute('pragma journal_mode=off')

            con.execute('''CREATE TABLE blocking_map
                           (block_key text, record_id {id_type})
                        '''.format(id_type=id_type))

            con.executemany("INSERT INTO blocking_map values (?, ?)",
                            self.fingerprinter(data.items()))

            self.fingerprinter.reset_indices()

            con.execute('''CREATE UNIQUE INDEX record_id_block_key_idx
                           ON blocking_map (record_id, block_key)''')
            con.execute('''CREATE INDEX block_key_idx
                           ON blocking_map (block_key)''')
            con.execute('''ANALYZE''')
            pairs = con.execute('''SELECT DISTINCT a.record_id, b.record_id
                                   FROM blocking_map a
                                   INNER JOIN blocking_map b
                                   USING (block_key)
                                   WHERE a.record_id < b.record_id''')

            for a_record_id, b_record_id in pairs:
                yield ((a_record_id, data[a_record_id]),
                       (b_record_id, data[b_record_id]))

            pairs.close()
            con.close()

    def cluster(self,
                scores: numpy.ndarray,
                threshold: float = 0.5) -> Clusters:
        r"""From the similarity scores of pairs of records, decide which groups
        of records are all referring to the same entity.

        Yields tuples containing a sequence of record ids and corresponding
        sequence of confidence score as a float between 0 and 1. The
        record_ids within each set should refer to the same entity and the
        confidence score is a measure of our confidence a particular entity
        belongs in the cluster.

        Each confidence scores is a measure of how similar the record is
        to the other records in the cluster. Let :math:`\phi(i,j)` be the pair-wise
        similarity between records :math:`i` and :math:`j`. Let :math:`N` be the number of records in the cluster.

        .. math::

           \text{confidence score}_i = 1 - \sqrt {\frac{\sum_{j}^N (1 - \phi(i,j))^2}{N -1}}

        This measure is similar to the average squared distance
        between the focal record and the other records in the
        cluster. These scores can be `combined to give a total score
        for the cluster
        <https://en.wikipedia.org/wiki/Variance#Discrete_random_variable>`_.

        .. math::

           \text{cluster score} = 1 - \sqrt { \frac{\sum_i^N(1 - \mathrm{score}_i)^2 \cdot (N - 1) } { 2 N^2}}

        Args:
            scores: a numpy `structured array <https://docs.scipy.org/doc/numpy/user/basics.rec.html>`_ with a dtype of `[('pairs', id_type, 2),
                    ('score', 'f4')]` where dtype is either a str
                    or int, and score is a number between 0 and
                    1. The 'pairs' column contains pairs of ids of
                    the records compared and the 'score' column
                    should contains the similarity score for that
                    pair of records.

                    For each pair, the smaller id should be first.

            threshold: Number between 0 and 1. We will only consider
                       put together records into clusters if the
                       `cophenetic similarity
                       <https://en.wikipedia.org/wiki/Cophenetic>`_ of
                       the cluster is greater than the threshold.

                       Lowering the number will increase recall,
                       raising it will increase precision

        .. code:: python

           > pairs = matcher.pairs(data)
           > scores = matcher.scores(pairs)
           > clusters = matcher.cluster(scores)
           > print(list(clusters))
           [((1, 2, 3), (0.790, 0.860, 0.790)),
            ((4, 5), (0.720, 0.720)),
            ((10, 11), (0.899, 0.899))]

        """

        logger.debug("matching done, begin clustering")

        yield from clustering.cluster(scores, threshold)

class RecordLinkMatching(IntegralMatching):
    """
    Class for Record Linkage, extends Matching.

    Use RecordLinkMatching when you have two datasets that you want to merge
    """

    def pairs(self, data_1: Data, data_2: Data) -> RecordPairs:
        """
        Yield pairs of records that share common fingerprints.

        Each pair will occur at most once. If you override this
        method, you need to take care to ensure that this remains
        true, as downstream methods, particularly :func:`one_to_one`,
        and :func:`many_to_one` assumes that every pair of records is
        compared no more than once.

        Args:
            data_1: Dictionary of records from first dataset, where the
                    keys are record_ids and the values are dictionaries
                    with the keys being field names
            data_2: Dictionary of records from second dataset, same
                    form as data_1

        .. code:: python

           > pairs = matcher.pairs(data_1, data_2)
           > print(list(pairs))
           [((1, {'name' : 'Pat', 'address' : '123 Main'}),
             (2, {'name' : 'Pat', 'address' : '123 Main'})),
            ((1, {'name' : 'Pat', 'address' : '123 Main'}),
             (3, {'name' : 'Sam', 'address' : '123 Main'}))
            ]
        """

        self.fingerprinter.index_all(data_2)

        id_type_a = core.sqlite_id_type(data_1)
        id_type_b = core.sqlite_id_type(data_2)

        # Blocking and pair generation are typically the first memory
        # bottlenecks, so we'll use sqlite3 to avoid doing them in memory
        with tempfile.TemporaryDirectory() as temp_dir:
            if self.in_memory:
                con = sqlite3.connect(':memory:')
            else:
                con = sqlite3.connect(temp_dir + '/blocks.db')

            con.execute('pragma journal_mode=off')

            con.executescript('''CREATE TABLE blocking_map_a
                                 (block_key text, record_id {id_type_a});

                                 CREATE TABLE blocking_map_b
                                 (block_key text, record_id {id_type_b});
                              '''.format(id_type_a=id_type_a,
                                         id_type_b=id_type_b))

            con.executemany("INSERT INTO blocking_map_a values (?, ?)",
                            self.fingerprinter(data_1.items()))

            con.executemany("INSERT INTO blocking_map_b values (?, ?)",
                            self.fingerprinter(data_2.items(), target=True))

            self.fingerprinter.reset_indices()

            con.executescript('''CREATE INDEX block_key_a_idx
                                 ON blocking_map_a (record_id, block_key);

                                 CREATE INDEX block_key_b_idx
                                 ON blocking_map_b (block_key, record_id);''')
            con.execute('''ANALYZE''')

            pairs = con.execute('''SELECT DISTINCT a.record_id, b.record_id
                                   FROM blocking_map_a a
                                   INNER JOIN blocking_map_b b
                                   USING (block_key)''')

            for a_record_id, b_record_id in pairs:
                yield ((a_record_id, data_1[a_record_id]),
                       (b_record_id, data_2[b_record_id]))

            pairs.close()
            con.close()

    def join(self,
             data_1: Data,
             data_2: Data,
             threshold: float = 0.5,
             constraint: JoinConstraint = "one-to-one") -> Links:
        """
        Identifies pairs of records that refer to the same entity.

        Returns pairs of record ids with a confidence score as a float
        between 0 and 1. The record_ids within the pair should refer to the
        same entity and the confidence score is the estimated probability that
        the records refer to the same entity.

        This method should only used for small to moderately sized
        datasets for larger data, you need may need to generate your
        own pairs of records and feed them to the :func:`~score`.

        Args:
            data_1: Dictionary of records from first dataset, where the
                    keys are record_ids and the values are dictionaries
                    with the keys being field names

            data_2: Dictionary of records from second dataset, same form
                    as data_1

            threshold: Number between 0 and 1. We
                       will consider records as potential
                       duplicates if the predicted probability of
                       being a duplicate is above the threshold.

                       Lowering the number will increase recall, raising it
                       will increase precision

            constraint: What type of constraint to put on a join.

                        'one-to-one'
                              Every record in data_1 can match at most
                              one record from data_2 and every record
                              from data_2 can match at most one record
                              from data_1. This is good for when both
                              data_1 and data_2 are from different
                              sources and you are interested in
                              matching across the sources. If,
                              individually, data_1 or data_2 have many
                              duplicates you will not get good
                              matches.
                        'many-to-one'
                              Every record in data_1 can match at most
                              one record from data_2, but more than
                              one record from data_1 can match to the
                              same record in data_2. This is good for
                              when data_2 is a lookup table and data_1
                              is messy, such as geocoding or matching
                              against golden records.
                        'many-to-many'
                              Every record in data_1 can match
                              multiple records in data_2 and vice
                              versa. This is like a SQL inner join.

        .. code:: python

           > links = matcher.join(data_1, data_2, threshold=0.5)
           > print(list(links))
           [((1, 2), 0.790),
            ((4, 5), 0.720),
            ((10, 11), 0.899)]

        """

        assert constraint in {'one-to-one', 'many-to-one', 'many-to-many'}, (
            '%s is an invalid constraint option. Valid options include '
            'one-to-one, many-to-one, or many-to-many' % constraint)

        pairs = self.pairs(data_1, data_2)
        pair_scores = self.score(pairs)

        if constraint == 'one-to-one':
            links = self.one_to_one(pair_scores, threshold)
        elif constraint == 'many-to-one':
            links = self.many_to_one(pair_scores, threshold)
        else:
            links = pair_scores[pair_scores['score'] > threshold]

        links = list(links)

        try:
            mmap_file = pair_scores.filename  # type: ignore
        except AttributeError:
            pass
        else:
            del pair_scores
            if mmap_file:
                os.remove(mmap_file)

        return links

    def one_to_one(self,
                   scores: numpy.ndarray,
                   threshold: float = 0.0) -> Links:
        """From the similarity scores of pairs of records, decide which
        pairs refer to the same entity.

        Every record in data_1 can match at most one record from
        data_2 and every record from data_2 can match at most one
        record from data_1. See
        https://en.wikipedia.org/wiki/Injective_function.

        This method is good for when both data_1 and data_2 are from
        different sources and you are interested in matching across
        the sources. If, individually, data_1 or data_2 have many duplicates
        you will not get good matches.

        Yields pairs of record ids with a confidence score as a float
        between 0 and 1. The record_ids within the pair should refer to the
        same entity and the confidence score is the estimated probability that
        the records refer to the same entity.

        Args:
            scores: a numpy `structured array <https://docs.scipy.org/doc/numpy/user/basics.rec.html>`_ with a dtype of `[('pairs', id_type, 2),
                    ('score', 'f4')]` where dtype is either a str
                    or int, and score is a number between 0 and
                    1. The 'pairs' column contains pairs of ids of
                    the records compared and the 'score' column
                    should contains the similarity score for that
                    pair of records.

            threshold: Number between 0 and 1. We
                       will consider records as potential
                       duplicates if the predicted probability of
                       being a duplicate is above the threshold.

                       Lowering the number will increase recall, raising it
                       will increase precision

        .. code:: python

           > pairs = matcher.pairs(data)
           > scores = matcher.scores(pairs, threshold=0.5)
           > links = matcher.one_to_one(scores)
           > print(list(links))
           [((1, 2), 0.790),
            ((4, 5), 0.720),
            ((10, 11), 0.899)]

        """
        if threshold:
            scores = scores[scores['score'] > threshold]

        logger.debug("matching done, begin clustering")

        yield from clustering.greedyMatching(scores)

    def many_to_one(self,
                    scores: numpy.ndarray,
                    threshold: float = 0.0) -> Links:
        """
        From the similarity scores of pairs of records, decide which
        pairs refer to the same entity.

        Every record in data_1 can match at most one record from
        data_2, but more than one record from data_1 can match to the same
        record in data_2. See
        https://en.wikipedia.org/wiki/Surjective_function

        This method is good for when data_2 is a lookup table and data_1
        is messy, such as geocoding or matching against golden records.

        Yields pairs of record ids with a confidence score as a float
        between 0 and 1. The record_ids within the pair should refer to the
        same entity and the confidence score is the estimated probability that
        the records refer to the same entity.

        Args:
            scores: a numpy `structured array <https://docs.scipy.org/doc/numpy/user/basics.rec.html>`_ with a dtype of `[('pairs', id_type, 2),
                    ('score', 'f4')]` where dtype is either a str
                    or int, and score is a number between 0 and
                    1. The 'pairs' column contains pairs of ids of
                    the records compared and the 'score' column
                    should contains the similarity score for that
                    pair of records.

            threshold: Number between 0 and 1. We
                       will consider records as potential
                       duplicates if the predicted probability of
                       being a duplicate is above the threshold.

                       Lowering the number will increase recall, raising it
                       will increase precision

        .. code:: python

           > pairs = matcher.pairs(data)
           > scores = matcher.scores(pairs, threshold=0.5)
           > links = matcher.many_to_one(scores)
           > print(list(links))
           [((1, 2), 0.790),
            ((4, 5), 0.720),
            ((7, 2), 0.623),
            ((10, 11), 0.899)]

        """

        logger.debug("matching done, begin clustering")

        yield from clustering.pair_gazette_matching(scores, threshold, 1)

class GazetteerMatching(Matching):

    def __init__(self,
                 num_cores: Optional[int],
                 in_memory: bool = False,
                 **kwargs) -> None:

        super().__init__(num_cores, in_memory, **kwargs)

        if self.in_memory:
            self.db = ':memory:'
        else:
            self.temp_dir = tempfile.TemporaryDirectory()
            self.db = self.temp_dir.name + '/blocks.db'

        self.indexed_data: Dict[RecordID, RecordDict] = {}

    def _close(self):
        if not self.in_memory:
            self.temp_dir.cleanup()

    def __del__(self):
        self._close()

    def index(self, data: Data) -> None:  # pragma: no cover
        """
        Add records to the index of records to match against. If a record in
        `canonical_data` has the same key as a previously indexed record, the
        old record will be replaced.

        Args:
            data: a dictionary of records where the keys
                  are record_ids and the values are
                  dictionaries with the keys being
                  field_names
        """

        self.fingerprinter.index_all(data)

        id_type = core.sqlite_id_type(data)

        con = sqlite3.connect(self.db)

        # Set journal mode to WAL.
        con.execute('pragma journal_mode=wal')

        con.execute('''CREATE TABLE IF NOT EXISTS indexed_records
                       (block_key text,
                        record_id {id_type},
                        UNIQUE(block_key, record_id))
                    '''.format(id_type=id_type))

        con.executemany("REPLACE INTO indexed_records VALUES (?, ?)",
                        self.fingerprinter(data.items(), target=True))

        con.execute('''CREATE INDEX IF NOT EXISTS
                       indexed_records_block_key_idx
                       ON indexed_records
                       (block_key, record_id)''')
        con.execute('''ANALYZE''')

        con.commit()
        con.close()

        self.indexed_data.update(data)

    def unindex(self, data: Data) -> None:  # pragma: no cover
        """
        Remove records from the index of records to match against.

        Args:
            data: a dictionary of records where the keys
                  are record_ids and the values are
                  dictionaries with the keys being
                  field_names
        """

        for field in self.fingerprinter.index_fields:
            self.fingerprinter.unindex({record[field]
                                        for record
                                        in data.values()},
                                       field)

        con = sqlite3.connect(self.db)
        con.executemany('''DELETE FROM indexed_records
                           WHERE record_id = ?''',
                        ((k, ) for k in data.keys()))

        con.commit()
        con.close()

        for k in data:
            del self.indexed_data[k]

    def blocks(self, data: Data) -> Blocks:
        """
        Yield groups of pairs of records that share fingerprints.

        Each group contains one record from data_1 paired with the records
        from the indexed records that data_1 shares a fingerprint with.

        Each pair within and among blocks will occur at most once. If
        you override this method, you need to take care to ensure that
        this remains true, as downstream methods, particularly
        :func:`many_to_n`, assumes that every pair of records is compared no
        more than once.

        Args:
            data: Dictionary of records, where the keys are record_ids
                  and the values are dictionaries with the keys being
                  field names

        .. code:: python

            > pairs = matcher.pairs(data)
            > print(list(pairs))
            [[((1, {'name' : 'Pat', 'address' : '123 Main'}),
               (8, {'name' : 'Pat', 'address' : '123 Main'})),
              ((1, {'name' : 'Pat', 'address' : '123 Main'}),
               (9, {'name' : 'Sam', 'address' : '123 Main'}))
              ],
             [((2, {'name' : 'Sam', 'address' : '2600 State'}),
               (5, {'name' : 'Pam', 'address' : '2600 Stat'})),
              ((2, {'name' : 'Sam', 'address' : '123 State'}),
               (7, {'name' : 'Sammy', 'address' : '123 Main'}))
             ]]
        """

        id_type = core.sqlite_id_type(data)

        con = sqlite3.connect(self.db, check_same_thread=False)

        con.execute('BEGIN')

        con.execute('''CREATE TEMPORARY TABLE blocking_map
                       (block_key text, record_id {id_type})
                    '''.format(id_type=id_type))
        con.executemany("INSERT INTO blocking_map VALUES (?, ?)",
                        self.fingerprinter(data.items()))

        pairs = con.execute('''SELECT DISTINCT a.record_id, b.record_id
                               FROM blocking_map a
                               INNER JOIN indexed_records b
                               USING (block_key)
                               ORDER BY a.record_id''')

        pair_blocks = itertools.groupby(pairs,
                                        lambda x: x[0])

        for _, pair_block in pair_blocks:

            yield [((a_record_id, data[a_record_id]),
                    (b_record_id, self.indexed_data[b_record_id]))
                   for a_record_id, b_record_id
                   in pair_block]

        pairs.close()
        con.execute("ROLLBACK")
        con.close()

    def score(self,
              blocks: Blocks) -> Generator[numpy.ndarray, None, None]:
        """
        Scores groups of pairs of records. Yields structured numpy arrays
        representing pairs of records in the group and the associated
        probability that the pair is a match.

        Args:
            blocks: Iterator of blocks of records

        """

        matches = core.scoreGazette(blocks,
                                    self.data_model,
                                    self.classifier,
                                    self.num_cores)

        return matches

    def many_to_n(self,
                  score_blocks: Iterable[numpy.ndarray],
                  threshold: float = 0.0,
                  n_matches: int = 1) -> Links:
        """
        For each group of scored pairs, yield the highest scoring N pairs

        Args:
            score_blocks: Iterator of numpy `structured arrays <https://docs.scipy.org/doc/numpy/user/basics.rec.html>`_,
                          each with a dtype of `[('pairs', id_type, 2),
                          ('score', 'f4')]` where dtype is either a str
                          or int, and score is a number between 0 and
                          1. The 'pairs' column contains pairs of ids of
                          the records compared and the 'score' column
                          should contains the similarity score for that
                          pair of records.

            threshold: Number between 0 and 1. We
                       will consider records as potential
                       duplicates if the predicted probability of
                       being a duplicate is above the threshold.

                       Lowering the number will increase recall, raising it
                       will increase precision

            n_matches: How many top scoring pairs to select per group

        """

        yield from clustering.gazetteMatching(score_blocks,
                                              threshold,
                                              n_matches)

    def search(self,
               data: Data,
               threshold: float = 0.0,
               n_matches: int = 1,
               generator: bool = False) -> LookupResults:  # pragma: no cover
        """
        Identifies pairs of records that could refer to the same entity,
        returns tuples containing tuples of possible matches, with a
        confidence score for each match. The record_ids within each
        tuple should refer to potential matches from a messy data
        record to canonical records. The confidence score is the
        estimated probability that the records refer to the same
        entity.

        Args:

            data: a dictionary of records from a messy
                  dataset, where the keys are record_ids and
                  the values are dictionaries with the keys
                  being field names.

            threshold: a number between 0 and 1. We will consider
                       records as potential duplicates if the predicted
                       probability of being a duplicate is
                       above the threshold.

                       Lowering the number will increase
                       recall, raising it will increase
                       precision
            n_matches: the maximum number of possible matches from
                       canonical_data to return for each record in
                       data. If set to `None` all possible
                       matches above the threshold will be
                       returned.
            generator: when `True`, match will generate a sequence of
                       possible matches, instead of a list.

        .. code:: python

            > matches = gazetteer.search(data, threshold=0.5, n_matches=2)
            > print(matches)
            [(((1, 6), 0.72),
              ((1, 8), 0.6)),
             (((2, 7), 0.72),),
             (((3, 6), 0.72),
              ((3, 8), 0.65)),
             (((4, 6), 0.96),
              ((4, 5), 0.63))]

        """
        blocks = self.blocks(data)
        pair_scores = self.score(blocks)
        search_results = self.many_to_n(pair_scores, threshold, n_matches)

        results = self._format_search_results(data, search_results)

        if generator:
            return results
        else:
            return list(results)

    def _format_search_results(self,
                               search_d: Data,
                               results: Links) -> LookupResults:

        seen: Set[RecordID] = set()

        for result in results:
            a: Optional[RecordID] = None
            b: RecordID
            score: float
            prepared_result: List[Tuple[RecordID, float]] = []
            for (a, b), score in result:  # type: ignore
                prepared_result.append((b, score))

            assert a is not None

            yield a, tuple(prepared_result)
            seen.add(a)

        for k in (search_d.keys() - seen):
            yield k, ()

class StaticMatching(Matching):
    """
    Class for initializing a dedupe object from a settings file.
    """

    def __init__(self,
                 settings_file: BinaryIO,
                 num_cores: Optional[int] = None,
                 in_memory: bool = False,
                 **kwargs) -> None:  # pragma: no cover
        """
        Args:
            settings_file: A file object containing settings
                           info produced from the
                           :func:`~dedupe.api.ActiveMatching.write_settings` method.

            num_cores: The number of cpus to use for parallel
                       processing, defaults to the number of cpus
                       available on the machine. If set to 0, then
                       multiprocessing will be disabled.

            in_memory: If True, :meth:`dedupe.Dedupe.pairs` will generate
                       pairs in RAM with the sqlite3 ':memory:' option
                       rather than writing to disk. May be faster if
                       sufficient memory is available.

        .. warning::

            If using multiprocessing on Windows or Mac OS X, then
            you must protect calls to the Dedupe methods with a
            `if __name__ == '__main__'` in your main module, see
            https://docs.python.org/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods

        """
        super().__init__(num_cores, in_memory, **kwargs)

        try:
            self.data_model = pickle.load(settings_file)
            self.classifier = pickle.load(settings_file)
            self.predicates = pickle.load(settings_file)
        except (KeyError, AttributeError):
            raise SettingsFileLoadingException(
                "This settings file is not compatible with "
                "the current version of dedupe. This can happen "
                "if you have recently upgraded dedupe.")
        except:  # noqa: E722
            raise SettingsFileLoadingException(
                "Something has gone wrong with loading the settings file. "
                "Try deleting the file")

        logger.info('Predicate set:')
        for predicate in self.predicates:
            logger.info(predicate)

        self._fingerprinter = blocking.Fingerprinter(self.predicates)

class ActiveMatching(Matching):
    """
    Class for training a matcher.
    """
    classifier = rlr.RegularizedLogisticRegression()

    def __init__(self,
                 variable_definition: Sequence[Mapping],
                 num_cores: Optional[int] = None,
                 in_memory: bool = False,
                 **kwargs) -> None:
        """
        Args:
            variable_definition: A list of dictionaries describing
                                 the variables will be used for
                                 training a model. See :ref:`variable_definitions`

            num_cores: The number of cpus to use for parallel
                       processing. If set to `None`, uses all cpus
                       available on the machine. If set to 0, then
                       multiprocessing will be disabled.

            in_memory: If True, :meth:`dedupe.Dedupe.pairs` will generate
                       pairs in RAM with the sqlite3 ':memory:' option
                       rather than writing to disk. May be faster if
                       sufficient memory is available.

        .. warning::

            If using multiprocessing on Windows or Mac OS X, then
            you must protect calls to the Dedupe methods with a
            `if __name__ == '__main__'` in your main module, see
            https://docs.python.org/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods

        """
        super().__init__(num_cores, in_memory, **kwargs)

        self.data_model = datamodel.DataModel(variable_definition)

        self.training_pairs: TrainingData
        self.training_pairs = {'distinct': [],
                               'match': []}
        self.active_learner: Optional[Union[labeler.DedupeDisagreementLearner,
                                            labeler.RecordLinkDisagreementLearner]]
        self.active_learner = None

    def cleanup_training(self) -> None:  # pragma: no cover
        '''
        Clean up data we used for training. Free up memory.
        '''
        del self.training_pairs
        del self.active_learner

    def _read_training(self, training_file: TextIO) -> None:
        '''
        Read training from previously built training data file object

        Args:
            training_file: file object containing the training data
        '''

        logger.info('reading training from file')
        training_pairs = json.load(training_file,
                                   object_hook=serializer._from_json)

        self.mark_pairs(training_pairs)

    def train(self,
              recall: float = 1.00,
              index_predicates: bool = True) -> None:  # pragma: no cover
        """
        Learn final pairwise classifier and fingerprinting rules. Requires that
        adequate training data has been already been provided.

        Args:
            recall: The proportion of true dupe pairs in our
                    training data that that the learned fingerprinting
                    rules must cover. If we lower the recall, there will
                    be pairs of true dupes that we will never
                    directly compare.

                    recall should be a float between 0.0 and 1.0.

            index_predicates: Should dedupe consider predicates
                              that rely upon indexing the
                              data. Index predicates can be slower
                              and take substantial memory. Without
                              index predicates, you may get lower
                              recall when true-dupes are not blocked
                              together.

        """
        assert self.active_learner is not None, \
               "Please initialize with the sample method"

        examples, y = flatten_training(self.training_pairs)
        self.classifier.fit(self.data_model.distances(examples), y)

        self.predicates = self.active_learner.learn_predicates(
            recall, index_predicates)
        self._fingerprinter = blocking.Fingerprinter(self.predicates)
        self.fingerprinter.reset_indices()

    def write_training(self, file_obj: TextIO) -> None:  # pragma: no cover
        """
        Write a JSON file that contains labeled examples

        :param file_obj: file object to write training data to

        .. code:: python

            with open('training.json', 'w') as f:
                matcher.write_training(f)

        """

        json.dump(self.training_pairs,
                  file_obj,
                  cls=serializer.TupleEncoder,
                  ensure_ascii=True)

    def write_settings(self,
                       file_obj: BinaryIO) -> None:  # pragma: no cover
        """
        Write a settings file containing the
        data model and predicates to a file object

        :param file_obj: file object to write settings data into

        .. code:: python

           with open('learned_settings', 'wb') as f:
               matcher.write_settings(f)

        """

        pickle.dump(self.data_model, file_obj)
        pickle.dump(self.classifier, file_obj)
        pickle.dump(self.predicates, file_obj)

    def uncertain_pairs(self) -> List[TrainingExample]:
        '''
        Returns a list of pairs of records from the sample of record pairs
        tuples that Dedupe is most curious to have labeled.

        This method is mainly useful for building a user interface for training
        a matching model.

       .. code:: python

          > pair = matcher.uncertain_pairs()
          > print(pair)
          [({'name' : 'Georgie Porgie'}, {'name' : 'Georgette Porgette'})]

        '''
        assert self.active_learner is not None, \
               "Please initialize with the sample method"
        return [self.active_learner.pop()]

    def mark_pairs(self, labeled_pairs: TrainingData) -> None:
        '''
        Add users labeled pairs of records to training data and update the
        matching model

        This method is useful for building a user interface for training a
        matching model or for adding training data from an existing source.

        Args:
            labeled_pairs: A dictionary with two keys, `match` and `distinct`
                           the values are lists that can contain pairs of
                           records

        .. code:: python

            labeled_examples = {'match'    : [],
                                'distinct' : [({'name' : 'Georgie Porgie'},
                                               {'name' : 'Georgette Porgette'})]
                                }
            matcher.mark_pairs(labeled_examples)

        .. note::
           `mark_pairs()` is primarily designed to be used with
           :func:`~uncertain_pairs` to incrementally build a training
           set.

           If you have existing training data, you should likely
           format the data into the right form and supply the training
           data to the :func:`~prepare_training` method with the
           `training_file` argument.

           If that is not possible or desirable, you can use
           `mark_pairs()` to train a linker with existing data.
           However, you must ensure that every record that
           appears in the `labeled_pairs` argument appears in either
           the data or training file supplied to the
           :func:`~prepare_training` method.

        '''
        self._checkTrainingPairs(labeled_pairs)

        self.training_pairs['match'].extend(labeled_pairs['match'])
        self.training_pairs['distinct'].extend(labeled_pairs['distinct'])

        if self.active_learner:
            examples, y = flatten_training(labeled_pairs)

            try:
                self.active_learner.mark(examples, y)
            except dedupe.predicates.NoIndexError as e:
                raise UserWarning(
                    ('The record\n'
                     '{unknown}\n'
                     'is not known to to the active learner. '
                     'Make sure all `labeled_pairs` '
                     'are in the data or training file '
                     'of the `prepare_training()` method').format(
                         unknown=e.failing_record))

    def _checkTrainingPairs(self, labeled_pairs: TrainingData) -> None:
        try:
            labeled_pairs.items()
            labeled_pairs['match']
            labeled_pairs['distinct']
        except (AttributeError, KeyError):
            raise ValueError('labeled_pairs must be a dictionary with keys '
                             '"distinct" and "match"')

        if labeled_pairs['match']:
            pair = labeled_pairs['match'][0]
            self._checkRecordPair(pair)

        if labeled_pairs['distinct']:
            pair = labeled_pairs['distinct'][0]
            self._checkRecordPair(pair)

        if not labeled_pairs['distinct'] and not labeled_pairs['match']:
            warnings.warn("Didn't return any labeled record pairs")

    def _checkRecordPair(self, record_pair: TrainingExample) -> None:
        try:
            a, b = record_pair
        except ValueError:
            raise ValueError("The elements of data_sample must be pairs "
                             "of record_pairs")
        try:
            record_pair[0].keys() and record_pair[1].keys()
        except AttributeError:
            raise ValueError("A pair of record_pairs must be made up of two "
                             "dictionaries ")

        self.data_model.check(record_pair[0])
        self.data_model.check(record_pair[1])

class StaticDedupe(StaticMatching, DedupeMatching):
    """
    Class for deduplication using saved settings. If you have already
    trained a :class:`Dedupe` object and saved the settings, you can
    load the saved settings with StaticDedupe.

    """

class Dedupe(ActiveMatching, DedupeMatching):
    """
    Class for active learning deduplication. Use deduplication when you have
    data that can contain multiple records that can all refer to the same
    entity.
    """

    canopies = True
    ActiveLearner = labeler.DedupeDisagreementLearner

    def prepare_training(self,
                         data: Data,
                         training_file: TextIO = None,
                         sample_size: int = 1500,
                         blocked_proportion: float = 0.9) -> None:
        '''
        Initialize the active learner with your data and, optionally,
        existing training data.

        Sets up the learner.

        Args:
            data: Dictionary of records, where the keys are
                  record_ids and the values are dictionaries
                  with the keys being field names
            training_file: file object containing training data
            sample_size: Size of the sample to draw
            blocked_proportion: The proportion of record pairs to be sampled from similar records, as opposed to randomly selected pairs.

        .. code:: python

           matcher.prepare_training(data_d, 150000, .5)

           # or
           with open('training_file.json') as f:
               matcher.prepare_training(data_d, training_file=f)

        '''

        # Reset active learner
        self.active_learner = None

        if training_file:
            self._read_training(training_file)
        self._sample(data, sample_size, blocked_proportion)

    def _sample(self,
                data: Data,
                sample_size: int = 15000,
                blocked_proportion: float = 0.5) -> None:
        '''Draw a sample of record pairs from the dataset
        (a mix of random pairs & pairs of similar records)
        and initialize active learning with this sample

        :param data: Dictionary of records, where the keys are
                     record_ids and the values are dictionaries
                     with the keys being field names

        :param sample_size: Size of the sample to draw

        :param blocked_proportion: Proportion of the sample that will be blocked

        '''
        self._checkData(data)

        # We need the active learner to know about all our
        # existing training data, so add them to data dictionary
        examples, y = flatten_training(self.training_pairs)

        self.active_learner = self.ActiveLearner(self.data_model,
                                                 data,
                                                 blocked_proportion,
                                                 sample_size,
                                                 index_include=examples)

        self.active_learner.mark(examples, y)

    def _checkData(self, data: Data) -> None:
        if len(data) == 0:
            raise ValueError(
                'Dictionary of records is empty.')

        self.data_model.check(next(iter(data.values())))

class Link(ActiveMatching):
    """
    Mixin Class for Active Learning Record Linkage
    """

    canopies = False
    ActiveLearner = labeler.RecordLinkDisagreementLearner

    def prepare_training(self,
                         data_1: Data,
                         data_2: Data,
                         training_file: Optional[TextIO] = None,
                         sample_size: int = 1500,
                         blocked_proportion: float = 0.9) -> None:
        '''
        Initialize the active learner with your data and, optionally,
        existing training data.

        Args:
            data_1: Dictionary of records from first dataset, where the
                    keys are record_ids and the values are dictionaries
                    with the keys being field names
            data_2: Dictionary of records from second dataset, same
                    form as data_1
            training_file: file object containing training data

            sample_size: The size of the sample to draw.

            blocked_proportion: The proportion of record pairs to
                                be sampled from similar records,
                                as opposed to randomly selected
                                pairs.

        .. code:: python

           matcher.prepare_training(data_1, data_2, 150000)

           with open('training_file.json') as f:
               matcher.prepare_training(data_1, data_2, training_file=f)

        '''
        # Reset active learner
        self.active_learner = None

        if training_file:
            self._read_training(training_file)
        self._sample(data_1,
                     data_2,
                     sample_size,
                     blocked_proportion)

    def _sample(self,
                data_1: Data,
                data_2: Data,
                sample_size: int = 15000,
                blocked_proportion: float = 0.5) -> None:
        '''
        Draws a random sample of combinations of records from
        the first and second datasets, and initializes active
        learning with this sample

        :param data_1: Dictionary of records from first dataset, where the
                       keys are record_ids and the values are dictionaries
                       with the keys being field names
        :param data_2: Dictionary of records from second dataset, same
                       form as data_1
        :param sample_size: Size of the sample to draw
        '''
        self._checkData(data_1, data_2)

        # We need the active learner to know about all our
        # existing training data, so add them to data dictionaries
        examples, y = flatten_training(self.training_pairs)

        self.active_learner = self.ActiveLearner(self.data_model,
                                                 data_1,
                                                 data_2,
                                                 blocked_proportion,
                                                 sample_size,
                                                 index_include=examples)

        self.active_learner.mark(examples, y)

    def _checkData(self, data_1: Data, data_2: Data) -> None:
        if len(data_1) == 0:
            raise ValueError(
                'Dictionary of records from first dataset is empty.')
        elif len(data_2) == 0:
            raise ValueError(
                'Dictionary of records from second dataset is empty.')

        self.data_model.check(next(iter(data_1.values())))
        self.data_model.check(next(iter(data_2.values())))

class RecordLink(Link, RecordLinkMatching):
    """
    Class for active learning record linkage.

    Use RecordLinkMatching when you have two datasets that you want to
    join.
    """

class StaticRecordLink(StaticMatching, RecordLinkMatching):
    """
    Class for record linkage using saved settings. If you have already
    trained a RecordLink instance, you can load the saved settings with
    StaticRecordLink.
    """

class Gazetteer(Link, GazetteerMatching):
    """
    Class for active learning gazetteer matching.

    Gazetteer matching is for matching a messy data set against a
    'canonical dataset'.  This class is useful for such tasks as matching messy
    addresses against a clean list
    """

class StaticGazetteer(StaticMatching, GazetteerMatching):
    """
    Class for gazetter matching using saved settings.

    If you have already trained a :class:`Gazetteer` instance, you can
    load the saved settings with StaticGazetteer.
    """

class EmptyTrainingException(Exception):
    pass

class SettingsFileLoadingException(Exception):
    pass

def flatten_training(training_pairs: TrainingData) -> Tuple[List[TrainingExample], numpy.ndarray]:
    examples: List[TrainingExample] = []
    y = []

    for label in ('match', 'distinct'):
        label = cast(Literal['match', 'distinct'], label)

        pairs = training_pairs[label]
        examples.extend(pairs)
        encoded_y = 1 if label == 'match' else 0
        y.extend([encoded_y] * len(pairs))

    return examples, numpy.array(y)


import faiss
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer

from tqdm.auto import tqdm
import pandas as pd
import scipy
import pickle
import json
import os
import numpy as np
from collections import Counter
import re

from datasets import (
    Dataset,
    load_from_disk,
    concatenate_datasets,
)
from konlpy.tag import Mecab

import numpy as np
from tqdm import tqdm, trange
import argparse
import random
import torch
import os
import torch.nn.functional as F
from transformers import BertModel, BertPreTrainedModel, AdamW, TrainingArguments, get_linear_schedule_with_warmup, AutoTokenizer
import pickle
from rank_bm25 import BM25Okapi, BM25Plus, BM25L, BM25
import time
from contextlib import contextmanager

import math
from multiprocessing import Pool, cpu_count

@contextmanager
def timer(name):
    t0 = time.time()
    yield
    print(f'[{name}] done in {time.time() - t0:.3f} s')
ban_words=("이따금","아마","절대로","무조건","한때","대략","오직",
          "오로지","감히","최소","아예","반드시","꼭","때때로","이미"
          ,"종종","졸곧","약간","기꺼이", "비록","꾸준히","일부러","어쩔", "문득", "어쨌든", "순전히", "필수")

mecab = Mecab()

class ES(BM25):
    def __init__(self, corpus, tokenizer=None, k1=1.2, b=0.75, delta=0):
        # Algorithm specific parameters
        self.k1 = k1
        self.b = b
        self.delta = delta
        super().__init__(corpus, tokenizer)

    def _calc_idf(self, nd):
        for word, freq in nd.items():
            idf = math.log(1 + ((self.corpus_size + 0.5 - freq) / (freq+0.5)))
            self.idf[word] = idf

    def get_scores(self, query):
        score = np.zeros(self.corpus_size)
        doc_len = np.array(self.doc_len)
        for q in query:
            q_freq = np.array([(doc.get(q) or 0) for doc in self.doc_freqs])
            score += (self.idf.get(q) or 0) * (self.delta + (q_freq) /
                                               (self.k1 * (1 - self.b + self.b * doc_len / self.avgdl) + q_freq))
        return score

def remove_q(query):
    stop = "|".join(
        "어느 무엇인가요 무엇 누가 누구인가요 누구인가 누구 어디에서 어디에 어디서 어디인가요 어디를 어디 언제 어떤 어떠한 몇 얼마 얼마나 뭐 어떻게 무슨 \?".split(
            " "
        )
    )
    rm = re.sub(stop, "", query).strip()
    return rm

class BertEncoder(BertPreTrainedModel):
  def __init__(self, config):
    super(BertEncoder, self).__init__(config)

    self.bert = BertModel(config)
    self.init_weights()

  def forward(self, input_ids, attention_mask=None, token_type_ids=None):
      outputs = self.bert(input_ids,attention_mask=attention_mask,token_type_ids=token_type_ids)

      pooled_output = outputs[1]

      return pooled_output
'''
tokenizer = AutoTokenizer.from_pretrained("bert-base-multilingual-cased")
model_checkpoint="bert-base-multilingual-cased"

# load pre-trained model on cuda (if available)
p_encoder = BertEncoder.from_pretrained(model_checkpoint).cuda()
q_encoder = BertEncoder.from_pretrained(model_checkpoint).cuda()

p_encoder.load_state_dict(torch.load("./p_encoder_fin.pth"))
q_encoder.load_state_dict(torch.load("./q_encoder_fin.pth"))

'''

def to_cuda(batch):
  return tuple(t.cuda() for t in batch)

class DenseRetrieval:
    def __init__(self, tokenize_fn, data_path="./data/", context_path="wikipedia_documents.json"):
        self.data_path = data_path
        self.wiki_embs = None
        with open(os.path.join(data_path, context_path), "r") as f:
            wiki = json.load(f)
        self.contexts = list(dict.fromkeys([v['text'] for v in wiki.values()]))
        pickle_name = f"./data/dense_embedding.bin"
        if os.path.isfile(pickle_name):
            with open(pickle_name,"rb") as file:
                self.wiki_embs=pickle.load(file)
                print("Pre")
        else:
            with torch.no_grad():
              self.wiki_embs = []
              for text in self.contexts:
                p = tokenizer(text, padding="max_length", truncation=True, return_tensors='pt').to('cuda')
                wiki_emb = p_encoder(**p).to('cpu').numpy()
                self.wiki_embs.append(wiki_emb)
            self.wiki_embs = torch.Tensor(self.wiki_embs).squeeze()  # (num_passage, emb_dim)
            with open(pickle_name,"wb") as file:
                pickle.dump(self.wiki_embs,file)
        self.f=open("./track.txt","w")

    def retrieve(self, query_or_dataset, topk=1):
        if isinstance(query_or_dataset, str):
            doc_scores, doc_indices = self.get_relevant_doc(query_or_dataset, k=topk)
            return doc_scores, doc_indices

        elif isinstance(query_or_dataset, Dataset):
            # make retrieved result as dataframe
            total = []
            super_count=0
            doc_scores, doc_indices = self.get_relevant_doc_bulk(query_or_dataset['question'], k=5)
            for idx, example in enumerate(tqdm(query_or_dataset, desc="Dense retrieval: ")):
                tmp = {
                    "question": example["question"],
                    "id": example['id'],
                    #"context_id": doc_indices[idx][0],  # retrieved id
                    "context": " ".join(self.contexts[doc_indices[idx][i]] for i in range(5))  # retrieved doument
                }
                if 'context' in example.keys() and 'answers' in example.keys():
                    tmp["original_context"] = example['context']  # original document
                    tmp["answers"] = example['answers']           # original answer
                total.append(tmp)
            cqas = pd.DataFrame(total)
            return cqas

    def get_relevant_doc(self, query, k=3):
        with torch.no_grad():
          q_seqs_val = tokenizer([query], padding="max_length", truncation=True, return_tensors='pt').to('cuda')
          q_emb = q_encoder(**q_seqs_val).to('cpu')  #(num_query, emb_dim)

        result = torch.matmul(q_emb, torch.transpose(self.wiki_embs, 0, 1))
        rank = torch.argsort(result, dim=1, descending=True).squeeze()
        self.f.write("=============\n")
        self.f.write("Query "+query+"\n")
        for idx,i in enumerate(range(k)):
            self.f.write(str(idx)+self.contexts[rank[i]][:75]+"\n")
        print(result.squeeze()[rank].tolist()[:k], rank.tolist()[:k])
        return result.squeeze()[rank].tolist()[:k], rank.tolist()[:k]

    def get_relevant_doc_bulk(self, queries, k=1):
        doc_scores = []
        doc_indices = []
        for query in queries:
            ret0,ret1=self.get_relevant_doc(query,k)
            doc_scores.append(ret0)
            doc_indices.append(ret1)
        self.f.close()
        return doc_scores, doc_indices

class SparseRetrieval:
    def __init__(self, tokenize_fn, data_path="./data/", context_path="wikipedia_documents.json"):
        self.data_path = data_path
        with open(os.path.join(data_path, context_path), "r") as f:
            wiki = json.load(f)

        self.contexts = list(dict.fromkeys([v['text'] for v in wiki.values()])) # set 은 매번 순서가 바뀌므로
        print(f"Lengths of unique contexts : {len(self.contexts)}")
        self.ids = list(range(len(self.contexts)))

        # Transform by vectorizer
        self.tfidfv = TfidfVectorizer(
            tokenizer=tokenize_fn,
            ngram_range=(1, 2),
            max_features=50000,
        )

        # should run get_sparse_embedding() or build_faiss() first.
        self.p_embedding = None
        self.indexer = None

    def get_sparse_embedding(self):
        # Pickle save.
        pickle_name = f"sparse_embedding.bin"
        tfidfv_name = f"tfidv.bin"
        emd_path = os.path.join(self.data_path, pickle_name)
        tfidfv_path = os.path.join(self.data_path, tfidfv_name)
        if os.path.isfile(emd_path) and os.path.isfile(tfidfv_path):
            with open(emd_path, "rb") as file:
                self.p_embedding = pickle.load(file)
            with open(tfidfv_path, "rb") as file:
                self.tfidfv = pickle.load(file)
            print("Embedding pickle load.")
        else:
            print("Build passage embedding")
            self.p_embedding = self.tfidfv.fit_transform(self.contexts)
            print(self.p_embedding.shape)
            with open(emd_path, "wb") as file:
                pickle.dump(self.p_embedding, file)
            with open(tfidfv_path, "wb") as file:
                pickle.dump(self.tfidfv, file)
            print("Embedding pickle saved.")

    def retrieve(self, query_or_dataset, topk=3):
        assert self.p_embedding is not None, "You must build faiss by self.get_sparse_embedding() before you run self.retrieve()."
        if isinstance(query_or_dataset, str):
            doc_scores, doc_indices = self.get_relevant_doc(query_or_dataset, k=topk)
            return doc_scores, doc_indices

        elif isinstance(query_or_dataset, Dataset):
            # make retrieved result as dataframe
            total = []
            doc_scores, doc_indices = self.get_relevant_doc_bulk(query_or_dataset['question'], k=20)
            for idx, example in enumerate(tqdm(query_or_dataset, desc="Sparse retrieval: ")):
                tmp = {
                    "question": example["question"],
                    "id": example['id'],
                    #"context_id": doc_indices[idx][0],  # retrieved id
                    "context": " ".join(self.contexts[doc_indices[idx][i]] for i in range(20))  # retrieved doument
                }
                if 'context' in example.keys() and 'answers' in example.keys():
                    tmp["original_context"] = example['context']  # original document
                    tmp["answers"] = example['answers']           # original answer
                total.append(tmp)
            cqas = pd.DataFrame(total)
            return cqas

    def get_relevant_doc(self, query, k=1):
        """
        참고: vocab 에 없는 이상한 단어로 query 하는 경우 assertion 발생 (예) 뙣뙇?
        """
        with timer("transform"):
            query_vec = self.tfidfv.transform([query])
        assert (
                np.sum(query_vec) != 0
        ), "오류가 발생했습니다. 이 오류는 보통 query에 vectorizer의 vocab에 없는 단어만 존재하는 경우 발생합니다."

        with timer("query ex search"):
            result = query_vec * self.p_embedding.T
        if not isinstance(result, np.ndarray):
            result = result.toarray()
        sorted_result = np.argsort(result.squeeze())[::-1]
        return result.squeeze()[sorted_result].tolist()[:k], sorted_result.tolist()[:k]

    def get_relevant_doc_bulk(self, queries, k=1):
        query_vec = self.tfidfv.transform(queries)
        assert (
                np.sum(query_vec) != 0
        ), "오류가 발생했습니다. 이 오류는 보통 query에 vectorizer의 vocab에 없는 단어만 존재하는 경우 발생합니다."

        result = query_vec * self.p_embedding.T
        if not isinstance(result, np.ndarray):
            result = result.toarray()
        doc_scores = []
        doc_indices = []
        for i in range(result.shape[0]):
            sorted_result = np.argsort(result[i, :])[::-1]
            doc_scores.append(result[i, :][sorted_result].tolist()[:k])
            doc_indices.append(sorted_result.tolist()[:k])
        return doc_scores, doc_indices

class BM25:
    def __init__(self, tokenize_fn, data_path="./data/", context_path="wikipedia_documents.json", k1=1.2, b=0.25):
        self.data_path = data_path
        with open(os.path.join(data_path, context_path), "r") as f:
            wiki = json.load(f)

        self.contexts = list(dict.fromkeys([v['text'] for v in wiki.values()])) # set 은 매번 순서가 바뀌므로
        self.l=[]
        for v in self.contexts:
            self.l.append(len(mecab.morphs(v)))
        self.l=np.array(self.l)
        print("Avg",np.average(self.l))
        self.l=self.l/np.average(self.l)
        self.l=(k1*(1-b+b*self.l)).astype(np.float32)
        print("L SHAP",self.l.shape)

        print(f"Lengths of unique contexts : {len(self.contexts)}")
        self.ids = list(range(len(self.contexts)))

        # Transform by vectorizer
        self.tfidfv = CountVectorizer(
            tokenizer=tokenize_fn,
            ngram_range=(1, 2),
            max_features=50000,
        )

        # should run get_sparse_embedding() or build_faiss() first.
        self.p_embedding = None
        self.a=None

    def get_sparse_embedding(self):
        # Pickle save.
        pickle_name = f"bm_embedding.bin"
        tfidfv_name = f"bm.bin"
        emd_path = os.path.join(self.data_path, pickle_name)
        tfidfv_path = os.path.join(self.data_path, tfidfv_name)
        if os.path.isfile(emd_path) and os.path.isfile(tfidfv_path):
            with open(emd_path, "rb") as file:
                self.p_embedding = pickle.load(file)
            with open(tfidfv_path, "rb") as file:
                self.tfidfv = pickle.load(file)
            print("Embedding pickle load.")
        else:
            print("Build passage embedding")
            self.p_embedding = self.tfidfv.fit_transform(self.contexts)
            print(self.p_embedding.shape)
            with open(emd_path, "wb") as file:
                pickle.dump(self.p_embedding, file)
            with open(tfidfv_path, "wb") as file:
                pickle.dump(self.tfidfv, file)
            print("Embedding pickle saved.")
        self.a=self.p_embedding
        z1=(len(self.contexts)-self.a.count_nonzero()+0.5)
        z2=(self.a.count_nonzero()+0.5)
        #z1=(len(self.contexts)-np.count_nonzero(self.a>0,axis=0)+0.5)
        #z2=(np.count_nonzero(self.a>0,axis=0)+0.5)
        z1=np.log(1+z1/z2)
        del(z2)
        self.a=(z1*((self.a)/(self.a+self.l[:,np.newaxis])))
        print(self.a)

    def retrieve(self, query_or_dataset, topk=1):
        assert self.p_embedding is not None, "You must build faiss by self.get_sparse_embedding() before you run self.retrieve()."
        if isinstance(query_or_dataset, str):
            doc_scores, doc_indices = self.get_relevant_doc(query_or_dataset, k=topk)
            print("[Search query]\n", query_or_dataset, "\n")

            for i in range(topk):
                print("Top-%d passage with score %.4f" % (i + 1, doc_scores[i]))
                print(self.contexts[doc_indices[i]])
            return doc_scores, [self.contexts[doc_indices[i]] for i in range(topk)]

        elif isinstance(query_or_dataset, Dataset):
            # make retrieved result as dataframe
            total = []
            with timer("query exhaustive search"):
                doc_scores, doc_indices = self.get_relevant_doc_bulk(query_or_dataset['question'], k=1)
            for idx, example in enumerate(tqdm(query_or_dataset, desc="Sparse retrieval: ")):
                # relev_doc_ids = [el for i, el in enumerate(self.ids) if i in doc_indices[idx]]
                tmp = {
                    "question": example["question"],
                    "id": example['id'],
                    "context_id": doc_indices[idx][0],  # retrieved id
                    "context": self.contexts[doc_indices[idx][0]]  # retrieved doument
                }
                if 'context' in example.keys() and 'answers' in example.keys():
                    tmp["original_context"] = example['context']  # original document
                    tmp["answers"] = example['answers']           # original answer
                total.append(tmp)

            cqas = pd.DataFrame(total)
            return cqas

    def get_relevant_doc(self, query, k=1):
        query_vec = self.tfidfv.transform([query])
        print("Query_vec",query_vec.toarray())
        result = query_vec * self.p_embedding.T
        if not isinstance(result, np.ndarray):
            result = result.toarray()
        sorted_result = np.argsort(result.squeeze())[::-1]
        return result.squeeze()[sorted_result].tolist()[:k], sorted_result.tolist()[:k]

    def get_relevant_doc_bulk(self, queries, k=1):
        query_vec = self.tfidfv.transform(queries)
        result = query_vec * self.p_embedding.T
        if not isinstance(result, np.ndarray):
            result = result.toarray()
        doc_scores = []
        doc_indices = []
        for i in range(result.shape[0]):
            sorted_result = np.argsort(result[i, :])[::-1]
            doc_scores.append(result[i, :][sorted_result].tolist()[:k])
            doc_indices.append(sorted_result.tolist()[:k])
        return doc_scores, doc_indices

class BM25Arti:
    def __init__(self, tokenize_fn, data_path="./data/", context_path="wikipedia_documents.json"):
        self.data_path = data_path
        with open(os.path.join(data_path, context_path), "r") as f:
            wiki = json.load(f)

        self.contexts = list(dict.fromkeys([v['text'] for v in wiki.values()])) # set 은 매번 순서가 바뀌므로
        print(f"Lengths of unique contexts : {len(self.contexts)}")
        self.ids = list(range(len(self.contexts)))

        # Transform by vectorizer
        pickle_name=f"token_corpus"
        emd_path = os.path.join(self.data_path, pickle_name)

        self.tokenized_corpus = None
        if os.path.isfile(emd_path):
            with open(emd_path, "rb") as file:
                self.tokenized_corpus = pickle.load(file)
        else:
            self.tokenized_corpus = [tokenize_fn(doc) for doc in self.contexts]
            with open(emd_path, "wb") as file:
                pickle.dump(self.tokenized_corpus, file)

        self.bm25 = None
        self.sCounter=Counter()

        pickle_name = f"bm25.bin"
        emd_path = os.path.join(self.data_path, pickle_name)
        if os.path.isfile(emd_path):
            with open(emd_path, "rb") as file:
                self.bm25 = pickle.load(file)
        else:
            self.bm25 = ES(self.tokenized_corpus)
            with open(emd_path, "wb") as file:
                pickle.dump(self.bm25, file)

    def retrieve(self, query_or_dataset, topk=20):
        if isinstance(query_or_dataset, str):
            doc_scores, doc_indices = self.get_relevant_doc(query_or_dataset, k=topk)
            return doc_scores, doc_indices

        elif isinstance(query_or_dataset, Dataset):
            # make retrieved result as dataframe
            f=open("./context.txt","w")
            total = []
            doc_scores, doc_indices = self.get_relevant_doc_bulk(query_or_dataset['question'], k=35)
            for idx, example in enumerate(tqdm(query_or_dataset, desc="Sparse retrieval: ")):
                standard=doc_scores[idx][0]
                self.sCounter+=Counter(doc_indices[idx])
                print("S-score",standard)
                standard_idx=0
                for idx2, doc_score in enumerate(doc_scores[idx]):
                    if doc_score>=standard*0.7:
                        standard_idx=idx2
                print("SIDX",standard_idx)
                print("DID",doc_indices[idx])
                temp_doc=" ".join(self.contexts[doc_indices[idx][i]] for i in range(standard_idx+1))
                f.write("\n------------\n")
                f.write(str(doc_indices[idx]))
                f.write("\n")
                f.write(example["question"]+"\n")
                for idid in range(standard_idx+1):
                    f.write("Doc id: ")
                    f.write(str(doc_indices[idx][idid]))
                    f.write("\n")
                    f.write(self.contexts[doc_indices[idx][idid]])

                print(temp_doc[:75])
                tmp = {
                    "question": example["question"],
                    "id": example['id'],
                    #"context_id": doc_indices[idx][0],  # retrieved id
                    "context":temp_doc   # retrieved doument
                }
                if 'context' in example.keys() and 'answers' in example.keys():
                    tmp["original_context"] = example['context']  # original document
                    tmp["answers"] = example['answers']           # original answer
                total.append(tmp)
            cqas = pd.DataFrame(total)
            f.write(str(self.sCounter.most_common(10)))
            f.close()
            return cqas

    def get_relevant_doc(self, query, k=3):

        result=self.bm25.get_scores(mecab.morphs(query))
        if not isinstance(result, np.ndarray):
            result = result.toarray()
        sorted_result = np.argsort(result.squeeze())[::-1]
        #print("Query "+query)
        #print(result.squeeze()[sorted_result].tolist()[:k], sorted_result.tolist()[:k])
        return result.squeeze()[sorted_result].tolist()[:k], sorted_result.tolist()[:k]

    def get_relevant_doc_bulk(self, queries, k=3):

        doc_scores = []
        doc_indices = []
        for query in queries:
            ret0,ret1=self.get_relevant_doc(query,k)
            doc_scores.append(ret0)
            doc_indices.append(ret1)

        return doc_scores, doc_indices

if __name__ == "__main__":
    # Test sparse
    org_dataset = load_from_disk("data/train_dataset")
    full_ds = concatenate_datasets(
        [
            org_dataset["train"].flatten_indices(),
            org_dataset["validation"].flatten_indices(),
        ]
    ) # train dev 를 합친 4192 개 질문에 대해 모두 테스트
    print("*"*40, "query dataset", "*"*40)
    print(full_ds)

    mecab = Mecab()
    def tokenize(text):
        # return text.split(" ")
        return mecab.morphs(text)

    wiki_path = "wikipedia_documents.json"
    retriever = BM25Arti(
        # tokenize_fn=tokenizer.tokenize,
        tokenize_fn=tokenize,
        data_path="data",
        context_path=wiki_path)
    query=" 1994년 FIFA 월드컵 당시 대한민국 축구 대표팀의 감독은 누구인가?"
    print(query)
    print(retriever.retrieve(query))


#!/usr/bin/env python

from __future__ import print_function
import numpy as np
import intera_interface
from random import randrange
#import intera_external_devices as kb
import keyboard as kb
import rospy
import rospkg

def printe(string):
    """
    Prints the string without adding a newline
    """
    print(string, end='')

def print_array(gameboard):
    """
    Prints an array in a comma separated format
    """
    print()
    for i in range(len(gameboard)):
        for j in range(len(gameboard[i])):
            printe(gameboard[i][j])
            printe(", ")
        print()

def findXYOffset(tuple):
    """
    Old implementation of implementing cell to cm conversion
    """
    homeX = -4.5
    homeY = 4.5
    baseX = -9
    baseY = 9
    xScale = 2-tuple[0]
    yScale = 2-tuple[1]
    xOffset = homeX + baseX * xScale
    yOffset = homeY + baseY * yScale

    unitScale = 1 #use this for changing units from cm
    return (xOffset*unitScale,yOffset*unitScale)

def displayFace(mode):
    """
    Displays a face on the sawyer robot corresponding to a mode:
    0: display a face for a Tie
    1: display a face for a Win
    2: display a face for a Loss
    11: display a face for Robot's turn
    22: display a face for the Player's turn
    """

    _head = intera_interface.Head()
    head_display = intera_interface.HeadDisplay()

    rospack = rospkg.RosPack()
    pkgpath = rospack.get_path('final-project-tic-tac-toe')

    if mode == 0:
        filepath = pkgpath + "/images/tie" + str(randrange(4)) + ".jpg"
    if mode == 1:
        filepath = pkgpath + "/images/win" + str(randrange(4)) + ".jpg"
    if mode == 2:
        filepath = pkgpath + "/images/lose.jpg"
    if mode == 11:
        filepath = pkgpath + "/images/board.jpg"
    if mode == 22:
        filepath = pkgpath + "/images/turn" + str(randrange(4)) + ".jpg"

    head_display.display_image(filepath)

def isGameOver(b):
    """
    Returns True if the game is over (either board is full or 3 in a row)
    Also returns the player that won if 3 in a row are detected
    """
    board = np.copy(b)
    moves = 0
    for i in range(3):
        for k in range(3):
            if board[i][k] == 2:
                board[i][k] = -1
            if board[i][k] != 0:
                moves += 1
    rows = np.copy(board)
    cols = np.copy(board).T
    diags = np.array([[0,0,0],[0,0,0]])
    for i in range(3):
        diags[0][i] = board[i][i]
        diags[1][i] = board[i][2-i]

    testWin = [0,0,0,0,0,0,0,0,0]
    for i in range(3):
        testWin[i] = np.sum(rows[i])
        testWin[i+3] = np.sum(cols[i])
    testWin[6] = np.sum(diags[0])
    testWin[7] = np.sum(diags[1])

    for i in range(8):
        if abs(testWin[i]) == 3:
            if testWin[i] < 0:
                #displayFace(2)
                return True,2
            #displayFace(1)
            return True,1

    if moves == 9:
        displayFace(0)
        return True, 0
    return False, 0

def findNextMove(gameboard, robot_player):
    """
    Given a 3x3 gameboard filled with 0 (no move), 1 (X), or 2 (O)
    Return a tuple corresponding to a cell on the tictactoe board
    Also return the shape to be drawn at that cell
    """
    nextMove = [-1, -1]
    rowSize = 3
    colSize = 3
    other_player = 1
    if robot_player == 1:
        other_player = 2

    #test for valid gameboard
    if len(gameboard) != rowSize:
        raise TypeError("Tic Tac Toe Board expected " + str(rowSize) + " rows, got " + str(len(gameboard)))
    for i in range(3):
        if len(gameboard[i]) != 3:
            raise TypeError("Tic Tac Toe Board expected " + str(colSize) + " columns, got " + str(len(gameboard[i])))
    ###

    print_array(gameboard)
    #test for gameover
    go,face = isGameOver(gameboard)
    if go:
        #print("GO")
        displayFace(face)
        return (0,0),-1

    #test to see if it is the robots turn
    currentMove = 0
    total_robot_moves = 0
    total_other_moves = 0
    for r in range(rowSize):
        for c in range(colSize):
            val = gameboard[r][c]
            if val == robot_player:
                total_robot_moves += 1
            if val == other_player:
                total_other_moves += 1
    currentMove = total_other_moves+total_robot_moves

    if(total_robot_moves > total_other_moves):
        print("Not the robots turn")
        return (-1,-1), 0
    ###

    #hard coded first move (go for center)
    if currentMove == 0:
        return (1,1), 1
    if currentMove == 1:
        if gameboard[1][1] == 0:
            return (1,1), 1
        return (0,0), 1
    ###

    weights = [[0,0,0],[0,0,0],[0,0,0]]

    for i in range(3):
        # get row weight
        pRobot = 0
        pOther = 0
        for k in range(3):
            cell = gameboard[i][k]
            if cell == robot_player:
                pRobot += 1
            if cell == other_player:
                pOther += 1
        w = pRobot - pOther
        if w > 1:
            w = 10
        if w < -1:
            w = 5

        # save row weight
        for k in range(3):
            if gameboard[i][k] != 0:
                weights[i][k] = -1
            else:
                weights[i][k] += abs(w)

        # get col weight
        pRobot = 0
        pOther = 0
        for k in range(3):
            cell = gameboard[k][i]
            if cell == robot_player:
                pRobot += 1
            if cell == other_player:
                pOther += 1
        w = pRobot - pOther
        if w > 1:
            w = 10
        if w < -1:
            w = 5

        # save col weight
        for k in range(3):
            if gameboard[k][i] != 0:
                weights[k][i] = -1
            else:
                weights[k][i] += abs(w)

        # end of loop

    # get left diagonal weight
    pRobot = 0
    pOther = 0
    for k in range(3):
        cell = gameboard[k][k]
        if cell == robot_player:
            pRobot += 1
        if cell == other_player:
            pOther += 1
    w = pRobot - pOther
    if w > 1:
        w = 10
    if w < -1:
        w = 5

    for k in range(3):
        if gameboard[k][k] != 0:
            weights[k][k] = -1
        else:
            weights[k][k] += abs(w)

    # get Right diagonal weight
    pRobot = 0
    pOther = 0
    for k in range(3):
        cell = gameboard[k][2-k]
        if cell == robot_player:
            pRobot += 1
        if cell == other_player:
            pOther += 1
    w = pRobot - pOther
    if w > 1:
        w = 10
    if w < -1:
        w = 5

    for k in range(3):
        if gameboard[k][2-k] != 0:
            weights[k][2-k] = -1
        else:
            weights[k][2-k] += abs(w)

    #print_array(weights)

    maxWeight = -1
    coor = (-1,-1)
    for i in range(3):
        for k in range(3):
            test_weight = weights[i][k]
            if test_weight > maxWeight:
                #print(test_weight)
                maxWeight = test_weight
                coor = (i,k)

    return coor, 1

def convertCoor(tuple):
    """
    Convert cell tuple from Bennett's gameboard numbering to Rico's numbering
    """
    return (-tuple[0]+1,-tuple[1]+1)

def convertShape2String(shapeInt):
    """
    Get the string corresponding to shape integer
    """
    if shapeInt == -1:
        return "end"
    if shapeInt == 0:
        return "idle"
    if shapeInt == 1:
        return "cross"
    if shapeInt == 2:
        return "circle"
    return "error"

def rotateBoard(board):
    """
    Rotate board from wrist camera
    (since wrist camera is upside down)
    """
    newBoard = np.copy(board)
    for i in range(3):
        for k in range(3):
            newBoard[2-i][2-k] = board[i][k]
    return newBoard

def waitKey():
    """
    Blocking function to wait for Enter key (or Escape key) to be pressed
    """
    kb.getch()
    print("CONTINUE KEY ACCEPTED")

def Update(camera):
    """
    This function takes the camera object from game_state.py
    It enables the camera, waits for a new picture to be taken, and then disables the camera
    It then uses the gameboard saved in the camera object
    (which has been updated now that the camera has taken a new picture)
    to find the cell where the next shape should be drawn.
    It returns a cell and a shape corresponding to that position and which shape the robot is playing as
    """
    rospy.sleep(1)
    camera.updateBoard()
    print(camera.gamestate)
    go, face = isGameOver(camera.gamestate)
    if go:
        displayFace(face)
        return (0,0),-1
    if np.sum(camera.gamestate) != 0:
        displayFace(22)
        print("Press Enter to continue...")
        waitKey()
    camera.cameras.start_streaming(camera.argsCamera)
    camera.once = False
    while(not camera.once):
        rospy.sleep(1)
    camera.cameras.stop_streaming(camera.argsCamera)
    rospy.sleep(1)
    go, face = isGameOver(camera.gamestate)
    if not go:
        displayFace(11)
    else:
        displayFace(face)
        return (0,0),-1
    move,shape = findNextMove(rotateBoard(camera.gamestate), 1)
    #print_array(rotateBoard(camera.gamestate))
    #print(move)

    print(":" + str(move))
    print(":" + str(shape))
    if move[0] == -1:
        return (0,0),0
    if move[1] == -1:
        return (0,0),0
    if shape == -1:
        return (0,0),-1
    return move,1

if __name__ == "__main__":
    """
    Testing for tttAI.py done through running this pyhton script directly
    """
    rospy.init_node('saw')
    board = [[1,1,2],[2,2,1],[1,1,1]]
    print_array(board)
    print()
    waitKey()

    bot = 1
    move,shape = findNextMove(board,bot)
    print("Move: " + str(move))
    print("Shape: " + str(shape))