import torch
import os

# Check GPU availability
use_cuda = torch.cuda.is_available()
gpu_ids = [0] if use_cuda else []
device = torch.device('cuda' if use_cuda else 'cpu')

dataset_name = 'all'  # DRIVE
#dataset_name = 'LES'  # LES
# dataset_name = 'hrf'  # HRF
# dataset_name = 'ukbb'  # UKBB
# dataset_name = 'all'
dataset = dataset_name
max_step = 30000  # 30000 for ukbb

batch_size = 8  # default: 4
print_iter = 100  # default: 100
display_iter = 100  # default: 100
save_iter = 5000  # default: 5000
first_display_metric_iter = max_step - save_iter  # default: 25000
lr = 0.0002  # if dataset_name!='LES' else 0.00005 # default: 0.0002
step_size = 7000  # 7000 for DRIVE
lr_decay_gamma = 0.5  # default: 0.5
use_SGD = False  # default:False

input_nc = 3
ndf = 32
netD_type = 'basic'
n_layers_D = 5
norm = 'instance'
no_lsgan = False
init_type = 'normal'
init_gain = 0.02
use_sigmoid = no_lsgan
use_noise_input_D = False
use_dropout_D = False
# torch.cuda.set_device(gpu_ids[0])
use_GAN = True  # default: True

# adam
beta1 = 0.5

# settings for GAN loss
num_classes_D = 1
lambda_GAN_D = 0.01
lambda_GAN_G = 0.01
lambda_GAN_gp = 100
lambda_BCE = 5
lambda_DICE = 5

input_nc_D = input_nc + 3

# settings for centerness
use_centerness = True  # default: True
lambda_centerness = 1
center_loss_type = 'centerness'
centerness_map_size = [128, 128]

# pretrained model
use_pretrained_G = True
use_pretrained_D = False
# model_path_pretrained_G = './log/patch_pretrain'
model_path_pretrained_G = ''
model_step_pretrained_G = 0
stride_height = 0
stride_width = 0
patch_size_list=[]
use_CAM = False

#use resize
use_resize = False
resize_w_h = (1920,512)
def set_dataset(name):
    global dataset_name, model_path_pretrained_G, model_step_pretrained_G
    global stride_height, stride_width,patch_size,patch_size_list,dataset,use_CAM,use_resize,resize_w_h
    dataset_name = name
    dataset = name
    if dataset_name == 'DRIVE':
        model_path_pretrained_G = './AV/log/DRIVE-2023_10_20_08_36_50(6500)'
        model_step_pretrained_G = 6500
    elif dataset_name == 'LES':
        model_path_pretrained_G = './AV/log/LES-2023_09_28_14_04_06(0)'
        model_step_pretrained_G = 0
    elif dataset_name == 'hrf':
        model_path_pretrained_G = './AV/log/HRF-2023_10_19_11_07_31(1500)'
        model_step_pretrained_G = 1500
    elif dataset_name == 'ukbb':
        model_path_pretrained_G = './AV/log/UKBB-2023_11_02_23_22_07(5000)'
        model_step_pretrained_G = 5000
    else:
        model_path_pretrained_G = './AV/log/ALL-2024_09_06_09_17_18(9000)'
        model_step_pretrained_G = 9000
    if dataset_name == 'DRIVE':
        patch_size_list = [64, 128, 256]
    elif dataset_name == 'LES':
        patch_size_list = [96, 384, 256]
    elif dataset_name == 'hrf':
        patch_size_list = [64, 384, 256]
    elif dataset_name == 'ukbb':
        patch_size_list = [96, 384, 256]
    else:
        patch_size_list = [96, 384, 512]
    patch_size = patch_size_list[2]

# path for dataset
    if dataset_name == 'DRIVE' or dataset_name == 'LES':
        stride_height = 50
        stride_width = 50
    elif dataset_name == 'ukbb' or dataset_name == 'hrf':
        use_CAM=False
        use_resize = True
        stride_height = 150
        stride_width = 150
    else:
        use_CAM=True
        use_resize = True
        stride_height = 150
        stride_width = 150
        

n_classes = 3

model_step = 0


# use av_cross
use_av_cross = False

use_high_semantic = False
lambda_high = 1  # A,V,Vessel

# use global semantic in local, huggingface set false
use_global_semantic = False
global_warmup_step = 0 if use_pretrained_G else 5000