from __future__ import unicode_literals
import re
import unicodedata
import torch
import streamlit as st
from transformers import T5ForConditionalGeneration, T5Tokenizer


def load_model():
    # 学習済みモデルをHugging Face model hubからダウンロードする
    model_dir_name = "sonoisa/t5-qiita-title-generation"

    # トークナイザー（SentencePiece）
    tokenizer = T5Tokenizer.from_pretrained(model_dir_name, is_fast=True)

    # 学習済みモデル
    trained_model = T5ForConditionalGeneration.from_pretrained(model_dir_name)

    # GPUの利用有無
    USE_GPU = torch.cuda.is_available()
    if USE_GPU:
        trained_model.cuda()

    return trained_model, tokenizer


def unicode_normalize(cls, s):
    pt = re.compile("([{}]+)".format(cls))

    def norm(c):
        return unicodedata.normalize("NFKC", c) if pt.match(c) else c

    s = "".join(norm(x) for x in re.split(pt, s))
    s = re.sub("－", "-", s)
    return s


def remove_extra_spaces(s):
    s = re.sub("[ 　]+", " ", s)
    blocks = "".join(
        (
            "\u4E00-\u9FFF",  # CJK UNIFIED IDEOGRAPHS
            "\u3040-\u309F",  # HIRAGANA
            "\u30A0-\u30FF",  # KATAKANA
            "\u3000-\u303F",  # CJK SYMBOLS AND PUNCTUATION
            "\uFF00-\uFFEF",  # HALFWIDTH AND FULLWIDTH FORMS
        )
    )
    basic_latin = "\u0000-\u007F"

    def remove_space_between(cls1, cls2, s):
        p = re.compile("([{}]) ([{}])".format(cls1, cls2))
        while p.search(s):
            s = p.sub(r"\1\2", s)
        return s

    s = remove_space_between(blocks, blocks, s)
    s = remove_space_between(blocks, basic_latin, s)
    s = remove_space_between(basic_latin, blocks, s)
    return s


def normalize_neologd(s):
    s = s.strip()
    s = unicode_normalize("０-９Ａ-Ｚａ-ｚ｡-ﾟ", s)

    def maketrans(f, t):
        return {ord(x): ord(y) for x, y in zip(f, t)}

    s = re.sub("[˗֊‐‑‒–⁃⁻₋−]+", "-", s)  # normalize hyphens
    s = re.sub("[﹣－ｰ—―─━ー]+", "ー", s)  # normalize choonpus
    s = re.sub("[~∼∾〜〰～]+", "〜", s)  # normalize tildes (modified by Isao Sonobe)
    s = s.translate(
        maketrans(
            "!\"#$%&'()*+,-./:;<=>?@[¥]^_`{|}~｡､･｢｣",
            "！”＃＄％＆’（）＊＋，－．／：；＜＝＞？＠［￥］＾＿｀｛｜｝〜。、・「」",
        )
    )

    s = remove_extra_spaces(s)
    s = unicode_normalize("！”＃＄％＆’（）＊＋，－．／：；＜＞？＠［￥］＾＿｀｛｜｝〜", s)  # keep ＝,・,「,」
    s = re.sub("[’]", "'", s)
    s = re.sub("[”]", '"', s)
    return s


CODE_PATTERN = re.compile(r"```.*?```", re.MULTILINE | re.DOTALL)
LINK_PATTERN = re.compile(r"!?\[([^\]\)]+)\]\([^\)]+\)")
IMG_PATTERN = re.compile(r"<img[^>]*>")
URL_PATTERN = re.compile(r"(http|ftp)s?://[^\s]+")
NEWLINES_PATTERN = re.compile(r"(\s*\n\s*)+")


def clean_markdown(markdown_text):
    markdown_text = CODE_PATTERN.sub(r"", markdown_text)
    markdown_text = LINK_PATTERN.sub(r"\1", markdown_text)
    markdown_text = IMG_PATTERN.sub(r"", markdown_text)
    markdown_text = URL_PATTERN.sub(r"", markdown_text)
    markdown_text = NEWLINES_PATTERN.sub(r"\n", markdown_text)
    markdown_text = markdown_text.replace("`", "")
    return markdown_text


def normalize_text(markdown_text):
    markdown_text = clean_markdown(markdown_text)
    markdown_text = markdown_text.replace("\t", " ")
    markdown_text = normalize_neologd(markdown_text).lower()
    markdown_text = markdown_text.replace("\n", " ")
    return markdown_text


def preprocess_qiita_body(markdown_text):
    return "body: " + normalize_text(markdown_text)[:4000]


def postprocess_title(title):
    return re.sub(r"^title: ", "", title)

st.title("Qiita記事タイトル案生成")

description_text = st.empty()

if "trained_model" not in st.session_state:
    description_text.text("...モデル読み込み中...")

    trained_model, tokenizer = load_model()
    trained_model.eval()
    
    st.session_state.trained_model = trained_model
    st.session_state.tokenizer = tokenizer

trained_model = st.session_state.trained_model
tokenizer = st.session_state.tokenizer

# GPUの利用有無
USE_GPU = torch.cuda.is_available()

description_text.text("記事の本文をコピペ入力して、タイトル生成ボタンを押すと、タイトル案が10個生成されます。\nGPUが使えないため生成に数十秒かかります。")
qiita_body = st.text_area(label="記事の本文", value="", height=300, max_chars=4000)
answer = st.button("タイトル生成")

if answer:
    title_fieids = st.empty()
    title_fieids.markdown("...生成中...")

    MAX_SOURCE_LENGTH = 512  # 入力される記事本文の最大トークン数
    MAX_TARGET_LENGTH = 64  # 生成されるタイトルの最大トークン数

    # 前処理とトークナイズを行う
    inputs = [preprocess_qiita_body(qiita_body)]
    batch = tokenizer.batch_encode_plus(
        inputs,
        max_length=MAX_SOURCE_LENGTH,
        truncation=True,
        padding="longest",
        return_tensors="pt",
    )

    input_ids = batch["input_ids"]
    input_mask = batch["attention_mask"]
    if USE_GPU:
        input_ids = input_ids.cuda()
        input_mask = input_mask.cuda()

    # 生成処理を行う
    outputs = trained_model.generate(
        input_ids=input_ids,
        attention_mask=input_mask,
        max_length=MAX_TARGET_LENGTH,
        return_dict_in_generate=True,
        output_scores=True,
        temperature=1.0,  # 生成にランダム性を入れる温度パラメータ
        num_beams=10,  # ビームサーチの探索幅
        diversity_penalty=1.0,  # 生成結果の多様性を生み出すためのペナルティ
        num_beam_groups=10,  # ビームサーチのグループ数
        num_return_sequences=10,  # 生成する文の数
        repetition_penalty=1.5,  # 同じ文の繰り返し（モード崩壊）へのペナルティ
    )

    # 生成されたトークン列を文字列に変換する
    generated_titles = [
        tokenizer.decode(
            ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )
        for ids in outputs.sequences
    ]

    # 生成されたタイトルを表示する
    titles = "## タイトル案:\n\n"

    for i, title in enumerate(generated_titles):
        titles += f"1. {postprocess_title(title)}\n"

    title_fieids.markdown(titles)