#!/usr/bin/env python
# coding: utf-8

# # Kuromoji completion を使用した日本語向けオートコンプリート機能の実装
# 

# ## 概要
# オートコンプリートは、あらかじめ登録されているデータに沿って、クエリの変換候補を提示する機能です。
# 
# 日本語検索プラグインである Kuromoji では、オートコンプリートを簡単に実装できる [Kuromoji completion][kuromoji-completion] を提供しています。
# 
# 本ラボでは、同機能を使用した日本語向けのサジェスタを実装します。
# 
# ### ラボの構成
# 
# 本ラボでは、ノートブック環境(JupyterLab) および Amazon OpenSearch Service を使用します。
# <img src="./img/architecture.png" width="1024"> 
# 
# [kuromoji-completion]: https://opensearch.org/docs/latest/analyzers/token-filters/kuromoji-completion/
# 
# ### 使用するデータセット
# 本ラボでは、 [Geolonia][geolonia] が CC BY 4.0 ライセンスの元で配布している [住所データツール v2 データセット][japanese-addresses-v2] を使用します。
# 
# [geolonia]: https://www.geolonia.com/
# [japanese-addresses-v2]: https://github.com/geolonia/japanese-addresses-v2

# ## 事前作業

# ### パッケージインストール

# In[1]:


get_ipython().system('pip install opensearch-py requests-aws4auth "awswrangler[opensearch]" --quiet')


# ### インポート

# In[2]:


import boto3
import json
import logging

import awswrangler as wr
import pandas as pd
from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth


# ### ヘルパー関数の定義
# 以降の処理を実行する際に必要なヘルパー関数を定義しておきます。

# In[3]:


def search_cloudformation_output(stackname, key):
    cloudformation_client = boto3.client("cloudformation", region_name=default_region)
    for output in cloudformation_client.describe_stacks(StackName=stackname)["Stacks"][0]["Outputs"]:
        if output["OutputKey"] == key:
            return output["OutputValue"]
    raise ValueError(f"{key} is not found in outputs of {stackname}.")


# ### 共通変数のセット

# In[4]:


default_region = boto3.Session().region_name
logging.getLogger().setLevel(logging.ERROR)


# ### OpenSearch クラスターへの接続確認
# 
# OpenSearch クラスターへのネットワーク接続性が確保されており、OpenSearch の Security 機能により API リクエストが許可されているかを確認します。
# 
# レスポンスに cluster_name や cluster_uuid が含まれていれば、接続確認が無事完了したと判断できます

# In[5]:


cloudformation_stack_name = "search-lab-jp"
opensearch_cluster_endpoint = search_cloudformation_output(cloudformation_stack_name, "OpenSearchDomainEndpoint")

credentials = boto3.Session().get_credentials()
service_code = "es"
auth = AWSV4SignerAuth(credentials=credentials, region=default_region, service=service_code)
opensearch_client = OpenSearch(
    hosts=[{"host": opensearch_cluster_endpoint, "port": 443}],
    http_compress=True, 
    http_auth=auth,
    use_ssl=True,
    verify_certs=True,
    connection_class = RequestsHttpConnection
)
opensearch_client.info()


# ## OpenSearch 関連リソースの作成
# テスト用のインデックスを作成します。今回はサジェスタとして Kuromoji Completion を使用します。
# 
# Kuromoji Completion は日本語をローマ字のタームに変換し、[completion][completion] タイプのフィールドに格納することでサジェストを実現しています。
# 
# [completion]: https://opensearch.org/docs/latest/field-types/supported-field-types/completion/

# ### データセットのダウンロードと変換

# In[6]:


dataset_dir = "./dataset/japanese-addresses-v2"
dataset_file_path = f"{dataset_dir}/ja.json"
get_ipython().run_line_magic('mkdir', '-p $dataset_dir')
get_ipython().system('curl -o {dataset_file_path} https://japanese-addresses-v2.geoloniamaps.com/api/ja.json')


# In[7]:


file = json.loads(open(dataset_file_path).read())
pref_df = pd.json_normalize(file, ["data"]).explode("cities").rename(columns={
    "code": "pref_code",
    "pref": "pref_name_kanji",
    "pref_k": "pref_name_kana",
    "pref_r": "pref_name_romaji",
    "point": "pref_point"
})
cities_df = pd.json_normalize(pref_df["cities"]).rename(columns={
    "code": "city_code",
    "point": "city_point",
    "city": "city_name_kanji",
    "city_k": "city_name_kana",
    "city_r": "city_name_romaji",
    "ward": "ward_name_kanji",
    "ward_k": "ward_name_kana",
    "ward_r": "ward_name_romaji",
    "county": "county_name_kanji",
    "county_k": "county_name_kana",
    "county_r": "county_name_romaji"
})
japanese_addresses_v2_df = pd.concat([pref_df.reset_index(drop=True).drop('cities', axis=1), cities_df], axis=1)
japanese_addresses_v2_df


# ### インデックスの作成

# In[8]:


index_name = "japanese-addresses-v2"
payload = {
  "mappings": {
    "dynamic_templates": [
      {
        "code": {
          "mapping": {
            "type": "keyword"
          },
          "match_mapping_type": "string",
          "path_match": "*_code"
        }
      },
      {
        "point": {
          "mapping": {
            "type": "geo_point"
          },
          "match_mapping_type": "string",
          "path_match": "*_point"
        }
      },
      {
        "name_kana_and_kanji": {
          "mapping": {
            "type": "completion",
            "analyzer": "kuromoji_completion_index",
            "search_analyzer": "kuromoji_completion_query",
            "preserve_separators": False,
            "preserve_position_increments": True,
            "max_input_length": 20,
            "fields": {
              "keyword": { 
                "type": "keyword"
              },
              "text": {
                "type": "text",
                "analyzer": "sudachi"
              },
            }
          },
          "match_mapping_type": "string",
          "match_pattern": "regex",
          "match": ".+_name_(kana|kanji)"
        }
      },
      {
        "name_romaji": {
          "mapping": {
            "type": "completion",
            "analyzer": "kuromoji_completion_index",
            "search_analyzer": "kuromoji_completion_query",
            "preserve_separators": False,
            "preserve_position_increments": True,
            "max_input_length": 20,
            "fields": {
              "keyword": { 
                "type":  "keyword"
              },
              "text": { 
                "type":  "text",
                "analyzer": "standard"
              }
            }
          },
          "match_mapping_type": "string",
          "path_match": "*_name_romaji"
        }
      }
    ]
  },
  "settings": {
    "index.number_of_shards": 1,
    "index.number_of_replicas": 0,
    "index.refresh_interval": -1,
    "analysis": {
      "analyzer": {
        "kuromoji_completion_index": {
          "mode": "index",
          "type": "kuromoji_completion"
        },
        "kuromoji_completion_query": {
          "mode": "query",
          "type": "kuromoji_completion"
        }
      }
    }
  }
}

try:
    # 既に同名のインデックスが存在する場合、いったん削除を行う
    print("# delete index")
    response = opensearch_client.indices.delete(index=index_name)
    print(json.dumps(response, indent=2))
except Exception as e:
    print(e)

response = opensearch_client.indices.create(index_name, body=payload)
response


# #### ドキュメントのロード
# ドキュメントのロードを行います。ドキュメントのロードは bulk API を使用することで効率よく進められますが、データ処理フレームワークを利用することでより簡単にデータを取り込むことも可能です。本ワークショップでは、[AWS SDK for Pandas][aws-sdk-pandas] を使用したデータ取り込みを行います。
# 
# [aws-sdk-pandas]: https://github.com/aws/aws-sdk-pandas
# 

# In[9]:


index_name = "japanese-addresses-v2"

response = wr.opensearch.index_df(
    client=opensearch_client,
    df=japanese_addresses_v2_df,
    use_threads=True,
    id_keys=["city_code"],
    index=index_name,
    bulk_size=1000,
    refresh=False
)


# response["success"] の値が DataFrame の件数と一致しているかを確認します。True が表示される場合は全件登録に成功していると判断できます。

# In[10]:


response["success"] == japanese_addresses_v2_df["pref_code"].count()


# 本ラボではデータ登録時に意図的に Refresh オプションを無効化しているため、念のため Refresh API を実行し、登録されたドキュメントが確実に検索可能となるようにします

# In[11]:


index_name = "japanese-addresses-v2"
response = opensearch_client.indices.refresh(index=index_name)
response = opensearch_client.indices.forcemerge(index=index_name)


# ## 入力補完のテスト
# 幾つかの候補クエリを用いてタ行の県の覧を取得します。各町ごとに県名を持っているため、サジェスト結果の重複を防ぐために `skip_duplicates` パラメーターをセットしています。

# In[12]:


index_name = "japanese-addresses-v2"

payload = {
  "suggest": {
    "suggest": {
      "prefix": "ｔ",
      "completion": {
        "field": "pref_name_kanji",
        "skip_duplicates": True,
        "size": 10
      }
    }
  }
}


try:
    response = opensearch_client.search(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

df = pd.json_normalize(response["suggest"]["suggest"]).explode("options")
pd.json_normalize(df["options"])


# In[13]:


index_name = "japanese-addresses-v2"

payload = {
  "suggest": {
    "suggest": {
      "prefix": "と",
      "completion": {
        "field": "pref_name_kanji",
        "skip_duplicates": True,
        "size": 10
      }
    }
  }
}


try:
    response = opensearch_client.search(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

df = pd.json_normalize(response["suggest"]["suggest"]).explode("options")
pd.json_normalize(df["options"])


# In[14]:


index_name = "japanese-addresses-v2"

payload = {
  "suggest": {
    "suggest": {
      "prefix": "とｙ",
      "completion": {
        "field": "pref_name_kanji",
        "skip_duplicates": True,
        "size": 10
      }
    }
  }
}


try:
    response = opensearch_client.search(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

df = pd.json_normalize(response["suggest"]["suggest"]).explode("options")
pd.json_normalize(df["options"])


# 漢字のフィールドに対してなぜ "とｙ" といった文字列で入力補完が成立するのでしょうか。
# completion フィールドタイプによるオートコンプリートでは、ドキュメントからアナライザーによって生成されたトークンを格納しておき、検索時に入力クエリから生成されたトークンと前方一致および正規表現によるマッチングを行います。
# 
# 以下は富山県を入力用のアナライザーで処理した結果です。

# In[15]:


index_name = "japanese-addresses-v2"

payload = {
  "text": "富山県",
  "analyzer": "kuromoji_completion_index"
}

try:
    response = opensearch_client.indices.analyze(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

pd.json_normalize(response["tokens"])


# 以下は検索用のクエリを検索用のアナライザーで処理した結果です。

# In[16]:


index_name = "japanese-addresses-v2"

payload = {
  "text": "とｙ",
  "analyzer": "kuromoji_completion_query"
}

try:
    response = opensearch_client.indices.analyze(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

pd.json_normalize(response["tokens"])


# データ格納時に生成された "toyama" と検索時に生成された "toy" が前方一致によりマッチングされたため、サジェストの結果として現れたということが確認できます。
# 
# 前方一致ではなく正規表現を使用する場合は、suggest API に regex 要素を指定します

# In[17]:


index_name = "japanese-addresses-v2"

payload = {
  "suggest": {
    "suggest": {
      "regex": ".*shima.*",
      "completion": {
        "field": "pref_name_kana",
        "skip_duplicates": True,
        "size": 10
      }
    }
  },
  "profile": True
}


try:
    response = opensearch_client.search(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

df = pd.json_normalize(response["suggest"]["suggest"]).explode("options")
pd.json_normalize(df["options"])


# 東御市(とうみし) のような難読地名については、漢字のフィールドに対して Kuromoji が正しく読み仮名を解釈できない可能性があります。以下は一例です。

# In[18]:


index_name = "japanese-addresses-v2"

payload = {
  "text": "東御市",
  "analyzer": "kuromoji_completion_index"
}

try:
    response = opensearch_client.indices.analyze(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

pd.json_normalize(response["tokens"])


# In[19]:


index_name = "japanese-addresses-v2"

payload = {
  "suggest": {
    "suggest": {
      "prefix": "toumi",
      "completion": {
        "field": "city_name_kanji"
      }
    }
  }
}


try:
    response = opensearch_client.search(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

response


# こういったケースでは Kuromoji Completion がインデクシング時に正しく東御市を "トウミシ" と解釈できるようにカスタム辞書をセットするのも手ではありますが、カスタム辞書を更新するたびに reindex が必要となります。代替アプローチとして本データセットのように、カナをあらかじめデータの方にセットしておき、カナフィールドでオートコンプリートを実行する方法もあります。

# In[20]:


index_name = "japanese-addresses-v2"

payload = {
  "suggest": {
    "suggest": {
      "prefix": "toumi",
      "completion": {
        "field": "city_name_kana"
      }
    }
  }
}


try:
    response = opensearch_client.search(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

df = pd.json_normalize(response["suggest"]["suggest"]).explode("options")
pd.json_normalize(df["options"])


# 漢字を使ったサジェストも可能です

# In[21]:


index_name = "japanese-addresses-v2"

payload = {
  "suggest": {
    "suggest": {
      "prefix": "東",
      "completion": {
        "field": "city_name_kanji"
      }
    }
  }
}


try:
    response = opensearch_client.search(
      index=index_name,
      body=payload
    )
except Exception as e:
    print(e)

df = pd.json_normalize(response["suggest"]["suggest"]).explode("options")
pd.json_normalize(df["options"])


# ## まとめ
# 本ラボでは、Kuromoji completion を使用したサジェスタの実装について学習しました。

# ## 後片付け

# ### インデックス削除
# 本ワークショップで使用したインデックスを削除します。インデックスの削除は Delete index API で行います。インデックスを削除するとインデックス内のドキュメントも削除されます。

# In[22]:


index_name = "japanese-addresses-v2"

try:
    response = opensearch_client.indices.delete(index=index_name)
    print(json.dumps(response, indent=2))
except Exception as e:
    print(e)


# ### データセット削除
# ダウンロードしたデータセットを削除します。./dataset ディレクトリ配下に何もない場合は、./dataset ディレクトリも合わせて削除します。

# In[23]:


get_ipython().run_line_magic('rm', '-rf {dataset_dir}')


# In[24]:


get_ipython().run_line_magic('rmdir', './dataset')