from enum import Enum

import typer


class ModelType(str, Enum):
    gpt2 = "gpt2"
    gpt2odin = "gpt2odin"
    gpt2background = "gpt2background"


app = typer.Typer()


@app.command()
def preparedata():
    from prepare_dataset import prepare_dataset

    prepare_dataset()


@app.command()
def train(
    model_name: ModelType,
    batch_size: int = 4,
    cuda: bool = False,
    epochs: int = 5,
    log_wandb: bool = False,
    gradient_accumulation: int = 4,
    use_amp: bool = False,
    train_ratio_perc: float = 5.0,
    eval_ratio_perc: float = 1.0,
):
    from train import train_model

    train_model(
        model_name=model_name,
        batch_size=batch_size,
        cuda=cuda,
        epochs=epochs,
        log_wandb=log_wandb,
        gradient_accumulation=gradient_accumulation,
        use_amp=use_amp,
        train_ratio_perc=train_ratio_perc,
        eval_ratio_perc=eval_ratio_perc,
    )


@app.command()
def eval(
    model_name: ModelType,
    model_checkpoint: str,
    cuda: bool = False,
    use_amp: bool = False,
):
    from evaluation import eval_model

    eval_model(
        model_name=model_name,
        model_checkpoint=model_checkpoint,
        cuda=cuda,
        use_amp=use_amp,
    )


@app.command()
def metrics():
    from metrics import calculate_metrics

    calculate_metrics()


if __name__ == "__main__":
    app()