import numpy as np
import os
import time
import requests
import json
import tempfile

DEBUG = False

def wait_flask_alive(server_full_address, check_interval = 0.05, timeout = 60):
    "Wait until flask server is alive"
    timeout_counter = timeout / check_interval
    while True:
        try:
            response = requests.get('{}/alive'.format(server_full_address))
            answer = response.content
            if answer == b'alive':
                break
        except:
            time.sleep(check_interval)
            timeout_counter -= 1

            if timeout_counter <= 0:
                raise TimeoutError("Timeout in waiting flask alive. Raising exception")

def send_flask_prediction(server_full_address, eval_step, input_dir, output_dir, debug = False):
    "Send prediction request to server"
    response = requests.post('{}/prediction'.format(server_full_address),
                             data=json.dumps({'eval_step': eval_step, "input_dir" : input_dir, "output_dir": output_dir}),
                             headers={"Content-type": "application/json"})
    if debug:
        print(response.content)
    answer = response.json()
    try:
        assert answer['status'] == 'prediction running' # Y a rien qui catch ca
    except:
        raise AssertionError('Expecting status "prediction running", instead got {}'.format(answer['status']))

def wait_flask_collect(server_full_address, check_interval = 0.05, debug = False, timeout=30):
    "Wait for prediction"
    collect_address = '{}/collect'.format(server_full_address)
    timeout_counter = timeout / check_interval
    while True:
        response = requests.get(collect_address)
        answer = response.content
        if answer == b'finished':
            return
        else:
            time.sleep(check_interval)
            timeout_counter -= 1
            if timeout_counter <= 0:
                raise TimeoutError("Timeout in waiting for prediction result. Raising exception")

def predict(server_full_address, eval_step, input_dir, output_dir, timeout=30):
    start_time = time.time()
    wait_flask_alive(server_full_address, timeout=timeout)
    send_flask_prediction(server_full_address, eval_step, input_dir, output_dir)
    remaining_time = int(timeout - (time.time() - start_time))
    wait_flask_collect(server_full_address, timeout=remaining_time)
    try:
        assert os.path.isfile(
            os.path.join(output_dir, 'Y{}.h5'.format(eval_step))
        )
    except:
        raise AssertionError('Server reports finished but prediction file not found')

def full_address(ip_address, port):
    return 'http://{}:{}'.format(ip_address, port)


import clean_data as data
import re

def clean_input_str(text: str) -> str:
    # remove leading and trailing whitespace
    text = text.lower().strip()
    # remove extra whitespace eg: '   ' + ' '
    text = re.sub(r'\s+', ' ', text)
    # remove non-alphanumeric characters
    text = re.sub(r'[^a-z0-9 ]', '', text)
    # remove leading and trailing whitespace that may have been the result of the string cleaning
    text = text.lower().strip()

    return text

def find_all(search_string: str) -> list[str]:
    results = []

    for key, value in data.data_dict.items():
        if search_string in key:
            for item in value:
                results.append(item)

    return sorted(results, key=len)

# Reduce the scope of the search so that the length of the translations are only twice as long as the input string to avoid long and random translations that may not be relevant.
def reduce_results(search_string: str, data_list: list[str]) -> list[str]:
    results = []

    for item in data_list:
        if len(item) <= (2*len(search_string)) :
            results.append(item)

    return sorted(results, key=len)

##### FILTER NOUN NOT WORKING #####
# Filter for only nouns
# def get_nouns(data_list: list[str]) -> list[str]:
#     results = []

#     for item in data_list:
#         # filter all posible articles which are always followed by nouns
#         if ('le ' in item) or ('la ' in item) or ('les ' in item) or ('l\'' in item) or ('un ' in item) or ('une ' in item) or ('de la ' in item) or ('du ' in item) or ('des ' in item) or ('à la ' in item) or ('au ' in item) or ('aux ' in item):
#             results.append(item)

#     return results

##### FILTER NOUN NOT WORKING #####
# get results for search text
# def get_result(search_text: str, show_less: bool, only_nouns: bool) -> list[str]:
#     # clean the input string
#     search_text = clean_input_str(search_text)

#     # find all possible translations
#     all_results = find_all(search_text)

#     if show_less and only_nouns:
#         # reduce the scope of the results and filter for nouns
#        return get_nouns(reduce_results(search_text, all_results))
#     elif show_less:
#         # reduce the scope of the results
#         return reduce_results(search_text, all_results)
#     elif only_nouns:
#         # filter all results for nouns
#         return get_nouns(all_results)
#     else:
#         return all_results

# get results for search text
def get_result(search_text: str, show_less: bool = False) -> list[str]:
    # clean the input string
    search_text = clean_input_str(search_text)

    # find all possible translations
    all_results = find_all(search_text)

    if show_less:
        # reduce the scope of the results
        return reduce_results(search_text, all_results)
    else:
        # return all results
        return all_results


from django.shortcuts import render

# Create your views here.
import json
from django.shortcuts import render
from django.views.generic.base import View
from search.models import ArticleType
from django.http import HttpResponse
from elasticsearch import Elasticsearch
from datetime import datetime
import redis

client = Elasticsearch(hosts=["127.0.0.1"])
redis_cli = redis.StrictRedis()

class IndexView(View):
    # 首页
    def get(self, request):
        topn_search = redis_cli.zrevrangebyscore("search_keywords_set", "+inf", "-inf", start=0, num=5)
        return render(request, "index.html", {"topn_search": topn_search})

class SearchSuggest(View):
    def get(self, request):
        key_words = request.GET.get('s', '')
        s_type = request.GET.get("s_type", 'article')
        if s_type == "article":
            s = ArticleType.search()

        re_datas = []
        if key_words:
            s = s.suggest('my_suggest', key_words, completion={
                "field": "suggest", "fuzzy": {
                    "fuzziness": 2
                },
                "size": 10
            })
            suggestions = s.execute_suggest()
            for match in suggestions.my_suggest[0].options:
                source = match._source
                re_datas.append(source["title"])
        return HttpResponse(json.dumps(re_datas), content_type="application/json")

class SearchView(View):
    def get(self, request):
        key_words = request.GET.get("q", "")
        s_type = request.GET.get("s_type", "article")
        jobbole_count = redis_cli.get("jobbole_count")
        start_time = datetime.now()
        hit_list = []
        # global response
        if s_type == "article":
            redis_cli.zincrby("search_keywords_set", 1, key_words)
            topn_search = redis_cli.zrevrangebyscore("search_keywords_set", "+inf", "-inf", start=0, num=5)
            page = request.GET.get("p","")
            try:
                page = int(page)
            except:
                page = 1
            response = client.search(
                index="pdfocr",
                body={
                    "query": {
                        "multi_match": {
                            "query": key_words,
                            "fields": ["tags", "pdfURL", "ocrText","jpgpath"]
                        }
                    },
                    "from": (page - 1) * 10,
                    "size": 10,
                    "highlight": {
                        "pre_tags": [''],
                        "post_tags": [''],
                        "fields": {
                            "pdfURL": {},
                            "ocrText": {},
                            "jpgpath": {},
                        }
                    }
                }
            )
            for hit in response["hits"]["hits"]:
                hit_dict = {}
                if "pdfURL" in hit["highlight"]:
                    hit_dict["pdfURL"] = "".join(hit["highlight"]["pdfURL"])
                else:
                    hit_dict["pdfURL"] = hit["_source"]["pdfURL"]
                if "ocrText" in hit["highlight"]:
                    hit_dict["ocrText"] = "".join(hit["highlight"]["ocrText"])
                else:
                    hit_dict["ocrText"] = hit["_source"]["ocrText"]

                if "jpgpath" in hit["highlight"]:
                    hit_dict["jpgpath"] = "".join(hit["highlight"]["jpgpath"])
                else:
                    hit_dict["jpgpath"] = hit["_source"]["jpgpath"]

                hit_dict["score"] = hit["_score"]
                hit_dict["id"] = hit["_id"]

                hit_list.append(hit_dict)

        end_time = datetime.now()
        last_seconds = (end_time - start_time).total_seconds()
        total_nums = response["hits"]["total"]['value']
        print(total_nums)
        if (page % 10) > 0:
            page_nums = int(total_nums / 10) + 1
        else:
            page_nums = int(total_nums / 10)

        return render(request, "result.html", {"page": page,
                                               "all_hits": hit_list,
                                               "key_words": key_words,
                                               "total_nums": total_nums,
                                               "page_nums": page_nums,
                                               "last_seconds": last_seconds,
                                               "jobbole_count": jobbole_count,
                                               "topn_search": topn_search
                                               })
class DetailView(View):
    # 首页
    def get(self, request):
        ocrText = request.GET.get("ocrText", "")
        jpgpath = request.GET.get("jpgpath","")

        return render(request, "detail.html", {"ocrText": ocrText,"jpgpath":jpgpath})


# Crie um programa que leia o ano de nascimento de sete pessoas. No final, mostre quantas pessoas ainda não atingiram a maioridade e quantas já são maiores.

from datetime import date

maioridade = 0
menoridade = 0
ano_atual = date.today().year

for i in range(1, 8):
    ano_nascimento = int(input(f'Em que ano a {i}ª pessoa nasceu? '))
    if ano_atual - ano_nascimento >= 21:
        maioridade += 1
    else:
        menoridade += 1

print(f'Ao todo, tivemos {maioridade} pessoas maiores de idade')
print(f'E também tivemos {menoridade} pessoas menores de idade')


import sys
input = sys.stdin.readline

N = int(input())
k = int(input())

start, end = 1, N*N #

while start <= end:
    mid = (start + end) // 2

    total = 0
    for i in range(1, N+1):
        total += min(mid // i, N)
        # mid보다 작은 수의 개수는 1x1, 1x2, 1x3, 2x1, 2x2, 3x1인데 그 수는 mid를 각 행으로 나눈 몫의 합에 해당한다
        # 또한 예를들어 mid가 5일 때 5보다 작은 수가 1x1, 1x2, 1x3, 1x4이더라도 4는 N*N 행렬을 벗어난 범위이기 때문에
        #한 행에서 mid보다 작은 수가 N개 보다 클 수는 없으므로 최대값을 N으로 제한한다.
    #total은 모든 행에서의 mid보다 작은 수의 개수의 합이며 total이 3이라면 mid는 적어도 4번째부터의 수에 해당한다.
    if total < k:
        start = mid+1
    elif total  >= k:
        end = mid-1
print(start)


import datetime
import time

from utils.dictionary import EntityDictionary
from utils.mention import extract_mention_and_plain_text_from_annotated_doc

def extract_mention_and_out_links_from_corpus(corpus_path):
    """
        只得到 mention_anchors 和 out_links，不需要同步生成 train_text
        由于中文 train_text 的生成需要分词，分词很耗时，可以先用这个函数生成一份 mention_anchors 和 out_links

    :param corpus_path:
    :return:
    """
    mention_anchors = dict()
    out_links = dict()
    self_links = dict()

    counter, mode_cnt = 0, 1000000
    start_time = int(time.time())
    last_update = start_time
    print("Extracting mention anchors and out links from corpus: \n\t{}".format(corpus_path))
    with open(corpus_path, "r", encoding="utf-8") as rf:
        for line in rf:
            counter += 1
            if counter % mode_cnt == 0:
                curr_update = int(time.time())
                print("\t#{}, batch_time: {}, total_time: {}".format(
                    counter,
                    str(datetime.timedelta(seconds=curr_update-last_update)),
                    str(datetime.timedelta(seconds=curr_update-start_time))
                ))
                last_update = curr_update
            try:
                instance_id, document = line.strip().split("\t\t")
                mention_anchor_list, _ = extract_mention_and_plain_text_from_annotated_doc(document)
                if out_links.get(instance_id) is None:
                    out_links[instance_id] = set()
                for mention, anchor, offset in mention_anchor_list:
                    mention = mention.lower()
                    if mention_anchors.get(mention) is None:
                        mention_anchors[mention] = dict()
                    if mention_anchors[mention].get(anchor) is None:
                        mention_anchors[mention][anchor] = 0
                    mention_anchors[mention][anchor] += 1
                    out_links[instance_id].add(anchor)

                    # 2020.10.28
                    if (instance_id == anchor):
                        self_links[mention] = self_links.get(mention, 0)+1

            except Exception as e:
                print(counter, e)
    ol = dict()
    for i in out_links:
        if len(out_links[i]) > 0:
            ol[i] = list(out_links[i])
    print("Extracted, total mentions: #{}, total time: {}".format(
        len(mention_anchors), str(datetime.timedelta(seconds=int(time.time())-start_time))))
    return mention_anchors, ol, self_links

def merge_mention_anchors(mention_anchors_list):
    """
        把多源的 mention_anchors 合并起来，例如合并分别从 abstract, article, infobox 中抽取的 mention_anchors

    :param mention_anchors_list:
    :return:
    """
    print("\nMerging mention anchors from {} sources".format(len(mention_anchors_list)))
    start_at = int(time.time())
    ma = dict()
    for mention_anchors in mention_anchors_list:
        for mention in mention_anchors:
            if len(mention) <= 1: continue
            if ma.get(mention) is None:
                ma[mention] = dict()
            for anchor in mention_anchors[mention]:
                if ma[mention].get(anchor) is None:
                    ma[mention][anchor] = 0
                ma[mention][anchor] += mention_anchors[mention][anchor]
    print("Merged, mentions: #{}, time: {}".format(
        len(ma), str(datetime.timedelta(seconds=int(time.time())-start_at))))
    return ma

def merge_out_links(out_links_list):
    """
        把多源的 out_links 合并起来，例如合并分别从 abstract, article, infobox 中抽取的 out_links

    :param out_links_list:
    :return:
    """
    print("\nMerging out links from {} sources".format(len(out_links_list)))
    start_at = int(time.time())
    ol = dict()
    for out_links in out_links_list:
        for inst in out_links:
            if len(out_links[inst]) > 0:
                if ol.get(inst) is None:
                    ol[inst] = set()
                for out_inst in out_links[inst]:
                    ol[inst].add(out_inst)
    for i in ol:
        if len(ol[i]) > 0:
            ol[i] = list(ol[i])
    print("Merged, out_links: #{}, time: {}".format(
        len(ol), str(datetime.timedelta(seconds=int(time.time()) - start_at))))
    return ol

def merge_self_links(self_links_list):
    """
        把多源的 self_links 合并起来，例如合并分别从 abstract, article, infobox 中抽取的 self_links

    :param self_links_list:
    :return:
    """
    print("\nMerging self links from {} sources".format(len(self_links_list)))
    start_at = int(time.time())
    sl = dict()
    for self_links in self_links_list:
        for inst in self_links:
            sl[inst] = sl.get(inst, 0) + self_links[inst]

    print("Merged, self_links: #{}, time: {}".format(
        len(sl), str(datetime.timedelta(seconds=int(time.time()) - start_at))))
    return sl

def expand_mention_anchors(source, mention_anchors):
    """
    从 mention_anchor.json 扩充词典
        a. 将满足以下条件的实体加入到全文统计的实体中，出现次数记为 1
            - 其 title 与 mention-anchor 字典中的某一 mention 相同
            - 该实体从未在语料中以 title 作为 mention 出现过
        b. 对于 title 没有作为 mention 出现过的实体
            - 以 title 作为 mention 构造 title-entity 字典

    :param source: string
    :param mention_anchors: dict
    :return: (dict, dict)
    """
    entity_dict = EntityDictionary.get_instance(source)

    title_entities = dict()

    print("\nExpanding mention anchors from entity dictionary...")
    start_at = int(time.time())
    for instance_id in entity_dict.entity_dict:
        mention = entity_dict.get_entity_by_id(instance_id).get_mention()
        if mention_anchors.get(mention) is not None:
            if mention_anchors[mention].get(instance_id) is None:
                mention_anchors[mention][instance_id] = 1
        else:
            title_entities[mention] = instance_id
    print("Expanded, title entities: #{}, mentions: #{}, time: {}".format(
        len(title_entities), len(mention_anchors), str(datetime.timedelta(seconds=int(time.time())-start_at))))
    return title_entities

def filter_mention_anchors(mention_anchors, link_m, freq_m, self_links, link_prob_th) -> dict:
    """
    1. filter out len(m) <= 1
    2. expand mention_anchors from entity dictionary
    3. filter out link(m) < 2
    4. filter out link_prob(m) < 0.0001
    :param mention_anchors:
    :return:
    """
    ma = dict()
    for m in mention_anchors:
        if len(m) > 1:
            ma[m] = mention_anchors[m]
    nma = dict()
    for m in ma:
        if m == '__all__':
            continue
        if link_m.get(m) is None or freq_m.get(m) is None or (link_m[m] - self_links.get(m, 0) < 2):
            continue
        if (float(link_m[m])/float(freq_m[m])) < link_prob_th:
            continue
        nma[m] = ma[m]
    return nma


from feature_process import *
from pose_cluster import *
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay, accuracy_score, roc_curve, RocCurveDisplay
from sklearn.inspection import permutation_importance
import tensorflow as tf
import csv

def train_balance(x,y):
    health = np.where(y==0)[0]
    pain = np.where(y==1)[0]
    sng = np.where(y==2)[0]
    mins = min([len(health),len(pain),len(sng)])
    health = np.random.choice(health, mins, replace=False)
    pain = np.random.choice(pain, mins, replace=False)
    sng = np.random.choice(sng, mins, replace=False)
    newidx = np.concatenate([health,pain,sng])
    return x[newidx],y[newidx]

class DataSet:
    '''
    Storing dataset to train/to test, root of related files, info of each single mice
    '''
    def __init__(self, dlc, bsoid=None, vidc=None, vids=None, dep=None, specific=[]):
        self.specific = specific
        self.all_treatment = ['Capbasal','Cap','pH5.2basal','pH5.2','pH7.4basal','pH7.4',
                                'pH5.2ASIC3KObasal','pH5.2ASIC3KO','CapTV1KObasal','CapTV1KO']
        self.files = {}
        self.files['dlc'] = self.load_paths(dlc, True)
        self.files['bsoid'] = self.load_paths(bsoid)
        self.files['vids'] = self.load_paths(vids)
        self.files['vidc'] = self.load_paths(vidc)
        self.files['dep'] = self.load_paths(dep)
        self.data_config()
        self.mclf=None

    def load_paths(self, root, sav_treat=False):
        if not root:
            return []
        files = os.listdir(root)
        sav_files = []
        treatments = []
        names = []
        for file in files:
            sav = True
            for sp in self.specific:
                if file.find(sp)==-1:
                    sav = False
                    break
            if not sav:
                continue
            if sav_treat:
                treatment = file.split('-')[0]
                name = file.split('-')[1]
                if file.find('basal')!=-1:
                    treatments.append(treatment+'basal')
                else:
                    treatments.append(treatment)
                names.append(name)
            sav_files.append(root+'/'+file)
        if sav_treat:
            self.names = names
            self.treatments = treatments
        return sav_files

    def data_config(self):
        self.ind={}
        self.ind['basal'] = np.array([i for i, j in enumerate(self.treatments) if j.find('basal')!=-1])
        for t in self.all_treatment:
            self.ind[t] = np.array([i for i, j in enumerate(self.treatments) if j == t])
        print('basal:',len(self.ind['basal']),' ,pain:',len(self.ind['Cap']),
                ' sng:',len(self.ind['pH5.2']),' pH7.4:',len(self.ind['pH7.4']),
                ' sngKO:',len(self.ind['pH5.2ASIC3KO']),' CapKO:',len(self.ind['CapTV1KO']))

    def sel_file(self, filetype='dlc', treatment='Cap'):
        if treatment == 'basal':
            return [self.files[filetype][i] for i, j in enumerate(self.treatments) if j.find('basal')!=-1]
        return [self.files[filetype][i] for i, j in enumerate(self.treatments) if j==treatment]

    def sel_feat(self, treatment='all'):
        if treatment == 'basal':
            return [self.mice_feat[i] for i, j in enumerate(self.treatments) if j.find('basal')!=-1]
        return [self.mice_feat[i] for i, j in enumerate(self.treatments) if j==treatment]

    def sel_data(self, treatment='all', sel_type='x_train'):
        if treatment == 'basal':
            return [self.data[sel_type][i] for i, j in enumerate(self.treatments) if j.find('basal')!=-1]
        return [self.data[sel_type][i] for i, j in enumerate(self.treatments) if j==treatment]

    def generate_feature(self, feat_type='frame'):
        self.mice_feat = []
        for i in range(len(self.files['dlc'])):
            if feat_type[:-1] == 'bs':
                tmp = miceFeature(self.treatments[i], bsoid=self.files['bsoid'][i], feat_type=feat_type)
            else:
                tmp = miceFeature(self.treatments[i], self.files['dlc'][i], feat_type=feat_type)#,self.files['vidc'][i],self.files['vids'][i],self.files['dep'][i])
            self.mice_feat.append(tmp)

    def generate_train_test(self, split=0.5, motion_del=False, k=1):
        '''
        validation setting as last k-th of each three treatment mice
        '''
        # config for mice_feat
        for miceF in self.mice_feat:
            if self.mclf:
                miceF.labeling(self.mclf,self.motion_score)
            else:
                miceF.labeling()
            miceF.train_config(split=split, motion_del=motion_del)

        # start
        all_sets = ['x_train','y_train','x_test','y_test','x_val','y_val']
        self.data = {}
        for s in all_sets:
            self.data[s] = []

        for t in self.all_treatment:
            inds = self.ind[t]
            k = k%len(inds)
            for i in range(len(inds)):
                if i==len(inds)-k-1:
                    continue
                ind = inds[i]
                self.data['x_train'].append(self.mice_feat[ind].x_train)
                self.data['y_train'].append(self.mice_feat[ind].y_train)
                self.data['x_test'].append(self.mice_feat[ind].x_test)
                self.data['y_test'].append(self.mice_feat[ind].y_test)
            ind = inds[len(inds)-k-1]
            self.data['x_val'].append(self.mice_feat[ind].feature)
            self.data['y_val'].append(self.mice_feat[ind].label)

    def pose_cls(self, sel=['random'], sel_num=20, embed=False, k=10, cls_type='km', clf_type='svm'):
        # miceF : miceFeature class object
        # get feature
        feat = []
        if sel[0]=='random':
            miceFs = self.mice_feat
            ind = np.random.choice(np.arange(len(miceFs)), sel_num, replace=False)
            for i in ind:
                feat.append(miceFs[i].feature)
        else:
            miceFs = []
            for s in sel:
                miceFs.extend(self.sel_feat(s))
            for miceF in miceFs:
                feat.append(miceF.feature)
        feat = np.concatenate(feat)
        # if is lstm => flatten to 2d feature
        if len(feat.shape)>2:
            feat = feat.reshape(len(feat), feat.shape[1]*feat.shape[2])
        # cluster
        if embed:
            embeder, embeddings = embedfeat(feat)
            motions, mclf = motion_cluster(embeddings, k, cls_type)
            self.embeder = embeder
        else:
            motions, mclf = motion_cluster(feat, k, cls_type)
        motion_num = len(np.unique(motions))
        if not mclf:
            mclf = motion_clf(feat, motions, clf_type=clf_type)
        # cluster predict and save result
        motionsB = [0]*motion_num
        motionsT = [0]*motion_num
        miceFsB, miceFsT = self.sel_feat('Capbasal'), self.sel_feat('Cap')
        for i in range(len(miceFsB)):
            miceFB = miceFsB[i]
            miceFT = miceFsT[i]
            if embed:
                motionB = motion_predict(miceFB.feature, mclf, embeder)
                motionT = motion_predict(miceFT.feature, mclf, embeder)
            else:
                motionB = motion_predict(miceFB.feature, mclf)
                motionT = motion_predict(miceFT.feature, mclf)
            for i in np.unique(motions):
                motionsB[i]+= len(np.where(motionB==i)[0])
                motionsT[i]+= len(np.where(motionT==i)[0])
        # motion score
        motion_num = len(motionsB)
        ratio = np.zeros((motion_num), dtype=float)
        for i in range(motion_num):
            if (motionsB[i]+motionsT[i])>0:
                ratio[i] = motionsT[i]/(motionsB[i]+motionsT[i])
        motion_score = np.zeros((motion_num), dtype=float)
        th = 0.4
        motion_score[(ratio<=th) | (ratio>=1-th)] = 1
        motion_score[(ratio>th) & (ratio<1-th)] = -1
        print("bad motions:", len(np.where(motion_score==-1)[0]))
        # plot
        x = np.arange(motion_num)
        width = 0.3
        plt.bar(x, motionsB, width, color='green', label='basal')
        plt.bar(x + width, motionsT, width, color='red', label='treat')
        plt.xticks(x + width / 2, x)
        plt.legend(bbox_to_anchor=(1,1), loc='upper left')
        #plt.show()
        #plt.savefig()
        self.mclf = mclf
        self.motionsB = motionsB
        self.motionsT = motionsT
        self.motion_score = motion_score

class miceFeature:
    '''
    Storing All data(file paths, landmarks, features ...) of single mice(file)
    '''
    def __init__(self, treatment, dlc=None, bsoid=None, vidc=None, vids=None, dep=None, feat_type='frame'):
        self.treatment = treatment
        if(dlc):
            self.dlcfile = dlc
            self.read_dlc()
        if(bsoid):
            self.bsoidfile = bsoid
        if(vidc):
            self.vidcfile = vidc
        if(vids):
            self.vidsfile = vids
        if(dep):
            self.depfile = dep

        self.count_feature(feat_type=feat_type)

    ### DLC functions #############################################################################
    def read_dlc(self):
        if not os.path.isfile(self.dlcfile):
            print("no file")
            return
        raw = np.genfromtxt(self.dlcfile, delimiter=",",dtype=int)[3:]
        getcol = tuple(np.arange(len(raw[0]))[np.arange(len(raw[0]))%3!=0])
        self.dlc_index = np.expand_dims(raw[:,0], axis=1)
        self.dlc_raw = raw[:,getcol]
        #remove nan
        notnan = ~np.isnan(self.dlc_raw).any(axis=1)
        self.dlc_raw = self.dlc_raw[notnan]
        self.dlc_index = self.dlc_index[notnan]
    def dlc_wrap(self):
        return np.resize(self.dlc_raw,(len(self.dlc_raw),int(self.dlc_raw.shape[1]/2),2))
    ###############################################################################################

    ### generate feature ##########################################################################
    def count_feature(self, feat_type='frame'):
        # config
        sel_dist=[[0,1],[0,2],[1,3],[2,3],[3,4],[3,5],[4,6],[5,6]]
        sel_ang=[[1,3,2],[0,3,6],[4,3,5]]
        sel_coord=[]
        normalize_range=(0,1)
        include_index = False
        window = 10
        if feat_type[-1]=='F':
            step = 10
        else:
            step = 5

########### frame ################################################
        if feat_type == 'frame':
            # config frame
            sel_dist=[[0,3],[3,4],[1,2]]
            sel_ang=[[0,1,3],[1,3,4]]
            # frame feature pre
            dist = count_dist(self.dlc_raw, sel_dist)
            ang = count_angle(self.dlc_raw, sel_ang)
            # disp = count_disp(self.dlc_raw, step=1, threshold=None)
            # frame feature
            feat = dist
            feat = np.hstack([feat, ang])
            # feat = np.hstack([feat, disp[:,0:1]])
            # normalize
            feat = feature_normalize(feat, normalize_range=normalize_range)
#########bsoid + cwt (full/half)##################################
        if feat_type[:-1] == 'bscwt':
            # frame feature pre
            dist = count_dist(self.dlc_raw, sel_dist)[1:]
            ang = count_angle(self.dlc_raw, sel_ang)[1:]
            disp = count_disp(self.dlc_raw, step=1, threshold=None)
            # frame feature
            feat = dist[:,5:6]
            feat = np.hstack([feat, dist[:,7:8]])
            feat = np.hstack([feat, ang[:,2:3]])
            feat = np.hstack([feat, disp[:,2:3]])
            # segment feature
            # seg = abs(fft_signal(feat, window=seg_window, flat=True))
            seg = cwt_signal(feat, window=window, step=step)
            # combine
            tmp = np.hstack([disp, ang])
            feat = np.hstack([seg, seg_statistic(tmp, count_types=['avg'], window=window, step=step)])
            feat = np.hstack([feat, seg_statistic(dist, count_types=['sum'], window=window, step=step)])
            # normalize
            feat = feature_normalize(feat, normalize_range=normalize_range)
########## bsoid ##########################################################
        if feat_type[:-1] == 'bs':
            savfile = joblib.load(self.bsoidfile)
            if len(savfile) > 10:
                feat = savfile
            else:
                feat = savfile[0]
########### bsoid LSTM ############################################
        if feat_type[:-1] == 'bsLSTM':
            # frame feature pre
            dist = count_dist(self.dlc_raw, sel_dist)[1:]
            ang = count_angle(self.dlc_raw, sel_ang)[1:]
            disp = count_disp(self.dlc_raw, step=1, threshold=None)
            # segment feature combine
            tmp = np.hstack([disp, ang, disp])
            tmp = feature_normalize(tmp, normalize_range=normalize_range)
            feat = generate_tmpfeat(tmp, window=window, step=step)
########### bsoid + cwt LSTM ############################################
        if feat_type[:-1] == 'bscwtLSTM':
            # frame feature pre
            dist = count_dist(self.dlc_raw, sel_dist)[1:]
            ang = count_angle(self.dlc_raw, sel_ang)[1:]
            disp = count_disp(self.dlc_raw, step=1, threshold=None)
            # frame feature
            feat = dist[:,5:6]
            feat = np.hstack([feat, dist[:,7:8]])
            feat = np.hstack([feat, ang[:,2:3]])
            feat = np.hstack([feat, disp[:,2:3]])
            # segment feature combine
            seg = cwt_signal(feat, window=window, step=step, flat=False)
            tmp = np.hstack([dist, ang, disp])
            tmp = feature_normalize(tmp, normalize_range=normalize_range)
            feat = generate_tmpfeat(tmp, window=window, step=step)
            feat = np.concatenate([feat, seg], axis=2)
#########################################################################
        self.feature = feat

    ### train test config ##########################################################################
    def labeling(self, mclf=None, motion_score=None):
        # pain:1 sng:2 health:0
        labels = np.zeros((self.feature.shape[0]), dtype=int)
        if self.treatment == 'pH5.2':
            labels[:] = 2
        elif self.treatment == 'pH7.4' or self.treatment.find('basal')!=-1 or \
                self.treatment.find('pH5.2ASIC3KO')!=-1 or self.treatment.find('CapTV1KO')!=-1:
            labels[:] = 0
        elif self.treatment == 'Cap':
            labels[:] = 1
        if mclf:
            motions = motion_predict(self.feature, mclf)
            for i in range(len(motion_score)):
                if motion_score[i] == -1:
                    labels[np.where(motions==i)] = 0 ## bad motion label
        self.label=labels

    def train_config(self, split=0.5, shuffle=True, motion_del=False):
        # select sample
        if motion_del:
            feat = self.feature[np.where(self.label!=-1)]
            label = self.label[np.where(self.label!=-1)]
        else:
            feat = self.feature
            label = self.label
        # shuffle
        ind = np.arange(len(feat))
        np.random.shuffle(ind)
        self.shuffle = ind
        # split
        if shuffle:
            feat = feat[ind]
            label = label[ind]
        if split==0:
            self.x_train = feat
            self.y_train = label
            self.x_test = []
            self.y_test = []
        else:
            # split : training portion
            sp = int(len(label)*split)
            self.x_train = feat[:sp,:]
            self.y_train = label[:sp]
            self.x_test = feat[sp:,:]
            self.y_test = label[sp:]

class Analysis:
    def __init__(self, model_type='svm', classes=3):
        if model_type == 'svm':
            self.model = SVC(kernel='rbf', C=1000)
        elif model_type == 'rf':
            self.model = RandomForestClassifier(random_state=42)
        elif model_type == 'dnn':
            self.model = DNN_model(classes)
        elif model_type == 'lstm':
            self.model = LSTM_model(classes)

    def train(self, x, y):
        self.model = self.model.fit(x,y)

    def test(self, x, y, show=False):
        sc = self.model.score(x, y)
        if show:
            print('accuracy = ',sc)
        return sc

    def analysis(self, x, y, seperate=False):
        pred = self.model.predict(x)

        # non-seperate
        tp = np.count_nonzero(((y==1) & (pred==1)) | ((y==2) & (pred==2)))
        tn = np.count_nonzero(((y==0) & (pred==0)) | ((y==-1) & (pred==-1)))
        fp = np.count_nonzero(((y==0) & ((pred==1)|(pred==2))) | ((y==-1) & ((pred==1)|(pred==2)))  |((y==1) & (pred==2)) | ((y==2) & (pred==1)) )  # mis postive
        fn = np.count_nonzero(((y==1) & ((pred==0)|(pred==-1)|(pred==2))) | ((y==2) & ((pred==0)|(pred==-1)|(pred==1))))  #|((y==0) & (pred==-1)) | ((y==-1) & (pred==0)) ) # mis negative
        toler = np.count_nonzero(((y==0) & (pred==-1)) | ((y==-1) & (pred==0))) # miss negative is useless
        if (fp+tn)==0:
            fa = 0
        else:
            fa = fp/(fp+tn)
        if (tp+fn)==0:
            dr = 0
        else:
            dr = tp/(tp+fn)
        acc = (tp+tn)/(tp+tn+fp+fn)
        print('accuracy = ', acc)
        print("false alarm: ", fa)
        print("detection rate: ", dr)
        return [acc,fa,dr]

    def analysis2(self, x, y):
        '''
        seperate detection rate of pain/sng
        '''
        pred = self.model.predict(x)
        tn = np.count_nonzero(((y==0) & (pred==0)) | ((y==-1) & (pred==-1)))
        fn_p = np.count_nonzero((y==1) & ((pred==0)|(pred==-1)|(pred==2)))
        tp_p = np.count_nonzero((y==1) & (pred==1))
        # fp_p = np.count_nonzero(((y==0) & (pred==1)) | ((y==-1) & (pred==1)) | ((y==2) & (pred==1)) )
        fn_s = np.count_nonzero((y==2) & ((pred==0)|(pred==-1)|(pred==1)))
        tp_s = np.count_nonzero((y==2) & (pred==2))
        # fp_s = np.count_nonzero(((y==0) & (pred==2)) | ((y==-1) & (pred==2)) | ((y==1) & (pred==2)) )
        if (tp_p+fn_p)==0:
            dr_p = 0
        else:
            dr_p = tp_p/(tp_p+fn_p)
        if (tp_s+fn_s)==0:
            dr_s = 0
        else:
            dr_s = tp_s/(tp_s+fn_s)
        return [dr_p, dr_s]

    def feat_importance(self, x, y, save_path=None):
        r = permutation_importance(self.model, x, y, n_repeats=10, random_state=0)
        feature_names = np.arange(len(x[0]))
        features = np.array(feature_names)
        # sorted_idx = r.importances_mean.argsort()
        plt.barh(features, r.importances_mean)
        plt.xlabel("Permutation Importance")
        if save_path:
            plt.savefig(save_path)
        return r.importances_mean

class LSTM_model:
    def __init__(self, classes):
        self.classes = classes
        self.build_model()

    def build_model(self):
        self.model = tf.keras.models.Sequential()
        self.model.add(tf.keras.layers.LSTM(32, return_sequences=False))
        self.model.add(tf.keras.layers.Dropout(0.2))
        self.model.add(tf.keras.layers.BatchNormalization())
        self.model.add(tf.keras.layers.Dense(self.classes, activation='softmax'))
        opt = tf.keras.optimizers.Adam(learning_rate=0.01)
        self.model.compile(optimizer=opt,
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    def fit(self, x, y):
        self.classes_ = np.arange(self.classes)
        if self.classes == 4:
            self.classes_ = self.classes_-1
        callbacks = [tf.keras.callbacks.EarlyStopping(monitor='accuracy', patience=20, mode='max')]
        self.model.fit(x, tf.one_hot(y,self.classes), epochs=200, batch_size=16) #,callbacks=callbacks)
        return self

    def predict(self, x):
        if self.classes == 4:
            return np.argmax(self.model.predict(x), axis=1)-1
        return np.argmax(self.model.predict(x), axis=1)

class DNN_model:
    def __init__(self, classes):
        self.classes = classes
        self.build_model()
    def build_model(self):
        self.model = tf.keras.models.Sequential()
        self.model.add(tf.keras.layers.Dense(32),activation='relu')
        self.model.add(tf.keras.layers.Dropout(0.2))
        self.model.add(tf.keras.layers.Dense(32),activation='relu')
        self.model.add(tf.keras.layers.Dropout(0.2))
        self.model.add(tf.keras.layers.BatchNormalization())
        self.model.add(tf.keras.layers.Dense(self.classes, activation='softmax'))
        opt = tf.keras.optimizers.Adam(learning_rate=0.01)
        self.model.compile(optimizer=opt,
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    def fit(self, x, y):
        self.classes_ = np.arange(self.classes)
        if self.classes == 4:
            self.classes_ = self.classes_-1
        callbacks = [tf.keras.callbacks.EarlyStopping(monitor='accuracy', patience=20, mode='max')]
        self.model.fit(x, tf.one_hot(y,self.classes), epochs=200, batch_size=16) #,callbacks=callbacks)
        return self

    def predict(self, x):
        if self.classes == 4:
            return np.argmax(self.model.predict(x), axis=1)-1
        return np.argmax(self.model.predict(x), axis=1)


import fileinput
import heapq
import math
import random
import time
import signal

#code venant d'un forum
class GracefulKiller:
  kill_now = False
  def __init__(self):
    signal.signal(signal.SIGINT, self.exit_gracefully)
    signal.signal(signal.SIGTERM, self.exit_gracefully)
    #code venant d'un forum
  def exit_gracefully(self,signum, frame):
    self.kill_now = True

"""
la fonction constructDict(Edges) prend en paramètre une liste d'arretes Edges et n le nombre de sommets
Elle retourne le graphe construit à l'aide de Edges
"""
def constructDict(Edges, n):
    G = {}

    for edge in Edges:
        key = int(edge[0])
        value = int(edge[1])
        if key not in G:
            G[key] = list()
        if value not in G:
            G[value] = list()
        G[key].append(value)
        G[value].append(key)
    for i in range(1, n+1):
        if i not in G.keys():
            G[i] = list()
    return G

"""
La fonction N(G,i) prend en paramètre le graphe G et le sommet i
La fonction retourne le nombre de voisin de i dans G
"""

def N(G, i):
    return len(G[i])

"""
La fonction getEdgesList(G) prend en paramètre le graphe G
Elle retourne la liste des arrêtes dans G
"""
def getEdgesList(G):
    Edges = set()
    for node in G:
        for voisin in G[node]:
            if (voisin,node) not in Edges:
                Edges.add((node,voisin))
    return Edges

"""
La fonction costP(i,G,C) prend en paramètre un sommet i,
un graphe G et C qui est un sous ensemble de G
Elle renvoie costP, le nombre de voisins de i dans C
"""
def costP(i, G, C):
    costP = 0
    killer=GracefulKiller()
    for node in C:
        if killer.kill_now:
          break
        if node in G[i]:
            costP = costP+1
    return costP

"""
La fonction costN(i,G,C) prend en paramètre un sommet i,
un graphe G et C qui est un sous ensemble de G
Elle renvoie costN, le nombre de sommets avec qui i n'est pas voisin dans C
"""
def costN(i, G, C):
    costN = len(C)-costP(i, G, C)
    return costN

"""
La fonction RN(i,G,C) prend en paramètre un sommet i,
un graphe G et C qui est un sous ensemble de G
Elle retourne costP-costN
"""
def RN(i, G, C):
    return costP(i, G, C)-costN(i, G, C)

"""
construction(G) prend en paramètre un graphe G, et renvoie une solution du clustering de G de facon a avoir un ensemble de clique
"""
def construction(G):
    killer=GracefulKiller()
    V = set()
    CL = {}
    # récupération et stockage du nombre de voisins de chaque sommet dans un tableau
    E = []
    for node in G:
        V.add(node)
        # on va mettre dans le tableau le couple (n,u) ou u est le sommet et n le nombre de voisins de u
        heapq.heappush(E, (N(G,node),node))
    heapq._heapify_max(E)
    n = len(V)

    Kbest = random.randint(1, n)
    Kmin = round(max(Kbest-math.sqrt(n), 1))
    Kmax = round(min(Kbest+math.sqrt(n), n))
    K = random.randint(Kmin, Kmax)

    # on choisis les k premiers elements de V, chaque sommet sélectionné est le premier élément d’un nouveau groupe CL[i]
    for i in range(K):
        CL[i] = list()
        node = (heapq._heappop_max(E))[1]
        CL[i].append(node)
        V.remove(node)

    # tant que V n'est pas vide on choisis au hasard un sommet j dans V et on le met dans un cluster CL[i] de sorte que CL[i] maximizes RN(j,Ci) = cost+(j,Ci)−cost−(j,Ci)
    while V and not killer.kill_now:
        node = V.pop()
        selectedCl = 0
        maxRn = RN(node,G,CL[selectedCl])
        lenCL = len(CL)

        for i in range(1, lenCL):
            if RN(node, G, CL[i]) > maxRn:
                maxRn = RN(node, G, CL[i])
                selectedCl = i
        # si d(node)=0 : node qui est le sommet en cours de traitement alors on met node dans un cluster à lui seul
        if len(G[node]) == 0:
            CL[lenCL] = list()
            CL[lenCL].append(node)
        else:
            CL[selectedCl].append(node)

    if not killer.kill_now :
        Edges = getEdgesList(G)

        for k in range(len(CL)):
            for node in CL[k]:
                # toutes les arretes reliant des sommets qui ne sont pas dans le même cluster sont supprimées
                for couple in set(Edges):
                    if killer.kill_now:
                        print("break")
                        break
                    u, v = couple
                    if (u == node and v not in CL[k]) or (v == node and u not in CL[k]):
                        Edges.remove((u,v))

                # on lie tous les sommets d'un cluster pour former une clique
                for voisin in CL[k]:
                    if killer.kill_now:
                        print("break")
                        break
                    if node != voisin and ((node,voisin) not in Edges and (voisin,node) not in Edges):
                        Edges.add((node,voisin))
            if killer.kill_now:
                print("break")
                break
    if not killer.kill_now:
        return Edges, CL
    else:
        return set(),list()

"""
La fonction sumCost(C,VpC) prend en paramètre un cluster C et un ensemble de sommets VpC (V privé de C)
Elle renvoie pour la somme des costP(i, G, VpC) pour tout les i appartenant à C
"""
def sumCost(C, VpC):
    killer=GracefulKiller
    sumCost = 0
    for i in C:
        if killer.kill_now:
            print("i")
            sumCost=math.inf
            break
        sumCost = sumCost+costP(i, G, VpC)
    return sumCost

"""
La fonction LocalSearch prend en paramètre un graphe G, et Une liste de Cluster CL
Elle essaye de trouver une meilleure solution à partir du graphe G et la renvoie.
"""
def LocalSearch(G, CL):
    killer=GracefulKiller()
    if not killer.kill_now:
        V = set()
        for node in G:
            V.add(node)
        n = len(V)
        selectedCLIndex = 0
        maxsumCost = sumCost(CL[selectedCLIndex], V-set(CL[selectedCLIndex]))
        for i in range(1, len(CL)):
            sumcost = sumCost(CL[i], V-set(CL[i]))
            if sumcost > maxsumCost:
                selectedCLIndex = i
                maxsumCost = sumcost

        # Local search empty cluster

        #on rentre dans le if que si le cluster selectionner
        if not killer.kill_now and (len(CL[selectedCLIndex]) > 1 or (len(CL[selectedCLIndex]) == 1 and len(G[CL[selectedCLIndex][0]]) != 0)):
            nodeSeen = []
            while not killer.kill_now and CL[selectedCLIndex] and len(nodeSeen) < len(CL[selectedCLIndex]):
                node = CL[selectedCLIndex].pop(0)

                if node not in nodeSeen:
                    nodeSeen.append(node)
                maxCostSelectedClIndex = 0
                maxCost = costP(node, G, CL[0])
                for i in range(1, len(CL)):
                    costp = costP(node, G, CL[i])
                    if costp > maxCost:
                        maxCostSelectedClIndex = i
                        maxCost = costp

                CL[maxCostSelectedClIndex].append(node)

        # End Local search empty cluster
            solution1=constructGraphWithCL(CL,G)

        # Local search Split cluster
        if not killer.kill_now :
            selectedCLIndex = 0
            maxsumCost = sumCost(CL[selectedCLIndex], V-set(CL[selectedCLIndex]))
            for i in range(1, len(CL)):
                if killer.kill_now:
                    break
                sumcost = sumCost(CL[i], V-set(CL[i]))
                if sumcost > maxsumCost:
                    selectedCLIndex = i
                    maxsumCost = sumcost

            Edges = getEdgesList(G)
            F = []
            if random.randint(0,1)==0:
                op=1
            else:
                op=-1
            if len(CL[selectedCLIndex]) > 2 and not killer.kill_now :
                seenCouple = set()
                for node in CL[selectedCLIndex]:
                    for voisin in CL[selectedCLIndex]:
                        if killer.kill_now:
                            break
                        if node != voisin and (voisin,node) not in seenCouple and ((node,voisin) not in Edges and (voisin,voisin) not in Edges):
                            # on calcule cos+(i,C)+cos+(j,C) ou cos+(i,C)-cos+(j,C) en sachant que i et j ne sont pas voisins dans G
                            costIJ = costP(node, G, CL[selectedCLIndex])+op*costP(voisin, G, CL[selectedCLIndex])
                            heapq.heappush(F, (abs(costIJ), (node,voisin)))
                            seenCouple.add((node,voisin))
                    if killer.kill_now:
                        break
                heapq._heapify_max(F)
                # on prend le coupe (i,j) tel que cost+(i,C) et cost+(j,C) est le plus grand dans C

                if len(F) > 0:
                    nodei, nodej = heapq._heappop_max(F)[1]
                    lenCL = len(CL)
                    CL[selectedCLIndex].remove(nodei)
                    CL[selectedCLIndex].remove(nodej)
                    CL[lenCL] = list()
                    CL[lenCL+1] = list()
                    CL[lenCL].append(nodei)
                    CL[lenCL+1].append(nodej)
                    while CL[selectedCLIndex] and not killer.kill_now:
                        node = CL[selectedCLIndex].pop()

                        if RN(node, G, CL[lenCL]) > RN(node, G, CL[lenCL+1]):
                            CL[lenCL].append(node)
                        else:
                            CL[lenCL+1].append(node)
        # END Local search Split cluster

        # on renvoie la meilleur solution entre le cluster split et le empty split
        if killer.kill_now:
            return dict(),True
        elif solutionCost(G,constructGraphWithCL(CL,G))[0]<solutionCost(G,solution1)[0]:
            return constructGraphWithCL(CL,G),False
        else:
            return solution1,False
    else:
        return dict(),True

"""
La fonction constructGraphWithCL(CL,G) prend en paramètre un graphe G et une lister de Cluster CL
Elle renvoie le graphe construit à l'aide du Cluster CL
"""
def constructGraphWithCL(CL,G):
    killer=GracefulKiller()
    Edges=getEdgesList(G)
    for k in range(len(CL)):
        for node in CL[k]:
            # toutes les arretes reliant des sommets qui ne sont pas dans le même cluster sont supprimées
            for couple in set(Edges):
                if killer.kill_now:
                    break
                u, v = couple
                if (u == node and v not in CL[k]) or (v == node and u not in CL[k]):
                    Edges.remove((u,v))

            # on lie tous les sommets d'un cluster pour former une clique
            for voisin in CL[k]:
                if killer.kill_now:
                    break
                if node != voisin and ((node,voisin) not in Edges and (voisin,node) not in Edges):
                    Edges.add((node,voisin))
            if killer.kill_now:
                break
    return constructDict(Edges, n)

    # fonction à implementer

"""
La fonction solutionCost(G, G1) prend en paramètre deux graphes : G le graphe de depart et G1 une solution de clustering de G
La fonction renvoie le cout de la solution G1 (le nombre d'arretes modifiées)
"""
def solutionCost(G, G1):

    EdgeG = getEdgesList(G)
    EdgeG1 = getEdgesList(G1)
    compareG1G = EdgeG.symmetric_difference(EdgeG1)
    for couple in EdgeG:
        node, voisin = couple
        if (node,voisin) in compareG1G and (voisin,node) in compareG1G:
            compareG1G.remove((node,voisin))
            compareG1G.remove((voisin,node))

    return len(compareG1G),compareG1G

#fonction venant d'un forum: elle renvoie l'heure actuelle en miliseconde
def current_time():
    return round(time.time() * 1000)

# GRASP FONCTION

def grasp(G, Tmax):
    killer = GracefulKiller()
    G1 = constructDict(construction(G)[0], len(G))
    timestart = current_time()
    while current_time()-timestart < Tmax and not killer.kill_now:
        G2,isKilled= LocalSearch(G, construction(G)[1])
        if not isKilled and solutionCost(G, G2)[0] < solutionCost(G, G1)[0]:
            G1 = G2

    return G1

if __name__ == '__main__':
    # 1. find a way to extract data from the input file
    G = {}
    for line in fileinput.input():
        line = line.split("\n")[0]
        l1 = line.split(" ")

        if l1[0] == "p" and l1[1] == "cep":
            n = int(l1[2])
        elif l1[0] != "c":
            key = int(l1[0])
            value = int(l1[1])
            if key not in G:
                G[key] = list()
            if value not in G:
                G[value] = list()
            G[key].append(value)
            G[value].append(key)
    # s'il existe un sommet dont le degrés est nul on l'ajoute quand meme dans le dict
    for i in range(1, n+1):
        if i not in G.keys():
            G[i] = list()

    # done
    if(len(G) > 0):
        # fonction pour laner le grasp:
        Tmax=300000  #5min
        G1 = grasp(G, Tmax)
        solution=solutionCost(G, G1)
        for coupl in solution[1]:
            print(coupl[0],coupl[1])

        # print(LocalSearch(G,construction(G)[1]))

    # 2. call the grasp fonction
    # 3. write the solution in a output file


import os
import pickle
from shutil import copyfile
from datetime import datetime
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from torch.cuda import get_device_name, current_device
import torch.nn.utils.rnn as rnn
# from tensorboardX import SummaryWriter
from vqa_dataset import VQADataSet
from model import CoattentionNet
from image_encoder import load_resnet

class ExperimentRunner():
    ''' Base class for runnung the experiment.
    '''

    def __init__(self, train_images_dir, train_questions_dir, train_annotations_dir,
                val_images_dir, val_questions_dir, val_annotations_dir , batch_size,
                num_of_epochs, num_of_workers, collate, train_results_dir, val_results_dir,
                saving_frequency, learning_rate, save_results=False):
        ''' Initializes the experiment runner class.
        Parameters:
            train_images_dir: string; Path to train images.
            train_questions_dir: string; Path to train questions.
            train_annotations_dir: string; Path to train annotations.
            val_images_dir: string; Path to validation images.
            val_questions_dir: string; Path to validation questions.
            val_annotations_dir: string; Path to validation annotations.
            batch_size: int; Batch size.
            num_of_epochs: int; Number of epochs.
            num_of_workers: int; Number of workers.
            collate: boolean; Flag to indicate that the results have already been preprocessed and saved. We will not need to do it again. Also images have been encoded.
            train_results_dir: string; Path to the saved train results.
            val_results_dir: string; Path to the saved validation results.
            saving_frequency: int; Frequency of saving in checkpoints.
            learning_rate: float; The learning rate of the algorithm.
            save_results: boolean; Flag for saving results such as vocabulary, top 1000 answers etc.
        '''

        self.train_dataset = VQADataSet("train", train_questions_dir, train_annotations_dir, train_images_dir,
                                  "COCO_train2014_", collate, save_results, train_results_dir)
        self.val_dataset = VQADataSet("val", val_questions_dir, val_annotations_dir, val_images_dir,
                                "COCO_val2014_", collate, save_results, val_results_dir)

        self.train_dataloader = DataLoader(self.train_dataset, batch_size=batch_size, shuffle=True,
                                          num_workers=num_of_workers, collate_fn=self.custom_collate_func)
        self.val_dataloader = DataLoader(self.val_dataset, batch_size=batch_size, shuffle=False,
                                        num_workers=num_of_workers, collate_fn=self.custom_collate_func)

        self.num_of_epochs = num_of_epochs
        self.num_of_workers = num_of_workers
        self.logging_frquency = 10  # every 10 training steps
        self.verbose = 50 # print every 50 steps
        self.batch_size = batch_size
        self.saving_frequency = saving_frequency  # epoch-wise
        self.learning_rate = learning_rate

        # Loading the model
        with open(os.path.join(train_results_dir, "question_vocabulary.pkl"), 'rb') as f:
            question_vocabulary = pickle.load(f)

        self.model = CoattentionNet(len(question_vocabulary), 1000).float()
        self.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
        if self.DEVICE == "cuda":
            print(f"Loading model with: {self.DEVICE} | {get_device_name(current_device())}...")
            self.model = self.model.cuda()
        else:
            print(f"Loading model with: {self.DEVICE}...")

        # Setting the optimizer
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate, weight_decay=1e-8)

        # Setting the loss function
        self.criterion = torch.nn.CrossEntropyLoss()

        # Weights initialization
        self.initialize_weights()

        # Logger
        # self.logger = SummaryWriter()  # rewrite this.

        # Image Encoder
        self.image_encoder = load_resnet(self.DEVICE)

        # Checkpoints
        self.checkpoints_dir = os.path.join("Data", "saved")
        if not os.path.isdir(self.checkpoints_dir):
            os.makedirs(self.checkpoints_dir)
        self.experiment_id = len(os.listdir(self.checkpoints_dir))

        if not os.path.isdir(os.path.join(self.checkpoints_dir, str(self.experiment_id))):
            print(f"Creatting {os.path.join(self.checkpoints_dir, str(self.experiment_id))}...")
            os.makedirs(os.path.join(self.checkpoints_dir, str(self.experiment_id)))

    def optimize(self, predicted_answers, true_answers):
        ''' Optimization step for the model.
        Parameters:
            predicted_answers: pytorch tensor object; Tensor of the model predictions.
            true_answers: pytorch tensor object; Ids of ground truth answers.
        Returns:
            loss: float; Loss for the model in current step.
        '''

        self.optimizer.zero_grad()
        loss = self.criterion(predicted_answers, true_answers)
        loss.backward()
        self.optimizer.step()

        return loss

    def validate(self):
        ''' Method for validating the model after training.
        Returns:
            accuracy: float; Accuracy of the model after validation.
        '''

        print("Validating the model...")
        accuracy = 0.0
        progress_bar = tqdm(total=len(self.val_dataloader))
        self.model.eval()
        with torch.no_grad():
            for batch_idx, (image_tensors, question_tensors, answer_tensors) in enumerate(self.val_dataloader):
                progress_bar.update(1)
                question_tensors = rnn.pack_sequence(question_tensors)
                image_tensors = image_tensors.to(self.DEVICE)
                image_tensors = self.image_encoder(image_tensors)
                image_tensors = image_tensors.view(image_tensors.size(0), image_tensors.size(1), -1)

                question_tensors = question_tensors.to(self.DEVICE)
                answer_tensors = answer_tensors.to(self.DEVICE)

                answer_tensors = torch.squeeze(answer_tensors)

                # predictions
                predicted_answer = self.model(image_tensors, question_tensors)

                for i in range(answer_tensors.shape[0]):
                    if torch.argmax(predicted_answer[i]).item() == answer_tensors[i]:
                        accuracy += 1.0

                # if (batch_idx + 1) % self.verbose == 0:
                #     print(f"validation accuracy: {round(accuracy / ((batch_idx + 1) * self.batch_size), 2)}")

            accuracy = round(accuracy / len(self.val_dataset), 2)

            return accuracy

    def train(self):
        ''' Method for training the model.
        '''

        print("\n\n--------------------------------------------------\nTraining the model...\n--------------------------------------------------\n\n")
        train_iteration = 0
        val_iteration = 0
        best_accuracy = 0.0
        loss_history = []
        accuracy_history = []

        for epoch in range(self.num_of_epochs):
            progress_bar = tqdm(total=len(self.train_dataloader))
            self.model.train()

            # if (epoch + 1) % 5 == 0:  # HYPER PARAMETER
            #     self.adjust_learning_rate()

            for batch_idx, (image_tensors, question_tensors, answer_tensors) in enumerate(self.train_dataloader):
                progress_bar.update(1)

                current_step = epoch * len(self.train_dataloader) + batch_idx

                question_tensors = rnn.pack_sequence(question_tensors)
                image_tensors = image_tensors.to(self.DEVICE)
                image_tensors = self.image_encoder(image_tensors)
                image_tensors = image_tensors.view(image_tensors.size(0), image_tensors.size(1), -1)

                question_tensors = question_tensors.to(self.DEVICE)
                answer_tensors = answer_tensors.to(self.DEVICE)

                answer_tensors = torch.squeeze(answer_tensors)

                # predictions
                predicted_answer = self.model(image_tensors, question_tensors)

                # Optimize the model according to the predictions
                loss = self.optimize(predicted_answer, answer_tensors)
                loss_history.append(loss)

                if (current_step + 1) % self.logging_frquency == 0:
                    # print(f"Epoch: {epoch}, Batch: {batch_idx}/{len(self.train_dataloader)} has loss: {loss}")

                    # self.logger.add_scalar("train/loss", loss.item(), train_iteration)
                    train_iteration += 1

            print(f"Epoch: {epoch} has loss: {loss}")

            # Validating
            if (epoch + 1) % self.saving_frequency == 0 or epoch == self.num_of_epochs - 1:
                validation_accuracy = self.validate()
                accuracy_history.append(validation_accuracy)
                print(f"Epoch: {epoch} has validation accuracy {validation_accuracy}")

                # self.logger.add_scalar("valid/accuracy", validation_accuracy, val_iteration)
                val_iteration += 1

                # Remember the best validation accuracy and save a checkpoint
                is_best = validation_accuracy > best_accuracy
                best_accuracy = max(best_accuracy, validation_accuracy)
                self.save_checkpoint(epoch, is_best, best_accuracy)

        # Closing tensorboard logger
        # logger_dir = os.path.join(self.checkpoints_dir, str(self.experiment_id), datetime.now().strftime("%d-%m-%y_%H-%M-%S"))
        # if not os.path.isdir(logger_dir):
        #     os.makedirs(logger_dir)
        # self.logger.export_scalars_to_json(logger_dir + 'tensorboard_summary.json')
        # self.logger.close()

    def custom_collate_func(self, seq_list):
        ''' Custome collate function to stack multiple images/questions/answers in a batch.
        Parameters:
            seq_list: *args; Whatever the dataloader returns.
        Returns:
            image_tensors: pytorch tensor object; Stacked image tensor.
            question_tensors: pytorch tensor object; Stacked question tensor.
            answer_tensors: pytorch tensor object; Stacked answer tensor.
        '''

        image_tensors, question_tensors, answer_tensors = zip(*seq_list)  # will recieve a batch of data points.
        lens = [len(question) for question in question_tensors]
        seq_order = sorted(range(len(lens)), key=lens.__getitem__, reverse=True)

        image_tensors = torch.stack([image_tensors[i] for i in seq_order])
        question_tensors = [question_tensors[i] for i in seq_order]
        answer_tensors = torch.stack([answer_tensors[i] for i in seq_order])

        return image_tensors, question_tensors, answer_tensors

    def initialize_weights(self):
        ''' Method to initialize model weights.
        '''

        for layer in self.model.modules():
            if isinstance(layer, (torch.nn.Conv1d, torch.nn.Linear)):
                try:
                    torch.nn.init.xavier_normal_(layer.weight)

                    try:
                        torch.nn.init.constant_(layer.bias.data, 0)
                    except:
                        pass
                except:
                    pass

    def adjust_learning_rate(self):
        ''' Sets the learning rate to the initial learning rate decayed by 10.
        '''

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = param_group['lr'] / 10

    def save_checkpoint(self, epoch, is_best, best_accuracy):
        ''' Method for sacving a checkpoint.
        Parameters:
            epoch: int; Epoch number for the checkpoint.
            is_best: boolean; Whether the checkpoint has the best accuracy or not.
            best_accuracy: float; Best accuracy value.
        '''

        state = {"epoch": epoch + 1,
                "model_state_dict": self.model.state_dict(),
                "optimizer_state_dict": self.optimizer.state_dict(),
                "best_accuracy": best_accuracy
                }

        checkpoint_path = os.path.join(self.checkpoints_dir, str(self.experiment_id), f"{epoch + 1}_checkpoint.pt")
        torch.save(state, checkpoint_path)

        if is_best:
            copyfile(checkpoint_path, os.path.join(self.checkpoints_dir, str(self.experiment_id), "best.pt"))  # create a separate copy of the best checkpoint

if __name__ == "__main__":
    train_image_dir = "Data/train/images/train10K"
    train_qjson = "Data/train/questions/train_quest_10K.json"
    train_ajson = "Data/train/annotations/train_ann_10K.json"
    collate = True
    train_results_dir = "Data/train/cache/"

    val_image_dir = "Data/val/images/val3K"
    val_qjson = "Data/val/questions/val_quest_3K.json"
    val_ajson = "Data/val/annotations/val_ann_3K.json"
    collate = True
    val_results_dir = "Data/val/cache/"

    saving_frequency = 10
    learning_rate = 0.001
    batch_size = 100
    num_epochs = 2
    num_workers = 5

    exp_runner = ExperimentRunner(train_image_dir, train_qjson, train_ajson, val_image_dir, val_qjson, val_ajson,
                                batch_size, num_epochs, num_workers, collate, train_results_dir, val_results_dir, saving_frequency, learning_rate)

    # exp_runner.save_checkpoint(100, True, 92.4)
    exp_runner.train()


import numpy as np
import pandas as pd
import inspect
import GPy
import scipy.optimize as opt
import bokeh.plotting as bp
from bokeh.embed import components
from bokeh.models import Legend
from bokeh.layouts import widgetbox
from bokeh.layouts import gridplot

import seaborn as sns

from acquisitions import LocalMove, SGD, Explore, Exploit, Phase, UCB, EI
from decisions import Propto, Softmax

"""
Get mean/variances of GP posterior predictive given a kernel and observations
Parameters:
    kernel: GPy kernel
    observed_x: list of observed actions
    observed_y: list of observed function values
    all_x: range over possible actions
"""

def gp(kernel, observed_x, observed_y, all_x):

    if len(observed_x) == 0:
        return np.zeros(len(all_x))[:,None], np.ones(len(all_x))[:,None]
    else:
        X = observed_x[:,None]
        Y = observed_y[:,None]
        m = GPy.models.GPRegression(X, Y, kernel)
        m.Gaussian_noise.variance = .00001
        mean, var = m.predict(all_x[:,None])
        return mean, var

"""
Get means and variances for each trial given a kernel and a sequence of actions
Parameters:
    kernel: GPy kernel to be used
    function: Function on which the given responses were made
    actions: sequence of actions given by participant
"""

def get_mean_var(kernel, function, actions):

    all_mean = []
    all_var = []
    all_x = np.arange(len(function))
    for i in range(len(actions)):
        observed_x = np.array(actions[:i])
        if len(observed_x) == 0:
            observed_y = []
        else:
            observed_y = function[observed_x]
        mean, var = gp(kernel, observed_x, observed_y, all_x)
        all_mean.append(mean.ravel())
        all_var.append(var.ravel())
    return np.array(all_mean), np.array(all_var)

"""
Get utilities for all actions on all trials given a sequence actions
Parameters:
    list actions: sequence of actions given by participant
    acquisition_type: acquisition function class
    list acq_params: list of paramaters to pass to acquisition function
    array all_means: trials x actions array of GP means. Pass None if acquisition
                     function does not use a GP
    array all_vars: trials x actions array of GP variances. Pass None if
                    acquisition function does not use a GP
"""

def get_utility(all_x, actions, rewards, acquisition_type, kernel, acq_params, all_means = [], all_vars = [], replace = True):

    utility_data = pd.DataFrame()
    for i in range(len(actions)):
        trial = 'trial_' + str(i + 1)
        if len(all_means) == 0:
            mean = None
        else:
            mean = all_means[i]
        if len(all_vars) == 0:
            var = None
        else:
            var = all_vars[i]

        if i > 0:
            last_x = actions[i - 1]
            last_y = rewards[i - 1]
        else:
            last_x = None
            last_y = None
        if i > 1:
            second_last_x = actions[i - 2]
            second_last_y = rewards[i - 2]
            idx = [j for j in range(len(actions[:i])) if actions[j] != last_x]
            if len(idx) > 0:
                unique_second_last_x = actions[idx[-1]]
                unique_second_last_y = rewards[idx[-1]]
            else:
                unique_second_last_x = None
                unique_second_last_y = None
                second_last_x = None
                second_last_y = None
        else:
            unique_second_last_x = None
            unique_second_last_y = None
            second_last_x = None
            second_last_y = None

        args = {'all_x': all_x, 'mean': mean, 'var': var, 'trial': i,
                'last_x': last_x, 'last_y': last_y,
                'second_last_x': second_last_x, 'second_last_y': second_last_y,
                'unique_second_last_x': unique_second_last_x, 'unique_second_last_y': unique_second_last_y,
                'ntrials': len(actions), 'trial': i, 'actions': actions[:i], 'rewards': rewards[:i],
                'kernel': kernel}
        acq_arg_names = list(inspect.signature(acquisition_type.__init__).parameters.keys())
        acq_args = {arg_name: args[arg_name] for arg_name in args.keys() if arg_name in acq_arg_names}
        acquisition = acquisition_type(**acq_args)
        utility = acquisition(*acq_params)
        if not replace:
            utility = [utility[j] if j not in actions[:i] else np.NaN for j in range(len(utility))]
        utility_data[trial] = utility
    return utility_data

"""
Get log likelihoods for all actions on all trials given a sequence actions
Parameters:
    DataFrame utility_data: action x trial data frame of utilities
    list actions: sequence of actions given by participant
    decision_type: decision function class
    list dec_params: list of paramaters to pass to acquisition function
"""

def get_likelihood(utility_data, actions, decision_type, dec_params):

    likelihood_data = pd.DataFrame()
    for i in range(len(actions)):
        trial = utility_data.columns[i]
        args = {'trial': i, 'x2': actions[i - 1] if i > 0 else None}
        dec_arg_names = list(inspect.signature(decision_type.__init__).parameters.keys())
        dec_args = {arg_name: args[arg_name] for arg_name in args.keys() if arg_name in dec_arg_names}
        decision = decision_type(**dec_args)
        utility = utility_data[trial]
        likelihood = np.log(decision(utility, *dec_params))

        likelihood_data[trial] = likelihood
    likelihood_data.index = utility_data.index
    return likelihood_data

"""
Get log likelihood given a set of actions and the log likelihoods associated with
a particular model
Parameters:
    list actions: sequence of actions given by participant
    DataFrame likelihood_data: action x trial data frame of log likelihoods
"""

def joint_log_likelihood(actions, likelihood_data):

    return np.nansum([likelihood_data[likelihood_data.columns[i]][actions[i]] for i in range(len(actions))])

"""
Find MLE acquisition and decision function parameters given a sequence of actions. Returns
a dictionary including all utilities/likelihoods across trials and values for model comparison.
Parameters:
    list actions: sequence of actions given by participant
    acquisition_type: acquisition function class
    decision_type: decision function class
    array all_means: trials x actions array of GP means. Pass None if acquisition
                     function does not use a GP
    array all_vars: trials x actions array of GP variances. Pass None if
                    acquisition function does not use a GP
"""

def fit_strategy(all_x, actions, rewards, acquisition_type, decision_type, replace, kernel, all_means = [], all_vars = [], method = 'DE', restarts = 5):

    init_acq_params = acquisition_type.init_params
    init_dec_params = decision_type.init_params
    nacq_params = len(init_acq_params)

    if nacq_params == 0:
        utility_data = get_utility(all_x, actions, rewards, acquisition_type, [], kernel, all_means = all_means, all_vars = all_vars, replace = replace)
        def obj(params):
            likelihood_data = get_likelihood(utility_data, actions, decision_type, params)
            jll = joint_log_likelihood(actions, likelihood_data)
            return -jll

    else:
        def obj(params):
            acq_params = params[:nacq_params]
            dec_params = params[nacq_params:]
            utility_data = get_utility(all_x, actions, rewards, acquisition_type, acq_params, kernel, all_means, all_vars, replace)
            likelihood_data = get_likelihood(utility_data, actions, decision_type, dec_params)
            jll = joint_log_likelihood(actions, likelihood_data)
            return -jll

    init_params = init_acq_params + init_dec_params
    bounds = acquisition_type.bounds + decision_type.bounds
    fun = np.inf
    final_x = None
    for i in range(restarts):
        if method == 'DE':
            x = opt.differential_evolution(obj, bounds)
        else:
            x = opt.minimize(obj, init_params, method = method, bounds = bounds)
        if x.fun < fun:
            fun = x.fun
            final_x = x
    print (final_x)
    params = final_x.x
    acq_params = params[:nacq_params]
    dec_params = params[nacq_params:]
    if nacq_params > 0:
        utility_data = get_utility(all_x, actions, rewards, acquisition_type, acq_params, kernel, all_means, all_vars, replace)
    likelihood_data = get_likelihood(utility_data, actions, decision_type, dec_params)
    n = len(actions)
    k = len(params)
    data = {'acquisition': acquisition_type.__name__, 'acquisition_params': acq_params.tolist(),
            'decision': decision_type.__name__, 'decision_params': dec_params.tolist(),
            'all_utilities': utility_data.to_dict(), 'all_likelihoods': likelihood_data.to_dict(),
            'MLE_likelihood': -fun, 'AIC': -2 * -fun + 2 * k, 'BIC': -2 * -fun + np.log(n) * k}
    random_ic = -2 * (np.log(1. / len(all_x)) * n)
    data['pseudo_r2'] = 1 - (data['AIC'] / random_ic)
    data['approx_bf'] = np.exp(random_ic - data['BIC'])
    return data

"""
fits combinations of kernels and acquisition/decision function strategies to the actions
of one participant
Parameters:
    list actions: sequence of actions given by participant
    Series function: Function on which the given responses were made
    list training_functions: List of functions on which to fit kernel parameters
    list kernels: List of GPy kernels to fit to actions
    list strategies: list of (acquisition_type, decision_type) tuples to fit to actions
"""

def fit_participant(participant_id, goal, actions, function, function_samples, kernels, strategies, method = 'DE', restarts = 5):

    all_x = np.arange(len(function))[:,None]
    stacked_function_samples_x = np.hstack([np.arange(len(f)) for f in function_samples])[:,None]
    stacked_function_samples_y = np.hstack(function_samples)[:,None]
    rewards = [function[a] for a in actions]
    if goal == 'find_max_last':
        replace = False
    else:
        replace = True
    all_data = []

    gp_strategies = [strategy for strategy in strategies if strategy[0].isGP]
    non_gp_strategies = [strategy for strategy in strategies if not strategy[0].isGP]

    for i in range(len(kernels)):
        kernel = kernels[i]
        m = GPy.models.GPRegression(X = stacked_function_samples_x, Y = stacked_function_samples_y, kernel = kernel)
        m.optimize()
        all_means, all_vars = get_mean_var(kernel, function, actions)
        for acquisition_type, decision_type in gp_strategies:
            data = fit_strategy(all_x, actions, rewards, acquisition_type, decision_type, replace, all_means = all_means, all_vars = all_vars, method = method, restarts = restarts)
            data['kernel'] = type(kernel).__name__
            data['all_means'] = {'trial_' + str(i + 1): all_means[i].ravel().tolist() for i in range(len(all_means))}
            data['all_vars'] = {'trial_' + str(i + 1): all_vars[i].ravel().tolist() for i in range(len(all_vars))}
            all_data.append(data)

    for acquisition_type, decision_type in non_gp_strategies:
        data = fit_strategy(all_x, actions, rewards, acquisition_type, decision_type, replace, all_means = [], all_vars = [], method = method, restarts = restarts)
        all_data.append(data)

    participant_data = {'id': participant_id, 'function': function.ravel().tolist(), 'actions': actions, 'rewards': rewards}
    participant_data['models'] = all_data
    return participant_data

def add_viz_data(data):

    for i in range(len(data)):
        participant = data[i]
        colormap = np.array(sns.color_palette("hls", len(participant['models']) + 1).as_hex())
        plots = []
        for trial in range(1, len(participant['models'][0]['all_utilities'].keys()) + 1):
            u_plot = bp.figure(title = "Utility of Next Action", plot_width = 700, plot_height = 400, tools = "")
            u_plot.toolbar.logo = None
            l_plot = bp.figure(title = "Log Likelihood of Next Action", plot_width = 700, plot_height = 400, tools = "")
            l_plot.toolbar.logo = None
            legend_items = []

            actions = participant['actions'][:trial - 1]
            rewards = participant['rewards'][:trial - 1]
            if len(rewards) == 0:
                min_reward = 0
                max_reward = 0
            else:
                min_reward = np.min(rewards)
                max_reward = np.max(rewards)
            next_action = participant['actions'][trial - 1]
            if len(actions) > 0:
                u_plot.circle(actions, rewards, color = colormap[-1])
                a = l_plot.circle(actions, rewards, color = colormap[-1])
                legend_items.append(('Actions', [a]))
            utilities = [[model['all_utilities']['trial_' + str(trial)][k] for k in range(len(model['all_utilities']['trial_' + str(trial)]))] for model in participant['models']]
            likelihoods = [[model['all_likelihoods']['trial_' + str(trial)][k] for k in range(len(model['all_likelihoods']['trial_' + str(trial)]))] for model in participant['models']]
            u_valmin = min(np.nanmin(np.array(utilities).ravel()), min_reward)
            u_valmax = max(np.nanmax(np.array(utilities).ravel()), max_reward)
            l_valmin = min(np.nanmin([l for l in np.array(likelihoods).ravel() if not np.isinf(l)]), min_reward) - .01
            l_valmax = max(np.nanmax(np.array(likelihoods).ravel()), max_reward) + .01

            u_plot.vbar(x = next_action, width = 0.1, bottom = u_valmin, top = u_valmax, color = "black")
            b = l_plot.vbar(x = next_action, width = 0.1, bottom = l_valmin, top = l_valmax, color = "black")
            legend_items.append(('Next Action             | p(Next)   | p(1,...,Next - 1)', [b]))
            for j in range(len(utilities)):
                utility = utilities[j]
                likelihood = likelihoods[j]
                index = [utility[k] for k in range(len(utility)) if not np.isnan(utility[k])]
                utility = [u if not np.isnan(u) and not np.isinf(u) else np.nanmin([ui for ui in utility if not np.isinf(ui)]) for u in utility]
                likelihood = [l if not np.isnan(l) and not np.isinf(l) else np.nanmin([li for li in likelihood if not np.isinf(li)]) for l in likelihood]
                acq = participant['models'][j]['acquisition']
                if 'kernel' in participant['models'][j]:
                    kern = '/' + participant['models'][j]['kernel']
                else:
                    kern = ''
                params = '(' + ','.join([str(np.round(p, 2)) for p in participant['models'][j]['acquisition_params'] + participant['models'][j]['decision_params']]) + ')'

                u_plot.line(list(range(len(utility))), utility, color = colormap[j])
                c = l_plot.line(list(range(len(likelihood))), likelihood, color = colormap[j])

                legend_string = acq + kern + params
                trial_l = np.round(likelihood[next_action], 2)
                likelihood_data = pd.DataFrame(participant['models'][j]['all_likelihoods'])
                total_l = np.round(joint_log_likelihood(actions[:trial], likelihood_data), 2)
                legend_string = legend_string + (' ' * (14 - len(legend_string))) + ' | ' + str(trial_l) + '      | ' + str(total_l)

                legend_items.append((legend_string, [c]))
            u_script, u_div = components(u_plot)
            l_script, l_div = components(l_plot)
            legend = Legend(items = legend_items)
            l_plot.add_layout(legend, 'right')

            gp_plot = bp.figure(title = "Expected Reward", plot_width = 400, plot_height = 400, tools = "")
            gp_plot.toolbar.logo = None
            gp_plot.circle(actions, rewards, color = colormap[-1], legend = "Actions")
            for j in range(len(data[i]['models'])):
                if 'kernel' in data[i]['models'][j].keys():
                    kernel = data[i]['models'][j]['kernel']
                    mean = data[i]['models'][j]['all_means']['trial_' + str(trial)]
                    var = data[i]['models'][j]['all_vars']['trial_' + str(trial)]
                    std = np.sqrt(np.array(var))
                    upper = np.array(mean) + 2 * std
                    lower = np.array(mean) - 2 * std
                    index = np.arange(len(mean))
                    gp_plot.line(index, mean, color = colormap[j], legend = kernel)

                    band_x = np.append(index, index[::-1])
                    band_y = np.append(lower, upper[::-1])
                    gp_plot.patch(band_x, band_y, color = colormap[j], fill_alpha = 0.2)
            gp_plot.legend.location = "top_right"
            f_script, f_div = components(gp_plot)

            p = gridplot([[gp_plot, u_plot], [l_plot]], toolbar_location=None)
            p_script, p_div = components(p)

            plots.append({'u_script': u_script, 'u_div': u_div, 'l_script': l_script, 'l_div': l_div,
                          'f_script': f_script, 'f_div': f_div, 'p_script': p_script, 'p_div': p_div})
        data[i]['plots'] = plots
    return data

def fit_all_participants(results, method = 'DE', restarts = 5):

    data = []
    for index, participant in results.iterrows():
        actions = participant['response']
        function = participant['function']
        function_n = ((function - np.mean(function)) / np.std(function))
        function_samples = list(participant['function_samples'].values())
        function_samples_n = np.array([((f - np.mean(f)) / np.std(f)) for f in function_samples])
        participant_id = participant['somataSessionId']
        function_name = participant['function_name']
        goal = participant['goal']
        score = participant['total_score']

        if function_name == 'pos_linear':
            kernels = [GPy.kern.RBF(1), GPy.kern.Linear(1) + GPy.kern.Bias(1)]
        elif function_name == 'neg_quad':
            kernels = [GPy.kern.RBF(1), GPy.kern.Poly(1, order = 2)]
        elif function_name == 'sinc_compressed':
            kernels = [GPy.kern.RBF(1), GPy.kern.StdPeriodic(1) + GPy.kern.RBF(1)]
        strategies = [(LocalMove, Softmax), (SGD, Softmax), (Phase, Softmax), (UCB, Softmax)]
        strategies = [(LocalMove, Softmax), (SGD, Softmax), (Explore, Softmax), (Exploit, Softmax), (Phase, Softmax), (UCB, Softmax)]
        kernels = [GPy.kern.RBF(1)]

        participant_data = fit_participant(participant_id, goal, actions, function_n, function_samples_n, kernels, strategies, method = method, restarts = restarts)
        participant_data['function_name'] = function_name
        participant_data['goal'] = goal
        participant_data['score'] = score
        data.append(participant_data)
    return data


import numpy as np
from utils import load_obj,split,textprocess
import os
import onnxruntime
from metrics import get_entities
import json

def get_attn_pad_mask(seq_q, seq_k):
    # print(seq_q)
    batch_size = 1
    len_q = len(seq_q[0])
    len_k = len(seq_k[0])
    # eq(zero) is PAD token
    pad_attn_mask = np.array(seq_k)==0  # (batch_size, 1, len_k/len_q) one is masking
    # pad_mask = np.repeat(pad_attn_mask,len_q,axis=0).reshape(1,len_q,len_k)
    return pad_attn_mask
    # return pad_attn_mask.repeat(batch_size, len_q, len_k)  # (batch_size, len_q, len_k)

def softmax(x):
    """
    Compute the softmax function for each row of the input x.

    Arguments:
    x -- A N dimensional vector or M x N dimensional numpy matrix.

    Return:
    x -- You are allowed to modify x in-place
    """
    orig_shape = x.shape

    if len(x.shape) > 1:
        # Matrix
        exp_minmax = lambda x: np.exp(x - np.max(x))
        denom = lambda x: 1.0 / np.sum(x)
        x = np.apply_along_axis(exp_minmax,1,x)
        denominator = np.apply_along_axis(denom,1,x)

        if len(denominator.shape) == 1:
            denominator = denominator.reshape((denominator.shape[0],1))

        x = x * denominator
    else:
        # Vector
        x_max = np.max(x)
        x = x - x_max
        numerator = np.exp(x)
        denominator =  1.0 / np.sum(numerator)
        x = numerator.dot(denominator)

    assert x.shape == orig_shape
    return x

def softmax_mask(logits_tgt, cls_idx, idx_mask):
    logits_tgt = logits_tgt[0,:,:]
    mask = idx_mask[str(cls_idx)]
    length, tgt_num = logits_tgt.shape[0], logits_tgt.shape[1]
    scores_exp = softmax(logits_tgt)
     # this step masks
    scores_exp = scores_exp * mask
    return scores_exp

class DataLoader_test(object):
    def __init__(self, save_dir):
        self.save_dir = save_dir
        self.word2idx = load_obj(self.save_dir + "dict.json")
        self.config = load_obj(self.save_dir + "Config.json")
        self.max_len = self.config["max_len"]

        self.WORD = {int(k):v for k,v in self.config["WORD"].items()}
        self.BOS = self.config["BOS"]
        self.UNK = self.config["UNK"]
        self.PAD = self.config["PAD"]

        assert self.BOS == self.word2idx[self.WORD[self.BOS]]
        assert self.UNK == self.word2idx[self.WORD[self.UNK]]
        assert self.PAD == self.word2idx[self.WORD[self.PAD]]

    def load_sentences(self, sent):
        """Loads sentences and tags from their corresponding files.
            Maps tokens and tags to their indices and stores them in the provided dict d.
        """
        sentence = []

        tokens = split(textprocess(sent))
        sentence.append(self.convert_tokens_to_ids(tokens))
        return tokens, sentence

    def convert_tokens_to_ids(self, tokens):
        sentence = []
        sentence.append(self.BOS)
        for tok in tokens:
            if tok in self.word2idx:
                sentence.append(self.word2idx[tok])
            else:
                sentence.append(self.UNK)
        pad = [self.PAD]*(self.max_len+1 - len(sentence))

        assert len(sentence + pad) == self.max_len+1
        return sentence + pad

class NLU_module():
    def __init__(self, save_dir = "./model_onnx/",model_nm = "transformer_mix.onnx"):
        self.save_dir = save_dir
        self.model_nm = model_nm
        self.Init_model()

    def Init_model(self):
        #init dataloader
        self.data_loader = DataLoader_test(self.save_dir)
        # init model
        self.ort_session = onnxruntime.InferenceSession(self.save_dir+ self.model_nm)
        # init dict
        self.idx2lbl = load_obj(self.save_dir + "idx2lbl.json")
        self.idx2cls  = load_obj(self.save_dir + "idx2cls.json")
        # get valid slot for a specific intent
        self.idx_mask = load_obj(self.save_dir + "idx_mask_onnx.json")

    def Inference(self, input_sentence):
        # read test_sentence
        # input_sentence = '导航到世纪大道一百一十八号'
        tokens, test_data = self.data_loader.load_sentences(input_sentence)

        # run inference
        pred_cls ,pred_lbls = self.test(test_data)

        # merge_slot
        slot = self.merged_slot(tokens, pred_lbls)

        ans = {}
        ans["Input_sentence"] = input_sentence.encode('utf-8').decode('utf-8')
        ans["Raw Labels"] = ' '.join(pred_lbls)
        ans["Intent"] = ''.join(pred_cls)
        ans["Megred Mentions"] = slot

        # return json.dumps(ans, ensure_ascii=False)
        return ans

    def test(self, enc):
        """Evaluate the model on `steps` batches."""
        enc_self_attn_mask = get_attn_pad_mask(enc, enc)
        x = (enc,enc_self_attn_mask)

        # compute ONNX Runtime output prediction
        ort_inputs = {self.ort_session.get_inputs()[i].name: x[i] for i in range(len(x))}
        logits_tgt, logits_clsf = self.ort_session.run(None, ort_inputs)

        pad_num = np.count_nonzero(enc)

        pred_cls = np.argmax(logits_clsf)
        score = logits_clsf[0,pred_cls]

        masked_logits_tgt= softmax_mask(logits_tgt, pred_cls, self.idx_mask)
        tgt_idx = np.argmax(masked_logits_tgt, axis = 1)
        score_tgt = 0

        pred_tags = tgt_idx[:pad_num]

        pred_lbls = []
        for idx in pred_tags:
            pred_lbls.append(self.idx2lbl[str(idx)])
        pred_cls = self.idx2cls[str(pred_cls)]

        return pred_cls ,pred_lbls

    def merged_slot(self,tokens, pred_lbls):
        chunks = get_entities(pred_lbls)
        slot_result = {}
        for chunk in chunks:
            tag, start, end = chunk[0], chunk[1], chunk[2]
            tok = ''.join(tokens[chunk[1]:chunk[2]+1])
            # string = '<{0}>: {1}'.format(tag, tok)
            while tag in slot_result:
                tag += '#'
            slot_result[tag]=tok
        return slot_result

if __name__ == '__main__':
    Module = NLU_module()
    # read test_sentence
    input_sentence = '找个北京的餐馆'
    results = Module.Inference(input_sentence)
    print(results)


input_products = input("Введите товар через запятую: ")
products_list = input_products.split(',')

print(f"Товары с нечётными индексами: {products_list[1::2]}")

"""
Задание 1: Инвентаризация магазина Романа

Описание:

Роман хочет провести инвентаризацию в своем магазине.
Помогите ему написать программу, которая выводит каждый товар с нечетным индексом из списка товаров.

"""


import os
import sys
from typing import List
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as ec
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import Select
from auth import login

def submit_logged_time(browser, duration: str) -> None:
    condition = ec.presence_of_element_located((By.ID, "time_entry_hours"))
    time_entry_field = WebDriverWait(browser, 10).until(condition)
    time_entry_field.send_keys(duration)
    activity_dropdown = Select(
        browser.find_element_by_id("time_entry_activity_id"))
    activity_dropdown.select_by_visible_text('Development')
    submit_btn = browser.find_element_by_name('commit')
    submit_btn.click()

def log_time(browser, issue_id: str, duration: str) -> None:
    url = f"https://apps.mohcc.gov.zw:8084/issues/{issue_id}/time_entries/new"
    login(browser, url)
    submit_logged_time(browser, duration)

def main(args: List[str]) -> None:
    if len(args) < 2:
        print("Expected issue ID and time taken")
        exit(1)
    issue_id, duration = args[1:]
    options = webdriver.FirefoxOptions()
    options.headless = True
    with webdriver.Firefox(service_log_path=os.devnull, options=options) as browser:
        log_time(browser, issue_id, duration)

if __name__ == "__main__":
    main(sys.argv)