#!/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
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 = "black" # 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}/trained-model.hdf5",
  "output_model_json_path": f"/asbdata/Philipp/HyphaeVsYeast/models/final_staining_models/model_{model_color}/trained-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)

predict_config = {
  #"input_dir": f"/home/jpraetor/Dev/HyphaeVsYeast/deep-learning/final_staining_models/model_{model_color}_fungus_patches.csv", # SegNet patches
  #"input_dir": f"/home/jpraetor/Dev/HyphaeVsYeast/deep-learning/final_staining_models/model_{model_color}_fungus_cutOuts.csv", # original 
  "input_dir": f"/home/jpraetor/Dev/HyphaeVsYeast/deep-learning/final_staining_models/model_{model_color}_fungus_cutOuts_single.csv", # original - single
  
  #"output_dir": f"/asbdata/Philipp/HyphaeVsYeast/images/ilastik_segmentation/prediction/SegNet/final_staining_models/probability_map/", # SegNet patches
  "output_dir": f"/asbdata/Philipp/HyphaeVsYeast/images/cutOuts/fungus/probability_map/", # original 
  "normalization": "zero_one",
  "extended": True
}

print(predict_config)

# Predict with model
# Here we have to set the model parameter as we disabled saving of the untrained model
SegNet.predict_data(model_config=model_config, config=predict_config)