# **About** : This notebook is used to train models.
#   - Use the `DEBUG` parameter to launch the code in debug mode (single fold, no logging)
#   - Specify the training parameters in the `Config` class. Feel free to experiment with the parameters, here are the main ones :
#     - `tile_size` : Tile size
#     - `reduce_factor` : Downscaling factor
#     - `on_spot_sampling` : Probability to accept a random tile with in the dataset
#     - `overlap_factor` : Tile overlapping during inference
#     - `selected_folds` : Folds to run computations for.
#     - `encoder` : Encoder as defined in [Segmentation Models PyTorch](https://github.com/qubvel/segmentation_models.pytorch)
#     - `decoder` : Decoders from [Segmentation Models PyTorch](https://github.com/qubvel/segmentation_models.pytorch)
#     - `num_classes` : Number of classes. Keep it at 2 to use the healthy and unhealthy classes
#     - `loss` : Loss function. We use the BCE but the lovasz is also interesting
#     - `optimizer` : Optimizer name
#     - `batch_size` : Training batch size, adapt the `BATCH_SIZES` dictionary to your gpu
#     - `val_bs` : Validation batch size
#     - `epochs` : Number of training epochs
#     - `iter_per_epoch` : Number of tiles to use per epoch
#     - `lr` : Learning rate. Will be decayed linearly
#     - `warmup_prop` : Proportion of steps to use for learning rate warmup
#     - `mix_proba` : Probability to apply MixUp with
#     - `mix_alpha` : Alpha parameter for MixUp
#     - `use_pl`: Probability to sample a tile from the pseudo-labeled images
#     - `use_external`: Probability to sample a tile from the external images
#     - `pl_path`: Path to pseudo labels generated by `notebooks/Inference_test.ipynb`
#     - `extra_path` : Path to extra labels generated by `notebooks/Json to Mask.ipynb`
#     - `rle_path` : Path to train labels downscaled by `notebooks/Image downscaling.ipynb`
# 
## Initialization
### Imports

import os
import sys
import torch
import warnings
import numpy as np
import pandas as pd

from tqdm.notebook import tqdm

sys.path.append("../code/")
# os.environ['CUDA_VISIBLE_DEVICES'] = "1,0"
warnings.simplefilter("ignore", UserWarning)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

import time
start = time.time()

def elapsed_time(start_time):
    return time.time() - start_time

transfer = False
if transfer == False:
    from training.main import k_fold
else:
    from training.main_transfer import k_fold

from utils.logger import (
    prepare_log_folder,
    save_config,
    create_logger,
    update_overall_logs,
)

from params import *

BATCH_SIZES = {
    "resnet18": 64,
    "resnet34": 32,
    "resnext50_32x4d": 32,
    "se_resnext50_32x4d": 32,
    "efficientnet-b0": 32,
    "efficientnet-b1": 32,
    "efficientnet-b2": 32,
    "efficientnet-b3": 16,
    "efficientnet-b4": 16,
    "efficientnet-b5": 16,
    "efficientnet-b6": 8,
}

class Config:
    """
    Parameters used for training
    """
    # General
    seed = 42
    verbose = 1

    device = "cuda" if torch.cuda.is_available() else "cpu"
    save_weights = True

    # Images
    tile_size = 512
    reduce_factor = 2
    on_spot_sampling = 1.9
    overlap_factor = 1.0

    img_dir = DATA_PATH + f"train_{tile_size}_red_{reduce_factor}"
    mask_dir = DATA_PATH + f"masks_{tile_size}_red_{reduce_factor}"

    # k-fold
    cv_column = "5fold"
    random_state = 0
    selected_folds = [0, 1, 2, 3, 4]

    # Model
    encoder = "efficientnet-b1"  # "resnet18" "resnext50_32x4d", "resnet34", "efficientnet-b5"
    decoder = "Unet"  # "Unet", "DeepLabV3Plus"
    encoder_weights = "imagenet"
    num_classes = 2

    # Training
    loss = "BCEWithLogitsLoss"  # "lovasz"
    activation = "none" if loss == "lovasz" else "sigmoid"

    optimizer = "Adam"

    batch_size = BATCH_SIZES[encoder]

    if tile_size == 512:
        batch_size = batch_size // 2

    if batch_size >= 32:
        epochs = 50 
    elif batch_size >= 16:
        epochs = 40
    elif batch_size >= 6:
        epochs = 30
    else:
        epochs = 25
    val_bs = batch_size * 2
    
    iter_per_epoch = 5000
    lr = 1e-3
    warmup_prop = 0.05

    first_epoch_eval = 0

    mix_proba = 0
    mix_alpha = 0.4
    if mix_proba > 0:
        epochs *= 3

    use_pl = 0.15
    use_external = 0.2

    pl_path = None  
    extra_path = None
#     [
#         "../input/train_extra.csv",  # Class healthy
#         "../input/train_extra_onlyfc.csv"  # Class unhealthy
#     ]
    rle_path = [
        f"../input/train_{reduce_factor}_fix.csv",  # Class healthy
        f"../input/train_{reduce_factor}_onlyfc.csv",  # Class unhealthy
    ]
    
class Config:
    """
    Parameters used for training
    """
    # General
    seed = 42
    verbose = 1

    device =  "cuda" if torch.cuda.is_available() else "cpu"
    save_weights = True

    # Images
    tile_size = 512
    reduce_factor = 2
    on_spot_sampling = 1.0 #----changed
    overlap_factor = 1.0 #----changed

    img_dir = DATA_PATH + f"train_{tile_size}_red_{reduce_factor}"
    mask_dir = DATA_PATH + f"masks_{tile_size}_red_{reduce_factor}"

    # k-fold
    cv_column = "5fold"
    random_state = 0
    selected_folds = [0 , 1, 2, 3, 4]

    # Model
    encoder = "efficientnet-b1"  # "resnet18" "resnext50_32x4d", "resnet34", "efficientnet-b5"
    decoder = "Unet"  # "Unet", "DeepLabV3Plus"
    encoder_weights = "imagenet"
    num_classes = 2

    # Training
    loss = "BCEWithLogitsLoss"  #"lovasz"
    activation = "none" if loss == "lovasz" else "sigmoid"

    optimizer = "Adam"

    batch_size = BATCH_SIZES[encoder]

    if tile_size == 512:
        batch_size = batch_size // 2

    if batch_size >= 32:
        epochs = 50
    elif batch_size >= 16:
        epochs = 40
    elif batch_size >= 6:
        epochs = 30
    else:
        epochs = 25 
    val_bs = batch_size * 2
    
    iter_per_epoch = 5000
    lr = 1e-3
    warmup_prop = 0.09

    first_epoch_eval = 0

    mix_proba = 0
    mix_alpha = 0.4
    if mix_proba > 0:
        epochs *= 3

    use_pl = 0 #0.15 - changed
    use_external = 0.2

    pl_path = None # - changed
    extra_path = None # - changed[
#         "../input/train_extra.csv",  # Class healthy
#         "../input/train_extra_onlyfc.csv"  # Class unhealthy
#     ]
    rle_path = [
        f"../input/train_{reduce_factor}_fix.csv",  # Class healthy
        f"../input/train_{reduce_factor}_onlyfc.csv",  # Class unhealthy
    ]

DEBUG = False
log_folder = None

if not DEBUG:
    log_folder = prepare_log_folder(LOG_PATH)
    print (log_folder)
    print(f"Logging results to {log_folder}")
    config_df = save_config(Config, log_folder + "config.json")
    create_logger(directory=log_folder, name="logs.txt")

metrics = k_fold(Config, log_folder=log_folder)

print ("Run time = ", elapsed_time(start))


# **About** : This notebook is used to train models.
#   - Use the `DEBUG` parameter to launch the code in debug mode (single fold, no logging)
#   - Specify the training parameters in the `Config` class. Feel free to experiment with the parameters, here are the main ones :
#     - `tile_size` : Tile size
#     - `reduce_factor` : Downscaling factor
#     - `on_spot_sampling` : Probability to accept a random tile with in the dataset
#     - `overlap_factor` : Tile overlapping during inference
#     - `selected_folds` : Folds to run computations for.
#     - `encoder` : Encoder as defined in [Segmentation Models PyTorch](https://github.com/qubvel/segmentation_models.pytorch)
#     - `decoder` : Decoders from [Segmentation Models PyTorch](https://github.com/qubvel/segmentation_models.pytorch)
#     - `num_classes` : Number of classes. Keep it at 2 to use the healthy and unhealthy classes
#     - `loss` : Loss function. We use the BCE but the lovasz is also interesting
#     - `optimizer` : Optimizer name
#     - `batch_size` : Training batch size, adapt the `BATCH_SIZES` dictionary to your gpu
#     - `val_bs` : Validation batch size
#     - `epochs` : Number of training epochs
#     - `iter_per_epoch` : Number of tiles to use per epoch
#     - `lr` : Learning rate. Will be decayed linearly
#     - `warmup_prop` : Proportion of steps to use for learning rate warmup
#     - `mix_proba` : Probability to apply MixUp with
#     - `mix_alpha` : Alpha parameter for MixUp
#     - `use_pl`: Probability to sample a tile from the pseudo-labeled images
#     - `use_external`: Probability to sample a tile from the external images
#     - `pl_path`: Path to pseudo labels generated by `notebooks/Inference_test.ipynb`
#     - `extra_path` : Path to extra labels generated by `notebooks/Json to Mask.ipynb`
#     - `rle_path` : Path to train labels downscaled by `notebooks/Image downscaling.ipynb`
# 
## Initialization
### Imports

import os
import sys
import torch
import warnings
import numpy as np
import pandas as pd

from tqdm.notebook import tqdm

sys.path.append("../code/")
# os.environ['CUDA_VISIBLE_DEVICES'] = "1,0"
warnings.simplefilter("ignore", UserWarning)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

import time
start = time.time()

def elapsed_time(start_time):
    return time.time() - start_time

transfer = False
if transfer == False:
    from training.main import k_fold
else:
    from training.main_transfer import k_fold

from utils.logger import (
    prepare_log_folder,
    save_config,
    create_logger,
    update_overall_logs,
)

from params import *

BATCH_SIZES = {
    "resnet18": 64,
    "resnet34": 32,
    "resnext50_32x4d": 32,
    "se_resnext50_32x4d": 32,
    "efficientnet-b0": 32,
    "efficientnet-b1": 32,
    "efficientnet-b2": 32,
    "efficientnet-b3": 16,
    "efficientnet-b4": 16,
    "efficientnet-b5": 16,
    "efficientnet-b6": 8,
}

class Config:
    """
    Parameters used for training
    """
    # General
    seed = 42
    verbose = 1

    device = "cuda" if torch.cuda.is_available() else "cpu"
    save_weights = True

    # Images
    tile_size = 512
    reduce_factor = 2
    on_spot_sampling = 1.9
    overlap_factor = 1.0

    img_dir = DATA_PATH + f"train_{tile_size}_red_{reduce_factor}"
    mask_dir = DATA_PATH + f"masks_{tile_size}_red_{reduce_factor}"

    # k-fold
    cv_column = "5fold"
    random_state = 0
    selected_folds = [0, 1, 2, 3, 4]

    # Model
    encoder = "efficientnet-b1"  # "resnet18" "resnext50_32x4d", "resnet34", "efficientnet-b5"
    decoder = "Unet"  # "Unet", "DeepLabV3Plus"
    encoder_weights = "imagenet"
    num_classes = 2

    # Training
    loss = "BCEWithLogitsLoss"  # "lovasz"
    activation = "none" if loss == "lovasz" else "sigmoid"

    optimizer = "Adam"

    batch_size = BATCH_SIZES[encoder]

    if tile_size == 512:
        batch_size = batch_size // 2

    if batch_size >= 32:
        epochs = 50 
    elif batch_size >= 16:
        epochs = 40
    elif batch_size >= 6:
        epochs = 30
    else:
        epochs = 25
    val_bs = batch_size * 2
    
    iter_per_epoch = 5000
    lr = 1e-3
    warmup_prop = 0.05

    first_epoch_eval = 0

    mix_proba = 0
    mix_alpha = 0.4
    if mix_proba > 0:
        epochs *= 3

    use_pl = 0.15
    use_external = 0.2

    pl_path = None  
    extra_path = None
#     [
#         "../input/train_extra.csv",  # Class healthy
#         "../input/train_extra_onlyfc.csv"  # Class unhealthy
#     ]
    rle_path = [
        f"../input/train_{reduce_factor}_fix.csv",  # Class healthy
        f"../input/train_{reduce_factor}_onlyfc.csv",  # Class unhealthy
    ]
    
class Config:
    """
    Parameters used for training
    """
    # General
    seed = 42
    verbose = 1

    device =  "cuda" if torch.cuda.is_available() else "cpu"
    save_weights = True

    # Images
    tile_size = 512
    reduce_factor = 2
    on_spot_sampling = 1.0 #----changed
    overlap_factor = 1.0 #----changed

    img_dir = DATA_PATH + f"train_{tile_size}_red_{reduce_factor}"
    mask_dir = DATA_PATH + f"masks_{tile_size}_red_{reduce_factor}"

    # k-fold
    cv_column = "5fold"
    random_state = 0
    selected_folds = [0 , 1, 2, 3, 4]

    # Model
    encoder = "efficientnet-b1"  # "resnet18" "resnext50_32x4d", "resnet34", "efficientnet-b5"
    decoder = "Unet"  # "Unet", "DeepLabV3Plus"
    encoder_weights = "imagenet"
    num_classes = 2

    # Training
    loss = "BCEWithLogitsLoss"  #"lovasz"
    activation = "none" if loss == "lovasz" else "sigmoid"

    optimizer = "Adam"

    batch_size = BATCH_SIZES[encoder]

    if tile_size == 512:
        batch_size = batch_size // 2

    if batch_size >= 32:
        epochs = 50
    elif batch_size >= 16:
        epochs = 40
    elif batch_size >= 6:
        epochs = 30
    else:
        epochs = 25 
    val_bs = batch_size * 2
    
    iter_per_epoch = 5000
    lr = 1e-3
    warmup_prop = 0.09

    first_epoch_eval = 0

    mix_proba = 0
    mix_alpha = 0.4
    if mix_proba > 0:
        epochs *= 3

    use_pl = 0 #0.15 - changed
    use_external = 0.2

    pl_path = None # - changed
    extra_path = None # - changed[
#         "../input/train_extra.csv",  # Class healthy
#         "../input/train_extra_onlyfc.csv"  # Class unhealthy
#     ]
    rle_path = [
        f"../input/train_{reduce_factor}_fix.csv",  # Class healthy
        f"../input/train_{reduce_factor}_onlyfc.csv",  # Class unhealthy
    ]

DEBUG = False
log_folder = None

if not DEBUG:
    log_folder = prepare_log_folder(LOG_PATH)
    print (log_folder)
    print(f"Logging results to {log_folder}")
    config_df = save_config(Config, log_folder + "config.json")
    create_logger(directory=log_folder, name="logs.txt")

metrics = k_fold(Config, log_folder=log_folder)

print ("Run time = ", elapsed_time(start))


# **About** : This notebook is used to train models. #   - Use the `DEBUG` parameter to launch the code in debug mode (single fold, no logging) #   - Specify the training parameters in the `Config` class. Feel free to experiment with the parameters, here are the main ones : #     - `tile_size` : Tile size #     - `reduce_factor` : Downscaling factor #     - `on_spot_sampling` : Probability to accept a random tile with in the dataset #     - `overlap_factor` : Tile overlapping during inference #     - `selected_folds` : Folds to run computations for. #     - `encoder` : Encoder as defined in [Segmentation Models PyTorch](https://github.com/qubvel/segmentation_models.pytorch) #     - `decoder` : Decoders from [Segmentation Models PyTorch](https://github.com/qubvel/segmentation_models.pytorch) #     - `num_classes` : Number of classes. Keep it at 2 to use the healthy and unhealthy classes #     - `loss` : Loss function. We use the BCE but the lovasz is also interesting #     - `optimizer` : Optimizer name #     - `batch_size` : Training batch size, adapt the `BATCH_SIZES` dictionary to your gpu #     - `val_bs` : Validation batch size #     - `epochs` : Number of training epochs #     - `iter_per_epoch` : Number of tiles to use per epoch #     - `lr` : Learning rate. Will be decayed linearly #     - `warmup_prop` : Proportion of steps to use for learning rate warmup #     - `mix_proba` : Probability to apply MixUp with #     - `mix_alpha` : Alpha parameter for MixUp #     - `use_pl`: Probability to sample a tile from the pseudo-labeled images #     - `use_external`: Probability to sample a tile from the external images #     - `pl_path`: Path to pseudo labels generated by `notebooks/Inference_test.ipynb` #     - `extra_path` : Path to extra labels generated by `notebooks/Json to Mask.ipynb` #     - `rle_path` : Path to train labels downscaled by `notebooks/Image downscaling.ipynb` # ## Initialization ### Imports # os.environ['CUDA_VISIBLE_DEVICES'] = "1,0" Parameters used for training # General # Images # k-fold # Model # "resnet18" "resnext50_32x4d", "resnet34", "efficientnet-b5" # "Unet", "DeepLabV3Plus" # Training # "lovasz" #     [ #         "../input/train_extra.csv",  # Class healthy #         "../input/train_extra_onlyfc.csv"  # Class unhealthy #     ] # Class healthy # Class unhealthy Parameters used for training # General # Images #----changed #----changed # k-fold # Model # "resnet18" "resnext50_32x4d", "resnet34", "efficientnet-b5" # "Unet", "DeepLabV3Plus" # Training #"lovasz" #0.15 - changed # - changed # - changed[ #         "../input/train_extra.csv",  # Class healthy #         "../input/train_extra_onlyfc.csv"  # Class unhealthy #     ] # Class healthy # Class unhealthy