import re
import json
from typing import Dict, Optional, Tuple

from root import ROOT_DIR
from scrc.data_classes.court_composition import CourtComposition
from scrc.data_classes.court_person import CourtPerson

from scrc.enums.court_role import CourtRole
from scrc.enums.gender import Gender
from scrc.enums.language import Language
from scrc.enums.political_party import PoliticalParty
from scrc.enums.section import Section

"""
This file is used to extract the judicial persons from decisions sorted by spiders.
The name of the functions should be equal to the spider! Otherwise, they won't be invocated!
Overview of spiders still todo: https://docs.google.com/spreadsheets/d/1FZmeUEW8in4iDxiIgixY4g0_Bbg342w-twqtiIu8eZo/edit#gid=0
"""

def XX_SPIDER(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract the court composition from decisions of the Baurekursgericht of Zurich
    :param header:      the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the court composition
    """
    # Seems to work for: CH_BVGer, mostly TI_Gerichte
    # Does not seem to work for: GE_Gerichte

    header = sections[Section.HEADER]
    language = namespace['language']

    role_regexes = {
        Language.DE:{
            Gender.MALE: {
                CourtRole.JUDGE: [r'Richter(?!in)', r'Einzelrichter(?!in)', r'Schiedsrichter(?!in)'],
                CourtRole.CLERK: [r'Gerichtsschreiber(?!in)']
            },
            Gender.FEMALE: {
                CourtRole.JUDGE: [r'Richterin(nen)?', r'Einzelrichterin(nen)?', r'Schiedsrichterin(nen)?'],
                CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?']
            }
        },
        Language.FR: {
            Gender.MALE: {
                CourtRole.JUDGE: [r'MM?\.(( et|,) Mmes?)? les? Juges?'],
                CourtRole.CLERK: [r'Greffier[^\w\s]*']
            },
            Gender.FEMALE: {
                CourtRole.JUDGE: [r'Mmes? l(a|es) Juges?', r'MMe et MM?\. les? Juges?'],
                CourtRole.CLERK: [r'Greffière.*Mme']
            }
        },
        Language.IT: {
            Gender.MALE: {
                CourtRole.JUDGE: [r'[Gg]iudici'],
                CourtRole.CLERK: [r'[Cc]ancelliere']
            },
            Gender.FEMALE: {
                CourtRole.JUDGE: [r'[Gg]iudice federal'],
                CourtRole.CLERK: [r'[Cc]ancelliera']
            }
        }
    }

    header = header.replace('U R T E I L', 'Urteil')
    header = header.replace('U R TE I L', 'Urteil')
    header = header.replace('URTEIL', 'Urteil')
    header = header.replace('Z W I S C H E N E N T S C H E I D', 'Zwischenentscheid')

    information_start_regex = r'Besetzung|Bundesrichter|Composition( de la Cour:)?|Composizione|Giudic[ie] federal|composta'

    start_pos = re.search(information_start_regex, header)
    if start_pos:
        # split off the first word
        header = header[start_pos.span()[1]:]
    end_pos = {
        Language.DE:
            re.search(r'.(?=(1.)?(Partei)|(Verfahrensbeteiligt))', header)
            or re.search('Urteil vom',header)
            or re.search(r'[Ii]n Sachen|Gegenstand', header)
            or re.search(r'\w{2,}\.', header),
        Language.FR:
            re.search(r'.(?=(Parties|Participant))', header)
            or re.search(r'Greffi[eè]re? M(\w)*\.\s\w*.', header),
        Language.IT:
            re.search(r'.(?=(Parti)|(Partecipant))', header)
            or re.search(r'[Cc]ancellier[ae]:?\s?\w*.', header) or re.search(r'\w{2,}\.', header),
    }
    end_pos = end_pos[language]
    if end_pos:
        header = header[:end_pos.span()[1]-1]

    composition = CourtComposition()
    composition = find_composition(header, role_regexes[language], namespace)
    if not composition: return None
    if len(composition.clerks)>5 or len(composition.judges) > 5:
        raise ValueError(f'Looks like the compostion of {namespace["id"]} was not correctly recognized and fulltext was included')
    return composition if composition else None

# check if court got assigned shortcut: SELECT count(*) from de WHERE lower_court is not null and lower_court <> 'null' and lower_court::json#>>'{court}'~'[A-Z0-9_]{2,}';
def CH_BGer(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract judicial persons from decisions of the Federal Supreme Court of Switzerland
    :param sections:    the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the sections dict
    """

    #ToDo: Use the functions extracted at the end of the file

    header = sections[Section.HEADER]

    information_start_regex = r'Besetzung|Bundesrichter|Composition( de la Cour:)?|Composizione|Giudic[ie] federal|composta'
    role_regexes = {
        Gender.MALE: {
            CourtRole.JUDGE: [r'Bundesrichter(?!in)', r'MM?\.(( et|,) Mmes?)? les? Juges?( fédéra(l|ux))?',
                              r'[Gg]iudici federali'],
            CourtRole.CLERK: [r'Gerichtsschreiber(?!in)', r'Greffier[^\w\s]*', r'[Cc]ancelliere']
        },
        Gender.FEMALE: {
            CourtRole.JUDGE: [r'Bundesrichterin(nen)?', r'Mmes? l(a|es) Juges? (fédérales?)?',
                              r'MMe et MM?\. les? Juges?( fédéra(l|ux))?', r'[Gg]iudice federal'],
            CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?', r'Greffière.*Mme', r'[Cc]ancelliera']
        }
    }

    skip_strings = get_skip_strings()

    start_pos = re.search(information_start_regex, header)
    if start_pos:
        header = header[start_pos.span()[0]:]
    end_pos = {
        Language.DE: re.search(r'.(?=(1.)?(Partei)|(Verfahrensbeteiligt))', header) or re.search('Urteil vom',
                                                                                header) or re.search(
            r'Gerichtsschreiber(in)?\s\w*.', header) or re.search(r'[Ii]n Sachen', header) or re.search(r'\w{2,}\.',
                                                                                header),
        Language.FR: re.search(r'.(?=(Parties|Participant))', header) or re.search(r'Greffi[eè]re? M(\w)*\.\s\w*.', header),
        Language.IT: re.search(r'.(?=(Parti)|(Partecipant))', header) or re.search(r'[Cc]ancellier[ae]:?\s\w*.',
                                                                            header) or re.search(r'\w{2,}\.', header),
    }
    end_pos = end_pos[namespace['language']]
    if end_pos:
        header = header[:end_pos.span()[1] - 1]

    header = header.replace(';', ',')
    header = header.replace('Th', '')
    header = header.replace(' und ', ', ')
    header = header.replace(' et ', ', ')
    header = header.replace(' e ', ', ')
    header = header.replace('MMe', 'Mme')
    header = re.sub(r'(?<!M)(?<!Mme)(?<!MM)(?<!\s\w)\.', ', ', header)
    header = re.sub(r'MM?\., Mme', 'M. et Mme', header)
    header = re.sub(r'Mmes?, MM?\.', 'MMe et M', header)
    header = header.replace('federali, ', 'federali')
    besetzungs_strings = header.split(',')

    def prepare_french_name_and_find_gender(person: CourtPerson) -> CourtPerson:
        """Removes the prefix from a french name and sets gender"""
        if person.name.find('M. ') > -1:
            person.name = person.name.replace('M. ', '')
            person.gender = Gender.MALE
        elif person.name.find('Mme') > -1:
            person.name = person.name.replace('Mme ', '')
            person.gender = Gender.FEMALE
        return person

    besetzung = CourtComposition()
    current_role = CourtRole.JUDGE
    last_person: CourtPerson = None
    last_gender = Gender.MALE

    for text in besetzungs_strings:
        text = text.strip()
        if len(text) == 0 or text in skip_strings[namespace['language']]:
            continue
        if re.search(r'(?<![Vv]ice-)[Pp]r[äée]sid',
                     text):  # Set president either to the current person or the last Person (case 1: Präsident Niklaus, case 2: Niklaus, Präsident)
            if last_person:
                besetzung.president = last_person
                continue
            else:
                text = text.split()[-1]
                president, _ = match_person_to_database(CourtPerson(text, last_gender, [], current_role), last_gender)
                besetzung.president = president
        has_role_in_string = False
        matched_gender_regex = False
        for gender in role_regexes:  # check for male and female all roles
            if matched_gender_regex:
                break
            role_regex = role_regexes[gender]
            for regex_key in role_regex:  # check each role
                regex = '|'.join(role_regex[regex_key])
                role_pos = re.search(regex, text)
                if role_pos:  # Found a role regex
                    last_role = current_role
                    current_role = regex_key
                    name_match = re.search(r'[A-Z][A-Za-z\-éèäöü\s]*(?= Urteil)|[A-Z][A-Za-z\-éèäöü\s]*(?= )',
                                           text[role_pos.span()[1] + 1:])
                    name = name_match.group() if name_match else text[role_pos.span()[1] + 1:]
                    if len(name.strip()) == 0:
                        if (last_role == CourtRole.CLERK and len(besetzung.clerks) == 0) or (last_role == CourtRole.JUDGE and len(besetzung.judges) == 0):
                            break

                        last_person_name = besetzung.clerks.pop().name if (last_role == CourtRole.CLERK) else besetzung.judges.pop().name # rematch in database with new role
                        last_person_new_match, _ = match_person_to_database(CourtPerson(last_person_name, court_role=current_role), gender)
                        if current_role == CourtRole.JUDGE:
                            besetzung.judges.append(last_person_new_match)
                        elif current_role == CourtRole.CLERK:
                            besetzung.clerks.append(last_person_new_match)
                    if namespace['language'] == Language.FR:
                        person = prepare_french_name_and_find_gender(CourtPerson(name))
                        gender = person.gender or gender
                        person.court_role = current_role
                    else:
                        person = CourtPerson(name, court_role=current_role)
                    matched_person, _ = match_person_to_database(person, gender)
                    if current_role == CourtRole.JUDGE:
                        besetzung.judges.append(matched_person)
                    elif current_role == CourtRole.CLERK:
                        besetzung.clerks.append(matched_person)
                    last_person = matched_person
                    last_gender = matched_person.gender
                    has_role_in_string = True
                    matched_gender_regex = True
                    break
        if not has_role_in_string:  # Current string has no role regex match
            if namespace['language'] == Language.FR:
                person = prepare_french_name_and_find_gender(CourtPerson(text))
                last_gender = person.gender or last_gender
            else:
                person = CourtPerson(text, last_gender, current_role)
            name_match = re.search(
                r'[A-Z][A-Za-z\-éèäöü\s]*(?= Urteil)|[A-Z][A-Za-z\-éèäöü\s]*(?= )|[A-Z][A-Za-z\-éèäöü\s]*', person.name)
            if not name_match:
                continue
            name = name_match.group()
            person.court_role = current_role
            matched_person, _ = match_person_to_database(person, last_gender)
            if current_role == CourtRole.JUDGE:
                besetzung.judges.append(matched_person)
            elif current_role == CourtRole.CLERK:
                besetzung.clerks.append(matched_person)
            last_person = person
    return besetzung

def ZG_Verwaltungsgericht(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract the court composition from decisions of the Verwaltungsgericht of Zug
    :param header:      the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the court composition of a decision
    """

    header = sections[Section.HEADER]

    role_regexes = {
        Gender.MALE: {
            CourtRole.JUDGE: [r'Richter(?!in)', r'Einzelrichter(?!in)', r'Schiedsrichter(?!in)'],
            CourtRole.CLERK: [r'Gerichtsschreiber(?!in)']
        },
        Gender.FEMALE: {
            CourtRole.JUDGE: [r'Richterin(nen)?', r'Einzelrichterin(nen)?', r'Schiedsrichterin(nen)?'],
            CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?']
        }
    }

    # regularize different forms of words
    header = header.replace('U R T E I L', 'Urteil')
    header = header.replace('U R TE I L', 'Urteil')
    header = header.replace('URTEIL', 'Urteil')
    header = header.replace('Z W I S C H E N E N T S C H E I D', 'Zwischenentscheid')

    information_start_regex = r'Mitwirkende|Einzelrichter'
    start_pos = re.search(information_start_regex, header)
    if start_pos:
        # split off the first word
        header = header[start_pos.span()[1]:]

    information_end_regex = r'Urteil|Zwischenentscheid'
    end_pos = re.search(information_end_regex, header)
    if end_pos:
        header = header[:end_pos.span()[1]]
        # split off the last word
        header = header.rsplit(' ', 1)[0]

    # this puts a comma before every judicial role
    for gender in role_regexes:
        for regex_key in role_regexes[gender]:
            for regex in role_regexes[gender][regex_key]:
                regex_group = r"(" + regex + r")"
                header = re.sub(regex_group, r', \1', header)

    composition = CourtComposition()
    composition = find_composition(header, role_regexes, namespace)
    return composition if composition else None

def ZH_Baurekurs(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract the court composition from decisions of the Baurekursgericht of Zurich
    :param header:      the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the court composition
    """

    header = sections[Section.HEADER]

    role_regexes = {
        Gender.MALE: {
            CourtRole.JUDGE: [r'Abteilungspräsident(?!in)', r'Baurichter(?!in)', r'Abteilungsvizepräsident(?!in)',
                              r'Ersatzrichter(?!in)'],
            CourtRole.CLERK: [r'Gerichtsschreiber(?!in)']
        },
        Gender.FEMALE: {
            CourtRole.JUDGE: [r'Abteilungspräsidentin(nen)?', r'Baurichterin(nen)?', r'Abteilungsvizepräsidentin(nen)?',
                              r'Ersatzrichterin(nen)?'],
            CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?']
        }
    }

    information_start_regex = r'Mitwirkende'
    start_pos = re.search(information_start_regex, header)
    if start_pos:
        # split off the first word
        header = header[start_pos.span()[1]:]

    information_end_regex = r'in Sachen'
    end_pos = re.search(information_end_regex, header)
    if end_pos:
        header = header[:end_pos.span()[1]]
        # split off the last two words
        header = header.rsplit(' ', 2)[0]

    composition = CourtComposition()
    composition = find_composition(header, role_regexes, namespace)
    return composition if composition else None

def ZH_Obergericht(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract the court composition from decisions of the Obergericht of Zurich
    :param header:      the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the court composition
    """

    header = sections[Section.HEADER]

    role_regexes = {
        Gender.MALE: {
            CourtRole.JUDGE: [r'Oberrichter(?!in)', r'Ersatzoberrichter(?!in)', r'Ersatzrichter(?!in)',
                              r'Kassationsrichter(?!in)', r'Vizepräsident(?!in)', r'Bezirksrichter(?!in)',
                              r'Handelsrichter(?!in)', r'Einzelrichter(?!in)'],
            CourtRole.CLERK: [r'Gerichtsschreiber(?!in)', r'Sekretär(?!in)']
        },
        Gender.FEMALE: {
            CourtRole.JUDGE: [r'Oberrichterin(nen)?', r'Ersatzoberrichterin(nen)?', r'Ersatzrichterin(nen)?',
                              r'Kassationsrichterin(nen)?', r'Vizepräsidentin(nen)?', r'Bezirksrichterin(nen)?',
                              r'Handelsrichterin(nen)?', r'Einzelrichterin(nen)?'],
            CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?', r'Sekretärin(nen)?']
        }
    }

    information_start_regex = r'Mitwirkend'
    start_pos = re.search(information_start_regex, header)
    if start_pos:
        # split off the first word
        header = header[start_pos.span()[1]:]

    information_end_regex = r'Zirkulationsbeschluss vom|Beschluss vom|Urteil vom|Verfügung vom|Beschluss und|in Sachen'
    end_pos = re.search(information_end_regex, header)
    if end_pos:
        header = header[:end_pos.span()[1]]
        # split off the last two words
        header = header.rsplit(' ', 2)[0]

    composition = CourtComposition()
    composition = find_composition(header, role_regexes, namespace)
    return composition if composition else None

def ZH_Sozialversicherungsgericht(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract the court composition from decisions of the Sozialversicherungsgericht of Zurich
    :param header:      the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the court composition
    """

    header = sections[Section.HEADER]

    role_regexes = {
        Gender.MALE: {
            CourtRole.JUDGE: [r'Sozialversicherungsrichter(?!in)', r'Ersatzrichter(?!in)', r'Schiedsrichter(?!in)'],
            CourtRole.CLERK: [r'Gerichtsschreiber(?!in)', r'Gerichtssekretär(?!in)']
        },
        Gender.FEMALE: {
            CourtRole.JUDGE: [r'Sozialversicherungsrichterin(nen)?', r'Ersatzrichterin(nen)?',
                              r'Schiedsrichterin(nen)?'],
            CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?', r'Gerichtssekretärin(nen)?']
        }
    }

    information_start_regex = r'Mitwirkende|Einzelrichter|Kammer'
    start_pos = re.search(information_start_regex, header)
    if start_pos:
        # split off the first word
        header = header[start_pos.span()[1]:]

    information_end_regex = r'Urteil vom|in Sachen'
    end_pos = re.search(information_end_regex, header)
    if end_pos:
        header = header[:end_pos.span()[1]]
        # split off the last two words
        header = header.rsplit(' ', 2)[0]

    # this puts a comma before every judicial role
    for gender in role_regexes:
        for regex_key in role_regexes[gender]:
            for regex in role_regexes[gender][regex_key]:
                regex_group = r"(" + regex + r")"
                header = re.sub(regex_group, r', \1', header)

    composition = CourtComposition()
    composition = find_composition(header, role_regexes, namespace)
    return composition if composition else None

def ZH_Steuerrekurs(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract the court composition from decisions of the Steuerrekursgericht of Zurich
    :param header:      the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the court composition
    """

    header = sections[Section.HEADER]

    role_regexes = {
        Gender.MALE: {
            CourtRole.JUDGE: [r'Abteilungspräsident(?!in)', r'Abteilungsvizepräsident(?!in)', r'Steuerrichter(?!in)',
                              r'Ersatzrichter(?!in)', r'Einzelrichter(?!in)'],
            CourtRole.CLERK: [r'Gerichtsschreiber(?!in)', r'Sekretär(?!in)']
        },
        Gender.FEMALE: {
            CourtRole.JUDGE: [r'Abteilungspräsidentin(nen)?', r'Abteilungsvizepräsidentin(nen)?',
                              r'Steuerrichterin(nen)?', r'Ersatzrichterin(nen)?', r'Einzelrichterin(nen)?'],
            CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?', r'Sekretärin(nen)?']
        }
        ,
        Gender.UNKNOWN: {
            CourtRole.JUDGE: [r'Ersatzmitglied(er)?', r'Mitglied(er)?']
        }
    }

    information_start_regex = r'Mitwirkend'
    start_pos = re.search(information_start_regex, header)
    if start_pos:
        # split off the first word
        header = header[start_pos.span()[1]:]

    information_end_regex = r'In Sachen|in Sachen'
    end_pos = re.search(information_end_regex, header)
    if end_pos:
        header = header[:end_pos.span()[1]]
        # split off the first word
        header = header.rsplit(' ', 2)[0]

    composition = CourtComposition()
    composition = find_composition(header, role_regexes, namespace)
    return composition if composition else None

def ZH_Verwaltungsgericht(sections: Dict[Section, str], namespace: dict) -> CourtComposition:
    """
    Extract the court composition from decisions of the Verwaltungsgericht of Zurich
    :param header:      the dict containing the sections per section key
    :param namespace:   the namespace containing some metadata of the court decision
    :return:            the court composition
    """

    header = sections[Section.HEADER]

    role_regexes = {
        Gender.MALE: {
            CourtRole.JUDGE: [r'Abteilungspräsident(?!in)', r'Abteilungsvizepräsident(?!in)',
                              r'Verwaltungsrichter(?!in)', r'Ersatzrichter(?!in)'],
            CourtRole.CLERK: [r'Gerichtsschreiber(?!in)', r'Gerichtssekretär(?!in)', r'Sekretär(?!in)']
        },
        Gender.FEMALE: {
            CourtRole.JUDGE: [r'Abteilungspräsidentin(nen)?', r'Abteilungsvizepräsidentin(nen)?',
                              r'Verwaltungsrichterin(nen)?', r'Ersatzrichterin(nen)?'],
            CourtRole.CLERK: [r'Gerichtsschreiberin(nen)?', r'Gerichtssekretärin(nen)?', r'Sekretärin(nen)?']
        }
    }

    information_start_regex = r'Mitwirkend'
    start_pos = re.search(information_start_regex, header)
    if start_pos:
        # split off the first word
        header = header[start_pos.span()[1]:]

    information_end_regex = r'In Sachen|in Sachen'
    end_pos = re.search(information_end_regex, header)
    if end_pos:
        header = header[:end_pos.span()[1]]
        # split off the last two words
        header = header.rsplit(' ', 2)[0]

    composition = CourtComposition()
    composition = find_composition(header, role_regexes, namespace)
    return composition if composition else None

def get_composition_strings(header: str) -> list:
    """
    Modifies the header of a decision and turns it into a list
    :param header:  the header of a decision
    :return:        a list of composition_strings
    """
    # repeating commas aren't necessary
    header = re.sub(r', *, *', ', ', header)
    # trying to join words that are split over two lines
    header = re.sub(r'- *, *', '', header)
    header = header.replace('- ', '')
    # regularize different forms to denote the Vorsitz
    header = header.replace('(Vorsitz)', 'Vorsitz')
    header = header.replace('Vorsitzender', 'Vorsitz')
    header = header.replace('Vorsitzende', 'Vorsitz')
    header = header.replace('Vorsitz', ', Vorsitz, ')
    # these word separators aren't relevant here
    header = header.replace(':', '')
    # a semicolon can be treated as a comma here
    header = header.replace(';', ',')
    # der & die aren't relevant for this task
    header = re.sub(r'\bder\b', '', header)
    header = re.sub(r'\bdie\b', '', header)
    # und & sowie separte different people
    header = header.replace(' und', ', ')
    header = header.replace(' et ', ', ')
    header = header.replace(' e ', ', ')
    header = header.replace(' sowie', ', ')
    # academic degrees presumably aren't relevant for this task
    header = header.replace('lic.', '')
    header = header.replace('iur.', '')
    header = header.replace('Dr.', '')
    header = header.replace('Prof.', '')
    header = header.replace('MLaw ', '')
    header = header.replace('M.A.', '')
    header = header.replace('HSG ', '')
    header = header.replace('PD ', '')
    header = header.replace('a.o.', '')
    header = header.replace('LL.M.', '')
    header = header.replace('LL. M.', '')
    header = header.replace('LLM ', '')
    # delete multiple spaces
    header = header.strip()
    header = re.sub(' +', ' ', header)
    # neither is this relevant
    header = header.replace(' als Einzelrichterin', '')
    header = header.replace(' als Einzelrichter', '')
    header = header.replace(';', ',')
    header = header.replace('Th', '')
    header = header.replace('MMe', 'Mme')
    header = re.sub(r'(?<!M)(?<!Mme)(?<!MM)(?<!\s\w)\.', ', ', header)
    header = re.sub(r'MM?\., Mme', 'M. et Mme', header)
    header = re.sub(r'Mmes?, MM?\.', 'MMe et M', header)
    header = header.replace('federali, ', 'federali')
    # uncomment to debug
    # print(header)
    return header.split(',')

def get_skip_strings() -> dict:
    """
    :return:    strings which should be skipped when extracting the court composition
    """
    return {
        Language.DE: ['Einzelrichter', 'Konkurskammer', 'Beschwerdeführerin', 'Beschwerdeführer', 'Kläger',
                      'Berufungskläger'],
        Language.FR: ['Juge suppléant', 'en qualité de juge unique'],
        Language.IT: ['Giudice supplente', 'supplente', 'statuendo']
    }

def match_person_to_database(person: CourtPerson, current_gender: Gender) -> Tuple[CourtPerson, bool]:
    """"Matches a name of a given role to a person from personal_information.json if found it adds the party and selects more accurate CourtRole"""
    personal_information_database = json.loads((ROOT_DIR / "legal_info" / "personal_information.json").read_text())

    results = []
    name = person.name.replace('.', '').strip()
    split_name = name.split()
    initial = False
    if len(split_name) > 1:
        initial = next((x for x in split_name if len(x) == 1), None)
        split_name = list(filter(lambda x: len(x) > 1, split_name))
    if person.court_role.value in personal_information_database: # The court_role of the person is in the database
        for subcategory in personal_information_database[person.court_role.value]:
            for cat_id in personal_information_database[person.court_role.value][subcategory]:
                for db_person in personal_information_database[person.court_role.value][subcategory][cat_id]:
                    if set(split_name).issubset(set(db_person['name'].split())): # The name of the person is in the database, then match it
                        if not initial or re.search(rf'\s{initial.upper()}\w*', db_person['name']):
                            person.name = db_person['name']
                            if db_person.get('gender'):
                                person.gender = Gender(db_person['gender'])
                            if db_person.get('party'):
                                person.party = PoliticalParty(db_person['party'])
                            results.append(person)
    else:
        for existing_role in personal_information_database: # Go through all roles in the database to find the person
            temp_person = CourtPerson(person.name, court_role=CourtRole(existing_role))
            db_person, match = match_person_to_database(temp_person, current_gender)
            if match:
                results.append(db_person)
    if len(results) == 1:
        if not results[0].gender:
            results[0].gender = current_gender
        return person, True
    return person, False

def find_composition(header: str, role_regexes: dict, namespace: dict) -> CourtComposition:
    """
    Find the court composition in the header of a decision
    :param header:          the string containing the header
    :param role_regexes:    the regexes for the court person roles
    :param namespace:       the namespace containing some metadata of the court decision
    :return:                the court composition
    """
    skip_strings = get_skip_strings()
    composition_strings = get_composition_strings(header)
    composition = CourtComposition()
    current_role = CourtRole.JUDGE
    last_person: CourtPerson = None
    person: CourtPerson = None
    last_gender = Gender.UNKNOWN

    # check if any of the role regexes can be found in the header.
    # if there are none, we can exit this function.
    def any_matches():
        for gender in role_regexes:
            role_regex = role_regexes[gender]
            for regex_key in role_regex:
                regex = '|'.join(role_regex[regex_key])
                if re.search(regex, header):
                    return True

    if not any_matches():
        return

    for text in composition_strings:
        text = text.strip()
        # delete the last character if it's a dot following a lower-case character
        if re.search(r'[a-z]\.$', text):
            text = text[:-1]
        if len(text) == 0 or text in skip_strings[namespace['language']]:
            continue
        if (re.search(r'Vorsitz', text) or re.search(r'(?<![Vv]ice-)[Pp]r[äée]sid', text)):
            # Set president either to the current person or the last Person (case 1: Präsident Niklaus, case 2: Niklaus, Präsident)
            if last_person:
                composition.president = last_person
                print('Set president to'+composition.president.name)
                continue
            else:
                pos = re.search(r'(?<![Vv]ice-)[Pp]r[äée]sid', text)
                if pos == None:
                    pos = re.search(r'Vorsitz[\w]*', text)
                # assign gender depending on the noun ending
                if re.search(r'räsidentin', text):
                    last_gender = Gender.FEMALE
                elif re.search(r'räsident\b', text):
                    last_gender = Gender.MALE
                text = text[pos.span()[1]:]
                text = text.strip()
                composition.president = CourtPerson(text, last_gender)
        has_role_in_string = False
        matched_gender_regex = False

        for gender in role_regexes:  # check for male and female all roles
            if matched_gender_regex:
                break
            role_regex = role_regexes[gender]
            for regex_key in role_regex:  # check each role
                regex = '|'.join(role_regex[regex_key])
                role_pos = re.search(regex, text)
                if role_pos:  # Found a role regex
                    last_role = current_role
                    current_role = regex_key
                    # a name can consist of letters, periods, dashes, and spaces but starts with a letter
                    name_match = re.search(r'[A-Za-zÀ-ž]+[A-Za-zÀ-ž\.\- ]*', text[role_pos.span()[1] + 1:])

                    name = name_match.group() if name_match else text[role_pos.span()[1] + 1:]
                    if len(name.strip()) == 0:
                        if (last_role == CourtRole.CLERK and len(composition.clerks) == 0) or (last_role == CourtRole.JUDGE and len(composition.judges) == 0):
                            break

                        if len(composition.clerks) != 0:
                            last_person_name = composition.clerks.pop().name if (last_role == CourtRole.CLERK) else composition.clerks.pop().name # rematch in database with new role
                            last_person_new_match = CourtPerson(last_person_name, gender, court_role=current_role)
                            if current_role == CourtRole.JUDGE:
                                composition.judges.append(last_person_new_match)
                            elif current_role == CourtRole.CLERK:
                                composition.clerks.append(last_person_new_match)
                    matched_person = CourtPerson(name, gender, court_role=current_role)
                    if current_role == CourtRole.JUDGE and len(name.strip()) != 0:
                        composition.judges.append(matched_person)
                    elif current_role == CourtRole.CLERK and len(name.strip()) != 0:
                        composition.clerks.append(matched_person)
                    last_person = matched_person
                    last_gender = matched_person.gender
                    has_role_in_string = True
                    matched_gender_regex = True
                    break
        if not has_role_in_string:  # Current string has no role regex match
            # a name can consist of letters, periods, dashes, and spaces but starts with a letter
            name_match = re.search(r'[A-Za-zÀ-ž]+[A-Za-zÀ-ž\.\- ]*', text)
            if not name_match:
                continue
            name = name_match.group()
            person = CourtPerson(name, last_gender, court_role=current_role)
            matched_person = person
            if current_role == CourtRole.JUDGE:
                composition.judges.append(matched_person)
            elif current_role == CourtRole.CLERK:
                composition.clerks.append(matched_person)
            last_person = person
    return composition if composition else None


from NN import *
from qm import *
import pandas as pd
import numpy as np
from plot_util import *

ms = 3000
n_iters = 10
delays = [10, 20, 40, 50]
pretune = 5
############## UNCOMMENT TO RUN NEW RESULTS #############
'''
args = get_args()
args['tuned'] = True
args['ms'] = ms
args['n_iters'] = n_iters
args['delay'] = True
args['pretune'] = pretune
args['eta'] = args['eta'] / 100

## run networks in parallel
total_time = 0
for d in delays:
    args['delay_ms'] = d
    print('Running ' + str(d) + ' ms delay ')
    f_stim = ['storing/delayed_pretrain:' + str(pretune) + '_delay' + str(d) +'.csv']
    initialize_file(f_stim[0], args)
    t = PQ(args, stim_l = f_stim)
    print('Time to run ' + str(d) + ' ms delay: ' + str(t/60) + ' minutes')
    total_time+=t

print('Total time to run ' + str(total_time/60) + ' minutes')
'''
### plot
dfs = list()
labels = list()
### read in files
for d in delays:
    f_stim = 'storing/delayed_pretrain:' + str(pretune) + '_delay' + str(d) +'.csv'
    stims = pd.read_csv(f_stim)
    stims = pd.DataFrame.transpose(stims)
    dfs.append(stims)
    labels.append(' '+ str(d) + ' ms')
pac = plot_args(n_curves = len(dfs)+2, rb=False)
pa = plot_args(n_curves = len(dfs))
pa['ylim'] = [0, 1.1]
pa['colors'] = pac['colors'][:len(dfs)]
pa['labels'] = labels
pa['xspot'] = 3000
pa['right_margins'] = True
pa['plt_errors'] = True
pa['yadds'] = [0, 0, .02, -.03]
pa['title'] = 'Delayed Plasticity Pre-Trained (' + str(pretune) + ')'
pa['title_size'] = pa['label_size']
plot_pqs(dfs, pa)


from experimentos.exp_base import exp
import json
import time

v_range = [10, 50, 100, 500, 1000, 5000, 10000]
p_range = [p/10 for p in range(0, 11)]

print("[*****] EXPERIMENT STARTED AT {} [*****]".format(time.ctime(time.time())))
experiment_dictionary = exp(v_range, p_range)
with open("experimentos/files/exp.json", 'w+') as fp:
    json.dump(experiment_dictionary, fp)
fp.close()
print("[*****] EXPERIMENT FINISHED AT {} [*****]".format(time.ctime(time.time())))


import time

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.select import Select
import pandas as pd

for i in range(18):
    value_xd = str(i)
    df = pd.read_csv(value_xd + '111_school_url.csv')

    driver = webdriver.Chrome('chromedriver.exe')
    driver.implicitly_wait(5)
    driver.get('https://www.google.com/')
    driver.maximize_window()

    school_list = list()
    a_class_list = list()

    chinese_list = list()
    math_list = list()
    science_list = list()
    social_list = list()
    art_list = list()
    tech_list = list()
    all_list = list()
    health_list = list()
    total_list = list()

    paper_report_list = list()
    doing_work_list = list()
    science_work_list = list()
    social_work_list = list()

    a = list()
    b = list()
    c = list()
    d = list()
    e = list()
    f = list()
    g = list()
    h = list()

    x = list()
    y = list()
    z = list()

    for url in df[value_xd]:
        driver.get(url)
        school_a_class = driver.find_element(By.XPATH, "//span").text
        school_list.append(school_a_class.split('-')[0])
        a_class_list.append(school_a_class.split('-')[1])

        class_record = driver.find_element(By.XPATH,
                                           "//td[contains(text(), '修課') and contains(@style, 'center')]/following-sibling::*[contains(@style, '910px')]").text

        chinese_list.append(True if '語文' in class_record else False)
        math_list.append(True if '數學' in class_record else False)
        science_list.append(True if '自然' in class_record else False)
        social_list.append(True if '社會' in class_record else False)
        art_list.append(True if '藝術' in class_record else False)
        tech_list.append(True if '科技' in class_record else False)
        all_list.append(True if '綜合' in class_record else False)
        health_list.append(True if '健康' in class_record else False)
        total_list.append(True if '總成績' in class_record else False)

        class_result = driver.find_element(By.XPATH,
                                           "//td[contains(text(), '課程') and contains(@style, 'center')]/following-sibling::*").text

        paper_report_list.append(True if '書面報告' in class_result else False)
        doing_work_list.append(True if '實作作品' in class_result else False)
        science_work_list.append(True if '自然' in class_result else False)
        social_work_list.append(True if '社會' in class_result else False)

        multi_result = driver.find_element(By.XPATH,
                                           "//td[contains(text(), '多元') and contains(@style, 'center')]/following-sibling::*").text

        a.append(True if '高中自主學習計畫與成果' in class_result else False)
        b.append(True if '社團活動經驗' in class_result else False)
        c.append(True if '擔任幹部經驗' in class_result else False)
        d.append(True if '服務學習經驗' in class_result else False)
        e.append(True if '競賽表現' in class_result else False)
        f.append(True if '非修課紀錄成果作品' in class_result else False)
        g.append(True if '檢定證照' in class_result else False)
        h.append(True if '特殊優良表現證明' in class_result else False)

        stu_his_result = driver.find_element(By.XPATH,
                                             "//td[contains(text(), '學習歷程') and contains(@style, 'center')]/following-sibling::*").text

        x.append(True if '高中學習歷程反思' in class_result else False)
        y.append(True if '就讀動機' in class_result else False)
        z.append(True if '未來學習計畫與生涯規劃' in class_result else False)

    df = pd.DataFrame({
        '學校': school_list,
        '系組': a_class_list,
        'chinese_list': chinese_list,
        'math_list': math_list,
        'science_list': science_list,
        'social_list': social_list,
        'art_list': art_list,
        'tech_list': tech_list,
        'all_list': all_list,
        'health_list': health_list,
        'total_list': total_list,
        'paper_report_list': paper_report_list,
        'doing_work_list': doing_work_list,
        'science_work_list': science_work_list,
        'social_work_list': social_work_list,
        'a': a,
        'b': b,
        'c': c,
        'd': d,
        'e': e,
        'f': f,
        'g': g,
        'h': h,
        'x': x,
        'y': y,
        'z': z,
    })

    df.to_csv(value_xd + 'result.csv')


class Solution:
    def licenseKeyFormatting(self, s: str, k: int) -> str:
        data = "".join(s.split("-"))[::-1]

        ret = []
        for i in range(0, len(data), k):
            ret.append(data[i:i+k].upper()[::-1])

        return "-".join(ret[::-1])


import os
import logging
from importlib import import_module
from argparse import ArgumentParser
from builder_utils import run_test_for_stage, init_logging

STAGE_INDEXING = ["preprocessing", "tokenisation",
                  "build_inverse", "build_right"]

INDEXING_CONSTANTS = {
    "data_dir": os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "data"),
    "stages": STAGE_INDEXING,
    "input_file_name": "data.in",
    "config_file_name": "config.yml",
    "size_of_record": 40,
    "size_of_record_right": 10
}

def _init_argparser() -> ArgumentParser:
    text = "Script for ruling indexation process." \
           "Data structure in config.yml, input file with name data.in"
    parser = ArgumentParser(description=text)
    parser.add_argument(
        "--stage", "-s",
        help=f"Stage of indexing, one of: {STAGE_INDEXING}",
        required=True)
    parser.add_argument(
        "--allstages", "-a",
        help=f"Process all stages of indexing, or only passed",
        action="store_true")
    parser.add_argument(
        "--withtests", "-w",
        help="Run tests after stage finishing",
        action="store_true")
    parser.add_argument(
        "--onlytests", "-o",
        help="Run tests after stage finishing",
        action="store_true")
    parser.add_argument("--loglvl", "-l", help=f"Logging level")
    return parser

def handle_stage_tests(stage: str):
    if args.withtests:
        test_success = run_test_for_stage(stage)
        if not test_success:
            logging.error(f"Tests failed for stage {stage}")
            exit(2)

if __name__ == "__main__":
    parser = _init_argparser()
    args = parser.parse_args()
    init_logging(args.loglvl)
    try:
        stage_target_id = STAGE_INDEXING.index(args.stage)
    except ValueError:
        logging.error(
            "Unknown stage of indexing: %s expected one of %s" %
            (args.stage, ', '.join(STAGE_INDEXING)), exc_info=True)
        exit(1)
    if args.allstages:
        stage_gen = enumerate(
            STAGE_INDEXING[:stage_target_id + 1])
        for stage_id, stage in stage_gen:
            if not args.onlytests:
                module = import_module(stage)
                module.run(INDEXING_CONSTANTS, stage_id)
            handle_stage_tests(stage)
    else:
        if not args.onlytests:
            module = import_module(args.stage)
            module.run(INDEXING_CONSTANTS, stage_target_id)
        handle_stage_tests(args.stage)


from __future__ import unicode_literals
from collections import namedtuple
from itertools import chain, count

class Position(namedtuple("Position", ["board", "whites_turn",
                                       "castling_white", "castling_black",
                                       "en_passant_target", "half_move_clock",
                                       "move_count"])):
    """Tuple that holds data about the state of a chess game. It somewhat
    matches Forsyth-Edwards Notation in both state contents and order.

    .. IMPORTANT::
        A chess position is a *state*, not a location.

    :param str board: 64-character string that holds the board pieces as
                      Unicode characters, and empty squares as spaces. The
                      first character is A8 on the chess board.
    :param bool whites_turn: Whether the current turn is white's turn.
    :param castling_white: Sequence of two boolean values that determine the
                           kingside and queenside castling rights of white.
    :type castling_white: :class:`CastlingRights`
    :param castling_black: Sequence of two boolean values that determine the
                           kingside and queenside castling rights of black.
    :type castling_black: :class:`CastlingRights`
    :param int en_passant_target: Index position of the spot behind the pawn
                                  that can be captured.
    :param int half_move_clock: Amount of half-moves since the last capture or
                                pawn move.
    :param int move_count: Total amount of full moves.
    """

class Move(namedtuple("Move", ["origin", "destination"])):
    """Tuple that holds data about a move from one index to another.

    :param int origin: From-index.
    :param int destination: To-index.
    """

class CastlingRights(namedtuple("CastlingRights", ["kingside", "queenside"])):
    """Tuple that holds data about the castling rights in either direction.

    :param bool kingside: Kingside castling rights.
    :param bool queenside: Queenside castling rights.
    """

#: White king piece.
KING_WHITE = '♔'
#: White queen piece.
QUEEN_WHITE = '♕'
#: White rook piece.
ROOK_WHITE = '♖'
#: White bishop piece.
BISHOP_WHITE = '♗'
#: White knight piece.
KNIGHT_WHITE = '♘'
#: White pawn piece.
PAWN_WHITE = '♙'
#: Set of white pieces.
WHITES = set([
    KING_WHITE, QUEEN_WHITE, ROOK_WHITE, BISHOP_WHITE, KNIGHT_WHITE, PAWN_WHITE
])

#: Black king piece.
KING_BLACK = '♚'
#: Black queen piece.
QUEEN_BLACK = '♛'
#: Black rook piece.
ROOK_BLACK = '♜'
#: Black bishop piece.
BISHOP_BLACK = '♝'
#: Black knight piece.
KNIGHT_BLACK = '♞'
#: Black pawn piece.
PAWN_BLACK = '♟'
#: Set of black pieces.
BLACKS = set([
    KING_BLACK, QUEEN_BLACK, ROOK_BLACK, BISHOP_BLACK, KNIGHT_BLACK, PAWN_BLACK
])

#: Set of kings.
KINGS = set([KING_WHITE, KING_BLACK])
#: Set of queens.
QUEENS = set([QUEEN_WHITE, QUEEN_BLACK])
#: Set of rooks.
ROOKS = set([ROOK_WHITE, ROOK_BLACK])
#: Set of bishops.
BISHOPS = set([BISHOP_WHITE, BISHOP_BLACK])
#: Set of knights.
KNIGHTS = set([KNIGHT_WHITE, KNIGHT_BLACK])
#: Set of pawns.
PAWNS = set([PAWN_WHITE, PAWN_BLACK])

#: Distance to north square in 1D array.
NORTH = -8
#: Distance to east square in 1D array.
EAST = 1
#: Distance to south square in 1D array.
SOUTH = 8
#: Distance to west square in 1D array.
WEST = -1

#: Dictionary that holds string identifiers as keys, and sets of associated
#: move offsets as values (White).
DIRECTIONS_WHITE = {
    KING_WHITE: set([NORTH, NORTH+EAST, EAST, EAST*2, SOUTH+EAST, SOUTH,
                     SOUTH+WEST, WEST*2, NORTH+WEST]),
    QUEEN_WHITE: set([NORTH, NORTH+EAST, EAST, SOUTH+EAST, SOUTH, SOUTH+WEST,
                      WEST, NORTH+WEST]),
    ROOK_WHITE: set([NORTH, EAST, SOUTH, WEST]),
    BISHOP_WHITE: set([NORTH+EAST, SOUTH+EAST, SOUTH+WEST, NORTH+WEST]),
    KNIGHT_WHITE: set([NORTH*2+EAST, NORTH+EAST*2, SOUTH+EAST*2, SOUTH*2+EAST,
                       SOUTH*2+WEST, SOUTH+WEST*2, NORTH+WEST*2,
                       NORTH*2+WEST]),
    PAWN_WHITE: set([NORTH, NORTH*2, NORTH+EAST, NORTH+WEST])
}

#: Dictionary that holds string identifiers as keys, and sets of associated
#: move offsets as values (Black).
DIRECTIONS_BLACK = {
    KING_BLACK: set([NORTH, NORTH+EAST, EAST, EAST*2, SOUTH+EAST, SOUTH,
                     SOUTH+WEST, WEST*2, NORTH+WEST]),
    QUEEN_BLACK: set([NORTH, NORTH+EAST, EAST, SOUTH+EAST, SOUTH, SOUTH+WEST,
                      WEST, NORTH+WEST]),
    ROOK_BLACK: set([NORTH, EAST, SOUTH, WEST]),
    BISHOP_BLACK: set([NORTH+EAST, SOUTH+EAST, SOUTH+WEST, NORTH+WEST]),
    KNIGHT_BLACK: set([NORTH*2+EAST, NORTH+EAST*2, SOUTH+EAST*2, SOUTH*2+EAST,
                       SOUTH*2+WEST, SOUTH+WEST*2, NORTH+WEST*2,
                       NORTH*2+WEST]),
    PAWN_BLACK: set([SOUTH+EAST, SOUTH, SOUTH*2, SOUTH+WEST])
}

#: Index locations of all squares of the north border.
BORDER_NORTH = set([i for i in range(8)])
#: Index locations of all squares of the east border.
BORDER_EAST = set([7+i*8 for i in range(8)])
#: Index locations of all squares of the south border.
BORDER_SOUTH = set([i for i in range(56, 64)])
#: Index locations of all squares of the west border.
BORDER_WEST = set([i*8 for i in range(8)])

#: 64-character representation of starting board.
INITIAL_BOARD = (
    "♜♞♝♛♚♝♞♜"  # 0-7
    "♟♟♟♟♟♟♟♟"  # 8-15
    "        "  # 16-23
    "        "  # 24-31
    "        "  # 32-39
    "        "  # 40-47
    "♙♙♙♙♙♙♙♙"  # 48-55
    "♖♘♗♕♔♗♘♖"  # 56-63
)

#: Beginning position of a standard chess game.
INITIAL_POSITION = Position(board=INITIAL_BOARD, whites_turn=True,
                            castling_white=CastlingRights(True, True),
                            castling_black=CastlingRights(True, True),
                            en_passant_target=None, half_move_clock=0,
                            move_count=0)

def get_piece(index, board):
    """Get the piece at the provided index on the board. Return an empty string
    if the index points to an empty square.

    :param int index: Index value from 0-63 pointing to a square on the board.
    :param str board: 64-character chess board representation.
    :return: Associated chess piece.
    :rtype: str
    """
    return board[index].strip()

def active_team(whites_turn):
    """Return the active team depending on whose turn it is.

    :param bool whites_turn: Boolean value signifying whose turn it is.
    :return: Active team.
    :rtype: set
    """
    return WHITES if whites_turn else BLACKS

def opponent(team):
    """Return the opposing team to the supplied team.

    :param set team: Set of chess pieces of a team.
    :return: Opposing team.
    :rtype: set
    """
    return BLACKS if team is WHITES else WHITES

def directions(team):
    """Return the directions associated with the team.

    :param set team: Set of chess pieces of a team.
    :return: Associated directions.
    :rtype: dict"""
    return DIRECTIONS_WHITE if team is WHITES else DIRECTIONS_BLACK

def castling_rights(team, position):
    """Return the castling rights embedded in :class:`Position` of the supplied
    team.

    :param set team: Set of chess pieces of a team.
    :param position: State of the game.
    :type position: :class:`Position`
    :return: Associated castling rights.
    :rtype: :class:`CastlingRights`
    """
    return (position.castling_white if team is WHITES
            else position.castling_black)

def legal_moves(origin, position):
    """Yield all legal moves of a piece at *index*.

    Takes into account:
      * Whether the move is actually possible.
      * Whether the king moves to a square that is under siege.
      * Whether the king tries to castle while under siege, or while the
        square the king has to pass through is under siege.
      * Whether the resulting move results in a check.

    :param int origin: Index position on the board.
    :param position: State of the game.
    :type position: :class:`Position`
    :return: All legal :class:`Moves <Move>`.
    :rtype: :class:`generator`
    """
    global NORTH, EAST, SOUTH, WEST

    piece = get_piece(origin, position.board)
    team = active_team(position.whites_turn)
    king = piece in KINGS
    if king:
        under_siege = set(besieged(team, position))
        current_check = origin in under_siege

    for move in _accessible_moves(origin, position):
        if king:
            offset = move.destination - move.origin
            half_offset = offset // 2
            if move.destination in under_siege:
                continue
            if (offset in (EAST*2, WEST*2)
                    and (current_check
                         or
                         half_offset in under_siege)):
                continue
        elif is_check(team, do_move(move, position, force=True)):
            continue

        yield move

def besieged(team, position):
    """Yield all squares that are besieged by enemy pieces.

    :param set team: Set of chess pieces of a team.
    :param position: State of the game.
    :type position: :class:`Position`
    :return: Integer values of besieged squares.
    :rtype: :class:`generator`
    """
    global NORTH, EAST, SOUTH, WEST, KINGS, PAWNS

    opposing_team = opponent(team)
    siege_set = set([])

    for i, piece in enumerate(position.board):
        if piece in opposing_team:
            for move in _accessible_moves(i, position, capture_moves=True):
                destination = move.destination
                if destination not in siege_set:
                    yield destination
                    siege_set.add(destination)

def is_check(team, position):
    """Determine whether the team is in check in the provided chess position.

    :param set team: Set of chess pieces of a team.
    :param position: State of the game.
    :type position: :class:`Position`
    :return: Whether the king of the provided team is in check.
    :rtype: bool
    """
    global KINGS

    under_siege = besieged(team, position)

    for i, piece in enumerate(position.board):
        if piece in team and piece in KINGS:
            if i in under_siege:
                return True
    return False

def is_move_legal(move, position):
    """Determine whether the supplied move is legal.

    :param move: Suggested move.
    :type move: :class:`Move`
    :param position: State of the game.
    :type position: :class:`Position`
    :return: Whether the move is legal.
    :rtype: bool
    """
    team = active_team(position.whites_turn)
    piece = get_piece(move.origin, position.board)
    destination_piece = get_piece(move.destination, position.board)

    if piece not in team or destination_piece in team:
        return False
    if move in legal_moves(move.origin, position):
        return True

    return False

def do_move(move, position, force=False, promotion_piece=None):
    """Apply a move to the board and return a new position based on the
    implications of the move.

    Also perform a check to confirm whether the supplied move is legal. This
    check is ignored if *force* is :const:`True`.

    An optional *promotion_piece* can be supplied in case a pawn is to be
    promoted. If no *promotion_piece* is supplied, the promotion will default
    to queen promotion.

    :param move: Move to be applied.
    :type move: :class:`Move`
    :param position: State of the game.
    :type position: :class:`Position`
    :keyword bool force: Whether to force the function regardless of whether
                         the move is legal.
    :keyword str promotion_piece: Identifier of piece that pawn should be
                                  promoted to.
    :return: New chess position.
    :rtype: :class:`Position`
    :raise ValueError: Move is illegal.
    """
    if force or is_move_legal(move, position):
        pass
    else:
        raise ValueError()  # TODO

    origin, destination = move
    team = active_team(position.whites_turn)
    piece = get_piece(origin, position.board)
    destination_piece = get_piece(destination, position.board)

    board = _apply_move(move, position.board)
    whites_turn = not position.whites_turn
    castling_white = position.castling_white
    castling_black = position.castling_black
    en_passant_target = None
    half_move_clock = position.half_move_clock + 1
    move_count = position.move_count
    if not position.whites_turn:
        move_count += 1

    if piece in KINGS:
        offset = destination - origin
        if offset == EAST*2:
            board = _apply_move(Move(origin+EAST*3, origin+EAST), board)
        elif offset == WEST*2:
            board = _apply_move(Move(origin+WEST*4, origin+WEST), board)
        if team is WHITES:
            castling_white = CastlingRights(False, False)
        else:
            castling_black = CastlingRights(False, False)
    elif piece in ROOKS:
        if team is WHITES:
            if origin % 8 == 0:
                castling_white = CastlingRights(castling_white.kingside, False)
            else:
                castling_white = CastlingRights(False,
                                                castling_white.queenside)
        else:
            if origin % 8 == 0:
                castling_black = CastlingRights(castling_black.kingside, False)
            else:
                castling_black = CastlingRights(False,
                                                castling_black.queenside)
    elif piece in PAWNS:
        offset = destination - origin
        half_move_clock = 0
        if offset == NORTH*2:
            en_passant_target = destination + SOUTH
        elif offset == SOUTH*2:
            en_passant_target = destination + NORTH
        elif offset in (NORTH+EAST, NORTH+WEST) and not destination_piece:
            board = _clear(destination + SOUTH, board)
        elif offset in (SOUTH+EAST, SOUTH+WEST) and not destination_piece:
            board = _clear(destination + NORTH, board)
        if destination in chain(BORDER_NORTH, BORDER_SOUTH):
            if not promotion_piece:
                if team is WHITES:
                    promotion_piece = QUEEN_WHITE
                else:
                    promotion_piece = QUEEN_BLACK
            board = _place_piece(destination, promotion_piece, board)
    if destination_piece:
        half_move_clock = 0

    return Position(board, whites_turn, castling_white, castling_black,
                    en_passant_target, half_move_clock, move_count)

def to_index(notation):
    """Return an index location for an algebraic notation (e.g., 'A8' -> 0)"""
    letter, number = notation[0].upper(), int(notation[1])
    return (ord(letter)-65) + (-(number-8) * 8)

def to_notation(index):
    """Return an algebraic notation for an index position (e.g., 63 -> 'H1')"""
    remainder, quotient = divmod(index, 8)
    return chr(quotient+65) + str(-(remainder-8))

def _accessible_moves(origin, position, capture_moves=False):
    """Generator that yields all plausible moves for a single piece."""
    global NORTH, BORDER_NORTH, EAST, SOUTH, BORDER_SOUTH, WEST, Move

    piece = get_piece(origin, position.board)
    if piece in WHITES:
        team = WHITES
    elif piece in BLACKS:
        team = BLACKS
    else:
        raise Exception()  # TODO
    pawn = piece in PAWNS
    king = piece in KINGS

    for offset in directions(team)[piece]:
        for destination in count(origin + offset, offset):
            # Piece in the bounds of the board?
            if not _in_bounds(destination - offset, destination):
                break

            destination_piece = get_piece(destination, position.board)

            # Destination occupied by piece of same colour?
            if destination_piece in team:
                break

            if king:
                # Offset is not a move that can capture?
                if capture_moves and offset in (EAST*2, WEST*2):
                    break
                # Offset is a castling move and castling not possible/allowed?
                if (offset == EAST*2
                        and get_piece(origin+EAST, position.board)
                        and not castling_rights(team, position).kingside
                        or
                        offset == WEST*2
                        and get_piece(origin+WEST, position.board)
                        and get_piece(origin+WEST*3, position.board)
                        and not castling_rights(team, position).queenside):
                    break
            elif pawn:
                # Offset is not a move that can capture?
                if capture_moves and offset in (NORTH, NORTH*2, SOUTH,
                                                SOUTH*2):
                    break
                # Offset is a vertical move forward and destination is an enemy
                # piece?
                if (destination_piece
                        and offset in (NORTH, NORTH*2, SOUTH, SOUTH*2)):
                    break
                # Offset is a two-squares move forward and pawn is not in its
                # initial spot?
                if (offset == NORTH*2
                        and (get_piece(origin+NORTH, position.board)
                             or
                             origin not in [index+NORTH for index in
                                            BORDER_SOUTH])
                        or
                        offset == SOUTH*2
                        and (get_piece(origin+SOUTH, position.board)
                             or
                             origin not in [index+SOUTH for index in
                                            BORDER_NORTH])):
                    break
                # Offset is a diagonal move and destination is not an enemy
                # piece or destination is not an en passant target?
                # Ignored if `capture_moves` is True.
                if (not capture_moves
                        and offset in (NORTH+EAST, SOUTH+EAST, SOUTH+WEST,
                                       NORTH+EAST)
                        and (not destination_piece
                             or
                             destination != position.en_passant_target)):
                    break

            # Move is legal.
            yield Move(origin, destination)

            # Piece cannot do multiple steps in the same direction?
            if piece in chain(KINGS, KNIGHTS, PAWNS):
                break

            # The destination was a capture move and the direction is therefore
            # further blocked?
            if destination_piece:
                break

def _in_bounds(origin, destination):
    """If a single move is applied to a piece, test whether it left the bounds
    of the board.

    This function is extremely limited in scope, and should be considered an
    implementation detail.
    """
    global EAST, BORDER_EAST, WEST, BORDER_WEST

    # North illegal?
    if destination < 0:
        return False
    # East illegal?
    if origin in BORDER_EAST and destination in BORDER_WEST:
        return False
    # Twice east illegal? (Knight)
    if (destination in BORDER_WEST
            and origin in [index+WEST for index in BORDER_EAST]
            or
            origin in BORDER_EAST
            and destination in [index+EAST for index in BORDER_WEST]):
        return False
    # South illegal?
    if destination > 63:
        return False
    # West illegal?
    if origin in BORDER_WEST and destination in BORDER_EAST:
        return False
    # Twice west illegal? (Knight)
    if (destination in BORDER_EAST
            and origin in [index+EAST for index in BORDER_WEST]
            or
            origin in BORDER_WEST
            and destination in [index+WEST for index in BORDER_EAST]):
        return False
    return True

def _apply_move(move, board):
    origin, destination = move
    board = _place_piece(destination, board[origin], board)
    if origin != destination:
        board = _clear(origin, board)
    return board

def _place_piece(origin, piece, board):
    if len(piece) != 1:
        raise Exception()  # TODO
    return "".join([board[:origin], piece, board[origin+1:]])

def _clear(origin, board):
    return "".join([board[:origin], " ", board[origin+1:]])


import cv2
import numpy as np

class RegionChooser:
    def __init__(self, frame):
        self._frame = frame
        self._clone = self._frame.copy()
        self._coords = []
        cv2.namedWindow("RegionChooser")
        cv2.setMouseCallback("RegionChooser", self._on_mouse)

    def _on_mouse(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            self._clone = self._frame.copy()
            self._coords += [(x,y)]
            print(self._coords)
        elif event == cv2.EVENT_LBUTTONUP:
            # draw to image
            self._shape = np.array([self._coords], dtype=np.int32)
            cv2.polylines(self._clone, pts=self._shape, isClosed=True, color=(55,243,46))
            cv2.imshow("RegionChooser", self._clone)

        elif event == cv2.EVENT_RBUTTONDOWN:
            # clear region
            self._clone = self._frame.copy()
            self._coords = []

    def set_shape(self, coords):
        self._shape = np.array([coords], dtype=np.int32)

    def cut(self, current_frame):
        mask = np.zeros(current_frame.shape, dtype=np.uint8)
        channel_count = current_frame.shape[2]
        ignore_mask_color = (0,) * channel_count
        cv2.fillPoly(mask, self._shape, ignore_mask_color)
        masked_image = cv2.bitwise_and(current_frame, mask)
        return masked_image

    def show(self):
        while True:
            cv2.imshow("RegionChooser", self._clone)
            key = cv2.waitKey(1)

            if key == ord('q'):
                cv2.destroyAllWindows()
                break


import requests
from flask import current_app

def confirm_subscription(confirmation_request):
    url = confirmation_request.get("SubscribeURL")
    if not url:
        current_app.logger.warning("SubscribeURL does not exist or empty")
        return

    response = requests.get(url)
    try:
        response.raise_for_status()
    except Exception as e:
        current_app.logger.warning("Response: %s", response.text)
        raise e

    return confirmation_request["TopicArn"]

def autoconfirm_subscription(req_json):
    if req_json.get("Type") == "SubscriptionConfirmation":
        current_app.logger.debug("SNS subscription confirmation url: %s", req_json["SubscribeURL"])
        subscribed_topic = confirm_subscription(req_json)
        return subscribed_topic


import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss, MultiLabelSoftMarginLoss
from .loss import FocalLoss, LabelSmoothingLoss
import torch.nn.functional as F

import sys
from transformers import BertPreTrainedModel, BertModel
from transformers import XLMRobertaConfig, RobertaConfig, RobertaModel

def compute_kl_loss( p, q, pad_mask = None):
    p_loss = F.kl_div(F.log_softmax(p, dim=-1), F.softmax(q, dim=-1), reduction='none')
    q_loss = F.kl_div(F.log_softmax(q, dim=-1), F.softmax(p, dim=-1), reduction='none')

# pad_mask is for seq-level tasks
    if pad_mask is not None:
        p_loss.masked_fill_(pad_mask, 0.)
        q_loss.masked_fill_(pad_mask, 0.)

# You can choose whether to use function "sum" and "mean" depending on your task
    p_loss = p_loss.sum()
    q_loss = q_loss.sum()

    loss = (p_loss + q_loss) / 2
    return loss

# class RobertaForSequenceClassification(BertPreTrainedModel):
#     config_class = RobertaConfig
#     base_model_prefix = "roberta"
#
#     def __init__(self, config):
#         super().__init__(config)
#         self.num_labels = config.num_labels
#         config.output_hidden_states=True
#
#         self.roberta = RobertaModel(config)
#         self.fc = nn.Linear(768, self.config.num_labels)
#
#         for param in self.roberta.parameters():
#             param.requires_grad = True
#
#         self.weights = nn.Parameter(torch.rand(13, 1))
#         self.dropouts = nn.ModuleList([
#             nn.Dropout(0.5) for _ in range(5)
#         ])
#         self.init_weights()
#
#     def forward(
#         self,
#         input_ids=None,
#         attention_mask=None,
#         token_type_ids=None,
#         position_ids=None,
#         head_mask=None,
#         inputs_embeds=None,
#         labels=None,
#     ):
#
#         outputs= self.roberta(
#             input_ids,
#             attention_mask=attention_mask,
#             token_type_ids=token_type_ids,
#             position_ids=position_ids,
#             head_mask=head_mask,
#             inputs_embeds=inputs_embeds,
#         )
#
#
#         batch_size = input_ids.shape[0]
#         ht_cls = torch.cat(outputs[2])[:, :1, :].view(
#             13, batch_size, 1, 768)
#
#         atten = torch.sum(ht_cls * self.weights.view(
#             13, 1, 1, 1), dim=[1, 3])
#         atten = F.softmax(atten.view(-1), dim=0)
#         feature = torch.sum(ht_cls * atten.view(13, 1, 1, 1), dim=[0, 2])
#
#         feature2 = torch.sum(ht_cls * atten.view(13, 1, 1, 1), dim=[0, 2])
#
#
#         if labels is not None:
#
#             for i, dropout in enumerate(self.dropouts):
#                 if i == 0:
#                     h = self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#
#                     loss_fct = CrossEntropyLoss()
#
#                     # carefully choose hyper-parameters
#                     loss = loss_fct(h.view(-1, self.num_labels), labels.view(-1))
#
#                 else:
#                     h += self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#                     loss_fct = CrossEntropyLoss()
#
#                     # carefully choose hyper-parameters
#                     loss = loss_fct(h.view(-1, self.num_labels), labels.view(-1))
#
#             loss = loss / len(self.dropouts)
#
#             outputs = (loss,) + outputs
#         else:
#             for i, dropout in enumerate(self.dropouts):
#                 if i == 0:
#                     h = self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#                 else:
#                     h += self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#
#             outputs = outputs
#         return outputs

class RobertaForSequenceClassification(BertPreTrainedModel):
    config_class = RobertaConfig
    base_model_prefix = "roberta"

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        config.output_hidden_states=True

        self.roberta = RobertaModel(config)
        self.fc = nn.Linear(768, self.config.num_labels)

        for param in self.roberta.parameters():
            param.requires_grad = True
        self.weights = nn.Parameter(torch.rand(13, 1))
        self.dropouts = nn.ModuleList([
            nn.Dropout(0.5) for _ in range(5)
        ])
        self.init_weights()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
    ):

        outputs= self.roberta(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )

        batch_size = input_ids.shape[0]
        ht_cls = torch.cat(outputs[2])[:, :1, :].view(
            13, batch_size, 1, 768)

        atten = torch.sum(ht_cls * self.weights.view(
            13, 1, 1, 1), dim=[1, 3])
        atten = F.softmax(atten.view(-1), dim=0)
        feature = torch.sum(ht_cls * atten.view(13, 1, 1, 1), dim=[0, 2])

        feature2 = torch.sum(ht_cls * atten.view(13, 1, 1, 1), dim=[0, 2])

        # if labels is not None:
        #
        #     for i, dropout in enumerate(self.dropouts):
        #         if i == 0:
        #             h = self.fc(dropout(feature))
        #             outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
        #
        #             h2 = self.fc(dropout(feature2))
        #             loss_fct = CrossEntropyLoss()
        #
        #             # carefully choose hyper-parameters
        #             loss = 0.5*(loss_fct(h.view(-1, self.num_labels), labels.view(-1)) + loss_fct(h2.view(-1, self.num_labels), labels.view(-1)))
        #
        #
        #             kl_loss = compute_kl_loss(feature, feature2)
        #             loss = loss + 4 * kl_loss
        #
        #         else:
        #             h += self.fc(dropout(feature))
        #             outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
        #             h2 += self.fc(dropout(feature))
        #
        #             loss_fct = CrossEntropyLoss()
        #
        #             # carefully choose hyper-parameters
        #             loss = 0.5*(loss_fct(h.view(-1, self.num_labels), labels.view(-1)) + loss_fct(h2.view(-1, self.num_labels), labels.view(-1)))
        #
        #             kl_loss = compute_kl_loss(feature, feature2)
        #             loss = loss + 4 * kl_loss
        #
        #     loss = loss / len(self.dropouts)
        #
        #     outputs = (loss,) + outputs
        # else:
        #     for i, dropout in enumerate(self.dropouts):
        #         if i == 0:
        #             h = self.fc(dropout(feature))
        #             outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
        #         else:
        #             h += self.fc(dropout(feature))
        #             outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
        #
        #     outputs = outputs
        # return outputs

        if labels is not None:

            for i, dropout in enumerate(self.dropouts):
                if i == 0:
                    h = self.fc(dropout(feature))
                    outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here

                    loss_fct = CrossEntropyLoss()

                    # carefully choose hyper-parameters
                    loss = loss_fct(h.view(-1, self.num_labels), labels.view(-1))

                    kl_loss = compute_kl_loss(feature, feature2)
                    loss = loss + 4 * kl_loss

                else:
                    h += self.fc(dropout(feature))
                    outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
                    loss_fct = CrossEntropyLoss()

                    # carefully choose hyper-parameters
                    loss = loss_fct(h.view(-1, self.num_labels), labels.view(-1))

                    kl_loss = compute_kl_loss(feature, feature2)
                    loss = loss + 4 * kl_loss

            loss = loss / len(self.dropouts)

            outputs = (loss,) + outputs
        else:
            for i, dropout in enumerate(self.dropouts):
                if i == 0:
                    h = self.fc(dropout(feature))
                    outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
                else:
                    h += self.fc(dropout(feature))
                    outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here

            outputs = outputs
        return outputs

import torch.nn.functional as F
class RobertaForSequenceClassification_cnn(BertPreTrainedModel):
    config_class = RobertaConfig
    base_model_prefix = "roberta"

    def __init__(self, config):
        super().__init__(config)

        # add cnn params
        config.output_hidden_states = True
        self.filter_sizes_ = (3, 4, 5)  # 卷积核尺寸
        self.num_filters_ = 256  # 卷积核数量(channels数)
        self.hidden_size_ = 768

        self.num_labels = config.num_labels

        self.roberta = RobertaModel(config)
        self.convs = nn.ModuleList(
            [nn.Conv2d(1, self.num_filters_, (k, self.hidden_size_)) for k in self.filter_sizes_])

        self.dropout = nn.Dropout(0.2)
        self.classifier = nn.Linear((self.num_filters_ * len(self.filter_sizes_))*3 + self.hidden_size_, config.num_labels)

        self.init_weights()

    def conv_and_pool(self, x, conv):
        x = F.relu(conv(x)).squeeze(3)
        x = F.max_pool1d(x, x.size(2)).squeeze(2)
        return x

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for computing the sequence classification/regression loss.
            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
            Classification (or regression if config.num_labels==1) loss.
        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Examples::

        from transformers import RobertaTokenizer, RobertaForSequenceClassification
        import torch

        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
        model = RobertaForSequenceClassification.from_pretrained('roberta-base')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, logits = outputs[:2]

        """
        outputs= self.roberta(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )
        # flat_input_ids = input_ids.view(-1, input_ids.size(-1))
        # flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
        # flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
        # flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
        #
        #
        # bert_output_a, pooled_output_a, hidden_output_a = self.roberta(input_ids=flat_input_ids, position_ids=flat_position_ids,
        #                     token_type_ids=flat_token_type_ids,
        #                     attention_mask=flat_attention_mask, head_mask=head_mask)

        # sequence_output = outputs[1]
        # last_cat = torch.cat(
        #     (sequence_output, outputs[2][-1][:,0], outputs[2][-2][:,0], outputs[2][-3][:,0]),
        #     1,
        # )
        sequence_output = outputs[2][-1]
        out1 = sequence_output.unsqueeze(1)
        out1 = torch.cat([self.conv_and_pool(out1, conv) for conv in self.convs], 1)

        sequence_output = outputs[2][-2]
        out2 = sequence_output.unsqueeze(1)
        out2 = torch.cat([self.conv_and_pool(out2, conv) for conv in self.convs], 1)

        sequence_output = outputs[2][-3]
        out3 = sequence_output.unsqueeze(1)
        out3 = torch.cat([self.conv_and_pool(out3, conv) for conv in self.convs], 1)

        out = torch.cat((out1, out2, out3, outputs[1]),1)
        pooled_output = self.dropout(out)
        logits = self.classifier(pooled_output)

        if labels is not None:
            if self.num_labels == 1:
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            outputs = [loss, ]
            outputs = outputs + [nn.functional.softmax(logits, -1)]
        else:
            outputs = nn.functional.softmax(logits, -1)

        return outputs  # (loss), logits, (hidden_states), (attentions)

class xlmroberta_model(RobertaForSequenceClassification):
    """
    This class overrides :class:`~transformers.RobertaForSequenceClassification`. Please check the
    superclass for the appropriate documentation alongside usage examples.
    """

    config_class = XLMRobertaConfig

class xlmroberta_model_cnn(RobertaForSequenceClassification_cnn):
    """
    This class overrides :class:`~transformers.RobertaForSequenceClassification`. Please check the
    superclass for the appropriate documentation alongside usage examples.
    """

    config_class = XLMRobertaConfig

# class bert_cls_model(BertPreTrainedModel):
#     def __init__(self, config):
#         super().__init__(config)
#         self.num_labels = config.num_labels
#         config.output_hidden_states = True
#         self.bert = BertModel(config)
#         for param in self.bert.parameters():
#             param.requires_grad = True
#
#         self.weights = nn.Parameter(torch.rand(13, 1))
#         self.dropouts = nn.ModuleList([
#             nn.Dropout(0.5) for _ in range(5)
#         ])
#         self.fc = nn.Linear(config.hidden_size, self.num_labels)
#
#         self.init_weights()
#
#     def forward(
#         self,
#         input_ids=None,
#         attention_mask=None,
#         token_type_ids=None,
#         position_ids=None,
#         head_mask=None,
#         inputs_embeds=None,
#         labels=None,
#     ):
#
#         outputs = self.bert(
#             input_ids,
#             attention_mask=attention_mask,
#             token_type_ids=token_type_ids,
#             position_ids=position_ids,
#             head_mask=head_mask,
#             inputs_embeds=inputs_embeds,
#         )
#
#         batch_size = input_ids.shape[0]
#         ht_cls = torch.cat(outputs[2])[:, :1, :].view(
#             13, batch_size, 1, 768)
#
#         atten = torch.sum(ht_cls * self.weights.view(
#             13, 1, 1, 1), dim=[1, 3])
#         atten = F.softmax(atten.view(-1), dim=0)
#         feature = torch.sum(ht_cls * atten.view(13, 1, 1, 1), dim=[0, 2])
#
#         if labels is not None:
#
#             for i, dropout in enumerate(self.dropouts):
#                 if i == 0:
#                     h = self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#
#                     loss_fct = CrossEntropyLoss()
#                     loss = loss_fct(h.view(-1, self.num_labels), labels.view(-1))
#                 else:
#                     h += self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#                     loss_fct = CrossEntropyLoss()
#                     loss = loss_fct(h.view(-1, self.num_labels), labels.view(-1))
#
#             loss = loss / len(self.dropouts)
#             outputs = (loss,) + outputs
#         else:
#             for i, dropout in enumerate(self.dropouts):
#                 if i == 0:
#                     h = self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#                 else:
#                     h += self.fc(dropout(feature))
#                     outputs = (h,) + outputs[2:]  # add hidden states and attention if they are here
#
#             outputs = outputs
#
#         return outputs


from db import db_client
import json
import os
import dialogflow
import requests
import json
from flask import jsonify, request
import re

menu_db = db_client.menu
menu_db = db_client.menu
review_db = db_client.review

def construct_custom_intent(type, items, text_message, found):
    message1 = {
    "fulfillmentText": "Your text response",
    "fulfillmentMessages": [
        {
            "text": {
                "text": [
                    text_message
                ]
            }
        },
        {
            "payload": {}
        }
    ]
        }
    message = {
        "richContent": [[]]
              }
    if(found == False): return jsonify(message1)
    if(type == 'menu_category'):
        for x in items['menu_items']:
            menu_item_rich = {
                "type": "list",
                "title": x['name'],
                "subtitle": f"Description: {x['description']}\nPrice: {x['price']}",
            }
            message['richContent'][0].append(menu_item_rich)
            message['richContent'][0].append({'type':'divider'})
    elif(type == 'menu_item'):
        menu_item_rich = {
            "type": "image",
            "rawUrl": items['image'],
        }
        menu_item_rich2 = {
           "type": "info",
           "title": items['name'],
           "subtitle": f"Menu-category: {items['category']}\nPrice: {items['price']}\nRating: {items['rating']}/5.0",
        }
        message['richContent'][0].append(menu_item_rich)
        message['richContent'][0].append(menu_item_rich2)
    elif(type == 'menu_category_general'):
        for x in items:
            # print(x['name']i)
            menu_category_rich = {
                "type": "list",
                "title": x['name'],
                }
            message['richContent'][0].append(menu_category_rich)
            message['richContent'][0].append({'type':'divider'})
    elif(type == 'review'):
        z=0
        for x in items:
            print(x)
            if(z==3):break
            review_item = {
                "type": "list",
                "title": x['review'],
                "subtitle": f"Rating: {x['rating']}/5\nDate Reviewed: {x['date_time']}"
            }
            message['richContent'][0].append(review_item)
            message['richContent'][0].append({'type':'divider'})
            z+=1
    elif(type == 'suggestion'):
        for x in items:
            for menu_item in x['menu_items']:
                if menu_item['chefs_pick'] == True:
                    menu_item_rich = {
                        "type": "list",
                        "title": menu_item['name'],
                        "subtitle": f"Description: {menu_item['description']}\nPrice: {menu_item['price']}\nRating: {menu_item['rating']}/5.0",
                    }
                    message['richContent'][0].append(menu_item_rich)
                    message['richContent'][0].append({'type':'divider'})
    elif(type == 'menu_review'):
        count = 0

        for z in items['review_list']:
            if(count == 3):break
            if('user' not in z): user = 'Anonymous'
            else: user = z['user']
            menu_review = {
                "type": "list",
                "title": f"{user}'s review of {items['name']}",
                "subtitle": f"Comment: {z['comment']}\nRating: {z['rating']}/5.0"
            }
            message['richContent'][0].append(menu_review)
            message['richContent'][0].append({'type':'divider'})
            count+=1
    message1['fulfillmentMessages'][1]['payload'] = message
    return jsonify(message1)
def handle_chats():
    data = request.get_json(silent=True)
    print(data['queryResult']['intent']['displayName'])
    intent = data['queryResult']['intent']['displayName']
    if ('menu_item' in intent):
        menu_item_raw = data['queryResult']['parameters']['menu_item']
        category = menu_db.find_one({'menu_items.name': re.compile(menu_item_raw, re.IGNORECASE)})
        menu_item_queried = menu_db.aggregate([
                {'$unwind': '$menu_items'},
                {'$match':{"menu_items.name": re.compile(menu_item_raw, re.IGNORECASE)}},
                {'$project': {
                            'name': '$menu_items.name',
                            'price': '$menu_items.price',
                            'rating': '$menu_items.rating',
                            'image': '$menu_items.media_urls',
                            'category': category['name']
                }}
        ])
        item = menu_item_queried.next()
        if not item:
            text_message = f"I'm sorry there are no {menu_item_raw} found"
            found = False
        else:
            found = True
            text_message = f"Here is the detail of {item['name']}"
        return construct_custom_intent('menu_item', item, text_message, found)
    elif('menu_category' in data['queryResult']['parameters']):
        menu_category = data['queryResult']['parameters']['menu_category']
        category = menu_db.find_one({'name': re.compile(menu_category, re.IGNORECASE)})
        if not category:
            text_message = f"I'm sorry there are no category of {menu_category}"
            found = False
        else:
            text_message = f"Here are the list of {menu_category}"
            found = True
        return construct_custom_intent('menu_category', category, text_message, found)
    elif('menu_category_general' in data['queryResult']['parameters']):
        category = menu_db.find({})
        text_message = "Here are the categories"
        return construct_custom_intent('menu_category_general',category, text_message, found=True)
    elif('review' in data['queryResult']['parameters']):
        reviews = review_db.find({})
        if not reviews:
            text_message = "I'm sorry there are currently no review for this restaurant"
            found = False
        else:
            text_message ="Here are some reviews for this restaurant"
            found = True
        return construct_custom_intent("review", reviews, text_message, found)
    elif('suggestion' in data['queryResult']['parameters']):
        categories = list(menu_db.find({}).sort("menu_items.rating"))
        print(categories)
        text_message="Here are the top rated dish in our menu"
        return construct_custom_intent('suggestion', categories, text_message, found = True)
    elif('menu_review' in intent):
        menu_item = data['queryResult']['parameters']['menu_item']
        menu = menu_db.find_one({'menu_items.name': re.compile(menu_item, re.IGNORECASE)})
        menu_item_queried = menu_db.aggregate([
                {'$unwind': '$menu_items'},
                {'$match':{"menu_items.name": re.compile(menu_item, re.IGNORECASE)}},
                {'$project': {
                            'name': '$menu_items.name',
                            'review_list': '$menu_items.review_list'
                }}
        ])
        item = menu_item_queried.next()
        if not menu:
            text_message = f"I'm sorry there are no review for  {item['name']}"
            found = False
        else:
            text_message = f"Here are the review for {item['name']}"
            found = True
        print(item)
        return construct_custom_intent('menu_review', item, text_message, found)
    else:
        text_message = "I'm sorry I couldn't understand what you said"
        message1 = {
            "fulfillmentMessages": [
                {"text": {"text": [text_message
                        ]}},{"payload": {}}]}
        return jsonify(message1)


"""

"""

import torch
from newspaper import Article
from sentence_transformers import SentenceTransformer, util
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
from typing import List, Tuple, Union
from torch import Tensor
import numpy as np

def retrieve(corpus_embeddings: Tensor,
             query_embedding: Tensor,
             top_k: int = 5
             ) -> List[int]:
    """Searches the corpus given the query and returns
    the top k results as indices."""
    cos_scores = util.pytorch_cos_sim(query_embedding, corpus_embeddings)[0]
    cos_scores = cos_scores.cpu()
    top_results = np.argpartition(-cos_scores, range(top_k))[0:top_k]
    return top_results.tolist()

def get_corpus_from_links(urls: List[str]) -> Tuple[List[str], List[str], Tensor]:
    # create and embed corpus
    corpus, titles = [], []
    for url in urls:
        article = Article(url)
        try:
            article.download()
            article.parse()
        except:
            continue
        corpus.append(article.text)
        titles.append(article.title)
    corpus_embeddings = model.encode(corpus, convert_to_tensor=True)
    return titles, corpus, corpus_embeddings

def get_embeddings_from_corpus(corpus: Union[List[str], str]) -> Tensor:
    if isinstance(corpus, str):
        return model.encode(corpus, convert_to_tensor=True).reshape(1, 768)
    else:
        return model.encode(corpus, convert_to_tensor=True)

def similarity_matrix(encoded_topics: List[Tensor]
                      ) -> List[List[float]]:
    topic_length = len(encoded_topics)
    m = [[0] * topic_length for x in range(topic_length)]
    for i, row in enumerate(encoded_topics):
        for j, col in enumerate(encoded_topics):
            m[i][j] = round(float(dist(row, col)), 4)
    return m

def dist(x, y):
    return util.pytorch_cos_sim(x,y)[0]

model = SentenceTransformer('distilbert-base-nli-stsb-mean-tokens')

urls = [
    "https://www.nytimes.com/2020/03/13/world/how-to-wash-your-hands-coronavirus.html",
    "https://www.nytimes.com/2020/06/11/well/live/the-hand-washing-wars.html",
    "https://www.nytimes.com/interactive/2020/04/15/burst/how-to-wash-your-hands.html",
    "https://www.nytimes.com/2016/04/21/health/washing-hands.html"
    ]

wiki = [
    "Hand washing (or handwashing), also known as hand hygiene, is the act of cleaning one's hands with soap (or equivalent materials) and water to remove viruses/bacteria/microorganisms, dirt, grease, or other harmful and unwanted substances stuck to the hands. Drying of the washed hands is part of the process as wet and moist hands are more easily recontaminated.", # hand washing -> first paragraph
    "Hand washing has many significant health benefits, including minimizing the spread of influenza, coronavirus, and other infectious diseases; preventing infectious causes of diarrhea; decreasing respiratory infections; and reducing infant mortality rate at home birth deliveries. A 2013 study showed that improved hand washing practices may lead to small improvements in the length growth in children under five years of age.", # hand washing -> public health -> healt benefits -> first paragraph
    "When not wearing a mask, the CDC, WHO, and NHS recommends covering the mouth and nose with a tissue when coughing or sneezing and recommends using the inside of the elbow if no tissue is available.[112][121][132] Proper hand hygiene after any cough or sneeze is encouraged.[112][121] The WHO also recommends that individuals wash hands often with soap and water for at least 20 seconds, especially after going to the toilet or when hands are visibly dirty, before eating and after blowing one's nose." # covid-19 pandemic -> hand washing -> first paragraph
    ]

paragraphs = [
    "Hand washing cleans the hand with soap to remove harmful viruses, bacteria and microorganisms." # compressed first wiki paragraph
    "Hand washing removes viruses from the hand." # more concise paragraph with emphasis on virus on hands
    "Hand washing is an effective hygiene measure in a pandemic" # more conside paragraph with emphasis is on measure in pandemic
    ]

topics = [
    "hand washing",
    "hand cleaning"
    "washing hands"
    "wash hands"
    "disinfect hands"
    ]

emb_corpus = get_corpus_from_links(urls)[2]
emb_wiki = get_embeddings_from_corpus(wiki)
emb_para = get_embeddings_from_corpus(paragraphs)
emb_topics = get_embeddings_from_corpus(topics)
embeddings = torch.cat((emb_corpus, emb_wiki, emb_para, emb_topics), dim=0).cpu()


import torch
from torch.autograd import Function
import torch.nn as nn

class MaxMinPoolingFunction(Function):
    """MinMax Layer for pooling the top and bottom instances of a given 1D embedding (section 2.3)"""
    def __init__(self, R):
        """
            Args:
                R: Number of top and bottom instances to pool after convolving
        """
        super().__init__()
        self.R = R

    def forward(self, x_input):
        """Propogates forward through the network given the 1D conv layer

            Args:
                x_input: Feature embedding from a 1D conv layer
                         of shape (batch_size, 1, n_regionss)

            Returns:
                output: top and bottom instances of shape (batch_size, 2 * R)
        """
        batch_size = x_input.shape[0]
        n_regions = x_input.shape[2]

        # sort the feature embedding and save the max & min indices for backprop
        x_sorted, x_indices = torch.sort(x_input.view(batch_size, n_regions),
                                         dim=1, descending=True)

        self.indices_max = x_indices.narrow(1, 0, self.R)
        self.indices_min = x_indices.narrow(1, -self.R, self.R)

        output_max = x_sorted.narrow(1, 0, self.R)
        output_min = x_sorted.narrow(1, -self.R, self.R)

        # concat the top & bottom instances for a MLP classifier (Figure 2)
        output = torch.cat((output_max, output_min), dim=1)

        self.save_for_backward(x_input)
        return output.view(batch_size, 2*self.R)

    def backward(self, grad_output):
        """Propogates back through the network returning the gradients only through
        the tiles selected as the top & bottom instances

            Args:
                grad_output: tensor of shape (batch_size, 2*R) storing the gradients
                             from the previous layer(s)

            Returns:
                grad_input: tensor of shape (batch_size, 1, n_regions) such that all
                            gradient are 0'd out except indices that stored the top &
                            bottom instances
        """

        x_input, = self.saved_tensors
        batch_size = x_input.shape[0]
        n_regions = x_input.shape[2]

        # store the gradients from the max * min R entries
        grad_output_max = grad_output.narrow(1, 0, self.R)
        grad_output_min = grad_output.narrow(1, -self.R, self.R)

        # 0 out all other gradients not used as a top/bottom instance
        grad_input_max = grad_output.new()
        grad_input_max.resize_(batch_size, n_regions).fill_(0)
        grad_input_max.scatter_(1, self.indices_max, grad_output_max)

        grad_input_min = grad_output.new()
        grad_input_min.resize_(batch_size, n_regions).fill_(0)
        grad_input_min.scatter_(1, self.indices_min, grad_output_min)

        # collapse both tensors by adding them along the tile axis
        grad_input = grad_input_max.add(grad_input_min)
        return grad_input.view(batch_size, 1, n_regions)

class MaxMinPooling(nn.Module):
    """MaxMin Module to add to our network such that pytorch can simply call forward"""
    def __init__(self, R):
        super().__init__()
        self.R = R

    def forward(self, x):
        x = MaxMinPoolingFunction(self.R)(x)
        return x

    def __repr__(self):
        return '{0} (R={1})'.format(self.__class__.__name__, self.R)