import os

from dotenv import load_dotenv

from .train_chat_brain import train_chat_brain
from .train_profile_brain import train_profile_brain

load_dotenv()


def main():
    processed_dir = os.getenv("PROCESSED_DIR", "data/processed")
    models_dir = os.getenv("MODELS_DIR", "outputs/models")
    artifacts_dir = os.getenv("ARTIFACT_DIR", "outputs/artifacts")
    reports_dir = os.getenv("REPORTS_DIR", "outputs/reports")

    print("\n================== TRAIN: CHAT BRAIN ==================")
    train_chat_brain(
        processed_dir=processed_dir,
        models_dir=models_dir,
        reports_dir=reports_dir,
        # min_recall=0.90,
    )

    # print("\n================== TRAIN: PROFILE BRAIN ==================")
    # train_profile_brain(
    #    processed_dir=processed_dir,
    #    artifacts_dir=artifacts_dir,
    #    gmm_components=3,
    #    kde_bandwidth=2.0,
    #    w_gmm=0.6,
    #    w_kde=0.4,
    # )


if __name__ == "__main__":
    main()