import albumentations as album
import cv2
import torch

from albumentations.pytorch import ToTensorV2

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
IMG_SHAPE = 416
SPLIT_SIZE = [IMG_SHAPE // 32, IMG_SHAPE // 16, IMG_SHAPE // 8]
NO_EPOCHS = 20
BATCH_SIZE = 7
WEIGHT_DECAY = 1e-3
NO_WORKERS = 4
NO_CLASSES = 9
LR = 1e-3
CONF_LIM = 0.05
MAP_IOU_LIM = 0.5
IOU_LIM = 0.45
SAVE_TRAINED_MODEL = True
MEM = True
LOAD_TRAINED_MODEL = True
TRAIN_MODEL_FILE = "trained_models/18_epoch_cars_kitti.pth.tar"
LABEL_DIR = r"data/kitti_single/training/label_2"
IMG_DIR = r"data/kitti_single/training/image_2"

ANCHORS = [
    [(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)],
    [(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)],
    [(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],
]

scale = 1.1
train_transforms = album.Compose(
    [
        album.LongestMaxSize(max_size=int(IMG_SHAPE * scale)),
        album.PadIfNeeded(
            min_height=int(IMG_SHAPE * scale),
            min_width=int(IMG_SHAPE * scale),
            border_mode=cv2.BORDER_CONSTANT,
        ),
        album.RandomCrop(width=IMG_SHAPE, height=IMG_SHAPE),
        album.ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4),
        album.OneOf(
            [
                album.ShiftScaleRotate(
                    rotate_limit=20, p=0.5, border_mode=cv2.BORDER_CONSTANT
                ),
                album.IAAAffine(shear=15, p=0.5, mode="constant"),
            ],
            p=1.0,
        ),
        album.HorizontalFlip(p=0.5),
        album.Blur(p=0.1),
        album.CLAHE(p=0.1),
        album.Posterize(p=0.1),
        album.ToGray(p=0.1),
        album.ChannelShuffle(p=0.05),
        album.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
        ToTensorV2(),
    ],
    bbox_params=album.BboxParams(format="yolo", min_visibility=0.4, label_fields=[],),
)
test_transforms = album.Compose(
    [
        album.LongestMaxSize(max_size=IMG_SHAPE),
        album.PadIfNeeded(
            min_height=IMG_SHAPE, min_width=IMG_SHAPE, border_mode=cv2.BORDER_CONSTANT
        ),
        album.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
        ToTensorV2(),
    ],
    bbox_params=album.BboxParams(format="yolo", min_visibility=0.4, label_fields=[]),
)

CLASSES = ['DontCare', 'Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'Misc']
