#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

@author: J-P Praetorius
@email: jan-philipp.praetorius@leibniz-hki.de or p.e.mueller07@gmail.com

Copyright by Jan-Philipp Praetorius

Research Group Applied Systems Biology - Head: Prof. Dr. Marc Thilo Figge
https://www.leibniz-hki.de/en/applied-systems-biology.html
HKI-Center for Systems Biology of Infection
Leibniz Institute for Natural Product Research and Infection Biology -
Hans Knöll Insitute (HKI)
Adolf-Reichwein-Straße 23, 07745 Jena, Germany

"""


import json
import tensorflow as tf

# own written scripts
import sys
from python import utils
from python import SegNet

device_config = {
  "cpus": "all",
  "gpus": [6],
  "log-device-placement": False
}

# Configure the devices
utils.setup_devices(device_config) # = Empty config = All CPUs, All GPUs

model_color = "red" # black , red

model_config = {
  "architecture": "SegNet", # MobileNetV2_pix2pix , SegNet
  "model_type": "segmentation",
  "learning_rate": 0.001,
  "output_model_path": f"/asbdata/Philipp/HyphaeVsYeast/models/final_staining_models/model_{model_color}/untrained-model.hdf5",
  "output_model_json_path": f"/asbdata/Philipp/HyphaeVsYeast/models/final_staining_models/model_{model_color}/untrained-model.hdf5",
  "regularization_method": "GaussianDropout",
  "regularization_lambda": 0.1,
  "weight_regularization": "l2",
  "weight_regularization_lambda": 0.001,
  "image_shape": [512, 512, 3],
  "n_classes": 2,
  "transfer_learning": False # False
}

print(model_config)

# Create the model
model = SegNet.build_model(model_config)

training_config = {
  "output_model_path": f"/asbdata/Philipp/HyphaeVsYeast/models/final_staining_models/model_{model_color}/trained-model.hdf5",
  "output_model_json_path": f"/asbdata/Philipp/HyphaeVsYeast/models/final_staining_models/model_{model_color}/trained-model.json",
  
  "input_dir": f"/asbdata/Philipp/HyphaeVsYeast/images/ilastik_segmentation/tiles/cross_validation_info/model_{model_color}_fungus.csv",
  "label_dir": f"/asbdata/Philipp/HyphaeVsYeast/images/ilastik_segmentation/tiles/cross_validation_info/model_{model_color}_fungus.csv",
  "input_validation_dir": "", 
  "label_validation_dir": "",
  
  "log_dir": "/asbdata/Philipp/HyphaeVsYeast/models/cross-validation/model_{}/log/".format(model_color),
  "max_epochs": 300,
  "batch_size": 32,
  "validation_split": 0.8,
  "augmentation_factor": 3,
  "class_weight_factor": 2,
  "use_elastic_transformation": True, # False,
  "normalization": "zero_one",
  "show_plots": True
}

print(training_config)

# Train the model
# Here we have to set the model parameter as we disabled saving of the untrained model
trained_model = SegNet.train_model(model_config=model_config, config=training_config, model=model)