import os
#os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import shutil
import sys

sys.path.append('./')
sys.path.append('../')
import time

import wandb
from wandb.keras import WandbCallback
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf

from deel.utils.yaml_to_params import load_yaml_config,getParams, getFunctionFromModules, dumdict2yaml
from deel.datasets.load_dataset import load_dataset
from deel.utils.binary_lip_trainer import BinaryLipTrainer,init_config



#print(tf.config.list_physical_devices('GPU'))
filename="./configs/toy/two_moons.yml"
if len(sys.argv)>=2:
        filename =sys.argv[1]
isExist = os.path.exists(filename)
full_config = load_yaml_config(filename)
#init_config(full_config)
trainer = BinaryLipTrainer(full_config)
print(trainer.xp_id)
trainer.init_training(full_config)
model = trainer.load_model_and_parameters(full_config)
model.summary()

dtset = load_dataset(getParams(full_config,'dataset'))
#x,y = next(dtset["train"].__iter__())
#print(x,y)
trainer.fit(dtset)
trainer.save_model()