import numpy as np

import torch

from monai.transforms import (

    Compose,

    LoadImaged,

    RandSpatialCropd,

    EnsureTyped,

    CastToTyped,

    NormalizeIntensityd,

    RandFlipd,

    CenterSpatialCropd,

    ScaleIntensityRanged,

    RandAffined,

    RandScaleIntensityd,

    RandShiftIntensityd,

    RandCoarseDropoutd,

    Rand2DElasticd,

    Lambdad,

    Resized,

    AddChanneld,

    RandGaussianNoised,

    RandGridDistortiond,

    RepeatChanneld,

    Transposed,

    OneOf,

    EnsureChannelFirstd,

    RandLambdad,

    Spacingd,

    FgBgToIndicesd,

    CropForegroundd,

    RandCropByPosNegLabeld,

    ToDeviced,

    SpatialPadd,

)

from default_config import basic_cfg

cfg = basic_cfg

# train

cfg.train = True

cfg.eval = False

cfg.start_eval_epoch = 5  # when use large lr, can set a large num

cfg.run_org_eval = False

cfg.run_tta_val = False

cfg.load_best_weights = False

cfg.amp = False

cfg.val_amp = False

cfg.num_workers = 8

cfg.model_type = "segres12"

# device

cfg.gpu = 0

cfg.device = "cuda:%d" % cfg.gpu

# lr

# warmup_restart, cosine

cfg.lr_mode = "warmup_restart"

cfg.lr = 5e-4

cfg.min_lr = 2e-4

cfg.weight_decay = 1e-6

cfg.epochs = 1000

cfg.restart_epoch = 100  # only for warmup_restart

cfg.eval_epochs = 10

cfg.finetune_lb = -1

# dataset

cfg.img_size = (160, 160, 80)

cfg.spacing = (1.5, 1.5, 1.5)

cfg.batch_size = 4

cfg.val_batch_size = 1

cfg.train_cache_rate = 0.0

cfg.val_cache_rate = 0.0

cfg.gpu_cache = False

cfg.val_gpu_cache = False

# val

cfg.roi_size = (224, 224, 80)

cfg.sw_batch_size = 4

# model

# loss

cfg.w_dice = 1.0

cfg.output_dir = "./output/segres12_all"

# transforms

cfg.train_transforms = Compose(

    [

        LoadImaged(keys=["image", "mask"]),

        EnsureChannelFirstd(keys=["image", "mask"]),

        # Spacingd(keys=["image", "mask"], pixdim=cfg.spacing, mode=("bilinear", "nearest")),

        RandSpatialCropd(

            keys=("image", "mask"),

            roi_size=cfg.img_size,

            random_size=False,

        ),

        Lambdad(keys="image", func=lambda x: x / x.max()),

        RandFlipd(keys=("image", "mask"), prob=0.5, spatial_axis=[0]),

        RandFlipd(keys=("image", "mask"), prob=0.5, spatial_axis=[1]),

        # RandFlipd(keys=("image", "mask"), prob=0.5, spatial_axis=[2]),

        RandAffined(

            keys=("image", "mask"),

            prob=0.5,

            rotate_range=np.pi / 12,

            translate_range=(cfg.img_size[0]*0.0625, cfg.img_size[1]*0.0625),

            scale_range=(0.1, 0.1),

            mode="nearest",

            padding_mode="reflection",

        ),

        OneOf(

            [

                RandGridDistortiond(keys=("image", "mask"), prob=0.5, distort_limit=(-0.05, 0.05), mode="nearest", padding_mode="reflection"),

                RandCoarseDropoutd(

                    keys=("image", "mask"),

                    holes=5,

                    max_holes=8,

                    spatial_size=(1, 1, 1),

                    max_spatial_size=(12, 12, 12),

                    fill_value=0.0,

                    prob=0.5,

                ),

            ]

        ),

        RandScaleIntensityd(keys="image", factors=(-0.2, 0.2), prob=0.5),

        RandShiftIntensityd(keys="image", offsets=(-0.1, 0.1), prob=0.5),

        EnsureTyped(keys=("image", "mask"), dtype=torch.float32),

    ]

)

cfg.val_transforms = Compose(

    [

        LoadImaged(keys=["image", "mask"]),

        EnsureChannelFirstd(keys=["image", "mask"]),

        # Spacingd(keys=["image", "mask"], pixdim=cfg.spacing, mode=("bilinear", "nearest")),

        Lambdad(keys="image", func=lambda x: x / x.max()),

        EnsureTyped(keys=("image", "mask"), dtype=torch.float32),

        # ToDeviced(keys=["image", "mask"], device="cuda:0"),

    ]

)

cfg.org_val_transforms = Compose(

    [

        LoadImaged(keys="image"),

        EnsureChannelFirstd(keys="image"),

        # Spacingd(keys="image", pixdim=cfg.spacing, mode="bilinear"),

        Lambdad(keys="image", func=lambda x: x / x.max()),

        # SpatialPadd(keys="image", spatial_size=cfg.img_size),

        EnsureTyped(keys="image", dtype=torch.float32),

    ]

)