from __future__ import annotations

from typing import Any, List

import numpy as np


def train_doc2vec(
    df: Any,
    token_col: str = "tagged_review",
    vector_size: int = 100,
    alpha: float = 0.025,
    min_alpha: float = 0.001,
    window: int = 8,
    epochs: int = 5,
    tag_prefix: str = "document",
):
    """Gensim Doc2Vec 모델을 학습하고 문서 벡터를 추출합니다.

    Parameters
    ----------
    df:
        pandas DataFrame(또는 유사 객체). `df[token_col]`에는
        **토큰 리스트(예: ['오늘', '날씨', '좋다'])**가 문서 단위로 들어있어야 합니다.
    token_col:
        토큰 리스트가 들어있는 컬럼명
    vector_size, alpha, min_alpha, window, epochs:
        Doc2Vec 하이퍼파라미터
    tag_prefix:
        문서 태그(prefix). 내부적으로 `f"{tag_prefix} {i}"` 형태로 태그가 생성됩니다.

    Returns
    -------
    model:
        학습된 `gensim.models.doc2vec.Doc2Vec` 모델
    vector_list:
        문서 순서대로 정렬된 문서 임베딩 벡터 리스트(List[np.ndarray])
    tagged_corpus_list:
        학습에 사용된 TaggedDocument 리스트
    """
    import gensim  # noqa: F401
    from gensim.models.doc2vec import TaggedDocument
    from gensim.models import doc2vec

    tagged_corpus_list = []
    for i, tokens in enumerate(df[token_col]):
        tag = f"{tag_prefix} {i}"
        tagged_corpus_list.append(TaggedDocument(tags=[tag], words=tokens))

    model = doc2vec.Doc2Vec(
        vector_size=vector_size,
        alpha=alpha,
        min_alpha=min_alpha,
        window=window,
    )
    model.build_vocab(tagged_corpus_list)
    model.train(tagged_corpus_list, total_examples=model.corpus_count, epochs=epochs)

    vector_list: List[np.ndarray] = []
    for i in range(len(df)):
        vector_list.append(model.dv[f"{tag_prefix} {i}"])

    return model, vector_list, tagged_corpus_list
