# -*- encoding:utf-8 -*-
import os
import torch
import codecs
import random
import pickle
import re
from tqdm import tqdm
from multiprocessing import Pool
from uer.utils.constants import *
from uer.utils.misc import count_lines
from uer.utils.seed import set_seed

# import pkuseg
# pku_seg = pkuseg.pkuseg(model_name="medicine",user_dict="../uer/utils/pku_seg_dict.txt")
# pku_seg_pos = pkuseg.pkuseg(model_name="medicine",user_dict="../uer/utils/pku_seg_dict.txt",postag=True)

# pos_dict = {}
# with open('../uer/utils/pos_tags_old.txt','r',encoding='utf-8') as f:
#     i = 0
#     for line in f.readlines():
#         if line:
#             pos_dict[line.strip().split()[0]] = i
#             i += 1

# # 获取本地术语表
# a = []
# term_dict = {}
# with open('../uer/utils/medical_terms/medical_terms(final).txt', 'r', encoding='utf-8') as f:
#     for line in f.readlines():
#         line = line.strip()
#         a.append(line)
#         term_dict[line.lower()] = 1

# term_set = set(a)

# max_num = max([len(line) for line in a])

def mask_seq(src, vocab_size):
    """
    mask input sequence for MLM task
    args:
        src: a list of tokens
        vocab_size: the vocabulary size
    """
    tgt_mlm = []
    for (i, token) in enumerate(src):
        if token == CLS_ID or token == SEP_ID:
            continue
        prob = random.random()
        if prob < 0.15:
            prob /= 0.15
            if prob < 0.8:
                src[i] = MASK_ID
            elif prob < 0.9:
                while True:
                    rdi = random.randint(1, vocab_size-1)
                    if rdi not in [CLS_ID, SEP_ID, MASK_ID]:
                        break
                src[i] = rdi
            tgt_mlm.append((i, token))
    return src, tgt_mlm

def merge_dataset(dataset_path, workers_num):
        # Merge datasets.
        f_writer = open(dataset_path, "wb")
        for i in range(workers_num):
            tmp_dataset_reader = open("dataset-tmp-"+str(i)+".pt", "rb")
            while True:
                tmp_data = tmp_dataset_reader.read(2^20)
                if tmp_data:
                    f_writer.write(tmp_data)
                else:
                    break
            tmp_dataset_reader.close()
            os.remove("dataset-tmp-"+str(i)+".pt")
        f_writer.close()

def seg_char(sent):
    """
    把句子按字分开，不破坏英文及数字结构
    """
    # 按中文汉字分割
    pattern = re.compile(r'([\u4e00-\u9fa5])')
    parts = pattern.split(sent)
    parts = [w for w in parts if len(w.strip())>0]

    # 按英文标点符号分割
    chars_list = []
    pattern = re.compile(r'([,?:;\'"!()<>，。；’‘、！])')
    for part in parts:
        chars = pattern.split(part)
        chars = [w for w in chars if len(w.strip())>0]
        chars_list += chars

    return chars_list

def max_match(txt, ano_dict, max_num):
    word_list = seg_char(txt) # 中文单字切割，保留英文和数字
    # print(word_list)
    new_word_list = []
    term_labels = []
    N = len(word_list)
    k = max_num
    i = 0
    while i < N:
        if i <= N - k:
            j = k
            while j > 0:
                token_tmp = ''.join(word_list[i:i + j])
                # print(token_tmp)
                if token_tmp.lower() in ano_dict.keys():
                    # print(token_tmp,'！!！!!!！!!!！!！!！!')
                    new_word_list.append(token_tmp)
                    term_labels.append(1)
                    i += j
                    break
                else:
                    j -= 1
            if j == 0:
                new_word_list.append(word_list[i])
                term_labels.append(0)
                i += 1
        else:
            j = N - i
            while j > 0:
                token_tmp = ''.join(word_list[i:i + j])
                # print(token_tmp)
                if token_tmp.lower() in ano_dict.keys():
                    # print(token_tmp, '！!！!!!！!!!！!！!！!')
                    new_word_list.append(token_tmp)
                    term_labels.append(1)
                    i += j
                    break
                else:
                    j -= 1
            if j == 0:
                new_word_list.append(word_list[i])
                term_labels.append(0)
                i += 1
    return new_word_list,term_labels

# ## 最大匹配分词
# txt = '我爱北京天安门，天安门上太阳升'
#
# ano_dict={'天安门':1,'北京':1,'太阳':1}
#
# max_num = 2
#
# print(max_match(txt,ano_dict,max_num))
# exit()

class Dataset(object):
    def __init__(self, args, vocab, tokenizer):
        self.vocab = vocab
        self.tokenizer = tokenizer
        self.corpus_path = args.corpus_path
        self.dataset_path = args.dataset_path
        self.seq_length = args.seq_length
        self.seed = args.seed
        self.stats_tokens = args.stats_tokens
        self.stats_instances = args.stats_instances
        self.stats_lines = args.stats_lines

    def build_and_save(self, workers_num):
        """
        Build dataset from the given corpus.
        Start workers_num processes and each process deals with a part of data.
        """

        if self.stats_instances:
            print('stats_instances...')
            ins_count = 0
            with open(self.dataset_path,'rb') as f:
                i = 0
                while True:
                    try:
                        instance = pickle.load(f)
                        ins_count += len(instance)
                    except EOFError:
                        break
                    i += 1
            print("instances: ", ins_count)
            exit()

        if self.stats_tokens:
            print('stats_tokens...')
            tokens_count = 0
            with open(self.corpus_path, mode="r", encoding="utf-8") as f:
                for line in tqdm(f.readlines()):
                    line = line.strip()
                    if line:
                        tokens_count += len(self.tokenizer.tokenize(line))
            print("tokens: ",tokens_count)
            exit()

        if self.stats_lines:
            lines_count = 0
            with open(self.corpus_path, mode="r", encoding="utf-8") as f:
                for line in tqdm(f.readlines()):
                    line = line.strip()
                    if line:
                        lines_count += 1
            print("lines: ",lines_count)
            exit()

        lines_num = count_lines(self.corpus_path)
        print("Total %d lines in courpus ... " % lines_num)
        print("Starting %d workers for building datasets ... " % workers_num)
        assert(workers_num >= 1)
        if workers_num == 1:
            self.worker(0, 0, lines_num)
        else:
            pool = Pool(workers_num)
            for i in range(workers_num):
                start = i * lines_num // workers_num
                end = (i+1) * lines_num // workers_num
                pool.apply_async(func=self.worker, args=[i, start, end])
            pool.close()
            pool.join()

        # Merge datasets.
        merge_dataset(self.dataset_path, workers_num)

    def worker(self, proc_id, start, end):
        raise NotImplementedError()

class DataLoader(object):
    def __init__(self, args, dataset_path, batch_size, proc_id, proc_num, shuffle=False):
        self.batch_size = batch_size
        self.instances_buffer_size = args.instances_buffer_size
        self.proc_id = proc_id
        self.proc_num = proc_num
        self.shuffle = shuffle
        self.f_read = open(dataset_path, "rb")
        self.read_count = 0
        self.start = 0
        self.end = 0
        self.buffer = []
        self.add_pos = args.add_pos

    def _fill_buf(self):
        try:
            self.buffer = []
            while True:
                instance = pickle.load(self.f_read)
                self.read_count += 1
                if (self.read_count - 1) % self.proc_num == self.proc_id:
                    self.buffer.append(instance)
                    if len(self.buffer) >= self.instances_buffer_size:
                        break
        except EOFError:
            # Reach file end.
            self.f_read.seek(0)

        if self.shuffle:
            random.shuffle(self.buffer)
        self.start = 0
        self.end = len(self.buffer)

    def _empty(self):
        return self.start >= self.end

    def __del__(self):
        self.f_read.close()

class BertDataset(Dataset):
    """
    Construct dataset for MLM and NSP tasks from the given corpus.
    Each document consists of multiple sentences,
    and each sentence occupies a single line.
    Documents in corpus must be separated by empty lines.
    """
    def __init__(self, args, vocab, tokenizer):
        super(BertDataset, self).__init__(args, vocab, tokenizer)
        self.docs_buffer_size = args.docs_buffer_size
        self.dup_factor = args.dup_factor
        self.short_seq_prob = args.short_seq_prob

    def worker(self, proc_id, start, end):
        print("Worker %d is building dataset ... " % proc_id)
        set_seed(self.seed)
        docs_buffer = []
        document = []
        pos = 0
        f_write = open("dataset-tmp-" + str(proc_id) + ".pt", "wb")
        with open(self.corpus_path, mode="r", encoding="utf-8") as f:
            while pos < start:
                try:
                    f.readline()
                except:
                    continue
                finally:
                    pos += 1
            while True:
                try:
                    line = f.readline()
                except:
                    continue
                finally:
                    pos += 1
                if not line.strip():
                    if len(document) >= 1:
                        docs_buffer.append(document)
                    document = []
                    if len(docs_buffer) == self.docs_buffer_size:
                        # Build instances from documents.
                        instances = self.build_instances(docs_buffer)
                        # Save instances.
                        for instance in instances:
                            pickle.dump(instance, f_write)
                        # Clear buffer.
                        docs_buffer = []
                        instances = []
                    continue
                sentence = [self.vocab.get(w) for w in self.tokenizer.tokenize(line)]
                if len(sentence) > 0:
                    document.append(sentence)

                if pos >= end - 1:
                    if len(docs_buffer) > 0:
                        instances = self.build_instances(docs_buffer)
                        for instance in instances:
                            pickle.dump(instance, f_write)
                    break
        f_write.close()

    def build_instances(self, all_documents):
        instances = []
        for _ in range(self.dup_factor):
            for doc_index in range(len(all_documents)):
                instances.extend(self.create_ins_from_doc(all_documents, doc_index))
        return instances

    def create_ins_from_doc(self, all_documents, document_index):
        document = all_documents[document_index]
        max_num_tokens = self.seq_length - 3
        target_seq_length = max_num_tokens
        if random.random() < self.short_seq_prob:
            target_seq_length = random.randint(2, max_num_tokens)
        instances = []
        current_chunk = []
        current_length = 0
        i = 0
        while i < len(document):
            segment = document[i]
            current_chunk.append(segment)
            current_length += len(segment)
            if i == len(document) - 1 or current_length >= target_seq_length:
                if current_chunk:
                    a_end = 1
                    if len(current_chunk) >= 2:
                        a_end = random.randint(1, len(current_chunk) - 1)

                    tokens_a = []
                    for j in range(a_end):
                        tokens_a.extend(current_chunk[j])

                    tokens_b = []
                    is_random_next = 0

                    # Random next
                    if len(current_chunk) == 1 or random.random() < 0.5:
                        is_random_next = 1
                        target_b_length = target_seq_length - len(tokens_a)

                        for _ in range(10):
                            random_document_index = random.randint(0, len(all_documents) - 1)
                            if random_document_index != document_index:
                                break

                        random_document = all_documents[random_document_index]
                        random_start = random.randint(0, len(random_document) - 1)
                        for j in range(random_start, len(random_document)):
                            tokens_b.extend(random_document[j])
                            if len(tokens_b) >= target_b_length:
                                break

                        num_unused_segments = len(current_chunk) - a_end
                        i -= num_unused_segments

                    # Actual next
                    else:
                        is_random_next = 0
                        for j in range(a_end, len(current_chunk)):
                            tokens_b.extend(current_chunk[j])

                    self.truncate_seq_pair(tokens_a, tokens_b, max_num_tokens)

                    # assert len(tokens_a) >= 1
                    # assert len(tokens_b) >= 1

                    src = []

                    src.append(CLS_ID)
                    for token in tokens_a:
                        src.append(token)

                    src.append(SEP_ID)

                    seg_pos = [len(src)]

                    for token in tokens_b:
                        src.append(token)

                    src.append(SEP_ID)

                    seg_pos.append(len(src))

                    src, tgt_mlm = mask_seq(src, len(self.vocab))

                    while len(src) != self.seq_length:
                        src.append(PAD_ID)

                    instance = (src, tgt_mlm, is_random_next, seg_pos)
                    instances.append(instance)
                current_chunk = []
                current_length = 0
            i += 1
        return instances

    def truncate_seq_pair(self, tokens_a, tokens_b, max_num_tokens):
        """ truncate sequence pair to specific length """
        while True:
            total_length = len(tokens_a) + len(tokens_b)
            if total_length <= max_num_tokens:
                break

            trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b
            # assert len(trunc_tokens) >= 1

            if random.random() < 0.5:
                del trunc_tokens[0]
            else:
                trunc_tokens.pop()

class BertDataLoader(DataLoader):
    def __iter__(self):
        while True:
            while self._empty():
                self._fill_buf()
            if self.start + self.batch_size >= self.end:
                instances = self.buffer[self.start:]
            else:
                instances = self.buffer[self.start: self.start + self.batch_size]

            self.start += self.batch_size

            src = []
            tgt_mlm = []
            is_next = []
            seg = []

            masked_words_num = 0
            for ins in instances:
                masked_words_num += len(ins[1])
            if masked_words_num == 0:
                continue

            for ins in instances:
                src.append(ins[0])
                tgt_mlm.append([0]*len(ins[0]))
                for mask in ins[1]:
                    tgt_mlm[-1][mask[0]] = mask[1]
                is_next.append(ins[2])
                seg.append([1]*ins[3][0] + [2]*(ins[3][1]-ins[3][0]) + [PAD_ID]*(len(ins[0])-ins[3][1]))

            yield torch.LongTensor(src), \
                torch.LongTensor(tgt_mlm), \
                torch.LongTensor(is_next), \
                torch.LongTensor(seg)

class LmDataset(Dataset):
    def worker(self, proc_id, start, end):
        print("Worker %d is building dataset ... " % proc_id)
        set_seed(self.seed)
        pos = 0
        f_write = open("dataset-tmp-" + str(proc_id) + ".pt", "wb")
        with open(self.corpus_path, mode="r", encoding="utf-8") as f:
            while pos < start:
                try:
                    f.readline()
                except:
                    continue
                finally:
                    pos += 1
            while True:
                try:
                    line = f.readline()
                except:
                    continue
                finally:
                    pos += 1

                src = [self.vocab.get(w) for w in self.tokenizer.tokenize(line)]
                tgt = src[1:]
                src = src[:-1]
                seg = [1] * len(src)
                if len(src) >= self.seq_length:
                    src = src[:self.seq_length]
                    tgt = tgt[:self.seq_length]
                    seg = seg[:self.seq_length]
                else:
                    while len(src) != self.seq_length:
                        src.append(PAD_ID)
                        tgt.append(PAD_ID)
                        seg.append(PAD_ID)

                pickle.dump((src, tgt, seg), f_write)

                if pos >= end - 1:
                    break

        f_write.close()

class LmDataLoader(DataLoader):
    def __iter__(self):
        while True:
            while self._empty():
                self._fill_buf()
            if self.start + self.batch_size >= self.end:
                instances = self.buffer[self.start:]
            else:
                instances = self.buffer[self.start: self.start + self.batch_size]

            self.start += self.batch_size

            src = []
            tgt = []
            seg = []

            for ins in instances:
                src.append(ins[0])
                tgt.append(ins[1])
                seg.append(ins[2])

            yield torch.LongTensor(src), \
                torch.LongTensor(tgt), \
                torch.LongTensor(seg)

class BilmDataset(Dataset):
    def worker(self, proc_id, start, end):
        print("Worker %d is building dataset ... " % proc_id)
        set_seed(self.seed)
        pos = 0
        f_write = open("dataset-tmp-" + str(proc_id) + ".pt", "wb")
        with open(self.corpus_path, mode="r", encoding="utf-8") as f:
            while pos < start:
                try:
                    f.readline()
                except:
                    continue
                finally:
                    pos += 1
            while True:
                try:
                    line = f.readline()
                except:
                    continue
                finally:
                    pos += 1

                src = [self.vocab.get(w) for w in self.tokenizer.tokenize(line)]
                if len(src) < 1:
                    continue
                tgt_forward = src[1:] + [SEP_ID]
                tgt_backward = [CLS_ID] + src[:-1]
                seg = [1] * len(src)
                if len(src) >= self.seq_length:
                    src = src[:self.seq_length]
                    tgt_forward = tgt_forward[:self.seq_length]
                    tgt_backward = tgt_backward[:self.seq_length]
                    seg = seg[:self.seq_length]
                else:
                    while len(src) != self.seq_length:
                        src.append(PAD_ID)
                        tgt_forward.append(PAD_ID)
                        tgt_backward.append(PAD_ID)
                        seg.append(PAD_ID)

                pickle.dump((src, tgt_forward, tgt_backward, seg), f_write)

                if pos >= end - 1:
                    break

        f_write.close()

class BilmDataLoader(DataLoader):
    def __iter__(self):
        while True:
            while self._empty():
                self._fill_buf()
            if self.start + self.batch_size >= self.end:
                instances = self.buffer[self.start:]
            else:
                instances = self.buffer[self.start: self.start + self.batch_size]

            self.start += self.batch_size

            src = []
            tgt_forward = []
            tgt_backward = []
            seg = []

            for ins in instances:
                src.append(ins[0])
                tgt_forward.append(ins[1])
                tgt_backward.append(ins[2])
                seg.append(ins[3])

            yield torch.LongTensor(src), \
                torch.LongTensor(tgt_forward), \
                torch.LongTensor(tgt_backward), \
                torch.LongTensor(seg)

class ClsDataset(Dataset):
    def worker(self, proc_id, start, end):
        print("Worker %d is building dataset ... " % proc_id)
        set_seed(self.seed)
        pos = 0
        f_write = open("dataset-tmp-" + str(proc_id) + ".pt", "wb")
        with open(self.corpus_path, mode="r", encoding="utf-8") as f:
            while pos < start:
                try:
                    f.readline()
                except:
                    continue
                finally:
                    pos += 1
            while True:
                try:
                    line = f.readline()
                except:
                    continue
                finally:
                    pos += 1

                line = line.strip().split('\t')
                if len(line) == 2:
                    label = int(line[0])
                    text = " ".join(line[1:])
                    src = [self.vocab.get(t) for t in self.tokenizer.tokenize(text)]
                    src = [CLS_ID] + src
                    tgt = label
                    seg = [1] * len(src)
                    if len(src) >= self.seq_length:
                        src = src[:self.seq_length]
                        seg = seg[:self.seq_length]
                    else:
                        while len(src) != self.seq_length:
                            src.append(PAD_ID)
                            seg.append(PAD_ID)
                    pickle.dump((src, tgt, seg), f_write)
                elif len(line) == 3: # For sentence pair input.
                    label = int(line[0])
                    text_a, text_b = line[1], line[2]

                    src_a = [self.vocab.get(t) for t in self.tokenizer.tokenize(text_a)]
                    src_a = [CLS_ID] + tokens_a + [SEP_ID]
                    src_b = [vocab.get(t) for t in tokenizer.tokenize(text_b)]
                    src_b = tokens_b + [SEP_ID]

                    src = src_a + src_b
                    seg = [1] * len(src_a) + [2] * len(src_b)

                    if len(src) >= self.seq_length:
                        src = src[:self.seq_length]
                        seg = seg[:self.seq_length]
                    else:
                        while len(src) != self.seq_length:
                            src.append(PAD_ID)
                            seg.append(PAD_ID)
                    pickle.dump((src, tgt, seg), f_write)
                else:
                    pass

                if pos >= end - 1:
                    break

        f_write.close()

class ClsDataLoader(DataLoader):
    def __iter__(self):
        while True:
            while self._empty():
                self._fill_buf()
            if self.start + self.batch_size >= self.end:
                instances = self.buffer[self.start:]
            else:
                instances = self.buffer[self.start: self.start + self.batch_size]

            self.start += self.batch_size

            src = []
            tgt = []
            seg = []

            for ins in instances:
                src.append(ins[0])
                tgt.append(ins[1])
                seg.append(ins[2])

            yield torch.LongTensor(src), \
                torch.LongTensor(tgt), \
                torch.LongTensor(seg)

class MlmDataset(Dataset):
    def __init__(self, args, vocab, tokenizer):
        super(MlmDataset, self).__init__(args, vocab, tokenizer)
        self.dup_factor = args.dup_factor

    def worker(self, proc_id, start, end):
        print("Worker %d is building dataset ... " % proc_id)
        set_seed(self.seed)
        f_write = open("dataset-tmp-" + str(proc_id) + ".pt", "wb")
        for _ in range(self.dup_factor):
            pos = 0
            with open(self.corpus_path, mode="r", encoding="utf-8") as f:
                while pos < start:
                    try:
                        f.readline()
                    except:
                        continue
                    finally:
                        pos += 1
                pbar = tqdm(total=end - start)
                pbar.set_description("Worker %d:" % proc_id)
                while True:
                    pbar.update(1)
                    try:
                        line = f.readline()
                    except:
                        continue
                    finally:
                        pos += 1

                    line = line.strip()
                    if not line:
                        continue

                    src = [self.vocab.get(w) for w in self.tokenizer.tokenize(line)]

                    if len(src) > self.seq_length:
                        src = src[:self.seq_length]
                    seg = [1] * len(src)

                    src, tgt = mask_seq(src, len(self.vocab))

                    while len(src) != self.seq_length:
                        src.append(PAD_ID)
                        seg.append(PAD_ID)

                    pickle.dump((src, tgt, seg), f_write)

                    if pos >= end - 1:
                        break
                pbar.close()
        f_write.close()

class MlmDataLoader(DataLoader):
    def __iter__(self):
        while True:
            while self._empty():
                self._fill_buf()
            if self.start + self.batch_size >= self.end:
                instances = self.buffer[self.start:]
            else:
                instances = self.buffer[self.start: self.start + self.batch_size]

            self.start += self.batch_size

            src = []
            tgt = []
            seg = []

            masked_words_num = 0
            for ins in instances:
                masked_words_num += len(ins[1])
            if masked_words_num == 0:
                continue

            for ins in instances:
                src.append(ins[0])
                seg.append(ins[2])
                tgt.append([0]*len(ins[0]))
                for mask in ins[1]:
                    tgt[-1][mask[0]] = mask[1]

            yield torch.LongTensor(src), \
                torch.LongTensor(tgt), \
                torch.LongTensor(seg)

class Csci_mlmDataset(Dataset):

    def __init__(self, args, vocab, tokenizer):
        super(Csci_mlmDataset, self).__init__(args, vocab, tokenizer)
        self.dup_factor = args.dup_factor

    def worker(self, proc_id, start, end):
        print("Worker %d is building dataset ... " % proc_id)
        set_seed(self.seed)
        f_write = open("dataset-tmp-" + str(proc_id) + ".pt", "wb")
        for _ in range(self.dup_factor):
            pos = 0
            with open(self.corpus_path, mode="r", encoding="utf-8") as f:
                while pos < start:
                    try:
                        f.readline()
                    except:
                        continue
                    finally:
                        pos += 1
                pbar = tqdm(total=end - start)
                pbar.set_description("Worker %d:" % proc_id)
                while True:
                    pbar.update(1)
                    try:
                        line = f.readline()
                    except:
                        continue
                    finally:
                        pos += 1

                    line = line.strip().lower()
                    if not line:
                        continue

                    tokens = []
                    for word in pku_seg.cut(line):
                        for w in self.tokenizer.tokenize(word):
                            tokens.append(w)

                    src_word = [self.vocab.get(w) for w in tokens]

                    # src_word = [self.vocab.get(w) for w in self.tokenizer.tokenize(line)]

                    if len(src_word) > self.seq_length:
                        src_word = src_word[:self.seq_length]
                    seg = [1] * len(src_word)

                    src_word, tgt = mask_seq(src_word, len(self.vocab))
                    # print('len(src_word)',len(src_word))

                    # tokens = [w for w in self.tokenizer.tokenize(line)]
                    # print([(i,a) for (i,a) in enumerate(tokens)])

                    src_pos = []
                    src_term = []
                    ## 加入pos 和terms
                    for (word, tag) in pku_seg_pos.cut(line):
                        piece_num = len(self.tokenizer.tokenize(word))
                        if word in term_set:
                            # print(word)
                            [src_term.append(1) for i in range(piece_num)]
                        else:
                            [src_term.append(0) for i in range(piece_num)]

                        [src_pos.append(pos_dict[tag]) for i in range(piece_num)]

                    if len(src_pos) > self.seq_length:
                        src_pos = src_pos[:self.seq_length]

                    # terms,labels = max_match(line.strip(),term_dict,max_num)
                    #
                    # for i,term in enumerate(terms):
                    #     if labels[i]:
                    #         for w in self.tokenizer.tokenize(term):
                    #             src_term.append(1)
                    #     else:
                    #         for w in self.tokenizer.tokenize(term):
                    #             src_term.append(0)

                    if len(src_term) > self.seq_length:
                        src_term = src_term[:self.seq_length]

                    # print('len(src_term)',len(src_term))
                    # print('src_term:',src_term)
                    # exit()

                    while len(src_word) != self.seq_length:
                        src_word.append(PAD_ID)
                        src_pos.append(pos_dict['[PAD]'])
                        src_term.append(2)
                        seg.append(PAD_ID)

                    if len(src_pos) != self.seq_length:
                        print('src_pos Problem~~~')
                        print(line)

                        tokens = [w for w in self.tokenizer.tokenize(line)]
                        print('tokens:\n', [(i, a) for (i, a) in enumerate(tokens)])

                        cut_to_tokens = []

                        for (word, tag) in pku_seg.cut(line.strip()):
                            # print(word,tag)
                            for w in self.tokenizer.tokenize(word):
                                cut_to_tokens.append(w)

                        print('cut to tokens:\n', [(i, a) for (i, a) in enumerate(cut_to_tokens)])

                        print('src_pos:\n', [(i, a) for (i, a) in enumerate(src_pos)])

                        exit()

                    if len(src_term) != self.seq_length:
                        print(line)

                        print('terms:\n',terms)
                        print('lables:\n',labels)

                        term_to_tokens = []
                        for i, term in enumerate(terms):
                            if labels[i]:
                                for w in self.tokenizer.tokenize(term):
                                    src_term.append(1)
                                    term_to_tokens.append(w)
                            else:
                                for w in self.tokenizer.tokenize(term):
                                    src_term.append(0)
                                    term_to_tokens.append(w)

                        print('term to tokens:\n',[(i,a) for (i,a) in enumerate(term_to_tokens)])

                        # tokens = [w for w in self.tokenizer.tokenize(line)]
                        print('tokens:\n',[(i,a) for (i,a) in enumerate(tokens)])
                        print('src_word:\n',[(i,a) for (i,a) in enumerate(src_word)])
                        print('src_pos:\n',[(i,a) for (i,a) in enumerate(src_pos)])
                        print('src_term:\n',[(i,a) for (i,a) in enumerate(src_term)])
                        print("tgt\n",[(i,a) for (i,a) in enumerate(tgt)])
                        print("seg\n",[(i,a) for (i,a) in enumerate(seg)])
                        exit()

                    # print(line)
                    # print('tokens:\n', [(i, a) for (i, a) in enumerate(tokens)])
                    # print('src_word:\n', [(i, a) for (i, a) in enumerate(src_word)])
                    # print('src_pos:\n', [(i, a) for (i, a) in enumerate(src_pos)])
                    # print('src_term:\n', [(i, a) for (i, a) in enumerate(src_term)])
                    # print("tgt\n", [(i, a) for (i, a) in enumerate(tgt)])
                    # print("seg\n", [(i, a) for (i, a) in enumerate(seg)])
                    # exit()

                    pickle.dump((src_word, src_pos, src_term, tgt, seg), f_write)

                    if pos >= end - 1:
                        break
                pbar.close()
        f_write.close()

class Csci_mlmDataLoader(DataLoader):
    def __iter__(self):
        while True:
            while self._empty():
                self._fill_buf()
            if self.start + self.batch_size >= self.end:
                instances = self.buffer[self.start:]
            else:
                instances = self.buffer[self.start: self.start + self.batch_size]

            self.start += self.batch_size

            src_word = []

            ## 加入pos和term
            src_pos = []
            src_term = []

            tgt = []
            seg = []

            masked_words_num = 0
            for ins in instances:
                masked_words_num += len(ins[3])
            if masked_words_num == 0:
                continue

            for ins in instances:
                src_word.append(ins[0])

                ## 加入pos和term
                src_pos.append(ins[1])
                src_term.append(ins[2])

                seg.append(ins[4])
                tgt.append([0]*len(ins[0]))
                for mask in ins[3]:
                    tgt[-1][mask[0]] = mask[1]

            yield torch.LongTensor(src_word), \
                torch.LongTensor(src_pos), \
                torch.LongTensor(src_term), \
                torch.LongTensor(tgt), \
                torch.LongTensor(seg)


import copy
from datetime import datetime

from abc import ABC, abstractmethod
from forensic.common.stream.stream import DataStruct, dynamic, BinaryReader, uint16le, uint32le, bytesbe

class CodesysV3FileFormatAbstract(ABC):
    @abstractmethod
    def read_tag(self, br: BinaryReader):
        pass

    @abstractmethod
    def read_list_of_tags(self, tag_data: bytes):
        pass

class CodesysBaseTag(DataStruct):
    tid = dynamic

    def _null_ending_bytes(self, data: bytes) -> bytes:
        return data[:data.find(b'\x00')]

class CodesysUnknownTag(CodesysBaseTag):
    data = dynamic

class CodesysListOfTags(CodesysBaseTag):
    tid = 0x81
    sub_tags = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.sub_tags = cff.read_list_of_tags(data)

class CodesysApplicationName(CodesysBaseTag):
    tid = 0x10
    name = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.name = self._null_ending_bytes(data)

class CodesysListOfCodeAreas(CodesysBaseTag):
    tid = 0x82
    sub_areas = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.sub_areas = cff.read_list_of_tags(data)

class CodesysCodeAreaHeader(CodesysBaseTag):
    tid = 0x21
    data = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.data = data

class CodesysCodeAreaData(CodesysBaseTag):
    tid = 0x22
    data = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.data = data

class CodesysCodeArea(CodesysBaseTag):
    tid = 0xA0
    area_type = dynamic
    pointer = dynamic
    code = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        br_code_area = BinaryReader(data)
        tag_header = cff.read_tag(br_code_area)

        if isinstance(tag_header, CodesysCodeAreaHeader):
            br_pointer = BinaryReader(tag_header.data)
            self.area_type = br_pointer.read(uint16le)
            self.pointer = br_pointer.read(uint32le)

            tag_code = cff.read_tag(br_code_area)

            if isinstance(tag_code, CodesysCodeAreaData):
                br_code = BinaryReader(tag_code.data)
                code_size = br_code.read(uint32le)
                self.code = br_code.read(bytesbe(code_size))

class CodesysProjectInfo(CodesysBaseTag):
    tid = 0x87
    attributes = dynamic
    unknown_tags = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.attributes = {}

        labels = {1: 'name',
                  2: 'project_version',
                  3: 'project_build_datetime',
                  4: 'author',
                  5: 'description',
                  6: 'ide_build_version',
                  8: 'codesys_version'}

        for sub_tag in cff.read_list_of_tags(data):
            if sub_tag.tid in labels:
                if sub_tag.tid == 3:
                    self.attributes[labels[sub_tag.tid]] = datetime.fromtimestamp(
                        int.from_bytes(sub_tag.data, byteorder='little'))
                else:
                    self.attributes[labels[sub_tag.tid]] = self._null_ending_bytes(sub_tag.data)

            else:
                self.attributes[f'unknown_tag {sub_tag.tid}'] = sub_tag.data

class CodesysFunctionSymbols(CodesysBaseTag):
    tid = 0x88
    symbols = dynamic
    unknown_tags = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.symbols = {}
        self.unknown_tags = []

        temp_symbol = None
        for sub_tag in cff.read_list_of_tags(data):
            if sub_tag.tid == 1:
                temp_symbol = {}
                br_sub_tag = BinaryReader(sub_tag.data)
                temp_symbol['type'] = br_sub_tag.read(uint16le)
                temp_symbol['pointer'] = br_sub_tag.read(uint32le)
                temp_symbol['raw'] = br_sub_tag.read(bytesbe(br_sub_tag.remaining_data()))
            elif sub_tag.tid == 2:
                if temp_symbol is None:
                    temp_symbol = {}
                temp_symbol['name'] = self.null_ending_bytes(sub_tag.data)
                self.symbols[temp_symbol['pointer']] = copy.deepcopy(temp_symbol)
                temp_symbol = None
            else:
                self.unknown_tags.append({'unknown tag': sub_tag.tid, 'raw': sub_tag.data})

class CodesysFooterAddresses(CodesysBaseTag):
    tid = 0x86
    sub_tags = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.sub_tags = cff.read_list_of_tags(data)

class CodesysInitJumpTable(CodesysBaseTag):
    tid = 0x60
    pointer_type = dynamic
    pointer = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        br = BinaryReader(data)
        self.pointer_type = br.read(uint16le)
        self.pointer = br.read(uint32le)

class CodesysSymbols(CodesysBaseTag):
    tid = 0x85
    symbols = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.symbols = cff.read_list_of_tags(data)

class CodesysSymbol(CodesysBaseTag):
    tid = 0x50
    pointer_type = dynamic
    pointer = dynamic
    zeros = dynamic
    name = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        br_symbol = BinaryReader(data)
        self.pointer_type = br_symbol.read(uint16le)
        self.pointer = br_symbol.read(uint32le)
        self.zeros = br_symbol.read(uint32le)
        self.name = self._null_ending_bytes(br_symbol.read(bytesbe(br_symbol.remaining_data())))

class CodesysSymbolExtended(CodesysBaseTag):
    tid = 0x88
    pointer_type = dynamic
    pointer = dynamic
    resolve_hash = dynamic
    index = dynamic
    raw = dynamic
    name = dynamic
    unknown_tags = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        self.unknown_tags = []

        for sub_tag in cff.read_list_of_tags(data):
            if sub_tag.tid == 1:
                br_sub_tag = BinaryReader(sub_tag.data)
                self.pointer_type = br_sub_tag.read(uint16le)
                self.pointer = br_sub_tag.read(uint32le)
                self.resolve_hash = br_sub_tag.read(uint32le)
                self.index = br_sub_tag.read(uint32le)
                self.raw = br_sub_tag.read(bytesbe(br_sub_tag.remaining_data()))
            elif sub_tag.tid == 2:
                self.name = self._null_ending_bytes(sub_tag.data)
            else:
                self.unknown_tags.append({'unknown tag': sub_tag.tid, 'raw': sub_tag.data})

class CodesysEntryPoint(CodesysBaseTag):
    tid = 0x61
    pointer_type = dynamic
    pointer = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        br = BinaryReader(data)
        self.pointer_type = br.read(uint16le)
        self.pointer = br.read(uint32le)

class CodesysInternalTableSymbolStart(CodesysBaseTag):
    tid = 0x6f
    pointer_type = dynamic
    pointer = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        br = BinaryReader(data)
        self.pointer_type = br.read(uint16le)
        self.pointer = br.read(uint32le)

class CodesysInternalTableSymbolEnd(CodesysBaseTag):
    tid = 0x71
    pointer_type = dynamic
    pointer = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        br = BinaryReader(data)
        self.pointer_type = br.read(uint16le)
        self.pointer = br.read(uint32le)

class CodesysInitApplication(CodesysBaseTag):
    tid = 0x63
    pointer_type = dynamic
    pointer = dynamic

    def __init__(self, cff: CodesysV3FileFormatAbstract, data: bytes):
        br = BinaryReader(data)
        self.pointer_type = br.read(uint16le)
        self.pointer = br.read(uint32le)


grid = [[1,1,1,1,1,1,1],
        [1,0,0,0,0,0,1],
        [1,0,1,1,1,0,1],
        [1,0,1,0,1,0,1],
        [1,0,1,1,1,0,1],
        [1,0,0,0,0,0,1],
        [1,1,1,1,1,1,1]]

rows = len(grid)
cols = len(grid[0])

closedIslands = 0
if rows == 0:
	print(closedIslands)
	exit(0)

def isclosed(grid,i,j,rows,cols):
	if grid[i][j] == -1 or grid[i][j] == 1:
		return True
	if i==0 or i==rows-1 or j==0 or j==cols-1:
		return False
	grid[i][j] = -1
	left = isclosed(grid,i-1,j,rows,cols)
	right = isclosed(grid,i+1,j,rows,cols)
	bottom = isclosed(grid,i,j+1,rows,cols)
	top = isclosed(grid,i,j-1,rows,cols)
	return left and right and bottom and top

for i in range(1,rows-1):
	for j in range(1,cols-1):
		if grid[i][j] == 0:
			if isclosed(grid,i,j,rows,cols):
				closedIslands+=1
print(closedIslands)


import logging
import cv2
import numpy as np
import pytesseract
import os
import time
import json
import re

from multiprocessing import Pool
from Levenshtein import distance

from .input_handler import InputHandler
from .grabscreen import grab_screen
from .utils import get_config, filter_mod

# This is a position of the inventory as fraction of the resolution
OWN_INVENTORY_ORIGIN = (0.6769531, 0.567361)

# These are the sockets positions as measured on 2560x1440 resolution
# with X_SCALE and Y_SCALE applied, i.e., scale * SOCKETS[i] is the i:th
# sockets absolute pixel position with origin in the middle of the skill tree
# I think the SCALE variables are in fact useless and a relics from the
# positions initially being measured at a view which wasn't zoomed out maximally
SOCKETS = {
    1: (-650.565, -376.013),
    2: (648.905, -396.45),
    3: (6.3354, 765.658),
    4: (-1700.9, 2424.17),
    5: (-2800.66, -215.34),
    6: (-1435.02, -2635.39),
    7: (1855.53, -2360.1),
    8: (2835.84, 230.5361),
    9: (1225.37, 2625.76),
    10: (-120.12471, 5195.44),
    11: (-3580.19, 5905.92),
    12: (-5395.86, 2120.42),
    13: (-6030.95, -115.7007),
    14: (-5400.59, -1985.18),
    15: (-3035.14, -5400.87),
    16: (160.10728, -5196.32),
    17: (3382.05, -5195.21),
    18: (5730.2, -1625.75),
    19: (6465.24, 190.3341),
    20: (5542.76, 1690.07),
    21: (3322.76, 6090.5),
}

# The offsets are specified in the same fashion as SOCKETS and are rough
# guesses which allow us to move to the general area and later refine the
# position of the socket through template matching
SOCKET_MOVE_OFFSET = {
    1: (0, 150),
    2: (0, 150),
    3: (0, 200),
    4: (0, 150),
    5: (-300, 200),
    6: (-100, 150),
    7: (-150, 0),
    8: (0, -150),
    9: (-100, -125),
    10: (170, 0),
    11: (-400, -900),
    12: (0, 300),
    13: (400, 200),
    14: (-250, -150),
    15: (-100, -150),
    16: (150, -150),
    17: (150, 500),  #
    18: (-300, 400),
    19: (-1000, -150),
    20: (-500, 500),
    21: (100, -1000),
}

# Scalers for the SOCKETS positions to convert them to 2560x1440 pixel positions
X_SCALE = 0.2
Y_SCALE = 0.2
CIRCLE_EFFECTIVE_RADIUS = 300

IMAGE_FOLDER = "data/images/"

# We're using a lot of template matching and all templates are defined here
# with matching thresholds (scores) and sizes per resolution
TEMPLATES = {
    "AmbidexterityCluster.png": {
        "1440p_size": (34, 34),
        "1440p_threshold": 0.95,
        "1080p_size": (26, 26),
        "1080p_threshold": 0.95,
    },
    "FreeSpace.png": {
        "1440p_size": (41, 41),
        "1440p_threshold": 0.98,
        "1080p_size": (30, 30),
        "1080p_threshold": 0.98,
    },
    "Notable.png": {
        "1440p_size": (30, 30),
        "1440p_threshold": 0.89,
        "1080p_size": (23, 23),
        "1080p_threshold": 0.85,
    },
    "NotableAllocated.png": {
        "1440p_size": (30, 30),
        "1440p_threshold": 0.93,
        "1080p_size": (23, 23),
        "1080p_threshold": 0.90,
    },
    "Jewel.png": {
        "1440p_size": (30, 30),
        "1440p_threshold": 0.92,
        "1080p_size": (23, 23),
        "1080p_threshold": 0.92,
    },
    "JewelSocketed.png": {
        "1440p_size": (30, 30),
        "1440p_threshold": 0.9,
        "1080p_size": (23, 23),
        "1080p_threshold": 0.9,
    },
    "LargeJewel.png": {
        "1440p_size": (39, 39),
        "1440p_threshold": 0.9,
        "1080p_size": (30, 30),
        "1080p_threshold": 0.88,
    },
    "LargeJewelSocketed.png": {
        "1440p_size": (39, 39),
        "1440p_threshold": 0.9,
        "1080p_size": (30, 30),
        "1080p_threshold": 0.88,
    },
    "Skill.png": {
        "1440p_size": (21, 21),
        "1440p_threshold": 0.87,
        "1080p_size": (15, 15),
        "1080p_threshold": 0.91,
    },
    "SkillAllocated.png": {
        "1440p_size": (21, 21),
        "1440p_threshold": 0.93,
        "1080p_size": (15, 15),
        "1080p_threshold": 0.91,
    },
}

# Defines the position of the text box which is cropped out and OCR'd per node
TXT_BOX = {"x": 32, "y": 0, "w": 900, "h": 320}

mod_files = {
    "passives": "data/passives.json",
    "passivesAlt": "data/passivesAlternatives.json",
    "passivesAdd": "data/passivesAdditions.json",
    "passivesVaalAdd": "data/passivesVaalAdditions.json",
}

class TreeNavigator:
    def __init__(self, resolution, halt_value):
        self.resolution = resolution
        self.input_handler = InputHandler(self.resolution)
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s %(message)s",
            datefmt="[%H:%M:%S %d-%m-%Y]",
        )
        self.log = logging.getLogger("tree_nav")
        self.config = get_config("tree_nav")
        self.find_mod_value_re = re.compile("(\(?(?:[0-9]*\.?[0-9]-?)+\)?)")
        self.nonalpha_re = re.compile("[^a-zA-Z]")
        self.origin_pos = (self.resolution[0] / 2, self.resolution[1] / 2)
        self.ingame_pos = [0, 0]
        self.px_multiplier = self.resolution[0] / 2560
        self.resolution_prefix = str(self.resolution[1]) + "p_"
        self.templates_and_masks = self.load_templates()
        self.passive_mods, self.passive_names = self.generate_good_strings(mod_files)
        self.passive_nodes = list(self.passive_mods.keys()) + list(
            self.passive_names.keys()
        )
        self.halt = halt_value
        self.first_run = True

    def _run(self):
        return not bool(self.halt.value)

    def eval_jewel(self, item_location):
        self.ingame_pos = [0, 0]
        item_name, item_desc = self._setup(item_location, copy=True)

        pool = Pool(self.config["ocr_threads"])
        jobs = {}

        if self.first_run:
            # We just initiated the module and not sure where we are
            # Thus, we better rectify our position estimate before starting
            self._refind_position(SOCKETS[1])
            self.first_run = False

        for socket_id in sorted(SOCKETS.keys()):
            if not self._run():
                return None, None, None
            found_socket = self._move_screen_to_socket(socket_id)
            if not found_socket and socket_id == 1:
                self.log.info("We are lost - trying to find known location")
                # We just initiated the search and have no clue where we are
                # Thus, we better rectify our position estimate before starting
                self._refind_position(SOCKETS[1])
            socket_nodes = self._analyze_nodes(socket_id)

            # Convert stats for the socket from image to lines in separate process
            self.log.info("Performing asynchronous OCR")
            jobs[socket_id] = pool.map_async(OCR.node_to_strings, socket_nodes)
            self.log.info("Analyzed socket %s" % socket_id)

        # Return to socket 1 to ease next search
        self._move_to_tree_pos_using_spaces(SOCKETS[1])

        self._setup(item_location)
        self.log.info("Waiting for last OCR to finish")
        item_stats = [
            {
                "socket_id": socket_id,
                "socket_nodes": self._filter_ocr_lines(
                    jobs[socket_id].get(timeout=300)
                ),
            }
            for socket_id in jobs
        ]
        pool.close()
        pool.join()
        return item_name, item_desc, item_stats

    def load_templates(self, threshold=128):
        templates_and_masks = {}
        for template_name in TEMPLATES.keys():
            template_path = os.path.join(IMAGE_FOLDER, template_name)
            img = cv2.imread(template_path, cv2.IMREAD_UNCHANGED)
            size = TEMPLATES[template_name][self.resolution_prefix + "size"]
            channels = cv2.split(img)
            mask = None
            if len(channels) > 3:
                mask = np.array(channels[3])
                mask[mask <= threshold] = 0
                mask[mask > threshold] = 255
                mask = cv2.resize(mask, size)

            img = cv2.imread(template_path, 0)
            img = cv2.resize(img, size)
            templates_and_masks[template_name] = {"image": img, "mask": mask}
        return templates_and_masks

    def _move_screen_to_socket(self, socket_id):
        self.log.debug("Moving close to socket %s" % socket_id)
        move_offset_tx, move_offset_ty = SOCKET_MOVE_OFFSET[socket_id]
        move_offset = self._tree_pos_to_xy(
            [move_offset_tx, move_offset_ty], offset=True
        )

        socket_tx, socket_ty = SOCKETS[socket_id]
        socket_xy = self._tree_pos_to_xy([socket_tx, socket_ty])
        compensation_offset = self._find_socket(socket_xy)
        if compensation_offset is None:
            found_socket = False
            compensation_offset = [0, 0]
        else:
            found_socket = True
            self.log.debug("Compensated navigation with %s" % compensation_offset)

        move_to = [
            socket_xy[0] + compensation_offset[0] + move_offset[0],
            socket_xy[1] + compensation_offset[1] + move_offset[1],
        ]

        x_offset = move_to[0] - self.resolution[0] / 2
        y_offset = move_to[1] - self.resolution[1] / 2

        self.input_handler.click(
            *move_to, *move_to, button=None, raw=True, speed_factor=1
        )
        self.input_handler.drag(self.origin_pos[0], self.origin_pos[1], speed_factor=1)
        self.input_handler.rnd_sleep(min=200, mean=300, sigma=100)
        self.ingame_pos = [socket_tx + move_offset_tx, socket_ty + move_offset_ty]
        return found_socket

    def _refind_position(self, desired_tree_pos):
        # If the current location has been determined to be incorrect
        # we can go to the bottom right corner and find a cluster close
        # to socket 21, namely the Ambidexterity cluster
        # This is a known location, which can then be used to calculate
        # our way to a desired position
        self.log.debug("Centering screen position")

        # Correct our tree position to a known value
        self._locate_screen_using_ambidexterity()

        # Find our way to the desired position
        self._move_to_tree_pos_using_spaces(desired_tree_pos)

    def _move_to_tree_pos_using_spaces(self, desired_tree_pos, max_position_error=5):
        dx = desired_tree_pos[0] - self.ingame_pos[0]
        dy = desired_tree_pos[1] - self.ingame_pos[1]
        self.log.debug("Moving to tree pos using spaces. Deltas: ({}, {})".format(dx, dy))
        while (abs(dx) + abs(dy)) > max_position_error:
            # Choose quadrant to find spaces in based on dx, dy
            right, bottom = dx >= 0, dy >= 0
            if right and not bottom:
                quadrant = 0
            elif not right and not bottom:
                quadrant = 1
            elif not right and bottom:
                quadrant = 2
            elif right and bottom:
                quadrant = 3

            # Find empty spaces that we can drag from
            spaces = self._find_empty_space(quadrant)
            if spaces is None:
                raise ValueError("Could not find an empty space, quitting.")

            # Choose a random empty space for maximum drag
            chosen_space = spaces[np.random.randint(spaces.shape[0])]

            # How far to drag the window to end up in the optimal place
            screen_move_x, screen_move_y = self._tree_pos_to_xy([dx, dy],
                                                                offset=True)

            # Calculate where our drag should end up to perform the move
            drag_x = chosen_space[0] - screen_move_x
            drag_y = chosen_space[1] - screen_move_y

            # We should only drag within the screen's resolution
            # Additionally, we use 100px margin to not trigger tree scroll
            drag_x = np.clip(drag_x, 100, self.resolution[0] - 100)
            drag_y = np.clip(drag_y, 100, self.resolution[1] - 100)

            # Drag
            self.input_handler.click(
                *chosen_space, *chosen_space, button=None, raw=True, speed_factor=1
            )
            self.input_handler.drag(drag_x, drag_y, speed_factor=1)
            self.input_handler.rnd_sleep(min=200, mean=300, sigma=100)

            # Calculate how far we've actually moved
            effective_move_x = chosen_space[0] - drag_x
            effective_move_y = chosen_space[1] - drag_y

            # Update our internal tree position
            self.ingame_pos = self._add_xy_offset_to_tree_pos(
                [effective_move_x, effective_move_y]
            )

            # Figure out how much we have left to move
            dx = desired_tree_pos[0] - self.ingame_pos[0]
            dy = desired_tree_pos[1] - self.ingame_pos[1]

    def _locate_screen_using_ambidexterity(self):
        # Essentially, this is _move_to_tree_pos_using_spaces but
        # only used to find the tree position by navigating to a known point
        self.log.debug("Moving to ambidexterity")
        ambidexterity_position = None
        assumed_ambidexterity_position = (0.25234375, 0.20555556)
        while ambidexterity_position is None:
            # Find empty spaces that we can drag from
            spaces = self._find_empty_space(3)
            if spaces is None:
                raise ValueError("Could not find an empty space, quitting.")

            # Choose the farthest empty space for maximum drag
            chosen_space = spaces[np.argmax(spaces.sum(axis=1))]

            # An arbitrary position in the top left region
            drag_location = (200, 200)

            # Drag
            self.input_handler.click(
                *chosen_space, *chosen_space, button=None, raw=True, speed_factor=1
            )
            self.input_handler.drag(drag_location[0], drag_location[1], speed_factor=1)
            self.input_handler.rnd_sleep(min=200, mean=300, sigma=100)

            # Are we there yet?
            # i.e., have we reached Ambidexterity, which in that case is at
            # roughly (646, 296) in absolute 1440p screen px position
            ambidexterity_position = self._find_icon(
                assumed_ambidexterity_position, "AmbidexterityCluster.png"
            )
        # Ambidexterity is located (-560, 850) from socket 21
        # Thus, this plus any (scaled) offset found by the template matcher is
        # our tree position
        self.ingame_pos = [
            SOCKETS[21][0]
            - 560
            + ambidexterity_position[0] / (X_SCALE * self.px_multiplier),
            SOCKETS[21][1]
            + 850
            + ambidexterity_position[1] / (Y_SCALE * self.px_multiplier),
        ]

    def _find_empty_space(self, quadrant):
        # Finds empty spaces that can be used to drag the screen
        # Used to recenter the screen
        # The quadrant argument is an int in [0, 1, 2, 3], corresponding to
        # [top-right, top-left, bottom-left, bottom-right]
        quadrant_translation = {0: [0.5, 0], 1: [0, 0], 2: [0, 0.5], 3: [0.5, 0.5]}
        fractional_lt = quadrant_translation[quadrant]
        lt = [
            int(fractional_lt[0] * self.resolution[0]),
            int(fractional_lt[1] * self.resolution[1]),
        ]
        rb = [int(lt[0] + self.resolution[0] / 2),
              int(lt[1] + self.resolution[1] / 2)]
        searched_area = grab_screen(tuple(lt + rb))
        searched_area = cv2.cvtColor(searched_area, cv2.COLOR_BGR2GRAY)

        locations = np.zeros_like(searched_area)

        centered_coordinates = self._match_image(searched_area, "FreeSpace.png")
        locations[tuple(centered_coordinates)] = 1

        rel_space_pos_yx = np.argwhere(locations == 1)
        rel_space_pos = rel_space_pos_yx.T[::-1].T
        if len(rel_space_pos) == 0:
            self.log.warning("Could not find any free spaces in tree!")
            return None
        screen_space_pos = rel_space_pos + lt

        # remove positions that are close to edges as these trigger scroll
        screen_space_pos = screen_space_pos[(screen_space_pos[:, 0] > 100) &
                            (screen_space_pos[:, 1] > 100) &
                            (screen_space_pos[:, 0] < self.resolution[0] - 100) &
                            (screen_space_pos[:, 1] < self.resolution[1] - 100)]
        return screen_space_pos

    def _find_icon(self, assumed_position, icon_name):
        # Finds the ambidexerity cluster icon in the region it sits in
        # if we are at the bottom-right corner of the tree
        # The exact location is used to refine our knowledge of our position
        abs_assumed_position = (
            assumed_position[0] * self.resolution[0],
            assumed_position[1] * self.resolution[1],
        )
        margin_side = int(0.05 * self.resolution[0])
        lt = [
            int(abs_assumed_position[0] - margin_side / 2),
            int(abs_assumed_position[1] - margin_side / 2),
        ]
        rb = [
            int(abs_assumed_position[0] + margin_side / 2),
            int(abs_assumed_position[1] + margin_side / 2),
        ]
        searched_area = grab_screen(tuple(lt + rb))
        searched_area = cv2.cvtColor(searched_area, cv2.COLOR_BGR2GRAY)
        locations = np.zeros((margin_side, margin_side))
        centered_coordinates = self._match_image(searched_area, icon_name)
        locations[tuple(centered_coordinates)] = 1
        rel_icon_pos_yx = np.argwhere(locations == 1)
        rel_icon_pos = rel_icon_pos_yx.T[::-1].T
        if len(rel_icon_pos) == 0:
            return None
        icon_offset = [
            int(rel_icon_pos[0][0] - margin_side / 2 + abs_assumed_position[0]),
            int(rel_icon_pos[0][1] - margin_side / 2 + abs_assumed_position[1]),
        ]
        return icon_offset

    def _click_socket(self, socket_pos, insert=True):
        self.log.debug("Clicking socket")
        xy = socket_pos
        lt = [xy[0] - 5 * self.px_multiplier, xy[1] - 5 * self.px_multiplier]
        rb = [xy[0] + 5 * self.px_multiplier, xy[1] + 5 * self.px_multiplier]
        if insert:
            self.input_handler.click(*lt, *rb, button="left", raw=True)
        else:
            self.input_handler.click(*lt, *rb, button="right", raw=True)
        self.input_handler.rnd_sleep(min=200, mean=300)

    def _tree_pos_to_xy(self, pos, offset=False):
        if offset:
            return [
                pos[0] * X_SCALE * self.px_multiplier,
                pos[1] * Y_SCALE * self.px_multiplier,
            ]
        uncentered_xy = [
            (pos[0] - self.ingame_pos[0]) * X_SCALE * self.px_multiplier,
            (pos[1] - self.ingame_pos[1]) * Y_SCALE * self.px_multiplier,
        ]
        xy = [
            int(uncentered_xy[0] + self.origin_pos[0]),
            int(uncentered_xy[1] + self.origin_pos[1]),
        ]
        return xy

    def _add_xy_offset_to_tree_pos(self, offset):
        tree_pos = [
            self.ingame_pos[0] + offset[0] / (X_SCALE * self.px_multiplier),
            self.ingame_pos[1] + offset[1] / (Y_SCALE * self.px_multiplier),
        ]
        return tree_pos

    def _analyze_nodes(self, socket_id):
        self.log.info("Analyzing nodes for socket id %s" % socket_id)
        nodes = []
        node_locations, socket_pos = self._find_nodes(socket_id)
        self.log.debug(
            "Found %s nodes for socket id %s" % (len(node_locations), socket_id)
        )
        self._click_socket(socket_pos)
        for location in node_locations:
            if not self._run():
                return
            node_stats = self._get_node_data(location)
            node = {
                "location": self._socket_offset_pos(socket_pos, location),
                "stats": node_stats,
            }
            nodes.append(node)
        self._click_socket(socket_pos, insert=False)
        return nodes

    def _socket_offset_pos(self, socket_pos, node_location):
        circle_radius = CIRCLE_EFFECTIVE_RADIUS * self.px_multiplier
        return [
            (node_location[0] - socket_pos[0]) / circle_radius,
            (node_location[1] - socket_pos[1]) / circle_radius,
        ]

    def _filter_ocr_lines(self, nodes_lines, max_dist=4):
        filtered_nodes = []
        for node in nodes_lines:
            names = []
            mods = []
            for line in node["stats"]:
                filtered_line = self._filter_nonalpha(line)
                if len(filtered_line) < 4 or filtered_line == "Unallocated":
                    continue
                if filtered_line in self.passive_names:
                    names.append(self.passive_names[filtered_line])
                elif filtered_line in self.passive_mods:
                    filtered_mod, value = filter_mod(line, regex=self.nonalpha_re)
                    new_mod = re.sub(
                        self.find_mod_value_re,
                        str(value),
                        self.passive_mods[filtered_line],
                        count=1,
                    )
                    mods.append(new_mod)
                else:
                    # Sometimes the OCR might return strange results. If so,
                    # as a last resort, check levenshtein distance to closest
                    # node. This shouldn't happen often.
                    best_distance = 99999999999
                    best_match = None
                    for possible_mod in self.passive_nodes:
                        d = distance(filtered_line, possible_mod)
                        if d < best_distance:
                            best_distance = d
                            best_match = possible_mod
                    if best_distance > max_dist:
                        continue
                    if best_match in self.passive_names:
                        names.append(self.passive_names[best_match])
                    elif best_match in self.passive_mods:
                        filtered_mod, value = filter_mod(line, regex=self.nonalpha_re)
                        new_mod = re.sub(
                            self.find_mod_value_re,
                            str(value),
                            self.passive_mods[best_match],
                            count=1,
                        )
                        mods.append(new_mod)

            if mods:
                filtered_nodes.append(
                    {"location": node["location"], "name": names, "mods": mods}
                )

        return filtered_nodes

    def _find_nodes(self, socket_id):
        self.input_handler.click(0.5, 0.07, 0.51, 0.083, button=None)
        socket_pos = self._tree_pos_to_xy(SOCKETS[socket_id])
        socket_offset = self._find_socket(socket_pos)
        if socket_offset is None:
            found_socket = False
            socket_offset = [0, 0]
        else:
            found_socket = True
            self.log.debug("Jewel socket offset correction: %s" % socket_offset)

        socket_pos[0] += socket_offset[0]
        socket_pos[1] += socket_offset[1]

        # Add some margin so that we dont accidentally cut any nodes off
        margin = 20 * self.px_multiplier

        x1 = int(socket_pos[0] - CIRCLE_EFFECTIVE_RADIUS * self.px_multiplier - margin)
        y1 = int(socket_pos[1] - CIRCLE_EFFECTIVE_RADIUS * self.px_multiplier - margin)
        x2 = int(x1 + 2 * CIRCLE_EFFECTIVE_RADIUS * self.px_multiplier + 2 * margin)
        y2 = int(y1 + 2 * CIRCLE_EFFECTIVE_RADIUS * self.px_multiplier + 2 * margin)

        nodes = self._get_node_locations_from_screen((x1, y1, x2, y2))
        nodes = self._filter_nodes(nodes, socket_pos)
        return nodes, socket_pos

    def _find_socket(self, socket_pos, side_len=100):
        lt = [int(socket_pos[0] - side_len / 2), int(socket_pos[1] - side_len / 2)]
        rb = [lt[0] + side_len, lt[1] + side_len]
        socket_area = grab_screen(tuple(lt + rb))
        socket_area = cv2.cvtColor(socket_area, cv2.COLOR_BGR2GRAY)

        locations = np.zeros((side_len, side_len))

        for template_name in [
            "Jewel.png",
            "JewelSocketed.png",
            "LargeJewel.png",
            "LargeJewelSocketed.png",
        ]:
            centered_coordinates = self._match_image(socket_area, template_name)
            locations[tuple(centered_coordinates)] = 1

        rel_node_pos_yx = np.argwhere(locations == 1)
        rel_node_pos = rel_node_pos_yx.T[::-1].T
        if len(rel_node_pos) == 0:
            self.log.warning("Could not find any jewel socket for compensating offset!")
            return None
        socket_offset = [
            int(rel_node_pos[0][0] - side_len / 2),
            int(rel_node_pos[0][1] - side_len / 2),
        ]
        return socket_offset

    def _filter_nodes(self, nodes, socket_pos, duplicate_min_dist=10):
        # filter duplicate nodes
        kept_node_indices = [len(nodes) - 1]
        z = np.array([[complex(c[0], c[1]) for c in nodes]])
        dist_matrix = abs(z.T - z)
        for node_idx in range(len(nodes) - 1):
            if np.min(dist_matrix[node_idx + 1 :, node_idx]) >= duplicate_min_dist:
                kept_node_indices.append(node_idx)
        nodes = np.array(nodes)
        nodes = nodes[kept_node_indices, :]

        # filter nodes outside jewel socket radius
        distances_to_socket = np.sqrt(np.sum((nodes - socket_pos) ** 2, axis=1))
        nodes = nodes[
            distances_to_socket <= CIRCLE_EFFECTIVE_RADIUS * self.px_multiplier
        ]
        return nodes

    def _get_node_locations_from_screen(self, box):
        jewel_area_bgr = grab_screen(box)
        jewel_area_gray = cv2.cvtColor(jewel_area_bgr, cv2.COLOR_BGR2GRAY)

        locations = np.zeros((box[2] - box[0], box[3] - box[1]))

        for template_name in [
            "Notable.png",
            "NotableAllocated.png",
            "Skill.png",
            "SkillAllocated.png",
        ]:
            centered_coordinates = self._match_image(jewel_area_gray, template_name)
            locations[tuple(centered_coordinates)] = 1

        rel_node_pos_yx = np.argwhere(locations == 1)
        rel_node_pos = rel_node_pos_yx.T[::-1].T
        abs_node_pos = rel_node_pos + [box[0], box[1]]
        return abs_node_pos

    def _match_image(self, screen, template_name):
        template = self.templates_and_masks[template_name]["image"]
        mask = self.templates_and_masks[template_name]["mask"]
        res = cv2.matchTemplate(screen, template, cv2.TM_CCORR_NORMED, mask=mask)
        coordinates = np.where(
            res >= TEMPLATES[template_name][self.resolution_prefix + "threshold"]
        )
        icon_size = (
            int(TEMPLATES[template_name][self.resolution_prefix + "size"][0]),
            int(TEMPLATES[template_name][self.resolution_prefix + "size"][1]),
        )
        icon_center_offset = [int(icon_size[0] / 2), int(icon_size[1] / 2)]
        centered_coordinates = [
            coordinates[0] + icon_center_offset[0],
            coordinates[1] + icon_center_offset[1],
        ]

        return centered_coordinates

    def _get_node_data(self, location):
        self.log.debug("Getting node stats at location %s" % location)
        lt = [
            location[0] - 7 * self.px_multiplier,
            location[1] - 7 * self.px_multiplier,
        ]
        rb = [
            location[0] + 7 * self.px_multiplier,
            location[1] + 7 * self.px_multiplier,
        ]
        self.input_handler.click(
            *lt,
            *rb,
            button=None,
            raw=True,
            speed_factor=self.config["node_search_speed_factor"]
        )
        textbox_lt = location + [TXT_BOX["x"], TXT_BOX["y"]]
        textbox_rb = textbox_lt + [
            int(TXT_BOX["w"] * self.px_multiplier),
            int(TXT_BOX["h"] * self.px_multiplier),
        ]

        jewel_area_bgr = grab_screen(tuple(np.concatenate([textbox_lt, textbox_rb])))
        return jewel_area_bgr

    def _setup(self, item_location, copy=False):
        item_desc = None
        item_name = None
        self.input_handler.click_hotkey("p")
        self.input_handler.rnd_sleep(min=150, mean=200, sigma=100)
        self.input_handler.click_hotkey("i")
        if copy:
            self.input_handler.rnd_sleep(min=150, mean=200, sigma=100)
            item = self.input_handler.inventory_copy(
                *item_location, OWN_INVENTORY_ORIGIN, speed_factor=2
            )
            item_desc = item.split("\n")[9].strip()
            item_name = item.split("\n")[1].strip()
        self.input_handler.rnd_sleep(min=150, mean=200, sigma=100)
        self.input_handler.inventory_click(*item_location, OWN_INVENTORY_ORIGIN)
        self.input_handler.rnd_sleep(min=150, mean=200, sigma=100)
        self.input_handler.click_hotkey("i")
        self.input_handler.rnd_sleep(min=150, mean=200, sigma=100)
        return item_name, item_desc

    def generate_good_strings(self, files):
        mods = {}
        names = {}
        for name in files:
            path = files[name]
            with open(path) as f:
                data = json.load(f)
                if isinstance(data, dict):
                    names.update({self._filter_nonalpha(k): k for k in data.keys()})
                    for key in data.keys():
                        if isinstance(data[key]["passives"][0], list):
                            mods.update(
                                {
                                    self._filter_nonalpha(e): e
                                    for e in data[key]["passives"][0]
                                }
                            )
                        else:
                            mods.update(
                                {
                                    self._filter_nonalpha(e): e
                                    for e in data[key]["passives"]
                                }
                            )
                else:
                    mods.update({self._filter_nonalpha(e): e for e in data})
        mods.pop("", None)
        return mods, names

    def _filter_nonalpha(self, value):
        return re.sub(self.nonalpha_re, "", value)

# Adapted from https://github.com/klayveR/python-poe-timeless-jewel
class OCR:
    @staticmethod
    def clahe(img, clip_limit=2.0, grid_size=(8, 8)):
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=grid_size)
        return clahe.apply(img)

    @staticmethod
    def getFilteredImage(src):
        srcH, srcW = src.shape[:2]
        src = cv2.resize(src, (int(srcW * 2), int(srcH * 2)))

        # HSV thresholding to get rid of as much background as possible
        src = cv2.cvtColor(src, cv2.COLOR_BGRA2BGR)
        hsv = cv2.cvtColor(src.copy(), cv2.COLOR_BGR2HSV)
        # Define 2 masks and combine them
        # mask1 for blue affix text
        # mask2 for yellow passive node name
        lower_blue = np.array([80, 10, 40])
        upper_blue = np.array([130, 180, 255])
        lower_yellow = np.array([10, 10, 190])
        upper_yellow = np.array([30, 200, 255])
        mask1 = cv2.inRange(hsv, lower_blue, upper_blue)
        mask2 = cv2.inRange(hsv, lower_yellow, upper_yellow)
        mask = cv2.bitwise_or(mask1, mask2)
        result = cv2.bitwise_and(src, src, mask=mask)
        b, g, r = cv2.split(result)
        b = OCR.clahe(b, 5, (5, 5))
        inverse = cv2.bitwise_not(b)
        return inverse

    @staticmethod
    def imageToStringArray(img):
        t = pytesseract.image_to_string(img, lang="eng", config="--oem 3 --psm 12 poe")
        t = t.replace("\n\n", "\n")
        lines = t.split("\n")
        return lines

    @staticmethod
    def node_to_strings(node):
        img = node["stats"]
        filt_img = OCR.getFilteredImage(img)
        text = OCR.imageToStringArray(filt_img)
        return {"location": node["location"], "stats": text}


import tensorflow as tf
import numpy as np

#from mpl_toolkits.mplot3d import Axes3D
#import matplotlib.pyplot as plt

import random

import signal
import sys
import time
import math
import csv
import os
import argparse
from datetime import datetime
from functools import reduce

print("\nPair Identification")
print("============================\n")

# Step 1: Input parameters
###################################################################################################

# Default parameters
# X, Y, Z bins
XBins = 32
YBins = 32
ZBins = 64

# File names
FileName = "PairIdentification.p1.sim.gz"
GeometryName = "$(MEGALIB)/resource/examples/geomega/GRIPS/GRIPS.geo.setup"

# Depends on GPU memory and layout

#Set later
BatchSize = 128

MaxEvents = 100000

# Determine derived parameters

OutputDataSpaceSize = ZBins

XMin = -43
XMax = 43

# XMin = -5
# XMax = +5

YMin = -43
YMax = 43

# YMin = -5
# YMax = +5

ZMin = 13
ZMax = 45

#Set in stone later
TestingTrainingSplit = 0.8

OutputDirectory = "Results"

parser = argparse.ArgumentParser(description='Perform training and/or testing of the pair identification machine learning tools.')
parser.add_argument('-f', '--filename', default='PairIdentification.p1.sim.gz', help='File name used for training/testing')
parser.add_argument('-m', '--maxevents', default='10000', help='Maximum number of events to use')
parser.add_argument('-s', '--testingtrainigsplit', default='0.1', help='Testing-training split')
parser.add_argument('-b', '--batchsize', default='128', help='Batch size')

args = parser.parse_args()

if args.filename != "":
  FileName = args.filename

if int(args.maxevents) > 1000:
  MaxEvents = int(args.maxevents)

if int(args.batchsize) >= 16:
  BatchSize = int(args.batchsize)

if float(args.testingtrainigsplit) >= 0.05:
  TestingTrainingSplit = float(args.testingtrainigsplit)

if os.path.exists(OutputDirectory):
  Now = datetime.now()
  OutputDirectory += Now.strftime("_%Y%m%d_%H%M%S")

os.makedirs(OutputDirectory)

###################################################################################################
# Step 2: Global functions
###################################################################################################

# Take care of Ctrl-C
Interrupted = False
NInterrupts = 0
def signal_handler(signal, frame):
  global Interrupted
  Interrupted = True
  global NInterrupts
  NInterrupts += 1
  if NInterrupts >= 2:
    print("Aborting!")
    sys.exit(0)
  print("You pressed Ctrl+C - waiting for graceful abort, or press  Ctrl-C again, for quick exit.")
signal.signal(signal.SIGINT, signal_handler)

# Everything ROOT related can only be loaded here otherwise it interferes with the argparse
from EventData import EventData

# Load MEGAlib into ROOT so that it is usable
import ROOT as M
M.gSystem.Load("$(MEGALIB)/lib/libMEGAlib.so")
M.PyConfig.IgnoreCommandLineOptions = True

###################################################################################################
# Step 3: Create some training, test & verification data sets
###################################################################################################

# Read the simulation file data:
DataSets = []

# Load geometry:
Geometry = M.MDGeometryQuest()
if Geometry.ScanSetupFile(M.MString(GeometryName)) == True:
  print("Geometry " + GeometryName + " loaded!")
else:
  print("Unable to load geometry " + GeometryName + " - Aborting!")
  quit()

Reader = M.MFileEventsSim(Geometry)
if Reader.Open(M.MString(FileName)) == False:
  print("Unable to open file " + FileName + ". Aborting!")
  quit()

print("\n\nStarted reading data sets")
NumberOfDataSets = 0
while NumberOfDataSets < MaxEvents:
  Event = Reader.GetNextEvent()
  if not Event:
    break

  if Event.GetNIAs() > 0:
    Data = EventData()
    if Data.parse(Event) == True:
      if Data.hasHitsOutside(XMin, XMax, YMin, YMax, ZMin, ZMax) == False:
        DataSets.append(Data)
        NumberOfDataSets += 1
        if NumberOfDataSets % 500 == 0:
            print("Data sets processed: {}".format(NumberOfDataSets))

print("Info: Parsed {} events".format(NumberOfDataSets))

# Split the data sets in training and testing data sets

TestingTrainingSplit = 0.75

numEvents = len(DataSets)

numTraining = int(numEvents * TestingTrainingSplit)

TrainingDataSets = DataSets[:numTraining]
ValidationDataSets = DataSets[numTraining:]

# For testing/validation split
# ValidationDataSets = TestingDataSets[:int(len(TestingDataSets)/2)]
# TestingDataSets = TestingDataSets[int(len(TestingDataSets)/2):]

print("###### Data Split ########")
print("Training/Testing Split: {}".format(TestingTrainingSplit))
print("Total Data: {}, Training Data: {}, Validation Data: {},Testing Data: {}".format(numEvents, len(TrainingDataSets), len(ValidationDataSets), len(TestingDataSets)))
print("##########################")

###################################################################################################
# Step 4: Setting up the neural network
###################################################################################################

#TODO: Tweak/optimize model
# Is there a better loss function?
#Make more efficient for larger data sets

print("Info: Setting up neural network...")

print("Info: Setting up 3D CNN...")
conv_model = tf.keras.models.Sequential(name='Pair Identification CNN')
conv_model.add(tf.keras.layers.Conv3D(filters=64, kernel_size=3, strides=2, input_shape=(XBins, YBins, ZBins, 1)))
# conv_model.add(tf.keras.layers.MaxPooling3D((2,2,1)))
conv_model.add(tf.keras.layers.LeakyReLU(alpha=0.25))
conv_model.add(tf.keras.layers.BatchNormalization())
conv_model.add(tf.keras.layers.Conv3D(filters=96, kernel_size=3, strides=1, activation='relu'))
conv_model.add(tf.keras.layers.BatchNormalization())
# conv_model.add(tf.keras.layers.MaxPooling3D((2,2,1)))
conv_model.add(tf.keras.layers.Flatten())
conv_model.add(tf.keras.layers.Dense(3*OutputDataSpaceSize, activation='relu'))
conv_model.add(tf.keras.layers.BatchNormalization())
print("Conv Model Summary: ")
print(conv_model.summary())

print("Info: Setting up Gamma Energy Model...")
base_model = tf.keras.models.Sequential(name='Gamma Model')
base_model.add(tf.keras.layers.Dense(int(1.5*OutputDataSpaceSize), activation='relu', input_shape=(1,)))
base_model.add(tf.keras.layers.BatchNormalization())
print("Base Model Summary: ")
print(base_model.summary())

print("Info: Setting up Combined NN...")
combinedInput = tf.keras.layers.concatenate([conv_model.output, base_model.output])
combinedLayer = tf.keras.layers.Dense(OutputDataSpaceSize, activation='softmax')(combinedInput)
combined_model = tf.keras.models.Model([conv_model.input, base_model.input], combinedLayer)
print("Combined Model Summary: ")
print(combined_model.summary())

combined_model.compile(optimizer = 'adam', loss = 'categorical_crossentropy', metrics = ['accuracy'])

# Session configuration
print("      ... configuration ...")
Config = tf.ConfigProto()
Config.gpu_options.allow_growth = True

# Create and initialize the session
print("      ... session ...")
Session = tf.Session(config=Config)
Session.run(tf.global_variables_initializer())

print("      ... listing uninitialized variables if there are any ...")
print(tf.report_uninitialized_variables())

print("      ... writer ...")
writer = tf.summary.FileWriter(OutputDirectory, Session.graph)
writer.close()

# Add ops to save and restore all the variables.
print("      ... saver ...")
Saver = tf.train.Saver()

K = tf.keras.backend
K.set_session(Session)

###################################################################################################
# Step 5: Training and evaluating the network
###################################################################################################

#TODO: attempt multiple models based on energy level
#TODO: Try a dual model setup for high and low energy setups: 10-20 and then 21+
# TODO: Add more robust model performance evaluation
#TODO: Modularize tensor set up and try sequence for training

BatchSize = 128

print("Initializing Tensors...")

class tensor_generator(tf.keras.utils.Sequence):

    def __init__(self, event_data, batch_size):
        self.event_data = event_data
        self.batch_size = batch_size

    def __len__(self):
        return int(len(self.event_data)/self.batch_size)

    def __getitem__(self, idx):
        pos_tensor = make_positional_tensor(self.event_data, idx, self.batch_size)
        gamma_tensor = make_gamma_tensor(self.event_data, idx, self.batch_size)
        label_tensor = make_label_tensor(self.event_data, idx, self.batch_size)

        if len(pos_tensor) != len(gamma_tensor) and len(pos_tensor) != len(label_tensor):
            raise Exception("Bad tensor sizes")

        return ([pos_tensor, gamma_tensor], label_tensor)

def make_positional_tensor(event_data, idx, batch_size):
    tensor = np.zeros(shape=(batch_size, XBins, YBins, ZBins, 1))
    for i in range(batch_size):
        Event = event_data[i + idx*batch_size]
        for j in range(len(Event.X)):
            XBin = int( (Event.X[j] - XMin) / ((XMax - XMin) / XBins) )
            YBin = int( (Event.Y[j] - YMin) / ((YMax - YMin) / YBins) )
            ZBin = int( (Event.Z[j] - ZMin) / ((ZMax - ZMin) / ZBins) )
            if XBin >= 0 and YBin >= 0 and ZBin >= 0 and XBin < XBins and YBin < YBins and ZBin < ZBins:
                tensor[i][XBin][YBin][ZBin][0] = Event.E[j]
    return tensor

def make_gamma_tensor(event_data, idx, batch_size):
    tensor = np.zeros(shape=(batch_size, 1))
    for i in range(batch_size):
        event = event_data[i + idx*batch_size]
        tensor[i][0] = event.GammaEnergy
    return tensor

def make_label_tensor(event_data, idx, batch_size):
    tensor = np.zeros(shape=(batch_size, OutputDataSpaceSize))
    for i in range(batch_size):
        Event = event_data[i + idx*BatchSize]
        # Set the layer in which the event happened
        if Event.OriginPositionZ > ZMin and Event.OriginPositionZ < ZMax:
            LayerBin = int ((Event.OriginPositionZ - ZMin) / ((ZMax- ZMin)/ ZBins) )
            tensor[i][LayerBin] = 1
        else:
            #TODO May need to reevaluate this line
            tensor[i][OutputDataSpaceSize-1] = 1
    return tensor

training_generator = tensor_generator(TrainingDataSets, BatchSize)
validation_generator = tensor_generator(ValidationDataSets, BatchSize)
# testing_generator = tensor_generator(TestingDataSets, BatchSize)

print("Training Model...")

history = combined_model.fit_generator(generator=training_generator, verbose=1, epochs=50, validation_data=validation_generator, shuffle=True)

print("Finished Training\nHistory is: \n")
print(history.history)

print("Evaluating Model...")

loss, acc = combined_model.evaluate_generator(generator=validation_generator, verbose=1)

print("Final Model has accuracy: {} with loss: {}".format(acc, loss))

#Non generated way of setting up data
# numBatches = int(len(TrainingDataSets)/BatchSize)
#
# #Elements are each one batch with [In, Out]--for running in batches
# tensors = []
# energy_tensors = []
#
# for i in range(numBatches):
#     if i % 10 == 0 and i > 0:
#         print("Created {} train tensors".format(i))
#
#     InputTensor = np.zeros(shape=(BatchSize, XBins, YBins, ZBins, 1))
#     OutputTensor = np.zeros(shape=(BatchSize, OutputDataSpaceSize))
#     InputEnergyTensor = np.zeros(shape=(BatchSize, 1))
#
#     for j in range(BatchSize):
#         Event = TrainingDataSets[j + i*BatchSize]
#         # Set the layer in which the event happened
#         if Event.OriginPositionZ > ZMin and Event.OriginPositionZ < ZMax:
#             LayerBin = int ((Event.OriginPositionZ - ZMin) / ((ZMax- ZMin)/ ZBins) )
#             OutputTensor[j][LayerBin] = 1
#         else:
#             #May need to reevaluate this line
#             OutputTensor[j][OutputDataSpaceSize-1] = 1
#
#       # Set all the hit locations and energies
#         for k in range(len(Event.X)):
#             XBin = int( (Event.X[k] - XMin) / ((XMax - XMin) / XBins) )
#             YBin = int( (Event.Y[k] - YMin) / ((YMax - YMin) / YBins) )
#             ZBin = int( (Event.Z[k] - ZMin) / ((ZMax - ZMin) / ZBins) )
#             if XBin >= 0 and YBin >= 0 and ZBin >= 0 and XBin < XBins and YBin < YBins and ZBin < ZBins:
#                 InputTensor[j][XBin][YBin][ZBin][0] = Event.E[k]
#
#         InputEnergyTensor[j][0] = Event.GammaEnergy
#
#     tensors.append([InputTensor, OutputTensor])
#     energy_tensors.append([InputEnergyTensor, OutputTensor])
#
#
# test_tensors = []
# test_energy_tensors = []
#
# for i in range(int(len(TestingDataSets)/BatchSize)):
#     if i % 10 == 0 and i > 0:
#         print("Created {} test tensors".format(i))
#
#     InputTensor = np.zeros(shape=(BatchSize, XBins, YBins, ZBins, 1))
#     OutputTensor = np.zeros(shape=(BatchSize, OutputDataSpaceSize))
#     InputEnergyTensor = np.zeros(shape=(BatchSize, 1))
#
#     for j in range(BatchSize):
#         Event = TestingDataSets[j + i*BatchSize]
#         # Set the layer in which the event happened
#         if Event.OriginPositionZ > ZMin and Event.OriginPositionZ < ZMax:
#             LayerBin = int ((Event.OriginPositionZ - ZMin) / ((ZMax- ZMin)/ ZBins) )
#             OutputTensor[j][LayerBin] = 1
#         else:
#             OutputTensor[j][OutputDataSpaceSize-1] = 1
#
#       # Set all the hit locations and energies
#         for k in range(len(Event.X)):
#             XBin = int( (Event.X[k] - XMin) / ((XMax - XMin) / XBins) )
#             YBin = int( (Event.Y[k] - YMin) / ((YMax - YMin) / YBins) )
#             ZBin = int( (Event.Z[k] - ZMin) / ((ZMax - ZMin) / ZBins) )
#             if XBin >= 0 and YBin >= 0 and ZBin >= 0 and XBin < XBins and YBin < YBins and ZBin < ZBins:
#                 InputTensor[j][XBin][YBin][ZBin][0] = Event.E[k]
#
#         InputEnergyTensor[j][0] = Event.GammaEnergy
#
#     test_tensors.append([InputTensor, OutputTensor])
#     test_energy_tensors.append([InputEnergyTensor, OutputTensor])
#
# if len(tensors) != len(energy_tensors):
#     print("ERROR two training inputs not of same size")
#
# if len(test_tensors) != len(test_energy_tensors):
#     print("ERROR two test inputs not of same size")
#
# print("Training Model...")
# history = []
# for i in range(len(tensors)):
#     history.append(combined_model.fit([tensors[i][0], energy_tensors[i][0]], tensors[i][1], epochs=5))
#
#
#
# print("Checking Performance...")
# acc_list = []
# loss_list = []
#
# accTestlist = []
# lossTestlist = []
#
# for i in range(len(tensors)):
#     loss, acc = combined_model.evaluate([tensors[i][0], energy_tensors[i][0]], tensors[i][1])
#     accTestlist.append(acc)
#     lossTestlist.append(loss)
#
# for i in range(len(test_tensors)):
#     loss, acc = combined_model.evaluate([test_tensors[i][0], test_energy_tensors[i][0]], test_tensors[i][1])
#     acc_list.append(acc)
#     loss_list.append(loss)
#
# for i in range(len(accTestlist)):
#     print('On training round {} the accuracy was {} with loss {}'.format(i, accTestlist[i], lossTestlist[i]))
#
# for i in range(len(acc_list)):
#     print('On test round {} the accuracy was {} with loss {}'.format(i, acc_list[i], loss_list[i]))


import pprint
import os
import re
import shutil
import sys
import tempfile
from io import StringIO

from copy import deepcopy
from unittest import TestCase as OrigTestCase
from vsc.install.headers import nicediff
from vsc.install.methodinspector import MethodInspector
from mock import patch

class TestCase(OrigTestCase):
    """Enhanced test case, provides extra functionality (e.g. an assertErrorRegex method)."""

    longMessage = True # print both standard messgae and custom message

    LOGCACHE = {}

    ASSERT_MAX_DIFF = 100
    DIFF_OFFSET = 5 # lines of text around changes

    # pylint: disable=arguments-differ
    def assertEqual(self, first, second, msg=None):
        """Make assertEqual always print useful messages"""

        try:
            super().assertEqual(first, second)
        except AssertionError as err:
            if msg is None:
                msg = str(err)
            else:
                msg = f"{msg}: {err}"

            if isinstance(first, str):
                txta = first
            else:
                txta = pprint.pformat(first)
            if isinstance(second, str):
                txtb = second
            else:
                txtb = pprint.pformat(second)

            diff = nicediff(txta, txtb, offset=self.DIFF_OFFSET)
            if len(diff) > self.ASSERT_MAX_DIFF:
                limit = f' (first {self.ASSERT_MAX_DIFF} lines)'
            else:
                limit = ''

            raise AssertionError(f"{msg}:\nDIFF{limit}:\n{''.join(diff[:self.ASSERT_MAX_DIFF])}")

    def setUp(self):
        """Prepare test case."""
        super().setUp()

        self.max_diff = None
        self.long_message = True

        self.orig_sys_stdout = sys.stdout
        self.orig_sys_stderr = sys.stderr

        self.orig_sys_argv = sys.argv
        sys.argv = deepcopy(self.orig_sys_argv)

        self.orig_workdir = os.getcwd()
        self.tmpdir = tempfile.mkdtemp()

    def convert_exception_to_str(self, err):
        """Convert an Exception instance to a string."""
        msg = err
        if hasattr(err, 'msg'):
            msg = err.msg
        elif hasattr(err, 'message'):
            msg = err.message
            if not msg:
                # rely on str(msg) in case err.message is empty
                msg = err
        elif hasattr(err, 'args'):  # KeyError in Python 2.4 only provides message via 'args' attribute
            msg = err.args[0]
        else:
            msg = err
        try:
            res = str(msg)
        except UnicodeEncodeError:
            res = msg.encode('utf8', 'replace')

        return res

    def assertErrorRegex(self, error, regex, call, *args, **kwargs):
        """
        Convenience method to match regex with the expected error message.
        Example: self.assertErrorRegex(OSError, "No such file or directory", os.remove, '/no/such/file')
        """
        try:
            call(*args, **kwargs)
            str_kwargs = ['='.join([k, str(v)]) for (k, v) in kwargs.items()]
            str_args = ', '.join(list(map(str, args)) + str_kwargs)
            self.assertTrue(False, f"Expected errors with {call.__name__}({str_args}) call should occur")
        except error as err:
            msg = self.convert_exception_to_str(err)
            if isinstance(regex, str):
                regex = re.compile(regex)
            self.assertTrue(regex.search(msg), f"Pattern '{regex.pattern}' is found in '{msg}'")

    def mock_stdout(self, enable):
        """Enable/disable mocking stdout."""
        sys.stdout.flush()
        if enable:
            sys.stdout = StringIO()
        else:
            sys.stdout = self.orig_sys_stdout

    def mock_stderr(self, enable):
        """Enable/disable mocking stdout."""
        sys.stderr.flush()
        if enable:
            sys.stderr = StringIO()
        else:
            sys.stderr = self.orig_sys_stderr

    def get_stdout(self):
        """Return output captured from stdout until now."""
        return sys.stdout.getvalue()

    def get_stderr(self):
        """Return output captured from stderr until now."""
        return sys.stderr.getvalue()

    def mock_logmethod(self, logmethod_func):
        """
        Intercept the logger logmethod. Use as
            mylogger = logging.getLogger
            mylogger.error = self.mock_logmethod(mylogger.error)
        """
        def logmethod(*args, **kwargs):
            if hasattr(logmethod_func, 'func_name'):
                funcname = logmethod_func.func_name
            elif hasattr(logmethod_func, 'im_func'):
                funcname = logmethod_func.im_func.__name__
            elif hasattr(logmethod_func, '__name__'):
                funcname = logmethod_func.__name__
            else:
                raise ValueError(f"Unknown logmethod {dir(logmethod_func)}")
            logcache = self.LOGCACHE.setdefault(funcname, [])
            logcache.append({'args': args, 'kwargs': kwargs})
            logmethod_func(*args, **kwargs)

        return logmethod

    def reset_logcache(self, funcname=None):
        """
        Reset the LOGCACHE
        @param: funcname: if set, only reset the cache for this log function
                (default is to reset the whole chache)
        """
        if funcname:
            self.LOGCACHE[funcname] = []
        else:
            self.LOGCACHE = {}

    def count_logcache(self, funcname):
        """
        Return the number of log messages for funcname in the logcache
        """
        return len(self.LOGCACHE.get(funcname, []))

    def gen_inspector(self, *args, **kwargs):
        """
        Convenience method to generate MethodInspector instance.
        All args are passed to the MethodInspector
        """
        return MethodInspector(*args, **kwargs)

    def create_patch(self, *args, **kwargs):
        """
        Create patch and return mocked whatever. Patch is added to tearDown
        All args are passed to patch
        """
        patcher = patch(*args, **kwargs)
        thing = patcher.start()
        self.addCleanup(patcher.stop)
        return thing

    def tearDown(self):
        """Cleanup after running a test."""
        self.mock_stdout(False)
        self.mock_stderr(False)
        self.reset_logcache()
        sys.argv = self.orig_sys_argv
        os.chdir(self.orig_workdir)
        shutil.rmtree(self.tmpdir)

        super().tearDown()

# backwards incompatible change
class VSCImportTest(TestCase):
    """
    catchall for old tests.
    """
    def test_deprecated_fail(self):
        """
        VSCImportTest is now deprecated and will always fail, use
            from vsc.install.commontest import CommonTest
        """
        self.assertTrue(False, msg='Use "from vsc.install.commontest import CommonTest" instead of VSCImportTest.')


from functools import cache
from typing import List

class Solution:
    """
    1. p选正号的数的和
    2. s所有数组的总和
    3. s-p选负号的数的和
    => p-(s-p)==t
    => 2p==s+t
    => p==(s+t)/2
    """

    def findTargetSumWays2(self, nums: List[int], target: int) -> int:
        n, s = len(nums), sum(nums)
        target += s
        if target < 0 or target % 2 != 0:
            return 0
        target //= 2

        @cache
        def dfs(i, c):
            if i < 0:
                return 1 if c == 0 else 0
            if c < nums[i]:
                return dfs(i - 1, c)
            return dfs(i - 1, c) + dfs(i - 1, c - nums[i])

        return dfs(n - 1, target)

    def findTargetSumWays3(self, nums: List[int], target: int) -> int:
        n, s = len(nums), sum(nums)
        target += s
        if target < 0 or target % 2 != 0:
            return 0
        target //= 2
        dp = [[0] * (target + 1) for _ in range(n + 1)]
        dp[0][0] = 1
        for i, x in enumerate(nums):
            for c in range(target + 1):
                if c < x:
                    dp[i + 1][c] = dp[i][c]
                else:
                    dp[i + 1][c] = dp[i][c] + dp[i][c - x]
        return dp[n][target]

    def findTargetSumWays4(self, nums: List[int], target: int) -> int:
        n, s = len(nums), sum(nums)
        target += s
        if target < 0 or target % 2 != 0:
            return 0
        target //= 2
        dp = [[0] * (target + 1) for _ in range(2)]
        dp[0][0] = 1
        for i, x in enumerate(nums):
            for c in range(target + 1):
                if c < x:
                    dp[(i + 1) % 2][c] = dp[i % 2][c]
                else:
                    dp[(i + 1) % 2][c] = dp[i % 2][c] + dp[i % 2][c - x]
        return dp[n % 2][target]

    def findTargetSumWays(self, nums: List[int], target: int) -> int:
        n, s = len(nums), sum(nums)
        target += s
        if target < 0 or target % 2 != 0:
            return 0
        target //= 2
        dp = [[0] * (target + 1)]
        dp[0] = 1
        for x in nums:
            for c in range(target, x - 1, -1):
                dp[c] = dp[c] + dp[c - x]
        return dp[target]