#!/usr/bin/env -S uv run
# /// script
# dependencies = ["fasttext", "matplotlib","numpy"]
# ///

# uv run fasttext_embeddings.py train ./corpus.txt --recursive --out ../models/kasturi_fasttext_v4
# uv run fasttext_embeddings.py test ../models/kasturi_fasttext_v4.bin --word ಕನ್ನಡ
# uv run fasttext_embeddings.py test ../models/kasturi_fasttext_v4.bin --word ರಾಜಧಾನಿ
# uv run fasttext_embeddings.py test ../models/kasturi_fasttext_v4.bin --sentence "ಬೆಂಗಳೂರು ಕರ್ನಾಟಕದ ರಾಜಧಾನಿ"
# uv run fasttext_embeddings.py compare ../models/kasturi_fasttext_v4.bin

import argparse
from pathlib import Path

import fasttext
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import font_manager


def setup_kannada_font():
    candidates = [
        "Noto Sans Kannada",
        "Noto Serif Kannada",
        "Lohit Kannada",
    ]

    available = {f.name for f in font_manager.fontManager.ttflist}

    for font in candidates:
        if font in available:
            plt.rcParams["font.family"] = font
            return

    print("Warning: Kannada font not found. Install fonts-noto-core.")


def cosine_similarity(v1, v2):
    denom = np.linalg.norm(v1) * np.linalg.norm(v2)
    if denom == 0:
        return 0.0
    return float(np.dot(v1, v2) / denom)


def compare_words_command(args):
    setup_kannada_font()

    model = fasttext.load_model(args.model)

    words = args.words or [
        "ಹಲಸು",
        "ಹಣ್ಣು",
        "ರೈತ",
        "ನದಿ",
        "ಮಳೆ",
        "ತೆರಿಗೆ",
        "ಅರಸ",
        "ಅರಸಿ",
        "ಹೆಣ್ಣು",
        "ಮಾನವ",
    ]

    matrix = []
    for w1 in words:
        row = []
        for w2 in words:
            row.append(
                cosine_similarity(
                    model.get_word_vector(w1),
                    model.get_word_vector(w2),
                )
            )
        matrix.append(row)

    matrix = np.array(matrix)

    fig, ax = plt.subplots(figsize=(12, 8))
    im = ax.imshow(matrix, vmin=0, vmax=1)

    ax.set_xticks(np.arange(len(words)))
    ax.set_yticks(np.arange(len(words)))
    ax.set_xticklabels(words, fontsize=12)
    ax.set_yticklabels(words, fontsize=12)

    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")

    for i in range(len(words)):
        for j in range(len(words)):
            ax.text(j, i, f"{matrix[i, j]:.3f}", ha="center", va="center", fontsize=9)

    ax.set_title("Kannada Word Similarity Matrix", fontsize=14)
    fig.colorbar(im, ax=ax)
    fig.tight_layout()

    plt.savefig(args.output, dpi=180, bbox_inches="tight")
    print(f"Saved matrix graph to {args.output}")


def train_command(args):
    corpus_path = Path(args.corpus)

    try:
        model = fasttext.train_unsupervised(
            input=str(corpus_path),
            model="skipgram",
            dim=args.dim,
            epoch=args.epoch,
            minCount=args.min_count,
            wordNgrams=args.word_ngrams,
            minn=args.minn,
            maxn=args.maxn,
            thread=args.threads,
        )
        print("training complete")
        model.save_model(f"{args.out}.bin")
        print(f"Saved {args.out}.bin")
    except Exception as e:
        print(f"Exception: {e}")


def test_command(args):
    model = fasttext.load_model(args.model)

    if args.word:
        for score, word in model.get_nearest_neighbors(args.word, k=args.topk):
            print(f"{word:<30} {score:.4f}")

    if args.sentence:
        vec = model.get_sentence_vector(args.sentence)
        print(f"Embedding dimensions: {len(vec)}")
        print(vec[:20])


def main():
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)

    t = sub.add_parser("train")
    t.add_argument("corpus", default="./corpus.txt")
    t.add_argument("--recursive", action="store_true")
    t.add_argument("--out", default="../models/kannada_fasttext")
    t.add_argument("--dim", type=int, default=300)
    t.add_argument("--epoch", type=int, default=20)
    t.add_argument("--min-count", type=int, default=3)
    t.add_argument("--word-ngrams", type=int, default=2)
    t.add_argument("--minn", type=int, default=3)
    t.add_argument("--maxn", type=int, default=8)
    t.add_argument("--threads", type=int, default=16)

    x = sub.add_parser("test")
    x.add_argument("model")
    x.add_argument("--word")
    x.add_argument("--sentence")
    x.add_argument("--topk", type=int, default=10)

    c = sub.add_parser("compare")
    c.add_argument("model")
    c.add_argument("--output", default="../models/compare_matrix.png")
    c.add_argument("--words", nargs="+")

    args = parser.parse_args()

    if args.command == "train":
        train_command(args)
    elif args.command == "test":
        test_command(args)
    elif args.command == "compare":
        compare_words_command(args)


if __name__ == "__main__":
    main()