import inquirer
import typer
import pyfiglet
from yaspin import yaspin
from train import train_model
from predict import make_predictions
import os

choice = ""

def main():
    choice = inquirer.list_input("What would you like to do?", choices=["Run tests", "Train a model", "Visualize training data"])
    
    if choice == "Run tests":
        predictions()
    elif choice == "Train a model":
        training()
    else:
        visualize()

def predictions():
    default_cats_path = "dataset/test_set/cats/"
    default_dogs_path = "dataset/test_set/dogs/"
    models_base_path = "models/"

    model_name = inquirer.list_input("Select the model to use", choices=os.listdir(models_base_path))
    model_path = os.path.join(models_base_path, model_name)

    dataset = inquirer.list_input("Select the testing data (default dataset)", choices=['Cats', 'Dogs'])
    
    if dataset == "Cats":
        with yaspin(text="Making predictions...", color="cyan") as sp:
            make_predictions(model_path, default_cats_path, "Cat", sp)
            sp.ok("DONE")
    else: 
        with yaspin(text="Making predictions...", color="cyan") as sp:
            make_predictions(model_path, default_dogs_path, "Dog", sp)
            sp.ok("DONE")

def training():
    text = inquirer.text(message="Enter the name of the new model")
    with yaspin(text="Training new model...", color="cyan") as sp:
        train_model(text, sp)
        sp.ok("DONE")

    
def visualize():
    print("Not available yet...\n")
    main()

if __name__ == "__main__":
    print(pyfiglet.figlet_format("Cats and Dogs"))
    print(pyfiglet.figlet_format("classification"))
    typer.run(main)