import gzip

def gzip_search(query: str, candidate_chunks: list[str], top_k: int=1):
    """
    文字列ベースで類似したテキストチャンクを推定するアルゴリズム.
    `query`, `chunk`, および`query + " " + chunk`をそれぞれgzipで圧縮し、編集距離のようなものをベースに評価する.

    Parameters:
        query (str): 検索クエリとして使用する文字列.
        top_k (int, optional): 返される類似チャンクの上位k個を指定する (default: 1).

    Returns:
        List[str]: 最も類似したテキストチャンクのリスト.

    ---
    Reference:
    -   “Low-Resource” Text Classification: A Parameter-Free Classification Method with Compressors (Jiang et al., Findings 2023)
        https://aclanthology.org/2023.findings-acl.426/
    """

    # 圧縮: query => Q
    Q = gzip.compress(query.encode())

    distance_from_Q = {}
    for chunk in candidate_chunks:
        # 圧縮: chunk => C (上記のquery => Qと同様)
        C = gzip.compress(chunk. encode())

        # queryとchunkを連結する
        query_chunk = query + " " + chunk
        # 共通事項をまとめて表現できるため、この値が小さいほど意味が近いということになる
        Q_C = gzip.compress(query_chunk. encode())

        # 編集距離の計算と似た形で、比較する文字列の長さを正規化する
        normalized_distance = (len(Q_C) - min(len(Q), len(C))) / max(len(Q), len(C))
        distance_from_Q[chunk] = normalized_distance

    # 最も近い`top_k`個までのチャンクを取得する
    return sorted(distance_from_Q, key=distance_from_Q.get)[:top_k]
    
# 疑似ローカルデータ
candidate_chunks = [
    "『北斗の拳』（ほくとのけん）は、武論尊（原作）、原哲夫（作画）による日本の漫画作品。",
    "『ちいかわ なんか小さくてかわいいやつ』（通称「ちいかわ」）は、稀代の天才漫画家ナガノによる作品。",
    "国会議事堂（こっかいぎじどう、英: National Diet Building）は、日本の国会が開催される議事堂。",
    "色覚異常（しきかくいじょう）とは、ヒトの色覚が正常色覚ではない事を示す診断名である。",
    "ライツアウトは、5×5の形に並んだライトをある法則にしたがってすべて消灯 (lights out) させることを目的としたパズル。"
]

query = "ちいかわって誰の作品ですか？"
closest = gzip_search(query, candidate_chunks, top_k=1)
print(closest)
# => ['『ちいかわ なんか小さくてかわいいやつ』（通称「ちいかわ」）は、稀代の天才漫画家ナガノによる作品。']
""" 表現「ちい」「かわ」「作品」などが共通して頻出するため、うまく検索できる """

query = "ルービックキューブとはどういうものですか？"
closest = gzip_search(query, candidate_chunks, top_k=1)
print(closest)
# => ['色覚異常（しきかくいじょう）とは、ヒトの色覚が正常色覚ではない事を示す診断名である。']
"""
最も関連しそうなのはパズルという同カテゴリにある「ライツアウト」に関するチャンクだが、
gzipによる検索は文字列ベースであり、潜在的な意味は拾えないため失敗する。
"""