"""
Feature tables for embeddings
"""

import re

# Mora list
MORAS = [
    # x
    "ア",
    "イ",
    "ウ",
    "エ",
    "オ",
    "ワ",
    # w
    "ウィ",
    "ウェ",
    "ウォ",
    # y
    "ヤ",
    "ユ",
    "ヨ",
    "イェ",
    # k
    "カ",
    "キ",
    "ク",
    "ケ",
    "コ",
    # kw
    "クァ",
    "クィ",
    "クェ",
    "クォ",
    "クヮ",
    # ky
    "キャ",
    "キュ",
    "キェ",
    "キョ",
    # g
    "ガ",
    "ギ",
    "グ",
    "ゲ",
    "ゴ",
    # gw
    "グァ",
    "グィ",
    "グェ",
    "グォ",
    "グヮ",
    # gy
    "ギャ",
    "ギュ",
    "ギェ",
    "ギョ",
    # s
    "サ",
    "スィ",
    "ス",
    "セ",
    "ソ",
    # sh
    "シ",
    "シェ",
    "シャ",
    "シュ",
    "ショ",
    # z
    "ザ",
    "ズィ",
    "ズ",
    "ゼ",
    "ゾ",
    "ヂェ",
    # j
    "ジ",
    "ジェ",
    "ジャ",
    "ジュ",
    "ジョ",
    "ヂャ",
    "ヂュ",
    "ヂョ",
    # t
    "タ",
    "ティ",
    "トゥ",
    "テ",
    "ト",
    # ty
    "テャ",
    "テュ",
    "テョ",
    # d
    "ダ",
    "ディ",
    "ドゥ",
    "デ",
    "ド",
    # dy
    "デャ",
    "デュ",
    "デョ",
    # ch
    "チ",
    "チェ",
    "チャ",
    "チュ",
    "チョ",
    # ts
    "ツァ",
    "ツィ",
    "ツ",
    "ツェ",
    "ツォ",
    # n
    "ナ",
    "ニ",
    "ヌ",
    "ネ",
    "ノ",
    # ny
    "ニャ",
    "ニュ",
    "ニェ",
    "ニョ",
    # h
    "ハ",
    "ヒ",
    "ヘ",
    "ホ",
    # hy
    "ヒャ",
    "ヒュ",
    "ヒェ",
    "ヒョ",
    # p
    "パ",
    "ピ",
    "プ",
    "ペ",
    "ポ",
    # pw
    "プァ",
    "プィ",
    "プェ",
    "プォ",
    # py
    "ピャ",
    "ピュ",
    "ピェ",
    "ピョ",
    "プャ",
    "プュ",
    "プョ",
    # b
    "バ",
    "ビ",
    "ブ",
    "ベ",
    "ボ",
    "ヴ",
    # v
    "ヴァ",
    "ヴィ",
    "ヴェ",
    "ヴォ",
    "ブァ",
    "ブィ",
    "ブェ",
    "ブォ",
    # by
    "ビャ",
    "ビュ",
    "ビェ",
    "ビョ",
    # vy
    "ヴャ",
    "ヴュ",
    "ヴョ",
    "ブャ",
    "ブュ",
    "ブョ",
    # f
    "ファ",
    "フィ",
    "フ",
    "フェ",
    "フォ",
    # fy
    "フャ",
    "フュ",
    "フョ",
    # m
    "マ",
    "ミ",
    "ム",
    "メ",
    "モ",
    # my
    "ミャ",
    "ミュ",
    "ミェ",
    "ミョ",
    # r
    "ラ",
    "リ",
    "ル",
    "レ",
    "ロ",
    # ry
    "リャ",
    "リュ",
    "リェ",
    "リョ",
    # N
    "ン",
    # cl
    "ッ",
    # Longvowel symbol
    "ー",
]
PUNCTUATIONS = [",", ".", "?", "!"]
MORAS += PUNCTUATIONS

# POS list defined in Unidic-CSJ
POS_UNIDIC_CSJ = [
    "代名詞:*:*:*",
    "副詞:*:*:*",
    "助動詞:*:*:*",
    # PARTICLE
    "助詞:係助詞:*:*",
    "助詞:副助詞:*:*",
    "助詞:接続助詞:*:*",
    "助詞:格助詞:*:*",
    "助詞:準体助詞:*:*",
    "助詞:終助詞:*:*",
    # VERB
    "動詞:一般:*:*",
    "動詞:非自立可能:*:*",
    # NOUN
    "名詞:助動詞語幹:*:*",
    "名詞:数詞:*:*",
    "名詞:普通名詞:サ変可能:*",
    "名詞:普通名詞:サ変形状詞可能:*",
    "名詞:普通名詞:一般:*",
    "名詞:普通名詞:副詞可能:*",
    "名詞:普通名詞:助数詞可能:*",
    "名詞:普通名詞:形状詞可能:*",
    # NOUN: NER
    "名詞:固有名詞:一般:*",
    "名詞:固有名詞:人名:一般",
    "名詞:固有名詞:人名:名",
    "名詞:固有名詞:人名:姓",
    "名詞:固有名詞:地名:一般",
    "名詞:固有名詞:地名:国",
    # ADJECTIVE
    "形容詞:一般:*:*",
    "形容詞:非自立可能:*:*",
    # NA-ADJECTIVE
    "形状詞:タリ:*:*",
    "形状詞:一般:*:*",
    "形状詞:助動詞語幹:*:*",
    # INTERJECTION
    "感動詞:フィラー:*:*",
    "感動詞:一般:*:*",
    # SUFFIX
    "接尾辞:動詞的:*:*",
    "接尾辞:名詞的:サ変可能:*",
    "接尾辞:名詞的:サ変形状詞可能:*",
    "接尾辞:名詞的:一般:*",
    "接尾辞:名詞的:副詞可能:*",
    "接尾辞:名詞的:助数詞:*",
    "接尾辞:名詞的:形状詞可能:*",
    "接尾辞:形容詞的:*:*",
    "接尾辞:形状詞的:*:*",
    # CONJUNCTION
    "接続詞:*:*:*",
    # PREFIX
    "接頭辞:*:*:*",
    # ADNOMINAL
    "連体詞:*:*:*",
    # SYMBOL
    "空白:*:*:*",
    "補助記号:一般:*:*",
    "補助記号:句点:*:*",
    "補助記号:括弧閉:*:*",
    "補助記号:括弧開:*:*",
    "補助記号:絵文字:一般:*",
    "補助記号:読点:*:*",
    "補助記号:ＡＡ:一般:*",
    "補助記号:ＡＡ:顔文字:*",
    "記号:一般:*:*",
    "記号:文字:*:*",
]

# POS list defined in OpenJTalk dictionary (NAIST-jdic)
POS_OPEN_JTALK = [
    "その他:間投:*:*",
    "フィラー:*:*:*",
    "感動詞:*:*:*",
    "記号:アルファベット:*:*",
    "記号:一般:*:*",
    "記号:括弧開:*:*",
    "記号:括弧閉:*:*",
    "記号:句点:*:*",
    "記号:空白:*:*",
    "記号:読点:*:*",
    "形容詞:自立:*:*",
    "形容詞:接尾:*:*",
    "形容詞:非自立:*:*",
    "助詞:格助詞:一般:*",
    "助詞:格助詞:引用:*",
    "助詞:格助詞:連語:*",
    "助詞:係助詞:*:*",
    "助詞:終助詞:*:*",
    "助詞:接続助詞:*:*",
    "助詞:特殊:*:*",
    "助詞:副詞化:*:*",
    "助詞:副助詞:*:*",
    "助詞:副助詞／並立助詞／終助詞:*:*",
    "助詞:並立助詞:*:*",
    "助詞:連体化:*:*",
    "助動詞:*:*:*",
    "接続詞:*:*:*",
    "接頭詞:形容詞接続:*:*",
    "接頭詞:数接続:*:*",
    "接頭詞:動詞接続:*:*",
    "接頭詞:名詞接続:*:*",
    "動詞:自立:*:*",
    "動詞:接尾:*:*",
    "動詞:非自立:*:*",
    "副詞:一般:*:*",
    "副詞:助詞類接続:*:*",
    "名詞:サ変接続:*:*",
    "名詞:ナイ形容詞語幹:*:*",
    "名詞:一般:*:*",
    "名詞:引用文字列:*:*",
    "名詞:形容動詞語幹:*:*",
    "名詞:固有名詞:一般:*",
    "名詞:固有名詞:人名:一般",
    "名詞:固有名詞:人名:姓",
    "名詞:固有名詞:人名:名",
    "名詞:固有名詞:組織:*",
    "名詞:固有名詞:地域:一般",
    "名詞:固有名詞:地域:国",
    "名詞:数:*:*",
    "名詞:接続詞的:*:*",
    "名詞:接尾:サ変接続:*",
    "名詞:接尾:一般:*",
    "名詞:接尾:形容動詞語幹:*",
    "名詞:接尾:助数詞:*",
    "名詞:接尾:助動詞語幹:*",
    "名詞:接尾:人名:*",
    "名詞:接尾:地域:*",
    "名詞:接尾:特殊:*",
    "名詞:接尾:副詞可能:*",
    "名詞:代名詞:一般:*",
    "名詞:代名詞:縮約:*",
    "名詞:動詞非自立的:*:*",
    "名詞:特殊:助動詞語幹:*",
    "名詞:非自立:一般:*",
    "名詞:非自立:形容動詞語幹:*",
    "名詞:非自立:助動詞語幹:*",
    "名詞:非自立:副詞可能:*",
    "名詞:副詞可能:*:*",
    "連体詞:*:*:*",
]

# Accent-con-type defined in Unidic-CSJ
ACCENT_CON_TYPES_UNIDIC_CSJ = [
    "1",
    "C1",
    "C2",
    "C3",
    "C4",
    "C5",
    "F1",
    "F2",
    "F3",
    "F4",
    "F6",
    "P1",
    "P2",
    "P4",
    "P6",
    "P13",
    "P14",
]

# Accent-con-type defined in OpenJTalk dictionary (NAIST-jdic)
ACCENT_CON_TYPES_OPEN_JTALK = [
    "1",
    "C1",
    "C2",
    "C3",
    "C4",
    "C5",
    "F1",
    "F2",
    "F3",
    "F4",
    "F6",
    "P1",
    "P2",
    "P4",
    "P6",
    "P13",
    "P14",
]

# C-type defined in Unidic-CSJ
C_TYPES_UNIDIC_CSJ = [
    "カ行変格",
    "サ行変格",
    "上一段-ア行",
    "上一段-カ行",
    "上一段-ガ行",
    "上一段-ザ行",
    "上一段-タ行",
    "上一段-ナ行",
    "上一段-ハ行",
    "上一段-バ行",
    "上一段-マ行",
    "上一段-ラ行",
    "下一段-ア行",
    "下一段-カ行",
    "下一段-ガ行",
    "下一段-サ行",
    "下一段-ザ行",
    "下一段-タ行",
    "下一段-ダ行",
    "下一段-ナ行",
    "下一段-ハ行",
    "下一段-バ行",
    "下一段-マ行",
    "下一段-ラ行",
    "五段-カ行",
    "五段-ガ行",
    "五段-サ行",
    "五段-タ行",
    "五段-ナ行",
    "五段-バ行",
    "五段-マ行",
    "五段-ラ行",
    "五段-ワア行",
    "助動詞-ジャ",
    "助動詞-タ",
    "助動詞-タイ",
    "助動詞-ダ",
    "助動詞-デス",
    "助動詞-ドス",
    "助動詞-ナイ",
    "助動詞-ナンダ",
    "助動詞-ヌ",
    "助動詞-ヒン",
    "助動詞-ヘン",
    "助動詞-マイ",
    "助動詞-マス",
    "助動詞-ヤ",
    "助動詞-ヤス",
    "助動詞-ヤン",
    "助動詞-ラシイ",
    "助動詞-レル",
    "助動詞-ンス",
    "形容詞",
    "文語カ行変格",
    "文語サ行変格",
    "文語ラ行変格",
    "文語上一段-マ行",
    "文語上一段-ワ行",
    "文語上二段-タ行",
    "文語上二段-バ行",
    "文語上二段-ラ行",
    "文語下二段-ア行",
    "文語下二段-カ行",
    "文語下二段-ガ行",
    "文語下二段-サ行",
    "文語下二段-タ行",
    "文語下二段-ダ行",
    "文語下二段-ナ行",
    "文語下二段-ハ行",
    "文語下二段-マ行",
    "文語下二段-ラ行",
    "文語下二段-ワ行",
    "文語助動詞-キ",
    "文語助動詞-ケリ",
    "文語助動詞-ゴトシ",
    "文語助動詞-ズ",
    "文語助動詞-タリ-完了",
    "文語助動詞-タリ-断定",
    "文語助動詞-ナリ-断定",
    "文語助動詞-ベシ",
    "文語助動詞-マジ",
    "文語助動詞-リ",
    "文語四段-カ行",
    "文語四段-ガ行",
    "文語四段-サ行",
    "文語四段-タ行",
    "文語四段-ハ行",
    "文語四段-バ行",
    "文語四段-マ行",
    "文語四段-ラ行",
    "文語形容詞-ク",
    "文語形容詞-シク",
    "無変化型",
    "特殊型",
]

# C-type defined in OpenJTalk dictionary (NAIST-jdic)
C_TYPES_OPEN_JTALK = [
    "カ変・クル",
    "カ変・来ル",
    "サ変・スル",
    "サ変・－スル",
    "サ変・－ズル",
    "ラ変",
    "一段",
    "一段・クレル",
    "一段・得ル",
    "上二・ダ行",
    "上二・ハ行",
    "下二・カ行",
    "下二・ガ行",
    "下二・タ行",
    "下二・ダ行",
    "下二・ハ行",
    "下二・マ行",
    "下二・得",
    "不変化型",
    "五段・カ行イ音便",
    "五段・カ行促音便",
    "五段・カ行促音便ユク",
    "五段・ガ行",
    "五段・サ行",
    "五段・タ行",
    "五段・ナ行",
    "五段・バ行",
    "五段・マ行",
    "五段・ラ行",
    "五段・ラ行アル",
    "五段・ラ行特殊",
    "五段・ワ行ウ音便",
    "五段・ワ行促音便",
    "四段・サ行",
    "四段・タ行",
    "四段・ハ行",
    "四段・バ行",
    "形容詞・アウオ段",
    "形容詞・イ段",
    "文語・キ",
    "文語・ケリ",
    "文語・ゴトシ",
    "文語・ナリ",
    "文語・ベシ",
    "文語・マジ",
    "文語・リ",
    "文語・ル",
    "特殊・ジャ",
    "特殊・タ",
    "特殊・タイ",
    "特殊・ダ",
    "特殊・デス",
    "特殊・ナイ",
    "特殊・ヌ",
    "特殊・マス",
    "特殊・ヤ",
]

# C-form defined in Unidic-CSJ
C_FORMS_UNIDIC_CSJ = [
    "ク語法",
    "仮定形-一般",
    "仮定形-融合",
    "命令形",
    "已然形-一般",
    "已然形-補助",
    "意志推量形",
    "未然形-サ",
    "未然形-セ",
    "未然形-一般",
    "未然形-撥音便",
    "未然形-補助",
    "終止形-一般",
    "終止形-促音便",
    "終止形-撥音便",
    "終止形-融合",
    "終止形-補助",
    "語幹-サ",
    "語幹-一般",
    "連体形-イ音便",
    "連体形-ウ音便",
    "連体形-一般",
    "連体形-促音便",
    "連体形-撥音便",
    "連体形-省略",
    "連体形-補助",
    "連用形-イ音便",
    "連用形-ウ音便",
    "連用形-ト",
    "連用形-ニ",
    "連用形-一般",
    "連用形-促音便",
    "連用形-撥音便",
    "連用形-省略",
    "連用形-融合",
    "連用形-補助",
]

# C-form defined in OpenJTalk dictionary (NAIST-jdic)
C_FORMS_OPEN_JTALK = [
    "ガル接続",
    "仮定形",
    "仮定縮約１",
    "仮定縮約２",
    "体言接続",
    "体言接続特殊",
    "体言接続特殊２",
    "命令ｅ",
    "命令ｉ",
    "命令ｒｏ",
    "命令ｙｏ",
    "基本形",
    "文語基本形",
    "未然ウ接続",
    "未然ヌ接続",
    "未然レル接続",
    "未然形",
    "未然特殊",
    "現代基本形",
    "連用ゴザイ接続",
    "連用タ接続",
    "連用テ接続",
    "連用デ接続",
    "連用ニ接続",
    "連用形",
    "音便基本形",
]

# Accent-type support Unidic-CSJ
ACCENT_TYPES_UNIDIC_CSJ = list(range(23))

# Accent-type support OpenJTalk
ACCENT_TYPES_OPEN_JTALK = list(range(50))

# Chain falg of OpenJTalk
CHAIN_FLAG_OPEN_JTALK = list(range(2))

# Word-type Table defined in Unidic-CSJ
WORD_TYPES_UNIDIC_CSJ = ["和", "漢", "外", "混", "固", "記号", "他"]

# Accent-mod-type defined in Unidic-CSJ
ACCENT_MOD_TYPES_UNIDIC_CSJ = ["M1@0", "M1@1", "M2@0", "M2@1", "M2@2", "M4@1"]

FEATURE_TABLES = {
    "unidic-csj": {
        "surface": None,  # surface will be loaded from vocab.pkl
        "mora": MORAS,
        "pos": POS_UNIDIC_CSJ,
        "c_type": C_TYPES_UNIDIC_CSJ,
        "c_form": C_FORMS_UNIDIC_CSJ,
        "accent_type": ACCENT_TYPES_UNIDIC_CSJ,
        "word_type": WORD_TYPES_UNIDIC_CSJ,
        "accent_con_type": ACCENT_CON_TYPES_UNIDIC_CSJ,
        "accent_mod_type": ACCENT_MOD_TYPES_UNIDIC_CSJ,
    },
    "open-jtalk": {
        "surface": None,  # surface will be loaded from vocab.pkl
        "mora": MORAS,
        "pos": POS_OPEN_JTALK,
        "c_type": C_TYPES_OPEN_JTALK,
        "c_form": C_FORMS_OPEN_JTALK,
        "accent_type": ACCENT_TYPES_OPEN_JTALK,
        "accent_con_type": ACCENT_CON_TYPES_OPEN_JTALK,
        "chain_flag": CHAIN_FLAG_OPEN_JTALK,
    },
}

RAW_FEATURE_KEYS = {
    "unidic-csj": [
        "surface",
        "pron",
        "pos",
        "c_type",
        "c_form",
        "accent_type",
        "word_type",
        "accent_con_type",
        "accent_mod_type",
    ],
    "open-jtalk": [
        "surface",
        "pron",
        "pos",
        "c_type",
        "c_form",
        "accent_type",
        "accent_con_type",
        "chain_flag",
    ],
}

_re_accent_con_f_type = re.compile(r"((形容詞|名詞|動詞)\%(F\d))")


def is_noun(pos_tag):
    return pos_tag.startswith("名詞")


def is_verb(pos_tag):
    return pos_tag.startswith("動詞")


def is_adjective(pos_tag):
    return pos_tag.startswith("形容詞")


def _pos_id_to_tag(pos_tag, unk_token):
    if is_noun(pos_tag):
        pos_key = "名詞"
    elif is_verb(pos_tag):
        pos_key = "動詞"
    elif is_adjective(pos_tag):
        pos_key = "形容詞"
    else:
        pos_key = unk_token

    return pos_key


def parse_accent_con_type(a_con_type, pos, unk_token="[UNK]"):
    if a_con_type is None:
        a_con_type = unk_token
    elif "%" in a_con_type:
        a_con_type = {
            con[1]: con[2] for con in _re_accent_con_f_type.findall(a_con_type)
        }
        pos_tag = _pos_id_to_tag(pos, unk_token)
        a_con_type = a_con_type.get(pos_tag, unk_token)

    return a_con_type
