#!/usr/bin/env python

import argparse
from models.svm import SVM_model
from models.gradient_boosting import Gradient_boosting_model
from models.logistic_regression import Logistic_model
from models.baseline import Baseline_model
from helper import default_labelling, not_reliable_labelling, only_fake_labelling, LIAR_labelling


def main() -> None:
    parser = argparse.ArgumentParser(
            prog="Fakenews detector",
            description="Train and test models",
            usage="The following is an example of training a logistic regression model on news_sample.parquet:\n"+
            "python main.py --model_type logistic --model_file logistic_news_sample.model --data_file news_sample.parquet --train",
    )
    parser.add_argument("--train", action="store_true", help="Whether model should be trained, if not set it will be tested instead")
    parser.add_argument("--validate", action="store_true", help="Wheter to use validation set when testing/validating")
    parser.add_argument("--model_type", "-t", required=True, choices=["baseline", "logistic", "svm", "gradient_boosting"], help="The type of model: baseline, logistic, ...")
    parser.add_argument("--model_file", "-f", required=True, default="", help="The model file to save to when training, or load from when testing")
    parser.add_argument("--data_file", "-d", required=True, help="The datafile used when training or testing")
    parser.add_argument("--label_translator", "-l", required=False, default = "", help="The translator function used by the model, such as \"not_reliable\", that only considers 'reliable' tagged news Real, ignored if not using --train.")
    parser.add_argument("--hyperparameters", "-p", required=False, nargs="+", default ="", help="The hyperparameters used when training the model, written like c=1")
    
    args = parser.parse_args()

    label_translator = default_labelling
    if "not_reliable" in args.label_translator.lower():
        label_translator = not_reliable_labelling
    if "only_fake" in args.label_translator.lower():
        label_translator = only_fake_labelling
    if "liar" in args.data_file.lower():
        label_translator = LIAR_labelling
        args.data_file = "LIAR.parquet"

    if args.model_type == "logistic":
        model = Logistic_model(label_translator=label_translator)
    elif args.model_type == "svm":
        model = SVM_model(label_translator=label_translator)
    elif args.model_type == "gradient_boosting":
        model = Gradient_boosting_model(label_translator=label_translator)
    else:
        model = Baseline_model(label_translator=label_translator)


    if args.train:
        hyperparameters:dict[str, float] = {}
        for parameter in args.hyperparameters:
            key, value = parameter.split("=")
            hyperparameters[key] = float(value)

        model.train(args.data_file, hyperparameters)
        model.save(args.model_file)

    if not args.train:
        model.load(args.model_file)
        if "liar" in args.data_file.lower():
            model.test("LIAR.parquet", validate=False)
        else:
            model.test(args.data_file, args.validate)

if __name__ == "__main__":
    main()
