from __future__ import unicode_literals
import re
import unicodedata
import torch
import streamlit as st
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import numpy as np
import scipy.spatial
from transformers import BertJapaneseTokenizer, BertModel
import pyminizip


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)
    s = s.lower()
    return s


def normalize_text(text):
    return normalize_neologd(text)


def normalize_title(title):
    title = title.strip()
    
    match = re.match(r"^「([^」]+)」$", title)
    if match:
        title = match.group(1)

    match = re.match(r"^POP素材「([^」]+)」$", title)
    if match:
        title = match.group(1)
    
    title = re.sub(r"(の?(?:イラスト|イラストの|イラストト|イ子のラスト|イラス|イラスト文字|「イラスト文字」|イラストPOP文字|ペンキ文字|タイトル文字|イラスト・メッセージ|イラスト文字・バナー|キャラクター(たち)?|マーク|アイコン|シルエット|シルエット素材|フレーム（枠）|フレーム|フレーム素材|テンプレート|パターン|パターン素材|ライン素材|コーナー素材|リボン型バナー|評価スタンプ|背景素材))+(\s*([0-9０-９]*|その[0-9０-９]+))(です。)?", "", title)
    
    title = normalize_text(title)
    
    if title.strip() == "":
        raise ValueError(title)
    
    return title


class SentenceBertJapanese:
    def __init__(self, model_name_or_path, device=None):
        self.tokenizer = BertJapaneseTokenizer.from_pretrained(model_name_or_path)
        self.model = BertModel.from_pretrained(model_name_or_path)
        self.model.eval()

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = torch.device(device)
        self.model.to(device)

    def _mean_pooling(self, model_output, attention_mask):
        token_embeddings = model_output[
            0
        ]  # First element of model_output contains all token embeddings
        input_mask_expanded = (
            attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        )
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
            input_mask_expanded.sum(1), min=1e-9
        )

    @torch.no_grad()
    def encode(self, sentences, batch_size=8):
        all_embeddings = []
        iterator = range(0, len(sentences), batch_size)
        for batch_idx in iterator:
            batch = sentences[batch_idx : batch_idx + batch_size]

            encoded_input = self.tokenizer.batch_encode_plus(
                batch, padding="longest", truncation=True, return_tensors="pt"
            ).to(self.device)
            model_output = self.model(**encoded_input)
            sentence_embeddings = self._mean_pooling(
                model_output, encoded_input["attention_mask"]
            ).to("cpu")

            all_embeddings.extend(sentence_embeddings)

        # return torch.stack(all_embeddings).numpy()
        return torch.stack(all_embeddings)


st.title("いらすと検索")
description_text = st.empty()

if "model" not in st.session_state:
    description_text.text("...モデル読み込み中...")
    model = SentenceBertJapanese("sonoisa/sentence-bert-base-ja-mean-tokens")
    st.session_state.model = model

    pyminizip.uncompress(
        "irasuto_items_20210224.pq.zip", st.secrets["ZIP_PASSWORD"], None, 1
    )
    
    df = pq.read_table("irasuto_items_20210224.parquet").to_pandas()
    st.session_state.df = df
    
    sentence_vectors = np.stack(df["sentence_vector"])
    st.session_state.sentence_vectors = sentence_vectors
    
model = st.session_state.model
df = st.session_state.df
sentence_vectors = st.session_state.sentence_vectors

description_text.text("説明文の意味が近い「いらすとや」画像を検索します。\nキーワードを列挙するよりも、自然な文章を入力した方が精度よく検索できます。\n画像は必ずリンク先の「いらすとや」さんのページを開き、そこからダウンロードしてください。")

def clear_result():
    result_text.text("")

prev_query = ""
query_input = st.text_input(label="説明文", value="", on_change=clear_result)

closest_n = st.number_input(label="検索数", min_value=1, value=10, max_value=100) 

search_buttion = st.button("検索")

result_text = st.empty()

if search_buttion or prev_query != query_input:
    query = normalize_text(query_input)
    prev_query = query_input
    query_embedding = model.encode([query]).numpy()

    distances = scipy.spatial.distance.cdist(
        query_embedding, sentence_vectors, metric="cosine"
    )[0]

    results = zip(range(len(distances)), distances)
    results = sorted(results, key=lambda x: x[1])

    md_content = ""
    for i, (idx, distance) in enumerate(results[0:closest_n]):
        page_url = df.iloc[idx]["page"]
        for img_url in df.iloc[idx]["images"]:
            md_content += f'1. <a href="{page_url}" target="_blank" rel="noopener noreferrer"><img src="{img_url}" width="100"></a>'
        md_content += f' [{distance / 2:.4f}: {df.iloc[idx]["description"]}]({page_url})\n'
        
    result_text.markdown(md_content, unsafe_allow_html=True)