from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# FastRCNNPredictor là đoạn model cuối, tương ứng với phần khoanh đỏ
# do bài toán chỉ có 2 class là wheat và non-wheat nên num_classes = 2

num_classes = 2
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, progress=False)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# import cv2, numpy....
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import SequentialSampler
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import albumentations as A
from albumentations.pytorch.transforms import ToTensor, ToTensorV2

BOX_COLOR = (0, 0, 255)
TEXT_COLOR = (255, 255, 255)
TRAIN_IMG_DIR = "./wheat-dataset/train"

#show 1 ảnh
def plot_img(img, size=(7,7), is_rgb=False):
    ....
    
#show nhiều ảnh
def plot_imgs(imgs, cols=5, size=7, is_rgb=False):
    .....
    
# vẽ bounding box lên ảnh
def visualize_bbox(img, boxes, thickness=3, color=BOX_COLOR):
    ...
    return img_copy

#load và tiền xử lí ảnh
def load_img(img_id, folder=TRAIN_IMG_DIR):
    ...
    return img

#chuyển đổi cặp [imgs, targets] sang dạng tensor
def data_to_device(images, targets, device=torch.device("cuda")):
    images = list(image.to(device) for image in images)
    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
    return images, targets


def expand_bbox(x):
    r = np.array(re.findall("([0-9]+[.]?[0-9]*)", x))
    if len(r) == 0:
        r = [-1, -1, -1, -1]
    return r


# đọc data từ file csv
# output là 1 list chứa thông tin về các ảnh
# mỗi phần tử bao gồm 1 image_id và 1 list các bounding box
def read_data_in_csv(csv_path="./wheat-dataset/train.csv"):
    df = pd.read_csv(csv_path)
    df['x'], df['y'],  df['w'], df['h'] = -1, -1, -1, -1
    df[['x', 'y', 'w', 'h']] = np.stack(df['bbox'].apply(lambda x: expand_bbox(x)))
    df.drop(columns=['bbox'], inplace=True)
    df['x'] = df['x'].astype(np.float)
    df['y'] = df['y'].astype(np.float)
    df['w'] = df['w'].astype(np.float)
    df['h'] = df['h'].astype(np.float)
    objs = []
    img_ids = set(df["image_id"])
    
    for img_id in tqdm(img_ids):
        records = df[df["image_id"] == img_id]
        boxes = records[['x', 'y', 'w', 'h']].values
        area = boxes[:,2]*boxes[:,3]
        boxes[:,2] = boxes[:,0] + boxes[:,2]
        boxes[:,3] = boxes[:,1] + boxes[:,3]

        obj = {
            "img_id": img_id,
            "boxes": boxes,
            "area":area
        }
        objs.append(obj)
    return objs


class WheatDataset(Dataset):
    def __init__(self, data, img_dir ,transform=None):
        self.data = data
        self.img_dir = img_dir
        self.transform = transform
        
    def __getitem__(self, idx):
        img_data = self.data[idx]
        bboxes = img_data["boxes"]
        box_nb = len(bboxes)
        labels = torch.ones((box_nb,), dtype=torch.int64)
        iscrowd = torch.zeros((box_nb,), dtype=torch.int64)
        img = load_img(img_data["img_id"], self.img_dir)
        area = img_data["area"]
        if self.transform is not None:
            sample = {
                "image":img,
                "bboxes": bboxes,
                "labels": labels,
                "area": area
            }
            sample = self.transform(**sample)
            img = sample['image']
            area = sample["area"]
            bboxes = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)

        target = {}
        target['boxes'] = bboxes.type(torch.float32)
        target['labels'] = labels
        target['area'] = torch.as_tensor(area, dtype=torch.float32)
        target['iscrowd'] = iscrowd
        target["image_id"] = torch.tensor([idx])
        return img, target
        
    def __len__(self):
        return len(self.data)
    

def collate_fn(batch):
    return tuple(zip(*batch))

#load data form csv file
data = read_data_in_csv()
shuffle(data)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# tạo transform cho dataset - các biến đổi để augmentation data
train_transform = A.Compose(
    [A.Flip(0.5), ToTensorV2(p=1.0)],
    bbox_params={
        "format":"pascal_voc",
        'label_fields': ['labels']
})

# khởi tạo Dataset và Dataloader
train_dataset = WheatDataset(train_data, img_dir=TRAIN_IMG_DIR, transform=train_transform)
train_loader = DataLoader(
    train_dataset,
    batch_size=8,
    shuffle=True,
    num_workers=2,
    collate_fn=collate_fn)


# Khởi tạo model
num_classes = 2
num_epochs = 5
iters = 1
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, progress=False)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0, weight_decay=0.0005)
model.to(device)


# tiến hành train model
for epoch in range(num_epochs):
    for images, targets in train_loader:
        images, targets = data_to_device(images, targets)
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()
        
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        iters += 1

        # show loss per 30 iteration
        if iters%30 == 0:
            print(f"Iteration #{iters} loss: {loss_value}")
            
        # để đơn giản, ta save model mỗi 90 iteration
        if iters%90 == 0:
            evaluate(model, val_loader, device=device)            
            model_path = f"./saved_model/model_{iters}_{round(loss_value, 2)}.pth"
            torch.save(model.state_dict(), model_path)
            model.train()