from tqdm import tqdm
from common import Sentence, Instance
from typing import List
import re
from transformers import BertTokenizer
import json
import copy
import random
type_map = {
    '质押':'zy',
    '股份股权转让': 'gfgqzr',
    '起诉': 'qs',
    '投资': 'tz',
    '减持': 'jc',
    '收购': 'sg',
    '判决': 'pj',
    '担保': 'db',
    '中标': 'zb',
    '签署合同': 'qsht'
}

class Reader:
    def __init__(self, digit2zero: bool=False):
        """
        Read the dataset into Instance
        :param digit2zero: convert the digits into 0, which is a common practice for LSTM-CRF.
        """
        self.digit2zero = digit2zero
        self.vocab = set()

    def read_txt(self, file_dir: str, number: int = -1, type: str = "", aug: bool = False) -> List[Instance]:
        count_0 = 0
        insts = []
        import os
        # print("Reading file: "+file_dir)

        if aug:
            mention_dict = {}
            aug_insts = []
            aug_examples = []
            Examples = []
            for f_temp in os.listdir(file_dir):
                file = os.path.join(file_dir, f_temp)
                in_file = open(file, 'r', encoding = 'utf-8')
                for line in in_file:
                    line = line.strip()
                    Examples.append(json.loads(line))
            for example in Examples:
                content = example['content']
                events = example['events']
                triggers = {}
                reverse_example = copy.deepcopy(example)
                aug_flag = False
                for event in events:
                    sub_flag = False
                    obj_flag = False
                    for mention in event['mentions']:
                        if not sub_flag and 'sub' in mention['role']:
                            sub_flag = True
                            sub_role = mention['role']
                            sub_word = mention['word']
                        if not obj_flag and 'obj' in mention['role']:
                            obj_flag = True
                            obj_role = mention['role']
                            obj_word = mention['word']
                        if mention['role'] not in mention_dict:
                            mention_dict[mention['role']] = [mention['word']]
                        else:
                            mention_dict[mention['role']].append(mention['word'])
                    if sub_flag and obj_flag and not aug_flag: # swap and replace
                        aug_flag = sub_role.replace('sub', 'aaa') == obj_role.replace('obj', 'aaa')
                if aug_flag:
                    reverse_example['content'] = swap(reverse_example['content'], sub_word, obj_word)
                    for event_id, event in enumerate(reverse_example['events']):
                        for mention_id, mention in enumerate(event['mentions']):
                            if mention['word'] == sub_word:
                                reverse_example['events'][event_id]['mentions'][mention_id]['word'] = obj_word
                            elif mention['word'] == obj_word:
                                reverse_example['events'][event_id]['mentions'][mention_id]['word'] = sub_word
                            temp_word = reverse_example['events'][event_id]['mentions'][mention_id]['word']
                            reverse_example['events'][event_id]['mentions'][mention_id]['span'] = get_new_span(content, [sub_word], [obj_word], mention['span'], temp_word)
                    aug_examples.append(reverse_example)
            aug_examples = aug_examples + copy.deepcopy(Examples)
            for example_id, example in enumerate(aug_examples):
                date_flag = True
                number_flag = True
                money_flag = True
                target_flag = True
                proportion_flag = True
                old_words = []
                new_words = []
                temp_content = example['content']
                for event_id, event in enumerate(example['events']):
                    for mention_id, mention in enumerate(event['mentions']):
                        if not date_flag and mention['role'] == 'date':
                            new_date = random.choice(mention_dict['date'])
                            new_words.append(new_date)
                            old_words.append(mention['word'])
                            date_flag = True
                            aug_examples[example_id]['content'] = swap(aug_examples[example_id]['content'], mention['word'], new_date)
                        if not number_flag and mention['role'] == 'number':
                            new_number = random.choice(mention_dict['number'])
                            new_words.append(new_number)
                            old_words.append(mention['word'])
                            number_flag = True
                            aug_examples[example_id]['content'] = swap(aug_examples[example_id]['content'], mention['word'], new_number)
                        if not money_flag and mention['role'] == 'money':
                            new_money = random.choice(mention_dict['money'])
                            new_words.append(new_money)
                            old_words.append(mention['word'])
                            money_flag = True
                            aug_examples[example_id]['content'] = swap(aug_examples[example_id]['content'], mention['word'], new_money)
                        if not target_flag and mention['role'] == 'target-company':
                            new_target = random.choice(mention_dict['sub-org'] + mention_dict['obj-org'])
                            new_words.append(new_target)
                            old_words.append(mention['word'])
                            target_flag = True
                            aug_examples[example_id]['content'] = swap(aug_examples[example_id]['content'], mention['word'], new_target)
                        if not proportion_flag and mention['role'] == 'proportion':
                            new_proportion = random.choice(mention_dict['proportion'] + mention_dict['share-per'])
                            new_words.append(new_proportion)
                            old_words.append(mention['word'])
                            proportion_flag = True
                            aug_examples[example_id]['content'] = swap(aug_examples[example_id]['content'], mention['word'], new_proportion)

                for event_id, event in enumerate(example['events']):
                    for mention_id, mention in enumerate(event['mentions']):
                        if mention['word'] in old_words:
                            idx = old_words.index(mention['word'])
                            replace_word = new_words[idx]
                            aug_examples[example_id]['events'][event_id]['mentions'][mention_id]['word'] = replace_word
                        replace_word = aug_examples[example_id]['events'][event_id]['mentions'][mention_id]['word']
                        aug_examples[example_id]['events'][event_id]['mentions'][mention_id]['span'] = get_new_span(temp_content, old_words, new_words, mention['span'], replace_word)
            aug_examples = check(aug_examples)
            for line in aug_examples:
                words = line['content']
                idx = line['id']
                if self.digit2zero:
                    words = re.sub('\d', '0', words)
                    count_0 += len(re.findall('0', words))
                words = list(words)
                events_dict = dict()
                for k in line['events']:
                    if not k['type'] in events_dict:
                        events_dict[k['type']] = [k]
                    else:
                        events_dict[k['type']].append(k)
                for t, k_list in events_dict.items():
                # for k in line['events']:
                    trigger = []
                    mentions = {}
                    # 不同事件需要单独标注，因为不同事件的实体会重合
                    labels = ['O'] * len(words)
                    evn_type = t
                    for k in k_list:
                        for i in k['mentions']:
                            start_span = i['span'][0]
                            end_span = i['span'][1]
                            role = i['role']
                            mentions[role] = line['content'][start_span:end_span]
                            if role == "trigger":
                                trigger.append(line['content'][start_span:end_span])
                            if end_span - start_span == 1:
                                labels[start_span] = "B-" + type_map[evn_type] + role
                                # labels[start_span] = "B-" + role
                            elif end_span - start_span == 2:
                                labels[start_span] = "B-" + type_map[evn_type] + role
                                labels[start_span+1] = "E-" + type_map[evn_type] + role
                                # labels[start_span] = "B-" + role
                                # labels[start_span+1] = "E-" + role
                            elif end_span - start_span > 2:
                                for i in range(start_span, end_span):
                                    if i == start_span:
                                        labels[i] = "B-" + type_map[evn_type] + role
                                        # labels[i] = "B-" + role
                                    elif i == end_span-1:
                                        labels[i] = "E-" + type_map[evn_type] + role
                                        # labels[i] = "E-" + role
                                    else:
                                        labels[i] = "I-" + type_map[evn_type] + role
                                        # labels[i] = "I-" + role
                            else:
                                print("Wrong span!")
                    aug_inst = Instance(Sentence(words), labels)
                    aug_inst.set_id(idx)
                    aug_inst.type = evn_type
                    aug_inst.trigger = trigger
                    aug_inst.content = line['content']
                    aug_inst.mentions = mentions
                    aug_insts.append(aug_inst)

        for f_temp in os.listdir(file_dir):
            file = os.path.join(file_dir, f_temp)
            print("Reading file: " + file)
            in_file = open(file, 'r', encoding = 'utf-8')
            for line in in_file:
                line = line.strip()
                line = json.loads(line)
                words = line['content']
                idx = line['id']
                if self.digit2zero:
                    words = re.sub('\d', '0', words)
                    count_0 += len(re.findall('0', words))
                words = list(words)
                events_dict = dict()
                for k in line['events']:
                    if not k['type'] in events_dict:
                        events_dict[k['type']] = [k]
                    else:
                        events_dict[k['type']].append(k)
                for t, k_list in events_dict.items():
                # for k in line['events']:
                    trigger = []
                    mentions = {}
                    # 不同事件需要单独标注，因为不同事件的实体会重合
                    labels = ['O'] * len(words)
                    evn_type = t
                    for k in k_list:
                        for i in k['mentions']:
                            start_span = i['span'][0]
                            end_span = i['span'][1]
                            role = i['role']
                            mentions[role] = line['content'][start_span:end_span]
                            if role == "trigger":
                                trigger.append(line['content'][start_span:end_span])
                            if end_span - start_span == 1:
                                labels[start_span] = "B-" + type_map[evn_type] + role
                                # labels[start_span] = "B-" + role
                            elif end_span - start_span == 2:
                                labels[start_span] = "B-" + type_map[evn_type] + role
                                labels[start_span+1] = "E-" + type_map[evn_type] + role
                                # labels[start_span] = "B-" + role
                                # labels[start_span+1] = "E-" + role
                            elif end_span - start_span > 2:
                                for i in range(start_span, end_span):
                                    if i == start_span:
                                        labels[i] = "B-" + type_map[evn_type] + role
                                        # labels[i] = "B-" + role
                                    elif i == end_span-1:
                                        labels[i] = "E-" + type_map[evn_type] + role
                                        # labels[i] = "E-" + role
                                    else:
                                        labels[i] = "I-" + type_map[evn_type] + role
                                        # labels[i] = "I-" + role
                            else:
                                print("Wrong span!")
                    inst = Instance(Sentence(words), labels)
                    inst.set_id(idx)
                    inst.type = evn_type
                    inst.trigger = trigger
                    inst.content = line['content']
                    inst.mentions = mentions
                    insts.append(inst)
                    if len(insts) == number:
                        break
        print("numbers being replaced by zero:", count_0)
        print("number of sentences: {}".format(len(insts)))
        if type == "":
            return insts
        else:
            if aug:
                return [inst for inst in insts if inst.type == type], [inst for inst in aug_insts if inst.type == type]
            else:
                return [inst for inst in insts if inst.type == type]

    def read_test_txt(self, file_dir: str, number: int = -1) -> List[Instance]:
        count_0 = 0
        insts = []
        import os
        # print("Reading file: "+file_dir)
        for f_temp in os.listdir(file_dir):
            file = os.path.join(file_dir, f_temp)
            print("Reading file: " + file)
            in_file = open(file, 'r', encoding = 'utf-8')
            for line in in_file:
                line = line.strip()
                line = json.loads(line)
                words = line['content']
                idx = line['id']
                if self.digit2zero:
                    words = re.sub('\d', '0', words)
                    count_0 += len(re.findall('0', words))
                words = list(words)
                labels = ['O'] * len(words)
                inst = Instance(Sentence(words), labels)
                inst.content = line['content']
                inst.set_test_id(idx)
                insts.append(inst)
                if len(insts) == number:
                    break
        print("numbers being replaced by zero:", count_0)
        print("number of sentences: {}".format(len(insts)))
        return insts

    def get_origin_query(self, file: str, number: int = -1):
        print("Reading file: " + file)
        with open(file, 'r') as f:
            f_dict = json.loads(f.read())
        query_list = f_dict.values()
        return query_list

def fix_trigger():
    # 修改valid_result
    id2content = dict()
    for dev in tqdm(devs):
        id2content[dev.id] = dev.content
    import json
    sub_data = open('./result_tmp.json', 'r', encoding='utf-8')
    fix_data = open('./result.json', 'w+', encoding='utf-8')
    cnt_zb = 0
    cnt_sg = 0
    cnt_qsht = 0
    cnt_db = 0
    cnt_pj = 0
    for i, line in tqdm(enumerate(sub_data.readlines())):
        line = line.strip()
        eval_dict = json.loads(line)
        # if i < 8132:
        # if i < 8267:
        for evn in eval_dict['events']:
            if evn['type'] == "中标":
                trigger_flag = False
                for mention in evn['mentions']:
                    if "trigger" in mention.values():
                        trigger_flag = True
                if not trigger_flag:
                    context = id2content[eval_dict['id']]
                    for candidate in ['中标']:
                        span_start = context.find(candidate)
                        if span_start != -1:
                            span_end = span_start + len(candidate)
                            evn['mentions'].append(
                                {"word": candidate, "span": [span_start, span_end], "role": "trigger"})
                            cnt_zb += 1
                            break
            elif evn['type'] == "收购":
                trigger_flag = False
                for mention in evn['mentions']:
                    if "trigger" in mention.values():
                        trigger_flag = True
                if not trigger_flag:
                    context = id2content[eval_dict['id']]
                    for candidate in ['收购', '间接控制', '控制权', '合并']:
                        span_start = context.find(candidate)
                        if span_start != -1:
                            span_end = span_start + len(candidate)
                            evn['mentions'].append(
                                {"word": candidate, "span": [span_start, span_end], "role": "trigger"})
                            cnt_sg += 1
                            break
            elif evn['type'] == "担保":
                trigger_flag = False
                for mention in evn['mentions']:
                    if "trigger" in mention.values():
                        trigger_flag = True
                if not trigger_flag:
                    context = id2content[eval_dict['id']]
                    for candidate in ['担保']:
                        span_start = context.find(candidate)
                        if span_start != -1:
                            span_end = span_start + len(candidate)
                            evn['mentions'].append(
                                {"word": candidate, "span": [span_start, span_end], "role": "trigger"})
                            cnt_db += 1
                            break
            elif evn['type'] == "签署合同":
                trigger_flag = False
                for mention in evn['mentions']:
                    if "trigger" in mention.values():
                        trigger_flag = True
                if not trigger_flag:
                    context = id2content[eval_dict['id']]
                    for candidate in ['签约', '签署合同', '签订', '订立', '合作']:
                        span_start = context.find(candidate)
                        if span_start != -1:
                            span_end = span_start + len(candidate)
                            evn['mentions'].append(
                                {"word": candidate, "span": [span_start, span_end], "role": "trigger"})
                            cnt_qsht += 1
                            break
            elif evn['type'] == "判决":
                trigger_flag = False
                for mention in evn['mentions']:
                    if "trigger" in mention.values():
                        trigger_flag = True
                if not trigger_flag:
                    context = id2content[eval_dict['id']]
                    for candidate in ['判决', '裁决', '审理', '宣判', '判令', '裁定', '判处', '仲裁']:
                        span_start = context.find(candidate)
                        if span_start != -1:
                            span_end = span_start + len(candidate)
                            evn['mentions'].append(
                                {"word": candidate, "span": [span_start, span_end], "role": "trigger"})
                            cnt_qsht += 1
                            break
            else:
                continue
        json.dump(eval_dict, fix_data, ensure_ascii=False)
        fix_data.write("\n")
    print("db: ", cnt_db)
    print("zb: ", cnt_zb)
    print("qsht: ", cnt_qsht)
    print("sg: ", cnt_sg)
    print("pj: ", cnt_pj)

def extract_by_reg():
    pattern_number = r'质押\D*(\d+\.?\d+[万亿])余?股'
    # pattern_number_2 = r'质押\D*(\d+,?\d+,?\d+\.?\d*)余?股'
    pattern_number_2 = r'质押\D*((?:\d+,){0,3}\d+\.?\d*)余?股'
    pattern_number_3 = r'(\d+\.?\d+[万亿])余?股\D*进行.*质押'
    pattern_number_4 = r'(\d+\.?\d+[万亿])余?股\D*质押(?:于|给)'
    pattern_number_5 = r'((?:\d+,){0,3}\d+\.?\d*)余?股\D*进行.*质押'
    pattern_number_6 = r'((?:\d+,){0,3}\d+\.?\d*)余?股\D*质押(于|给)'
    pattern_number_7 = r'(\d+\.?\d+[万亿])余?股\D*(?:(?:被)|(?:处于))质押'
    pattern_number_8 = r'((?:\d+,){0,3}\d+\.?\d*)余?股\D*(?:(?:被)|(?:处于))质押'
    pattern_list = [pattern_number, pattern_number_2, pattern_number_3,
                    pattern_number_4, pattern_number_5, pattern_number_6,
                    pattern_number_7, pattern_number_8]
    p = 0
    t = 0
    err_list = []
    for content, number in [(t.content, t.mentions['number']) for t in trains if
                            t.type == '质押' and 'number' in t.mentions]:
        t += 1
        for pattern in pattern_list:
            if re.search(pattern, content) and re.search(pattern, content).group(1) == number:
                p += 1
                break
        else:
            err_list.append((content, number))
    print(p / t)

def fix_NUM():
    pattern_number = r'质押\D*(\d+\.?\d+[万亿])余?股'
    # pattern_number_2 = r'质押\D*(\d+,?\d+,?\d+\.?\d*)余?股'
    pattern_number_2 = r'质押\D*((?:\d+,){0,3}\d+\.?\d*)余?股'
    pattern_number_3 = r'(\d+\.?\d+[万亿])余?股\D*进行.*质押'
    pattern_number_4 = r'(\d+\.?\d+[万亿])余?股\D*质押(?:于|给)'
    pattern_number_5 = r'((?:\d+,){0,3}\d+\.?\d*)余?股\D*进行.*质押'
    pattern_number_6 = r'((?:\d+,){0,3}\d+\.?\d*)余?股\D*质押(于|给)'
    pattern_number_7 = r'(\d+\.?\d+[万亿])余?股\D*(?:(?:被)|(?:处于))质押'
    pattern_number_8 = r'((?:\d+,){0,3}\d+\.?\d*)余?股\D*(?:(?:被)|(?:处于))质押'
    pattern_list = [pattern_number, pattern_number_2, pattern_number_3,
                    pattern_number_4, pattern_number_5, pattern_number_6,
                    pattern_number_7, pattern_number_8]
    # 修改valid_result
    id2content = dict()
    for dev in tqdm(devs):
        id2content[dev.id] = dev.content
    import json
    sub_data = open('../valid_result_15_fix.json', 'r', encoding='utf-8')
    fix_data = open('../valid_result_15_fix_num.json', 'w+', encoding='utf-8')
    cnt_add = 0
    cnt_fix = 0
    for i, line in tqdm(enumerate(sub_data.readlines())):
        line = line.strip()
        eval_dict = json.loads(line)
        # if i < 8132:
        # if i < 8267:
        for evn in eval_dict['events']:
            if evn['type'] == "质押":
                number_flag = False
                for mention in evn['mentions']:
                    if "number" in mention.values():
                        number_flag = True
                        mention_num = mention
                context = id2content[eval_dict['id']]
                for pattern in pattern_list:
                    if re.search(pattern, context):
                        number = re.search(pattern, context).group(1)
                        span_start = context.find(number)
                        span_end = span_start + len(number)
                        if number_flag:
                            if mention_num != {"word": number, "span": [span_start, span_end], "role": "number"}:
                                evn['mentions'].remove(mention_num)
                                evn['mentions'].append({"word": number, "span": [span_start, span_end], "role": "number"})
                                cnt_fix+=1
                                print(eval_dict['id'])
                        else:
                            evn['mentions'].append(
                                {"word": number, "span": [span_start, span_end], "role": "number"})
                            cnt_add += 1
                        break
            else:
                continue
        json.dump(eval_dict, fix_data, ensure_ascii=False)
        fix_data.write("\n")
    print("add: ", cnt_add)
    print("fix: ", cnt_fix)

def error_output():
    error_data = open('../data/error_info.json', 'w+', encoding='utf8')
    type_map = {
        'zy': '质押',
        'gfgqzr': '股份股权转让',
        'qs': '起诉',
        'tz': '投资',
        'jc': '减持',
        'sg': '收购',
        'pj': '判决'
    }
    for suffix in ['zy', 'gfgqzr', 'qs', 'tz', 'jc', 'sg', 'pj']:
        type = type_map[suffix]
        pres = reader.read_txt(file_dir_err, -1, type)
        for pre in pres:
            err_dict = dict()
            err_dict['type'] = type
            ids = pre.id
            err_dict['id'] = ids
            mentions_pre = pre.mentions
            mentions_true = [t.mentions for t in trains if t.id==ids and t.type==type][0]
            for k, v in mentions_true.items():
                if k in mentions_pre.keys() and v != mentions_pre[k]:
                    err_dict[k] = (v, mentions_pre[k])
                elif k not in mentions_pre.keys():
                    err_dict[k] = (v, '')
            err_dict['content'] = pre.content
            import json
            if len(err_dict.keys()) == 3:
                continue
            json.dump(err_dict, error_data, ensure_ascii=False)
            error_data.write('\n')

def error_output_csv():
    import pandas as pd
    import numpy as np
    error_list = []
    type_map = {
        'zy': '质押',
        'gfgqzr': '股份股权转让',
        'qs': '起诉',
        'tz': '投资',
        'jc': '减持',
        'sg': '收购',
        'pj': '判决'
    }
    for suffix in ['zy', 'gfgqzr', 'qs', 'tz', 'jc', 'sg', 'pj']:
        type = type_map[suffix]
        pres = reader.read_txt(file_dir_err, -1, type)
        for pre in pres:
            ids = pre.id
            mentions_pre = pre.mentions
            mentions_true = [t.mentions for t in trains if t.id == ids and t.type == type][0]
            for k, v in mentions_true.items():
                if k in mentions_pre.keys() and v != mentions_pre[k]:
                    error_list.append([ids, type, k, v, mentions_pre[k], pre.content])
                elif k not in mentions_pre.keys():
                    error_list.append([ids, type, k, v, '', pre.content])
    error_array = np.asarray(error_list)
    df = pd.DataFrame(error_array, index=None)
    df.columns = ['id', 'type', 'role', 'answer', 'predict', 'content']
    df.to_csv('../data/error_info.csv', index=False)

def error_output_csv_2(file_pre, file_true):
    import pandas as pd
    import numpy as np
    error_list = []
    content_dict = dict()
    pre_dict = dict()
    true_dict = dict()
    for line in open(file_pre,'r'):
        line = line.strip()
        line = json.loads(line)
        ids = line['id']
        pre_dict[ids] = line['events']
        content_dict[ids] = line['contents']
    for line in open(file_true,'r'):
        line = line.strip()
        line = json.loads(line)
        ids = line['id']
        true_dict[ids] = line['events']

    for pre_id in pre_dict.keys():
        ids = pre_id
        content = content_dict[ids]
        pre_events =  pre_dict[pre_id]
        for evn_pre in pre_events:
            type = evn_pre['type']
            true_events = true_dict[pre_id]
            mentions_pre = dict()
            for m in evn_pre['mentions']:
                mentions_pre[m['role']] = m['word']
            for evn_true in true_events:
                if evn_true['type'] == type:
                    mentions_true = dict()
                    for m in evn_true['mentions']:
                        mentions_true[m['role']] = m['word']
                    for k, v in mentions_true.items():
                        if k in mentions_pre.keys() and v != mentions_pre[k]:
                            error_list.append([ids, type, k, v, mentions_pre[k], content])
                        elif k not in mentions_pre.keys():
                            error_list.append([ids, type, k, v, '', content])
    error_array = np.asarray(error_list)
    df = pd.DataFrame(error_array, index=None)
    df.columns = ['id', 'type', 'role', 'answer', 'predict', 'content']
    df.to_csv('../data/error_info.csv', index=False)

def fix_classify():
    print('fixing classify...')
    classify_dict = {
        'base': ['质押', '股份股权转让', '起诉', '投资', '减持'],
        'trans': ['收购', '判决']
    }
    import json
    dev_base = open('../data/dev/dev_base.json', 'r', encoding='utf8')
    dev_trans = open('../data/dev/trans_dev.json', 'r', encoding='utf8')
    base_ids = []
    trans_ids = []
    for line in tqdm(dev_base.readlines()):
        line = line.strip()
        base_dict = json.loads(line)
        base_ids.append(base_dict['id'])
    for line in tqdm(dev_trans.readlines()):
        line = line.strip()
        trans_dict = json.loads(line)
        trans_ids.append(trans_dict['id'])
    sub_data = open('../valid_result_16_fix.json', 'r', encoding='utf-8')
    fix_data = open('../valid_result_16_fix_classify.json', 'w+', encoding='utf-8')
    fix_num = 0
    fix_num_base = 0
    fix_num_trans = 0
    for i, line in tqdm(enumerate(sub_data.readlines())):
        line = line.strip()
        eval_dict = json.loads(line)
        ids = eval_dict['id']
        remove_evns = []
        if ids in base_ids:
            type_evn = 'base'
        else:
            type_evn = 'trans'
        for evn in eval_dict['events']:
            if evn['type'] not in classify_dict[type_evn]:
                remove_evns.append(evn)
                if type_evn == 'base':
                    fix_num_base += 1
                else:
                    fix_num_trans += 1
        for e in remove_evns:
            fix_num += 1
            eval_dict['events'].remove(e)
        json.dump(eval_dict, fix_data, ensure_ascii=False)
        fix_data.write("\n")
    print('total_fix: '+str(fix_num))
    print('base_fix: '+str(fix_num_base))
    print('trans_fix: '+str(fix_num_trans))

def add_trigger():
    sub_data = open('../valid_result_16_fix.json', 'r', encoding='utf-8')
    fix_data = open('../valid_result_16_fix_add_trigger.json', 'w+', encoding='utf-8')

    id_list = []
    add_cnt = 0
    for i, line in tqdm(enumerate(sub_data.readlines())):
        line = line.strip()
        eval_dict = json.loads(line)
        ids = eval_dict['id']
        sub_org_cnt = 0
        trigger_cnt = 0
        for evn in eval_dict['events']:
            if evn['type'] == '股份股权转让':
                for mention in evn['mentions']:
                    if mention['role'] == 'sub-org':
                        sub_org_cnt += 1
                    if mention['role'] == 'trigger':
                        mention_trigger = mention
                        trigger_cnt += 1
                if trigger_cnt != 0:
                    for i in range(sub_org_cnt - trigger_cnt):
                        add_cnt += 1
                        evn['mentions'].append(mention_trigger)
        json.dump(eval_dict, fix_data, ensure_ascii=False)
        fix_data.write("\n")
    print(add_cnt)

# def fix_gfgqzr_tz_pairs():
#     print('fixing pairs...')
#     fix_map = {
#         'target-company': 'obj',
#         'trigger': 'trigger',
#         'money': 'money',
#         'date': 'date',
#         'obj-org': 'sub'
#     }
#     import json
#     sub_data = open('../valid_result_16_fix.json', 'r', encoding='utf-8')
#     fix_data = open('../valid_result_16_fix_gfgqzr_tz_pairs.json', 'w+', encoding='utf-8')
#     for i, line in tqdm(enumerate(sub_data.readlines())):
#         line = line.strip()
#         eval_dict = json.loads(line)
#         origin_evn = dict()
#         fix_evn = dict()
#         for evn in eval_dict['events']:
#             if evn['type']=='股份股权转让' or evn['type']=='投资':
#                 origin_evn[evn['type']] = evn
#         if len(origin_evn)==2:
#             for mention in origin_evn['股份股权转让']['mentions']:
#                 if mention['role'] in fix_map.keys():

def get_dev_ids(file_dir):
    import pandas as pd
    type_map = {
        'zy': '质押',
        'gfgqzr': '股份股权转让',
        'qs': '起诉',
        'tz': '投资',
        'jc': '减持',
        'sg': '收购',
        'pj': '判决'
    }
    dev_ids = []
    for suffix in ['zy', 'gfgqzr', 'qs', 'tz', 'jc', 'sg', 'pj']:
        # read trains
        trains_add_devs = reader.read_txt(file_dir, -1, type_map[suffix])
        trains = trains_add_devs[:int(0.8 * len(trains_add_devs))]
        devs = trains_add_devs[int(0.8 * len(trains_add_devs)):]
        dev_ids.extend([d.id for d in devs])
    dev_ids = list(set(dev_ids))
    df = pd.DataFrame(dev_ids,index=None,columns=['dev_id'])
    df.to_csv('./dev_ids.csv',)
    return df

def add_content2result(devs):
    id2content = dict()
    for d in devs:
        id2content[d.id] = d.content
    sub_data = open('./valid_result.json', 'r', encoding='utf-8')
    fix_data = open('./valid_result_add_content.json', 'w+', encoding='utf-8')
    for i, line in tqdm(enumerate(sub_data.readlines())):
        line = line.strip()
        eval_dict = json.loads(line)
        ids = eval_dict['id']
        eval_dict['contents'] = id2content[ids]
        json.dump(eval_dict, fix_data, ensure_ascii=False)
        fix_data.write("\n")

def check(examples):
    ID_list = []
    good_examples = []
    cnt = 0
    for idx, example in enumerate(examples):
        if example['id'] not in ID_list:
            ID_list.append(example['id'])
        else:
            continue
        flag = True
        content = example['content']
        for event in example['events']:
            for mention in event['mentions']:
                if content[mention['span'][0]:mention['span'][1]] != mention['word']:
                    flag = False
                    #print(mention['span'])
                    #print(mention['word'])
                    #print(content[mention['span'][0]:mention['span'][1]])
        if flag:
            good_examples.append(example)
        else:
            #print(example, idx)
            cnt += 1
    print("Total %d error augmented examples !" %(cnt))
    return good_examples

def swap(sentence, swap1, swap2):
    if swap2 == swap1:
        return sentence
    temp = '````||||'
    sentence = sentence.replace(swap1, temp)
    sentence = sentence.replace(swap2, swap1)
    sentence = sentence.replace(temp, swap2)
    return sentence

def get_new_span(sentence, swaps1, swaps2, old_span, temp_word):
    if old_span[0] == 0:
        return [0, len(temp_word)]
    old_sentence = sentence[:old_span[0]]
    for i in range(len(swaps1)):
        old_sentence = swap(old_sentence, swaps1[i], swaps2[i])
    new_span = [len(old_sentence), len(old_sentence) + len(temp_word)]
    return new_span

if __name__ == "__main__":
    reader = Reader(True)
    file_dir = "./data/train"
    file_dir_test = "./data/dev"
    file_dir_err = './data/error'
    trains = reader.read_txt(file_dir, -1,"")
    devs = reader.read_test_txt(file_dir_test, -1)
    # get_dev_ids(file_dir)
    # add_content2result(devs)
    # error_output()
    # error_output_csv()
    fix_trigger()
    # extract_by_reg()
    # fix_NUM()
    # fix_classify()


from setuptools import setup, find_packages
import sys, os

version = '0.0'

setup(
	name='ckanext-webstorer',
	version=version,
	description="Tasks that upload data to the webstore.",
	long_description="""\
	""",
	classifiers=[], # Get strings from http://pypi.python.org/pypi?%3Aaction=list_classifiers
	keywords='',
	author='CKAN team.',
	author_email='ckan-dev@okfn.org',
	url='https://github.com/okfn/ckenext-webstorer',
	license='AGPL',
	packages=find_packages(exclude=['ez_setup', 'examples', 'tests']),
	namespace_packages=['ckanext', 'ckanext.webstorer'],
	include_package_data=True,
	zip_safe=False,
    install_requires = [
        'messytables>=0.1.1',
    ],
	entry_points=\
	"""
    [paste.paster_command]
    webstorer = ckanext.webstorer.commands:Webstorer

    [ckan.plugins]
    webstorer = ckanext.webstorer.plugin:WebstorerPlugin
	""",
)


import config # indirect use of config

db: 'AsyncSQLPool' # type hinting
version: 'Version' # once again, type hinting

packets: dict['Packets', 'BanchoPacket']

# cache some things like bcrypt for speeeeeeeeeeeeeeeeeed hours
cache = {
    'bcrypt': {}, # store bcrypt pws for speed hours
    'user': {} # temporarily used to cache tokens until we have a player object
}

players = [] # player list
queue: 'SimpleQueue'


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy
import numpy as np
from tensorflow.contrib.learn.python.learn.datasets import base
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import random_seed

import scipy.stats as stats
import pylab as pl

import helper_functions

def dense_to_one_hot(labels_dense, num_classes):
    """Convert class labels from scalars to one-hot vectors."""
    num_labels = labels_dense.shape[0]
    index_offset = numpy.arange(num_labels) * num_classes
    labels_one_hot = numpy.zeros((num_labels, num_classes))
    labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1
    return labels_one_hot

def get_one_hot(sentence, mean_length, num_unique_chars,one_hot_dict):
    """
    Given a sentences (for example 'abcdserer'), this function returns
    a one hot matrix of shape (mean_length rows x num_unique chars)
    based on the one_hot_dict argument which looks like:
    {a:[1,0,0,...,0], b[0,1,0,0....0], ...}
    :param sentence:
    :param mean_length:
    :param num_unique_chars:
    :param one_hot_dict:
    :return:
    """
    sentence_array = np.empty([mean_length, num_unique_chars], dtype=float)
    for i in range(mean_length):
        if(i < len(sentence)):
            sentence_array[i] = one_hot_dict[sentence[i]]
        else:
            sentence_array[i] = one_hot_dict['1']
    return sentence_array

class DataSet(object):
    def __init__(self,
                 sentences,
                 labels=None,
                 num_classes=None,
                 uniqe_chars=None,
                 num_unique_char=None,
                 mean_len=None,
                 fake_data=False,
                 one_hot=False,
                 dtype=dtypes.float32,
                 reshape=True,
                 seed=None,
                 normalize=False,
                 mode=None,
                 one_hot_dict=None):
        """Construct a DataSet.
        """
        seed1, seed2 = random_seed.get_seed(seed)
        # If op level seed is not set, use whatever graph level seed is returned
        numpy.random.seed(seed1 if seed is None else seed2)
        dtype = dtypes.as_dtype(dtype).base_dtype

        if mode == 'train':
            # mean_length is the most common length among sentenes
            self._mean_length = mean_len
            # unique_chars is the dictionary of unique characters which
            # looks like {a:0, b:1, ...., z:25}
            self._uniqe_chars = uniqe_chars

            # we add '0' and '1' in unique char dictionary to denote
            # beginning and end of sentences
            uniqe_chars['0'] = num_unique_char
            uniqe_chars['1'] = num_unique_char+1

            # getting values, i.e, numbers from unique characters dict
            el_dense = [i for i in uniqe_chars.values()]
            self._num_unique_chars = num_unique_char+2 # +2 is for beginning and ending of sentences.

            # change char indices to one_hot
            el_one_hot = dense_to_one_hot(np.array(el_dense), self.num_unique_chars)

            # create a dict of char->one_hot
            self._one_hot_dict = {}
            for k,v in zip(uniqe_chars.keys(), el_one_hot):
                self._one_hot_dict[k] = v

        else: # if valid/test
            self._one_hot_dict = one_hot_dict
            self._num_unique_chars = num_unique_char
            self._mean_length = mean_len

        self._epochs_completed = 0
        self._index_in_epoch = 0
        self._sentences = sentences
        self._labels = dense_to_one_hot(labels, num_classes)

        if dtype not in (dtypes.uint8, dtypes.float32):
            raise TypeError('Invalid image dtype %r, expected uint8 or float32' %
                            dtype)
        if fake_data:
            self._num_examples = 10000
            self.one_hot = one_hot
        else:
            assert len(sentences) == labels.shape[0], (
                'sentences.shape: %s labels.shape: %s' % (len(sentences), labels.shape))
            self._num_examples = len(sentences)

            # Convert shape from [num examples, rows, columns, depth]
            # to [num examples, rows*columns] (assuming depth == 1)
            # if reshape:
            #     #assert images.shape[3] == 1
            #     images = images.reshape(images.shape[0],
            #                             images.shape[1] * images.shape[2])
            # if dtype == dtypes.float32 and normalize:
            #      # Convert from [0, 255] -> [0.0, 1.0].
            #     images = images.astype(numpy.float32)
            #     images = numpy.multiply(images, 1.0 / 255.0)

    # labels in one hot format
    @property
    def one_hot_dict(self):
        return self._one_hot_dict

    @property
    def labels(self):
        return self._labels

    @property
    def mean_length(self):
        return self._mean_length

    @property
    def num_unique_chars(self):
        return self._num_unique_chars

    @property
    def num_examples(self):
        return self._num_examples

    @property
    def epochs_completed(self):
        return self._epochs_completed

    # this function changes sentences in string format into
    # one hots based on the characters
    def sentences_to_one_hot(self, sents):
        sentence_array = np.empty([len(sents),self._mean_length,self._num_unique_chars ], dtype=float)
        for i in range(len(sents)):
            sentence = sents[i]
            one_hot_sentence = get_one_hot(sentence,self.mean_length, self._num_unique_chars,self.one_hot_dict)
            sentence_array[i] = one_hot_sentence
        return sentence_array.reshape(sentence_array.shape[0], sentence_array.shape[1] * sentence_array.shape[2])

    # this returns the sentences in one_hot format
    @property
    def one_hot_sentences(self):
        return self.sentences_to_one_hot(self._sentences)

    # this is for adding test_sentences to data so we
    # could test our model
    def add_test_sentences(self,test_sentences):
        self.test_sentences = test_sentences

    # this function returns one hots of the test sentences we just added
    @property
    def one_hot_test_sentences(self):
        return self.sentences_to_one_hot(self.test_sentences)

    # this returns sentences in string format
    @property
    def sentences(self):
        return self._sentences

    def next_batch(self, batch_size, fake_data=False, shuffle=False):
        """Return the next `batch_size` examples from this data set."""
        start = self._index_in_epoch
        # Shuffle for the first epoch
        if self._epochs_completed == 0 and start == 0 and shuffle:
            perm0 = numpy.arange(self._num_examples)
            numpy.random.shuffle(perm0)
            #self._sentences = self.sentences[perm0]

            self._sentences = [self.sentences[i] for i in perm0]
            self._labels = self.labels[perm0]

        # Go to the next epoch
        if start + batch_size > self._num_examples:
            # Finished epoch
            self._epochs_completed += 1
            # Get the rest examples in this epoch
            rest_num_examples = self._num_examples - start
            sentences_rest_part = self._sentences[start:self._num_examples]
            labels_rest_part = self._labels[start:self._num_examples]
            # Shuffle the data
            if shuffle:
                perm = numpy.arange(self._num_examples)
                numpy.random.shuffle(perm)
                #self._sentences = self.sentences[perm]
                self._sentences = [self.sentences[i] for i in perm]
                self._labels = self.labels[perm]
            # Start next epoch
            start = 0
            self._index_in_epoch = batch_size - rest_num_examples
            end = self._index_in_epoch
            sentences_new_part = self._sentences[start:end]
            labels_new_part = self._labels[start:end]
            return self.sentences_to_one_hot(numpy.concatenate((sentences_rest_part, sentences_new_part), axis=0)), numpy.concatenate(
                (labels_rest_part, labels_new_part), axis=0)
        else:
            self._index_in_epoch += batch_size
            end = self._index_in_epoch
            return self.sentences_to_one_hot(self._sentences[start:end]), self._labels[start:end]

def get_mean_len(data):
    """
    This function calculates the distribution of length
    of sentences and returns the most common
    length among sentences. You can also see the plot
    of the normal distribution if you uncomment the
    commented lines.
    :param data: sentences
    :return: most common length among sentences
    """
    all_len = []
    for sentence in data:
        length = len(sentence)
        all_len.append(length)

    h = sorted(all_len)
    #fit = stats.norm.pdf(h, np.mean(h), np.std(h))
    #pl.plot(h, fit, '-o')

    y, x, _ = pl.hist(h, normed=True)
    #print (x.max())
    #print (y.max())

    #pl.show()
    return int(x.max())

def get_data_set(data_file, labels_file,test=False):

    if test:
        data = helper_functions.read_test_file(data_file)
        labels = np.zeros(len(data)) # dummy labels
        return data, labels.astype(int), None, None, None
    # read sentences to python list without the ending newline char
    data, unique_chars = helper_functions.read_text_file(data_file)

    # read labels to numpy array (float)
    labels = np.genfromtxt(labels_file)

    # get the most common sentence length
    mean_sentence_len = get_mean_len(data)

    # trim sentences to equal length and add '0', and '1' to indicate
    # beginning and end of sentences
    #data, unique_chars = pad_sentences(data, avg_sentence_len, unique_chars)

    # get classes (authors)
    classes = np.unique(labels)

    return data, labels.astype(int), unique_chars, mean_sentence_len, classes

class Data(object):

    def __init__(self,data_dir,
                   classes=None,
                   fake_data=False,
                   one_hot=False,
                   dtype=dtypes.float32,
                   reshape=True,
                   seed=None,
                   shape=None):
        # directories should have ending slash, otherwise won't work

        train_file = 'xtrain_obfuscated_small.txt'
        train_labels_file = 'ytrain_small.txt'
        test_file = 'xtest_obfuscated.txt'

#        train_file = 'subset_xtrain_obfuscated.txt'
#        train_labels_file = 'subset_ytrain.txt'
#        test_file = 'subset_xtest_obfuscated.txt'
        valid_file = 'xvalidation_obfuscated.txt'
        valid_labels_file = 'yvalidation.txt'

        train_sentences, train_labels, unique_chars, mean_sentence_len, classes \
            = get_data_set(data_dir+train_file, data_dir+train_labels_file)

        validation_sentences, validation_labels, _, _, _ \
            = get_data_set(data_dir + valid_file, data_dir + valid_labels_file)

        test_sentences, test_labels, _, _, _ \
            = get_data_set(data_dir + test_file, data_dir + test_file, test=True)

        options_train = dict(labels=train_labels, num_classes=len(classes), \
                       uniqe_chars=unique_chars,num_unique_char=len(unique_chars),\
                       mean_len=mean_sentence_len,\
                       dtype=dtype, reshape=reshape, seed=seed, normalize=True,mode='train' )

        # get a DataSet instance for train
        self._train = DataSet(train_sentences, **options_train)

        options_valid = dict(labels=validation_labels, num_classes=len(classes), \
                             uniqe_chars=unique_chars, num_unique_char=len(unique_chars), \
                             mean_len=mean_sentence_len, \
                             dtype=dtype, reshape=reshape, seed=seed, normalize=True, mode='validation',
                             one_hot_dict=self.train.one_hot_dict)

        self._validation = DataSet(validation_sentences, **options_valid)

        options_test = dict(labels=test_labels, num_classes=len(classes), \
                            uniqe_chars=unique_chars, num_unique_char=len(unique_chars), \
                            mean_len=mean_sentence_len, \
                            dtype=dtype, reshape=reshape, seed=seed, normalize=True, mode='test',
                            one_hot_dict=self.train.one_hot_dict)

        self._test = DataSet(test_sentences, **options_test)

    @property
    def train(self):
        return self._train
    @property
    def test(self):
        return self._test
    @property
    def validation(self):
        return self._validation


# Hack Server

import json
import time
from datetime import datetime
import socketio

from aiohttp import web

from PIL import Image, ImageMode
import io

connected = False

sio = socketio.AsyncServer(cors_allowed_origins='*')

app = web.Application()

# bugbug disable show_index before we start uploading more files
app.add_routes([web.static('/app', 'www', show_index=True)])

sio.attach(app)

@sio.event
async def connect(sid, environ):

    global connected

    print('connect ', sid)
    connected = True

@sio.event
def disconnect(sid):

    global connected

    print('disconnect ', sid)
    connected = False

@sio.on('message')
async def print_message(sid, message):

    print("Socket ID: " , sid)

    timestamp = datetime.timestamp(datetime.now())

    image = Image.open(io.BytesIO(message))

    file_name = str(timestamp) + '.jpg'

    print('received file: ', file_name)

    # bugbug save images outside of the web folder structure - should not be accessible
    image.save('www/data/' + file_name, "JPEG")

    await sio.emit('image_path', 'http://ec2-34-251-228-120.eu-west-1.compute.amazonaws.com:8080/app/data/' + file_name)

if __name__ == '__main__':
    web.run_app(app)


#!/usr/bin/env python
# Copyright 2017 ARC Centre of Excellence for Climate Systems Science
# author: Scott Wales <scott.wales@unimelb.edu.au>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import print_function
import iris
import argparse
import os

iris.FUTURE.netcdf_no_unlimited = True

def main():
    parser = argparse.ArgumentParser(description="Convert a file to netcdf using Iris. Multiple inputs will be merged into a single file")
    parser.add_argument('--output','-o', help="Output file name")
    parser.add_argument('--compression','-C',help="Compression level",choices=range(0,10),default=4,metavar='{0-9}',type=int)
    parser.add_argument('input', help="Input file name [UM/GRIB/NetCDF format]", nargs="+")
    args = parser.parse_args()

    if args.output:
        outfile = args.output
    else:
        basename = os.path.basename(args.input[0])
        (root, ext) = os.path.splitext(basename)
        outfile = root + '.nc'

    try:
        convert(args.input, outfile, compression_level=args.compression)
        print("Created %s"%outfile)
    except Exception:
        print("Unable to convert %s to NetCDF format"%args.input)
        raise

def convert(infiles, outfile, compression_level):
    # Do we need to compress?
    compress = compression_level != 0
    cubes = iris.load(infiles)
    iris.fileformats.netcdf.save(cubes, outfile, zlib=compress, complevel=compression_level)

if __name__ == '__main__':
    main()


is_debug = 0
is_down = 0
learning_rate = 0.01
import os
import re
import gc
import sys
import glob
import pandas as pd
import numpy as np
import datetime
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import KFold

#========================================================================
# original library
HOME = os.path.expanduser('~')
sys.path.append(f"{HOME}/kaggle/data_analysis/library/")
sys.path.append(f"{HOME}/kaggle/data_analysis/model/")
import MS_utils
from params_MS import params_lgb
import utils, ml_utils
from utils import logger_func
try:
    if not logger:
        logger=logger_func()
except NameError:
    logger=logger_func()
#========================================================================

"""
argv[1]: comment
argv[2]: feature_key
"""
# Basic Args
seed = 1208
set_type = 'all'
fold_n = 5
key, target, ignore_list = MS_utils.get_basic_var()
ignore_list = [key, target, 'country_group', 'down_flg']
comment = sys.argv[1]

# Base
base_path = '../input/base_exclude*'
base_path = '../input/base_Av*'
base_path = '../input/base_group*'
base = utils.read_df_pkl(base_path)[[key, target, 'country_group']]

#========================================================================
# Make Validation
#  vi_col = 'f000_AvSigVersion'
#  ignore_list.append(vi_col)

#  cv = KFold(n_splits=fold_n, shuffle=False, random_state=seed)
#  if is_debug:
#      base_train = base[~base[target].isnull()].head(10000).sort_values(by=vi_col)
#  else:
#      base_train = base[~base[target].isnull()].sort_values(by=vi_col)

#  kfold = list(cv.split(base_train[[key, 'country_group']], base_train[target]))
#  del base_train
#  gc.collect()
#  base = base[[key, target, 'country_group']]

if is_down:
    kfold = MS_utils.get_kfold(base=base[base['down_flg']==0].reset_index(drop=True))
else:
    kfold = MS_utils.get_kfold(base=base)
kfold = zip(*kfold)

if sys.argv[2].count('f'):
    train, test = MS_utils.get_feature_set(feat_key=sys.argv[2], base=base)
else:
    train, test = MS_utils.get_dataset(base=base)
print(train.shape, test.shape)

if is_debug:
    train = train.head(10000)
    test = test.head(500)

#========================================================================
# Categorical Encode
cat_cols = utils.get_categorical_features(df=train, ignore_list=ignore_list)
print(f"Categorical: {cat_cols}")

#Fit LabelEncoder
for col in cat_cols:
    # 最も頻度の多いカテゴリでimpute
    max_freq = train[~train[col].isnull()][col].mode().values[0]
    train[col].fillna(max_freq, inplace=True)
    test[col].fillna(max_freq, inplace=True)
    le = LabelEncoder().fit(pd.concat([train[col], test[col]], axis=0).value_counts().index.tolist())
    train[col] = le.transform(train[col])
    test[col]  = le.transform(test[col])
#========================================================================

#========================================================================
# Down
if is_down:
    train = train[train['down_flg']==0].reset_index(drop=True)
    down_cols = ['down_flg', 'country_group']
    train.drop(down_cols, axis=1, inplace=True)
    test.drop( down_cols, axis=1 , inplace=True)
#========================================================================

Y = train[target]

model_type_list = ['lgb', 'rmf', 'lgr']
model_type_list = ['lgb']
model_type = 'lgb'
metric = 'auc'

feim_list = []
score_list = []
oof_pred = np.zeros(len(train))
y_test = np.zeros(len(test))

use_cols = [col for col in train.columns if col not in ignore_list]
#  use_cols = [col for col in train.columns if col not in ignore_list and col.count('Smart')]
x_test = test[use_cols]
start_time = "{0:%Y%m%d_%H%M%S}".format(datetime.datetime.now())

for model_type in model_type_list:

    if model_type=='lgb':
        params = params_lgb()
        #  params['num_leaves'] = 4
        params['num_threads'] = 36
        #  params['num_threads'] = 60
        params['learning_rate'] = learning_rate
    else:
        params = {}
    logger.info(f"{model_type} Train Start!!")
    for num_fold, (trn_idx, val_idx) in enumerate(kfold):
        x_train, y_train = train[use_cols].iloc[trn_idx, :], Y.iloc[trn_idx]
        x_val, y_val = train[use_cols].iloc[val_idx, :], Y.iloc[val_idx]

        #  x_train = train.set_index(key)[use_cols + [target]].loc[trn_idx, :]
        #  x_val = train.set_index(key)[use_cols + [target]].loc[val_idx, :]
        #  y_train = x_train[target]
        #  y_val = x_val[target]
        #  x_train = x_train[use_cols]
        #  x_val = x_val[use_cols]

        logger.info(f"Fold{num_fold} | Train:{x_train.shape} | Valid:{x_val.shape}")

        score, tmp_oof, tmp_pred, feim, _ = ml_utils.Classifier(
            model_type=model_type
            , x_train=x_train
            , y_train=y_train
            , x_val=x_val
            , y_val=y_val
            , x_test=x_test
            , params=params
            , seed=seed
            , get_score=metric
        )
        feim_list.append(feim.set_index('feature').rename(columns={'importance':f'imp_{num_fold}'}))

        logger.info(f"Fold{num_fold} CV: {score}")
        score_list.append(score)
        oof_pred[val_idx] = tmp_oof
        y_test += tmp_pred

    feim = pd.concat(feim_list, axis=1)
    feim_cols = [col for col in feim.columns if col.count('imp_')]
    feim['importance'] = feim[feim_cols].mean(axis=1)
    feim.drop(feim_cols, axis=1, inplace=True)
    feim.sort_values(by='importance', ascending=False, inplace=True)
    feim['rank'] = np.arange(len(feim))+1

    cv_score = np.mean(score_list)
    logger.info(f'''
    #========================================================================
    # Model: {model_type}
    # CV   : {cv_score}
    #========================================================================''')

    y_test /= (num_fold+1)

    pred_col = 'prediction'
    train[pred_col] = oof_pred
    test[pred_col] = y_test
    stack_cols = [key, target, pred_col]

    df_stack = pd.concat([train[stack_cols], test[stack_cols]], ignore_index=True, axis=0)

    #========================================================================
    # Saving
    feim.to_csv(f'../valid/{start_time[4:12]}_valid_{model_type}_SET-{set_type}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB.csv', index=True)
    utils.to_pkl_gzip(obj=df_stack, path=f'../stack/{start_time[4:12]}_stack_{model_type}_SET-{set_type}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB')

    submit = pd.read_csv('../input/sample_submission.csv').set_index(key)
    submit[target] = test[pred_col].values
    submit.to_csv(f'../submit/{start_time[4:12]}_submit_{model_type}_SET-{set_type}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB.csv', index=True)
    #========================================================================


import os

cmd01 = "git clone https://github.com/h2oai/h2ogpt.git"
cmd02 = "cd h2ogpt"
cmd00 = "pip install numpy==1.23 pandas==1.5.3"
cmd03 = "cd .."
cmd04 = "xcopy h2ogpt src_demo /E /H /C /I"
cmd05 = "pip install xformers"

cms = [cmd01, cmd00]
cm00 = [cmd04, cmd05]

# group command 1
if(os.path.exists('h2ogpt') == False):
    for cm in cms:
        print('- Execute =', cm)
        os.system(cm)

# Install packages
fnames = ['requirements.txt', 'reqs_optional/requirements_optional_langchain.txt',
          'reqs_optional/requirements_optional_langchain.urls.txt']
cmd_pip = "pip install -r "
os.chdir('h2ogpt')
for fn in fnames:
    full_cmd = cmd_pip + fn
    print('- Execute =', full_cmd)
    os.system(full_cmd)

# group command 2
os.chdir('..')
if(os.path.exists('src_demo') == False):
    for cm in cm00:
        print('- Execute =', cm)
        os.system(cm)


from __future__ import print_function
from __future__ import division
import torch, torchvision
import torchvision.models as models
import os
import sklearn
from sklearn.model_selection import KFold
from sklearn.metrics import confusion_matrix
from matplotlib import pyplot as plt
from matplotlib import colors
from PIL import Image, ImageDraw
from tqdm import tqdm
from imgaug import augmenters as iaa
from sklearn.metrics import confusion_matrix
import pandas as pd
import numpy as np
import cv2
import random
import time
import math
from glob import glob
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
import copy
import sys
import argparse
import csv
import torch_optimizer as torch_optimizer
from shutil import copyfile
import json
from bisect import bisect_left

np.set_printoptions(threshold=np.inf)
torch.set_printoptions(threshold=np.inf)

TRAIN = 0
TEST = 1

# Set random seed
np.random.seed(0)

# Number of classes in the dataset
img_size = 224

# Set number of folds
num_folds = 10

def set_parameter_requires_grad(model):
    modify_pretrained_params = True
    for param in model.parameters():
        param.requires_grad = modify_pretrained_params

def weights_init(layer):
    if isinstance(layer, nn.Linear):
        nn.init.xavier_uniform_(layer.weight.data)

class GlobalAvgPool2d(nn.Module):
    def forward (self, x):
        return torch.mean(x.view(x.size(0), x.size(1), -1), dim=2)

def initialize_model(arch, num_classes):
    # Initialize these variables which will be set in this if statement. Each of these
    #   variables is model specific.
    model_ft = None

    if arch == "resnet":
        """ Resnet101 """
        model_ft = models.resnet101()
        num_ftrs = model_ft.fc.in_features
        model_ft.fc = nn.Linear(num_ftrs, num_classes)

    elif arch == "mobilenet":
        """ Mobilenet """
        model_ft = models.mobilenet_v2()
        num_ftrs = model_ft.classifier[1].in_features

    elif arch == "densenet":
        """ Densenet """
        model_ft = models.densenet201() #DenseNet201
        num_ftrs = model_ft.classifier.in_features
    else:
        print(f"Unknown model name {arch}. Choose from resnet, mobilenet, or densenet")
        quit()

    model_ft.classifier =  nn.Sequential(
                GlobalAvgPool2d(), #Equivalent to GlobalAvgPooling in Keras
#                nn.Linear(1920, 1024),
                nn.Linear(num_ftrs, 1024),
                nn.ReLU(),
                nn.Linear(1024, 1024),
                nn.ReLU(),
                nn.Linear(1024, 512),
                nn.ReLU(),
                nn.Linear(512, num_classes))

    return model_ft

# Initialize the model for this run# Print the model we just instantiated
#print(model_ft)

def f1_loss(y_pred, y_true, is_training=False):

    tp = (y_true * y_pred).sum().to(torch.float32)
    tn = ((1 - y_true) * (1 - y_pred)).sum().to(torch.float32)
    fp = ((1 - y_true) * y_pred).sum().to(torch.float32)
    fn = (y_true * (1 - y_pred)).sum().to(torch.float32)

    epsilon = 1e-7

    precision = tp / (tp + fp + epsilon)
    recall = tp / (tp + fn + epsilon)

    f1 = 2 * (precision * recall) / (precision + recall + epsilon)
    f1.requires_grad = is_training
    return f1

def plot(metric, train, val, epoch, fold, data_name, arch):

    max_y = max(train + val)

    #plt.title(f"{round(max(val), 2)} - {metric} - {data_name} - {xp_description} - {model_name} - {fold}-fold")
    plt.title(f"{metric} - {data_name} - {arch} - {fold}-fold")
    plt.plot(range(epoch + 1), train, label="train")
    plt.plot(range(epoch + 1), val, label="val")
    plt.axis([0, len(train)-1, 0, max_y + 0.1])
    plt.xlabel("Epoch")
    plt.ylabel(f"{metric}")
    plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0)
    plt.savefig(f"models/{data_name}/{arch}/plot_{metric}_{fold}.jpg", bbox_inches="tight")
    plt.clf()

def train_model(model, dataloaders_dict, device, criterion, optimizer, fold, scheduler, data_name, epochs, arch):

    # Initialize training
    since = time.time()
    train_acc_history, val_acc_history = [], []
    train_loss_history, val_loss_history = [], []

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    # Iterate over number of epochs
    for epoch in range(epochs):
        print('Epoch {}/{}'.format(epoch, epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                # Set model to training mode
                model.train()
            else:
                # Set model to evaluate mode
                model.eval()

            # Initialze epoch metrics
            running_loss, running_f1, running_corrects = 0.0, 0.0, 0

            # Iterate over data.
            # Get inputs, labels and path from Custom DataLoader
            for inputs, labels, path in tqdm(dataloaders_dict[phase]):

                # Send to GPU (if using GPU)
                inputs, labels = inputs.to(device), labels.to(device).long()

                # Zero the parameter gradients
                optimizer.zero_grad()

                # forward
                with torch.set_grad_enabled(phase == 'train'):

                    # Get model outputs
                    outputs = model(inputs)

                    # Get model loss
                    loss = criterion(outputs, labels)

                    # Get model predictions from soft max
                    _, preds = torch.max(outputs, 1)

                    # Optimize if Train
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # Update Running Statistics During Epoch

                # Loss
                running_loss += loss.item() * inputs.size(0)

                # F1
                running_f1 += f1_loss(preds, labels.data)

                # Accuracy
                running_corrects += torch.sum(preds == labels.data)

            # Summarize Epoch Statistics

            # Loss
            epoch_loss = running_loss / len(dataloaders_dict[phase].dataset)

            # F1
            epoch_f1 = running_f1 / len(dataloaders_dict[phase].dataset)

            # Accuracy
            epoch_acc = (running_corrects.double() / len(dataloaders_dict[phase].dataset)).cpu()

            print('\n{} Loss: {:.4f} F1 {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_f1, epoch_acc))

            # If Train, Update Training Summary Statistics
            if phase == 'train':
                train_acc_history.append(epoch_acc)
                train_loss_history.append(epoch_loss)

            # If Val, Update Validation Summary Statistics
            if phase == 'val':
                val_acc_history.append(epoch_acc)
                val_loss_history.append(epoch_loss)
                scheduler.step(epoch_loss)

            # If Best Performing model on validation set
            if phase == 'val' and epoch_acc > best_acc:

                # Set new best accuracy
                best_acc = epoch_acc

                # Copy model. Because PyTorch
                best_model_wts = copy.deepcopy(model.state_dict())

                # Remove previous best model. Saves space to keep only the best model
                if "prev_best_name" in locals():
                    os.remove(prev_best_name)

                # Save best model
                fname = f"models/{data_name}/{arch}/{fold}_{best_acc}_{epoch}.pt"
                torch.save(best_model_wts, fname)

                # Save name of save model as previous best for future iterative deletion
                prev_best_name = fname
                print(f'model succesfully saved with val_acc = {best_acc}')

        plot("loss", train_loss_history, val_loss_history, epoch, fold, data_name, arch)
        plot("acc", train_acc_history, val_acc_history, epoch, fold, data_name, arch)

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))
    print('Best train Acc: {:4f}'.format(train_acc_history[-1]))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

def test_model(model, dataloaders_dict, device, fold, test_img_paths, data_name, xp_description, model_name, mixed, is_inception=False):
    since = time.time()

    class_to_acc = {}
    class_to_count = {}
    preds_array = np.array([])
    labels_array = np.array([])

    global num_to_class
    global class_to_num
    global report_csv

    # Each epoch has a training and validation phase
    for phase in ['test']:
        model.eval()   # Set model to evaluate mode

        running_f1 = 0.0
        running_corrects = 0
        running_corrects_5 = 0
        running_corrects_10 = 0

        # Iterate over data.
        for inputs, labels, paths in tqdm(dataloaders_dict[phase]):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            preds_cpu = preds.cpu().numpy()
            labels_cpu = labels.cpu().numpy()

            preds_array = np.append(preds_array, preds_cpu)
            labels_array = np.append(labels_array, labels_cpu)

            # forward
            # track history if only in train
            # statistics
            running_f1 += f1_loss(preds, labels.data)
            running_corrects += torch.sum(preds == labels.data)

            outputs_cpu = outputs.detach().cpu().numpy()
            for i in range(outputs_cpu.shape[0]):
                top_10_indices = outputs_cpu[i].argsort()[-10:][::-1]
                for index, k in enumerate(top_10_indices):
                    if int(labels_cpu[i]) == k and index < 5:
                        running_corrects_5 += 1
                        running_corrects_10 += 1
                    elif int(labels_cpu[i]) == k:
                        running_corrects_10 += 1

            #print(preds, labels.data, preds == labels.data)

            #running_corrects_5 += torch.sum(preds == labels.data)

            #print("\nAccuracy", torch.sum(preds == labels.data).double() / batch_size)
            #print("\nModel Prediction", preds_cpu)
            #print("\nLabels", labels_cpu)

            for i, pred in enumerate(preds_cpu):
                label = int(labels_cpu[i])
                img_name = paths[i].split("/")[-1]
                if not num_to_class[label] in class_to_acc:
                    class_to_acc[num_to_class[label]] = 0
                    class_to_count[num_to_class[label]] = 0
                if not img_name in democracy_dict:
                    democracy_dict[img_name] = [pred]
                else:
                    democracy_dict[img_name].append(pred)
                if pred == label:
                    class_to_acc[num_to_class[label]] += 1
                elif mixed.lower() != "true":
                    report_csv.append([img_name, num_to_class[label], num_to_class[pred]])
                else:
                    report_csv.append([img_name, "test_sample", num_to_class[pred]])
                class_to_count[num_to_class[label]] += 1

        epoch_f1 = running_f1 / len(dataloaders[phase].dataset)
        epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
        epoch_acc_top5 = running_corrects_5 / len(dataloaders[phase].dataset)
        epoch_acc_top10 = running_corrects_10 / len(dataloaders[phase].dataset)

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    #print('Confusion Matrix', confusion_matrix(labels_array, preds_array))

    if mixed.lower() != "true":
        print("Class Acc", class_to_acc)

        for i, key in enumerate(class_to_acc):
            print(key,
            class_to_acc[key],
            class_to_count[key],
            round(class_to_acc[key]/class_to_count[key], 3),
            class_to_count[key] - class_to_acc[key])

        print("\n", sum(class_to_acc.values()), " / ", sum(class_to_count.values()),
            '\nAcc: {:4f}'.format(epoch_acc))

        print("\nTop 5 accuracy", epoch_acc_top5)
        print("\nTop 10 accuracy", epoch_acc_top10)

        global test_iter_accuracy
        test_iter_accuracy += (epoch_acc.item() / num_folds)

        print("\n Test Iter Accuracy ", test_iter_accuracy)

        with open(f"Models/{data_name}/{xp_description}/{model_name}/performance_report_{fold}.csv", "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerows(report_csv)

        for pred_error in report_csv[1:]:
            dest_dir = f"Models/{data_name}/{xp_description}/{model_name}/pred_error_{fold}/"
            ## Vector

            if "mist" in os.uname()[1]:
                if data_name == "alus":
                    print("ALUS PATH UNIMPLEMENTED")
                elif data_name == "PC":
                    src_path = f"../../alex/ParksCanada/{pred_error[1]}/{pred_error[0]}"
                    file_name = f"{pred_error[1]}_{pred_error[2]}_{pred_error[0]}"
            else:
                print("HOST NAME NOT IMPLEMENTED")
                if data_name == "alus":
                    src_path = f"../../../../../scratch/gobi2/schnei/Alus/ALUS_Full_Data/{pred_error[0]}"
                    file_name = f"{pred_error[0].split('.')[0]}_{pred_error[2]}.jpg"
                else:
                    print("PC image path not implemented")

            if not os.path.exists(dest_dir):
                os.makedirs(dest_dir)

            copyfile(src_path, os.path.join(dest_dir, file_name))

        build_confusion_matrix(labels_array, preds_array, [*class_to_acc], fold, epoch_acc, data_name, xp_description, model_name)

    else:
        create_report_csv(data_name, xp_description, model_name, fold)

def create_report_csv(data_name, xp_description, model_name, fold):
    global mixed_dict
    global report_csv

    classifications = list(num_to_class.values())

    for sample in report_csv[1:]:
        photo_id = "_".join(sample[0].split("_")[:5])

        if photo_id not in mixed_dict:
            mixed_dict[photo_id] = {k:0 for k in classifications}
        mixed_dict[photo_id][sample[2]] += 1

    sorted_mixed_dict = dict(sorted(mixed_dict.items()))
    mixed_pred_csv = []
    headers = ["Photo_ID"]
    classifications = list(num_to_class.values())
    headers += classifications
    mixed_pred_csv.append(headers)
    for photo_id, preds in sorted_mixed_dict.items():
        preds_list = list(preds.values())
        row = [photo_id] + preds_list
        mixed_pred_csv.append(row)

    csv_path = f"Models/{data_name}/{xp_description}/{model_name}/pred_mixed_{fold}.csv"

    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerows(mixed_pred_csv)

    mixed_dict = {}

def build_confusion_matrix(labels_array, preds_array, labels, fold, epoch_acc, data_name, xp_description, model_name):

    labels_class_array = np.array([num_to_class[int(i)] for i in labels_array])
    preds_class_array = np.array([num_to_class[int(i)] for i in preds_array])

    conf_mtrx = confusion_matrix(labels_class_array, preds_class_array, labels)

    conf_mtrx_norm = conf_mtrx / conf_mtrx.max(axis=0)
    conf_mtrx_norm = conf_mtrx / conf_mtrx.max(axis=1)

    cmap = colors.ListedColormap(['#ffffff', '#e8ffff', '#d9f1ff', '#bfe6ff', '#8cd3ff',
                                    '#59bfff', '#26abff', '#0da2ff', '#009dff', '#0055ff'])

    bounds=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]
    norm = colors.BoundaryNorm(bounds, cmap.N)

    plot = plt.imshow(conf_mtrx_norm, interpolation='nearest', cmap = cmap, norm=norm)

    plt.xticks([i for i in range(len(labels))], labels)
    plt.xticks(rotation=80)
    plt.yticks([i for i in range(len(labels))], labels)
    plt.xlabel("Ground Truth")
    plt.ylabel("Prediciton")
    plt.title(f"Confusion Matrix - {fold} - {round(epoch_acc.item(), 4)}acc")

    for i in range(conf_mtrx.shape[0]):
        for k in range(conf_mtrx.shape[1]):
            plt.text(k, i, str(conf_mtrx[i][k]), horizontalalignment='center', verticalalignment='center')

    plt.gcf().subplots_adjust(bottom=0.3)
    plt.savefig(f"Models/{data_name}/{xp_description}/{model_name}/plot_confusion_matrix_{fold}.jpg")
    plt.clf()

def split_data(X, n:int = num_folds):
    print('\nSplitting data')
    kf = KFold(n_splits=n, shuffle=False)
    kf.get_n_splits(X)
    kfolds = []
    for train_index, test_index in kf.split(X):
        #print("TRAIN:", train_index, "VAL:", val_index)
        kfolds.append((train_index, test_index))
    return kfolds

def yield_fold(X, Y, X_paths, train_val_or_test, chosen_fold, kfolds):
    #print("TRAIN LABELS", len(Y[kfolds[chosen_fold][TRAIN]]))
    #print("TEST LABELS", len(Y[kfolds[chosen_fold][VAL]]))

    val_index = int(len(X[kfolds[chosen_fold][TRAIN]]) * 0.9)

    if train_val_or_test == "train":
        train_X = X[kfolds[chosen_fold][TRAIN]][:val_index]
        train_Y = Y[kfolds[chosen_fold][TRAIN]][:val_index]
        train_paths = X_paths[kfolds[chosen_fold][TRAIN]][:val_index]
        return (train_X, train_Y, train_paths)
    elif train_val_or_test == "val":
        val_X = X[kfolds[chosen_fold][TRAIN]][:-val_index]
        val_Y = Y[kfolds[chosen_fold][TRAIN]][:-val_index]
        val_paths = X_paths[kfolds[chosen_fold][TRAIN]][:-val_index]
        return (val_X, val_Y, val_paths)
    else:
        test_X = X[kfolds[chosen_fold][TEST]]
        test_Y = Y[kfolds[chosen_fold][TEST]]
        test_paths = X_paths[kfolds[chosen_fold][TEST]]
        return (test_X, test_Y, test_paths)

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, X_images, targets, X_paths, img_aug, transform):
        self.X_images = X_images
        self.target = targets
        self.img_aug = img_aug
        self.transform = transform
        self.X_paths = X_paths
    def __len__(self):
        return len(self.X_images)
    def __getitem__(self, idx):
        sample = self.X_images[idx]
        sample = self.img_aug.augment_image(sample).copy()
        sample = sample.astype("float32") / 255.0
        sample = self.transform(sample)
        #print("TRANSFORM", sample)
        return (sample, self.target[idx], self.X_paths[idx])

img_aug = {
    'train': iaa.Sequential([
                iaa.Fliplr(p=0.5),
                iaa.SomeOf((0, 1), [
                    iaa.Affine(scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},  # Affine: Scale/zoom,
                           translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},  # Translate/move
                           rotate=(-90, 90), shear=(-4, 4)),  # Rotate and Shear
                    iaa.PiecewiseAffine(scale=(0, 0.05)),  # Distort Image similar water droplet  1.76
                    ]),
                iaa.SomeOf((0, 2),[
                    iaa.Add((-3, 3)),  # Overall Brightness
                    iaa.Multiply((0.95, 1.05), per_channel=0.2),  # Brightness multiplier per channel
                    iaa.Sharpen(alpha=(0.1, 0.75), lightness=(0.85, 1.15)),  # Sharpness
                    iaa.WithColorspace(to_colorspace='HSV', from_colorspace='RGB',  # Random HSV increase
                                       children=iaa.WithChannels(0, iaa.Add((-20, 20)))),
                    iaa.WithColorspace(to_colorspace='HSV', from_colorspace='RGB',
                                       children=iaa.WithChannels(1, iaa.Add((-20, 20)))),
                    iaa.WithColorspace(to_colorspace='HSV', from_colorspace='RGB',
                                       children=iaa.WithChannels(2, iaa.Add((-20, 20)))),
                    iaa.AddElementwise((-5, 5)),  # Per pixel addition
                    iaa.CoarseDropout((0.0, 0.015), size_percent=(0.02, 0.25)),  # Add large black squares
                    iaa.GaussianBlur(sigma=(0.1, 0.75)),  # GaussianBlur
                    iaa.Grayscale(alpha=(0.1, 1.0)),  # Random Grayscale conversion
                    iaa.Dropout(p=(0, 0.08), per_channel=0.2),  # Add small black squares
                    iaa.AdditiveGaussianNoise(scale=(0.0, 0.03 * 255), per_channel=0.5), # Add Gaussian per pixel noise
                    iaa.ElasticTransformation(alpha=(0, 1.0), sigma=0.15),  # Distort image by rearranging pixels
                    #iaa.weather.Clouds(),
                    iaa.weather.Fog(),
                    #iaa.weather.Snowflakes()

                ])

            ]),
    'val': iaa.Sequential([
                iaa.Fliplr(p=0)
            ]),
    'test': iaa.Sequential([
                iaa.Fliplr(p=0)
            ])
}

data_transforms ={
    'train': transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]),
    'test': transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]),
}

def buildImageAspectRatio(X_path):
    img = cv2.imread(X_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    resize_x = int(img.shape[1] * img_size / max(img.shape))
    resize_y = int(img.shape[0] * img_size / max(img.shape))

    push_x = (img_size - resize_x) // 2
    push_y = (img_size - resize_y) // 2

    resized_img = cv2.resize(img, (resize_x, resize_y))

    canvas = np.zeros((img_size, img_size, 3)).astype("uint8") + 255
    canvas[push_y:resized_img.shape[0] + push_y, push_x:resized_img.shape[1] + push_x, :] = resized_img
    return canvas

def createData(X_paths, Y, data_name):
    reset = True

    if not os.path.exists("metadata/Arrays_Batches"):
        os.makedirs("metadata/Arrays_Batches")

    if not os.path.exists("metadata/Arrays_Data"):
        os.makedirs("metadata/Arrays_Data")

    data_batch = 0
    for i, X_path in enumerate(tqdm(X_paths)):
        if reset == True:
            reset = False
            X = np.expand_dims(buildImageAspectRatio(X_path), axis=0)
        else:
            X = np.vstack((X, np.expand_dims(buildImageAspectRatio(X_path), axis=0)))
        if not i == 0 and i % 999 == 0:
            reset = True
            np.save(f"Arrays_Batches/{data_name}_Input_{data_batch}_{len(X)}.npy", X)
            np.save(f"Arrays_Batches/{data_name}_Labels_{data_batch}_{len(Y)}.npy", Y)
            data_batch += 1
        if i == len(X_paths) - 1:
            np.save(f"Arrays_Batches/{data_name}_Input_{data_batch}_{len(X)}.npy", X)
            np.save(f"Arrays_Batches/{data_name}_Labels_{data_batch}_{len(Y)}.npy", Y)
            data_batch += 1

    data_paths = []
    for batch in range(data_batch):
        data_paths.append(glob(f'metadata/Arrays_Batches/{data_name}_Input_{batch}_*')[0])

    for i, data_path in enumerate(tqdm(data_paths)):
        data = np.load(data_path)
        labels = np.load(data_path.replace("Input", "Labels"))

        if i == 0:
            X = data
        else:
            X = np.vstack((X, data))
        print(X.shape, Y.shape)

    np.save(f'metadata/Arrays_Data/{data_name}_Input_{len(X)}.npy', X)
    np.save(f'metadata/Arrays_Data/{data_name}_Labels_{len(Y)}.npy', Y)

def createCSV(csv_path, data_name):

    if not os.path.exists("CSVs"):
        os.makedirs("CSVs")

    csv_lst = []
    csv_lst.append(["filepath", "species"])

    for file_path in tqdm(glob(f"extracted/{data_name}/*")):
        print(file_path)
        species = "_".join(file_path.split("/")[-1].split("_")[:2])
        print(species)
        csv_lst.append([file_path, species])

    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerows(csv_lst)

def main(data_path, data_name, arch, batch_size, fold, xp, epochs):

    if not os.path.exists(f"models/{data_name}/{arch}"):
        os.makedirs(f"models/{data_name}/{arch}")

    if not os.path.exists(f"metadata"):
        os.makedirs(f"metadata")

    # Create Input CSV
    csv_path = os.path.join("CSVs", data_name) + '.csv'

    if not os.path.exists(csv_path):
        createCSV(csv_path, data_name)

    # Read data from input csv
    data = pd.read_csv(csv_path)
    num_to_class = {idx : species for idx, species in enumerate(data['species'].unique())}
    class_to_num = {species : idx for idx, species in enumerate(data['species'].unique())}
    num_classes = len(num_to_class)

    # Save num_classes for testing
    with open(f"metadata/{data_name}_num_to_class.json", "w") as fp:
        json.dump(num_to_class, fp)

    print(data)
    print("Num Classes: ", num_classes)

    # Create & Load data
    X_paths = np.array(data['filepath'])
    Y = np.array([class_to_num[each] for each in data['species']])

    input_file_path = f"metadata/Arrays_Data/{data_name}_Input_{len(X_paths)}.npy"
    label_file_path = f"metadata/Arrays_Data/{data_name}_Labels_{len(X_paths)}.npy"

    if not os.path.exists(input_file_path) or not os.path.exists(label_file_path):
        createData(X_paths, Y, data_name)
    X = np.load(input_file_path)
    Y = np.load(label_file_path)

    # Shuffle data
    permutation = np.random.permutation(len(Y))
    X = X[permutation]
    Y = Y[permutation]
    X_paths = X_paths[permutation]

    # Get data based on Fold
    kfolds = split_data(X)

    # Create train, val, and test image datasets
    image_datasets = {phase: CustomDataset(*yield_fold(X, Y, X_paths, phase, fold, kfolds), img_aug[phase], data_transforms[phase]) for phase in ['train', 'val', 'test']}

    # Initialize model
    model = initialize_model(arch, num_classes)

    # Detect if we have a GPU available
    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = nn.DataParallel(model).to(device)

    if xp.lower() == "train":
        # Build train and val pyTorch dataloaders
        dataloaders_dict = {phase: torch.utils.data.DataLoader(image_datasets[phase], batch_size=batch_size, shuffle=True, num_workers=4) for phase in ['train', 'val']}

        # Set Adam optimizer
        optimizer = optim.Adam(model.parameters(), lr=0.001)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=10, min_lr =0.000001, verbose=True)

        # Assign weights per class during training. Helps with overfitting to common classes
        unique, counts = np.unique(Y, return_counts=True)
        class_weights = [(max(counts)/ count)**0.5 for count in counts]

        # Setup the loss function
        class_weights = torch.FloatTensor(class_weights).to(device)
        criterion = nn.CrossEntropyLoss(weight=class_weights)

        # Train and evaluate model
        model_ft = train_model(model, dataloaders_dict, device, criterion, optimizer, fold, scheduler, data_name, epochs, arch)

    # elif xp.lower() == "test":
    #     dataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['test']}
    #     if not ckpt:
    #         try:
    #             model = sorted(glob(f"Models/{data_name}/{xp_description}/{model_name}/{fold}_*.pt"))[-1]
    #         except IndexError as e:
    #             print(e, data_name, xp_description, model_name, fold)
    #             quit()
    #         model_ft.load_state_dict(torch.load(model))
    #     else:
    #         model_ft.load_state_dict(torch.load(ckpt))
    #     print(f"\nLoading model {model}")
    #     test_img_paths = X_paths[kfolds[fold][TEST]]
    #     test_model(model_ft, dataloaders_dict, device, fold, test_img_paths, data_name, xp_description, model_name, mixed, is_inception=(model_name=="inception"))
    # else:
    #     print("Select xp as either 'train' or 'test'")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_path", default="images")
    parser.add_argument("--data_name", default="Alus")
    parser.add_argument("--arch", default="mobilenet")          #densenet,  resnet, mobilenet
    parser.add_argument("--batch_size", default=32, type=int)
    parser.add_argument("--fold", default=0, type=int)
    parser.add_argument("--xp", default="train")                #train or test
    parser.add_argument("--epochs", default=25, type=int)
    args = parser.parse_args()

    main(args.data_path, args.data_name, args.arch, args.batch_size, args.fold, args.xp, args.epochs)


#!/usr/bin/env python

# Copyright (c) Microsoft. All rights reserved.
# Licensed under the MIT license. See LICENSE file in the project root for
# full license information.

import os
import random
import time
import sys
import iothub_client
from iothub_client import IoTHubClient, IoTHubClientError, IoTHubTransportProvider
from iothub_client import IoTHubMessage, IoTHubMessageDispositionResult, IoTHubError, DeviceMethodReturnValue
from iothub_client_args import get_iothub_opt, OptionError

# HTTP options
# Because it can poll "after 9 seconds" polls will happen effectively
# at ~10 seconds.
# Note that for scalabilty, the default value of minimumPollingTime
# is 25 minutes. For more information, see:
# https://azure.microsoft.com/documentation/articles/iot-hub-devguide/#messaging
TIMEOUT = 241000
MINIMUM_POLLING_TIME = 9

# messageTimeout - the maximum time in milliseconds until a message times out.
# The timeout period starts at IoTHubClient.send_event_to_output.
# By default, messages do not expire.
MESSAGE_TIMEOUT = 10000

RECEIVE_CONTEXT = 0
AVG_WIND_SPEED = 10.0
MIN_TEMPERATURE = 20.0
MIN_HUMIDITY = 60.0
MESSAGE_COUNT = 5
RECEIVED_COUNT = 0
TWIN_CONTEXT = 0
METHOD_CONTEXT = 0

# global counters
SEND_CALLBACKS = 0

# Choose HTTP, AMQP or MQTT as transport protocol.  Currently only MQTT is supported.
PROTOCOL = IoTHubTransportProvider.MQTT

# String containing Hostname, Device Id & Device Key in the format:
# "HostName=<host_name>;DeviceId=<device_id>;SharedAccessKey=<device_key>;ModuleId=<module_id>;GatewayHostName=<gateway>"
CONNECTION_STRING = "[Device Connection String]"

MSG_TXT = "{\"deviceId\": \"myPythonDevice\",\"windSpeed\": %.2f,\"temperature\": %.2f,\"humidity\": %.2f}"

def send_confirmation_callback(message, result, user_context):
    global SEND_CALLBACKS
    print ( "Confirmation[%d] received for message with result = %s" % (user_context, result) )
    map_properties = message.properties()
    key_value_pair = map_properties.get_internals()
    print ( "    Properties: %s" % key_value_pair )
    SEND_CALLBACKS += 1
    print ( "    Total calls confirmed: %d" % SEND_CALLBACKS )

class HubManager(object):

    def __init__(
            self,
            connection_string,
            protocol=IoTHubTransportProvider.MQTT):
        self.client_protocol = protocol
        self.client = IoTHubClient(connection_string, protocol)
        if protocol == IoTHubTransportProvider.HTTP:
            self.client.set_option("timeout", TIMEOUT)
            self.client.set_option("MinimumPollingTime", MINIMUM_POLLING_TIME)
        # set the time until a message times out
        self.client.set_option("messageTimeout", MESSAGE_TIMEOUT)
        # some embedded platforms need certificate information
        self.set_certificates()

    def set_certificates(self):
        from iothub_client_cert import CERTIFICATES
        try:
            CERT_FILE = os.environ['EdgeModuleCACertificateFile']
            file = open(CERT_FILE)
            self.client.set_option("TrustedCerts", file.read())
            print ( "set_option TrustedCerts successful" )
        except IoTHubClientError as iothub_client_error:
            print ( "set_option TrustedCerts failed (%s)" % iothub_client_error )

        file.close()

    # Sends a message to the queue with outputQueueName, "temperatureOutput" in the case of the sample.
    def send_event_to_output(self, outputQueueName, event, properties, send_context):
        if not isinstance(event, IoTHubMessage):
            event = IoTHubMessage(bytearray(event, 'utf8'))

        if len(properties) > 0:
            prop_map = event.properties()
            for key in properties:
                prop_map.add_or_update(key, properties[key])

        self.client.send_event_async(
            outputQueueName, event, send_confirmation_callback, send_context)

def main(connection_string, protocol):
    try:
        print ( "\nPython %s\n" % sys.version )
        print ( "IoT Hub Client for Python" )

        hub_manager = HubManager(connection_string, protocol)

        print ( "Starting the IoT Hub Python sample using protocol %s..." % hub_manager.client_protocol )

        content = "Hello World from Python Blob APi"

        while True:
            # send a few messages every minute
            print ( "IoTHubClient sending %d messages" % MESSAGE_COUNT )

            for message_counter in range(0, MESSAGE_COUNT):
                temperature = MIN_TEMPERATURE + (random.random() * 10)
                humidity = MIN_HUMIDITY + (random.random() * 20)
                msg_txt_formatted = MSG_TXT % (
                    AVG_WIND_SPEED + (random.random() * 4 + 2),
                    temperature,
                    humidity)

                msg_properties = {
                    "temperatureAlert": 'true' if temperature > 28 else 'false'
                }
                hub_manager.send_event_to_output("temperatureOutput", msg_txt_formatted, msg_properties, message_counter)
                print ( "IoTHubClient.send_event_to_output accepted message [%d] for transmission to IoT Hub." % message_counter )

            # Wait for Commands or exit
            print ( "IoTHubClient waiting for commands, press Ctrl-C to exit." )

            status_counter = 0
            while status_counter < 6:
                status = hub_manager.client.get_send_status()
                print ( "Send status: %s" % status )
                time.sleep(10)
                status_counter += 1

    except IoTHubError as iothub_error:
        print ( "Unexpected error %s from IoTHub" % iothub_error )
        return
    except KeyboardInterrupt:
        print ( "IoTHubClient sample stopped" )

def usage():
    print ( "Usage: iothub_client_sample_module_sender.py -p <protocol> -c <connectionstring>" )
    print ( "    protocol        : <amqp, http, mqtt>" )
    print ( "    connectionstring: <HostName=<host_name>;DeviceId=<device_id>;SharedAccessKey=<device_key>;ModuleId=<module_id>;GatewayHostName=<gateway>" )

if __name__ == '__main__':
    try:
        # (CONNECTION_STRING, PROTOCOL) = get_iothub_opt(sys.argv[1:], CONNECTION_STRING)
        CONNECTION_STRING = os.environ["EdgeHubConnectionString"]
    except OptionError as option_error:
        print ( option_error )
        usage()
        sys.exit(1)

    if PROTOCOL != IoTHubTransportProvider.MQTT:
        protocol("Only the MQTT protocol is currently supported")
        sys.exit(1)

    main(CONNECTION_STRING, PROTOCOL)