#!/usr/bin/env -S uv run
# /// script
# dependencies = ["gensim"]
# ///

# uv run word2vec_embeddings.py train ./corpus.txt --out ../models/kannada_kasturi_word2vec.model --model skipgram
# uv run word2vec_embeddings.py test ../models/kannada_kasturi_word2vec.model --word ಕನ್ನಡ
# uv run word2vec_embeddings.py test ../models/kannada_kasturi_word2vec.model --word ರಾಜಧಾನಿ

import argparse
from pathlib import Path

from gensim.models import Word2Vec
from gensim.utils import simple_preprocess


def read_sentences(corpus_path: Path):
    with corpus_path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                # For Kannada, whitespace tokenization is often better than simple_preprocess
                yield line.split()


def train_word2vec(args):
    sentences = list(read_sentences(Path(args.input)))

    model = Word2Vec(
        sentences=sentences,
        vector_size=args.dim,
        window=args.window,
        min_count=args.min_count,
        workers=args.threads,
        epochs=args.epoch,
        sg=1 if args.model == "skipgram" else 0,
    )

    model.save(args.out)
    model.wv.save_word2vec_format(args.out + ".vec")

    print(f"Saved model: {args.out}")
    print(f"Saved vectors: {args.out}.vec")


def test_word(args):
    model = Word2Vec.load(args.model_path)

    print(model.wv[args.word])

    print("\nNearest:")
    for word, score in model.wv.most_similar(args.word, topn=10):
        print(word, score)


def main():
    parser = argparse.ArgumentParser()

    sub = parser.add_subparsers(dest="cmd", required=True)

    train = sub.add_parser("train")
    train.add_argument("input", default="./corpus.txt")
    train.add_argument("--out", default="../models/kannada_kasturi_word2vec.model")
    train.add_argument("--model", choices=["skipgram", "cbow"], default="skipgram")
    train.add_argument("--dim", type=int, default=300)
    train.add_argument("--window", type=int, default=5)
    train.add_argument("--epoch", type=int, default=20)
    train.add_argument("--min-count", type=int, default=2)
    train.add_argument("--threads", type=int, default=16)

    test = sub.add_parser("test")
    test.add_argument("model_path")
    test.add_argument("--word", required=True)

    args = parser.parse_args()

    if args.cmd == "train":
        train_word2vec(args)
    elif args.cmd == "test":
        test_word(args)


if __name__ == "__main__":
    main()