import argparse

from kartezio.callback import CallbackSave, CallbackVerbose
from kartezio.core.base import ModelDraft
from kartezio.core.helpers import singleton
from kartezio.data import read_dataset
from kartezio.enums import CSV_DATASET
from kartezio.utils.io import pack_one_directory


@singleton
class TrainingArgs:
    def __init__(self):
        self.parser = argparse.ArgumentParser()
        self.set_arguments()

    def parse(self):
        return self.parser.parse_args()

    def set_arguments(self):
        self.parser.add_argument("output_directory", help="path to output directory")
        self.parser.add_argument(
            "dataset_directory", help="path to the dataset directory"
        )
        self.parser.add_argument(
            "--indices",
            help="list of indices to select among dataset for the training",
            nargs="+",
            type=int,
            default=None,
        )
        self.parser.add_argument(
            "--dataset_filename",
            help=f"name of the dataset file, default is {CSV_DATASET}",
            default=CSV_DATASET,
        )
        self.parser.add_argument(
            "--genome",
            help="initial genome to instantiate population",
            default=None,
        )
        self.parser.add_argument(
            "--generations",
            help="Number of generations",
            default=100,
            type=int,
        )


kartezio_parser = TrainingArgs()


def get_args():
    return kartezio_parser.parse()


class KartezioTraining:
    def __init__(self, model: ModelDraft, reformat_x=None, frequency=1, preview=False):
        self.args = get_args()
        self.model = model
        self.dataset = read_dataset(
            self.args.dataset_directory,
            counting=True,
            preview=preview,
            indices=self.args.indices,
        )
        if frequency < 1:
            frequency = 1
        self.callbacks = [
            CallbackVerbose(frequency=frequency),
            CallbackSave(self.args.output_directory, self.dataset, frequency=frequency),
        ]
        self.reformat_x = reformat_x

    def run(self):
        train_x, train_y = self.dataset.train_xy
        if self.reformat_x:
            train_x = self.reformat_x(train_x)

        if self.callbacks:
            for callback in self.callbacks:
                callback.set_decoder(self.model.parser)
                self.model.attach(callback)
        elite, history = self.model.fit(train_x, train_y)
        return elite


def train_model(
    model: ModelDraft,
    dataset,
    output_directory,
    preprocessing=None,
    callbacks="default",
    callback_frequency=1,
    pack=True,
):
    if callbacks == "default":
        verbose = CallbackVerbose(frequency=callback_frequency)
        save = CallbackSave(output_directory, dataset, frequency=callback_frequency)
        callbacks = [verbose, save]
        workdir = str(save.workdir._path)
        print(f"Files will be saved under {workdir}.")
    if callbacks:
        for callback in callbacks:
            callback.set_decoder(model.parser)
            model.attach(callback)

    train_x, train_y = dataset.train_xy
    if preprocessing:
        train_x = preprocessing(train_x)

    res = model.fit(train_x, train_y)
    if pack:
        pack_one_directory(workdir)

    return res
