import os
from tensorflow.python.platform import build_info as build
import tensorflow as tf
import gc
from meow_ml.src.training.train import train
from meow_ml.src.hyperparam_optimization.tune import tune
import numpy as np


def split_runs_across_gpus(runs, num_gpus):
    return np.array_split(runs, num_gpus)  

def init_gpu(gpu_id):
    tf.keras.backend.clear_session()
    gc.collect()  
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    os.environ["CUDA_VISIBLE_DEVICES"] = "2"
    os.environ["TF_GPU_ALLOCATOR"] = "cuda_malloc_async"
    os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)  
    os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
    physical_devices = tf.config.list_physical_devices('GPU')
    tf.config.experimental.set_memory_growth(physical_devices[0], enable=True)

def init_cpu():
    tf.keras.backend.clear_session()
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    gc.collect()  
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    # physical_devices = tf.device('/cpu:0')
    # tf.config.set_visible_devices(physical_devices[0], 'CPU')
  
def sub_process_entry_point(runs_config, data_generator_config, mlflow_config, gpu_id=None, is_optuna=False):
    init_gpu(gpu_id)
    if gpu_id:
        strategy = tf.distribute.MirroredStrategy()
        with strategy.scope():
            for run_config in runs_config:
                if is_optuna:
                    tune(study_variables=run_config, data_generator_config=data_generator_config, mlflow_config=mlflow_config)
                else:
                    train(model_config=run_config, data_generator_config=data_generator_config, mlflow_config=mlflow_config)
    else:
        for run_config in runs_config:
            if is_optuna:
                tune(study_variables=run_config, data_generator_config=data_generator_config, mlflow_config=mlflow_config)
            else:
                    train(model_config=run_config, data_generator_config=data_generator_config, mlflow_config=mlflow_config)
