from chia import containers, instrumentation
from chia.components import classifiers
from chia import helpers
from chillax import chillax_classifier, chillax_extrapolator

import config as pcfg
import argparse


class CheapObserver(instrumentation.Observer):
    def update(self, message: instrumentation.Message):
        print(f"Message: {message}")


def main(config_files):
    configs = [
        pcfg.config_from_json(config_file, read_from_file=True)
        for config_file in config_files
    ]
    config = pcfg.ConfigurationSet(*configs)

    classifiers.ClassifierFactory.name_to_class_mapping.update(
        {"chillax": chillax_classifier.CHILLAXKerasHC}
    )

    # This shouldn't be necessary, but...
    helpers.setup_environment()
    obs = instrumentation.NamedObservable("Experiment")

    experiment_container = containers.ExperimentContainer(config, outer_observable=obs)

    with experiment_container.exception_shroud:
        obs.log_info("Hello!")

        # Now, build the extrapolator
        if "extrapolator" in config.keys():
            extrapolator = chillax_extrapolator.CHILLAXExtrapolatorFactory.create(
                config["extrapolator"],
                knowledge_base=experiment_container.knowledge_base,
                observers=experiment_container.observers,
            )
            experiment_container.classifier.extrapolator = extrapolator

        experiment_container.runner.run()

    # Make sure all the data is saved
    obs.send_shutdown(successful=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog="chillax.experiment_selfsupervised")
    parser.add_argument("config_file", type=str, nargs="+")
    args = parser.parse_args()
    main(config_files=args.config_file)