import logging
import os
import subprocess
import datetime

from .const import DEFAULT_DIM, DEFAULT_LEARNING_RATE, DEFAULT_NUM_EPOCHS, PKG_NAME

_LOGGER = logging.getLogger(PKG_NAME)


def main(
    path_to_starspace: str,
    input: str,
    output: str,
    num_epochs: str = DEFAULT_NUM_EPOCHS,
    dim: str = DEFAULT_DIM,
    learning_rate: str = DEFAULT_LEARNING_RATE,
):
    """Train a starspace model.

    Args:
        path_to_starspace (str): Path to starspace folder (must be prebuilt).
        input (str): Path to input file generated by preprocess.
        num_epochs (str): Number of epochs.
        dim (str): Dimension of embeddings.
        learning_rate (str): Learning rate.
        output (str): Path to output folder.
    """
    _LOGGER.info("Running training...")
    _LOGGER.info(f"Start: {datetime.datetime.now()}")

    # PLACE TRAINING CODE HERE
    model = os.path.join(output, "starspace_trained_model")
    path_to_starspace = os.path.join(path_to_starspace, "starspace")

    star_command = [
        path_to_starspace,
        "train",
        "-trainFile",
        input,
        "-model",
        model,
        "-trainMode",
        "0",
        "-dim",
        dim,
        "-epoch",
        num_epochs,
        "-negSearchLimit",
        "5",
        "-thread",
        "20",
        "-lr",
        learning_rate,
    ]

    if os.path.exists(model):
        star_command.extend(["-initModel", model])

    # Start the subprocess
    process = subprocess.Popen(star_command)

    # Wait for the subprocess to complete
    process.wait()

    _LOGGER.info("Train bedembed done.")
    _LOGGER.info(f"End: {datetime.datetime.now()}")
