import pandas as pd
import numpy as np
import asyncio
from angle_emb import AnglE, Prompts
from angle_emb.utils import cosine_similarity

angle = AnglE.from_pretrained(
    "hon9kon9ize/bert-large-cantonese-sts", pooling_strategy='cls')
df = None

def get_embedding(text):
    response = angle.encode(text, to_numpy=True)

    return response[0]

def get_dataframe():
    global df

    if df is not None:
        return df

    df = pd.read_json("emb.json")
    df = df.drop_duplicates(subset=["artist_name", "track_name"])

    return df


def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def find_songs(text):
    df = get_dataframe()
    text_embedding = get_embedding(text)
    df["similarity"] = df["lyrics_embedding"].apply(
        lambda x: cosine_similarity(x, text_embedding)
    )
    df = df.sort_values(by="similarity", ascending=False)
    top_5 = df.head(5)

    return "### 以下係你嘅歌名推介：\n" + "\n".join(
        [
            f"- {row['artist_name']} 嘅 **[「{row['track_name']}」](https://open.spotify.com/track/{row['track_id']})**（相似度：{row['similarity']:.2f}）"
            for _, row in top_5.iterrows()
        ]
    )

# Create Gradio application
import gradio as gr


async def create_demo():
    example1 = """故事開始於一個悲傷的雨夜，主角站在街頭，淋雨中凝視著遠方。心裡充滿了對一段失去的愛情的痛苦和迷惘。回想起往事，他感到愛情並非他當初所想像的那麼美好，無法找到回到對方身邊的路，更不用說如何忘記過往。

淚水在眼底打轉，他感到迷失，不知道該往哪裡去，只是心中不斷地呼喚著對方的名字，渴望重新找回失去的愛。他開始思考，是應該安靜地離開這段過往，還是勇敢地留下來，面對愛情的種種無奈。

在迷惘中，他決定給自己一個機會。或許他應該離開，或者他應該在原地等待，等待對方明白他所付出的愛是永遠不會離開的。這段故事充滿了對愛情的掙扎、遺憾和無奈，卻也帶著一絲希望和堅持。結局如何，只有時間能給予答案。"""
    example2 = "跌咗嘢試下唔好搵"
    example3 = "香港有國安法"
    example4 = "雞！全部都係雞！"
    description = """呢個 space 利用咗廣東話 Bert 語言模型，將歌詞轉換成向量，再用 cosine similarity 計算輸入嘅文字同 2394 首粵語歌唧歌詞之間嘅相似度，嚟畀出個歌名推介。

⚠️ 要註意!呢個唔係關鍵字搜尋，而係用文字嘅意思去比對每首歌嘅歌詞內容，所以有時候會有啲奇怪嘅結果。
"""
    css = """
.output {
  padding: 10px;
  min-height: 200px;
}
"""

    demo = gr.Interface(
        fn=find_songs,
        css=css,
        inputs=[
            gr.Textbox(
                label="求其打啲嘢去搵歌名",
                lines=5,
                placeholder="請輸入歌詞",
            ),
        ],
        outputs=[
            gr.Markdown(elem_classes="output"),
        ],
        examples=[example1, example2, example3, example4],
        title="粵語流行歌相似度",
        description=description,
        analytics_enabled=False,
        allow_flagging=False,
    )

    return demo


# Run the application
if __name__ == "__main__":
    demo = asyncio.run(create_demo())
    demo.launch()