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 = 'DRIVE'  # DRIVE
#dataset_name = 'LES'  # LES
#dataset_name = 'hrf'  # HRF
dataset = dataset_name

max_step = 30000  # 30000 for ukbb
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]
patch_size = patch_size_list[2]
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 = r"../RIP/weight"

model_step_pretrained_G = 'best_drive'


# path for dataset
stride_height = 50
stride_width = 50


n_classes = 3

model_step = 0

# use CAM
use_CAM = False

#use resize
use_resize = False
resize_w_h = (256,256)

#use av_cross
use_av_cross = False

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

# use global semantic
use_global_semantic = True
global_warmup_step = 0 if use_pretrained_G else 5000

# use network
use_network = 'convnext_tiny' # swin_t,convnext_tiny

dataset_path = {'DRIVE': './data/AV_DRIVE/training/',

                'hrf': './data/hrf/training/',

                'LES': './data/LES_AV/training/',

                }
trainset_path = dataset_path[dataset_name]


print("Dataset:")
print(trainset_path)
print(use_network)