"""

PCB.py - Main script for PCB Mask R-CNN network

Usage:

- Train the network
python PCB.py -m train -e 100 -s 100 -l 0.001 -b resnet101 -r y

- Perform the segmentation
python PCB.py -m segmentation -r y

- Perform the metrics calculation
python PCB.py -m metrics -r y

-m      training mode       train/segmentation/metrics
-e      number of epochs    number
-s      steps in epoch      number
-r      generate html file  y/n
-l      learning rate       number
-b      backbone network    ResNet50/ResNet101
-a      augmentation        no / light / heavy

"""

# Basic imports -------------------------------------------------------------------------------------------------------#

import os
import sys
import json
import glob
from array import array
import datetime
import numpy as np
import skimage.draw
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.image as mpimg
import tensorflow as tf
from timeit import default_timer as timer
from datetime import timedelta

import PCB_results
import PCB_yolo
import PCB_metrics_func

import warnings
from imgaug import augmenters as iaa
warnings.filterwarnings("ignore")

# Mask R-CNN import ---------------------------------------------------------------------------------------------------#

from mrcnn.config import Config
from mrcnn import visualize
from mrcnn import model as modellib, utils
from mrcnn import utils
from mrcnn import visualize
from mrcnn.visualize import display_images
import mrcnn.model as modellib
from mrcnn.model import log
import final

# Default configuration -----------------------------------------------------------------------------------------------#

# Root folder of the project
PATH_ROOT = "/content/drive/My Drive/DIP21/"

# Path to dataset folder
PATH_DATASET = PATH_ROOT + "main/dataset"

# Path to pre-trained ImageNet weights
PATH_IN_WEIGHTS = "/content/drive/MyDrive/DIP21/resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5"

# Path to images to segment
PATH_IMAGES_TO_SEGMENT = PATH_ROOT + "test_files/"

# Path to the weights used for segmentation

# ResNet50/101, 120 epochs, 300 steps, learning rate 0.002, light augmentation
PATH_WEIGHTS_50  = '/content/drive/MyDrive/DIP21/log_files/RN50_e120_s300_l002_lightaug/mask_rcnn_object_0119.h5'
PATH_WEIGHTS_101 = '/content/drive/MyDrive/DIP21/log_files/RN101_e200_s300_l0015_lightaug/mask_rcnn_object_0199.h5'


# Path to log folder
PATH_LOGS = PATH_ROOT + "log_files"

# Results path
PATH_RESULTS = PATH_ROOT + "results"

# Number of epochs in training
EPOCHS = 100

# Number of steps in training
EPOCH_STEPS = 100

# Learning rate speed
LEARNING_RATE = 0.001

# Backbone used (resnet50 / resnet101)
BB_NET = "resnet101"

DEVICE = '/gpu:0'


# Print functions -----------------------------------------------------------------------------------------------------#
def pr_div():
    for z in range(60):
        print("*", end="")
    print("")


def pr_prog(a, b):
    print("Progress " + str(round((a * 100) / b)) + "%...")

# Parsing and checking atributes --------------------------------------------------------------------------------------#

if __name__ == '__main__':
    import argparse

    commandParser = argparse.ArgumentParser()
    # -m    Mode                        train / segmentation / metrics
    commandParser.add_argument("-m", required=True)
    # -e    Number of epochs            [number] / 10 default
    commandParser.add_argument("-e", required=False)
    # -s    Number of steps             [number] / 100 default
    commandParser.add_argument("-s", required=False)
    # -r    Generate HTML results file  y / n
    commandParser.add_argument("-r", required=False)
    # -l    Learning rate               [number] / 0.001 default
    commandParser.add_argument("-l", required=False)
    # -b    Backbone net                resnet50 / resnet101
    commandParser.add_argument("-b", required=False)
    # -a    Augmentation                 no / light / heavy
    commandParser.add_argument("-a", required=False)

    arguments_list = commandParser.parse_args()

    # Verify mode
    if arguments_list.m not in ['train', 'segmentation', 'metrics']:
        print("This mode is not supported!\nUse 'train' or 'segmentation' or 'metrics' instead!")
        exit()
    
    # Verify number of epochs
    if arguments_list.e and arguments_list.e.isnumeric() and 0 < int(arguments_list.e) <= 500:
        EPOCHS = int(arguments_list.e)
    
    # Verify number of epoch steps
    if arguments_list.s and arguments_list.s.isnumeric() and 10 <= int(arguments_list.s) <= 10000:
        EPOCH_STEPS = int(arguments_list.s)
    
    # Verify learning rate
    if arguments_list.l and float(arguments_list.l) and 0.0001 < float(arguments_list.l) <= 0.1:
        LEARNING_RATE = float(arguments_list.l)

    # Verify backbone network
    if arguments_list.b not in ['resnet50', 'resnet101']:
        print("This backbone network is not supported!\nUse 'resnet50' or 'resnet101' as backbone instead!")
        exit()
    elif arguments_list.b == "resnet50":
        BB_NET = "resnet50"
        PATH_WEIGHTS = PATH_WEIGHTS_50
    else:
        PATH_WEIGHTS = PATH_WEIGHTS_101

    # Verify augmentation
    if arguments_list.m == "train":
      if arguments_list.a not in ['no', 'heavy', 'light']:
        print("This mode of augmentation is not supported!\nUse 'no' or 'light' or 'heavy' instead!")



# Formatters ----------------------------------------------------------------------------------------------------------#

def datetime_to_str(input):
    input = str(input)
    time_h = input.split(":")[0]
    time_m = input.split(":")[1]
    time_m_s = input.split(":")[2]
    time_s = time_m_s[:2]
    time_ms = time_m_s[3:5] if (len(time_m_s) > 2) else time_m_s[3:]
    return str(time_h + "h " + time_m + "m " + time_s + "s " + time_ms + "0ms ")

# Custom dataset class ------------------------------------------------------------------------------------------------#

class DatasetPCB(utils.Dataset):

    def get_PCB_dataset(self, path_dataset, subset):

        # Allow training or validation dataset only

        assert subset in ["training", "validation", "testing"]

        # PCB dataset classes

        self.add_class("object", 1, "relay")
        self.add_class("object", 2, "capacitor")
        self.add_class("object", 3, "ic")
        self.add_class("object", 4, "led")

        names = {"relay": 1, "capacitor": 2, "ic": 3, "led": 4}

        # Palm on another category when using YOLO hand detection
        if arguments_list.m == 'segmentation':
          self.add_class("object", 5, "hand")
          names = {"relay": 1, "capacitor": 2, "ic": 3, "led": 4, "hand": 5}
        

        # Loads annotation file
        path_dataset = os.path.join(path_dataset, subset)
        an_file = json.load(open(os.path.join(path_dataset, "annotations.json")))                            
        an = list(an_file.values())

        # Skip images without regions defined
        an = [a for a in an if a['regions']]

        # Loop trough annotations
        for an_item in an:

            # Get annotations shapes
            polygons = [r['shape_attributes'] for r in an_item['regions']]

            # Get annotated objects names
            objects = [s['region_attributes']['name'] for s in an_item['regions']]

            num_ids = [names[o] for o in objects]

            # Get image size
            image_path = os.path.join(path_dataset, an_item['filename'])
            image = skimage.io.imread(image_path)
            height, width = image.shape[:2]

            self.add_image(
                "object",
                image_id=an_item['filename'],
                path=image_path,
                width=width, height=height,
                polygons=polygons,
                num_ids=num_ids)

    def load_mask(self, image_id):

        image_info = self.image_info[image_id]
        if image_info["source"] != "object":
            return super(self.__class__, self).load_mask(image_id)

        info = self.image_info[image_id]
        if info["source"] != "object":
            return super(self.__class__, self).load_mask(image_id)
        num_ids = info['num_ids']
        mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
                        dtype=np.uint8)
        for i, p in enumerate(info["polygons"]):
            rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])
            mask[rr, cc, i] = 1

        num_ids = np.array(num_ids, dtype=np.int32)
        return mask, num_ids

    def get_img_path(self, image_id):
        info = self.image_info[image_id]
        if info["source"] == "object":
            print("Image reference " + info["path"])
            return info["path"]
            exit()
        else:
            super(self.__class__, self).get_img_path(image_id)


# Prepare datasets ----------------------------------------------------------------------------------------------------#

# Training dataset
training_dataset = DatasetPCB()
training_dataset.get_PCB_dataset(PATH_DATASET, "training")
training_dataset.prepare()

# Validation dataset
validation_dataset = DatasetPCB()
validation_dataset.get_PCB_dataset(PATH_DATASET, "validation")
validation_dataset.prepare()

# Segmentation
seg_dataset = DatasetPCB()
seg_dataset.get_PCB_dataset(PATH_DATASET, "validation")
seg_dataset.prepare()

# Test dataset
test_dataset = DatasetPCB()
test_dataset.get_PCB_dataset(PATH_DATASET, "testing")
test_dataset.prepare()

# Setup model configurations ------------------------------------------------------------------------------------------#

class TrainConf(Config):
    NAME = "object"
    IMAGES_PER_GPU = 1
    STEPS_PER_EPOCH = EPOCH_STEPS
    DETECTION_MIN_CONFIDENCE = 0.85
    BACKBONE = BB_NET
    NUM_CLASSES = 5
    DETECTION_MAX_INSTANCES = 100

conf1 = TrainConf()

class InferenceConf(Config):
    NAME = "object"
    IMAGES_PER_GPU = 1
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    NUM_CLASSES = 5
    BACKBONE = BB_NET
    DETECTION_MIN_CONFIDENCE = 0.85

conf2 = InferenceConf()


# Segmentation ------------------------------------------------------------------------------------------#

def do_segmentation():
    
    print("Preparing image segmentation...")
    pr_div()

    # Perform YOLO hand detection
    print("Performing YOLO hand detection ...")
    yl = PCB_yolo.PCByolo(PATH_IMAGES_TO_SEGMENT, PATH_ROOT + 'main/YOLO/models/')
    yl.perform_detection()

    print("Total YOLO detection time: "+ str(round(yl.get_time(), 2)) + "s")
    print("Processed "+str(yl.get_images()) + " images, found "+str(len(yl.get_results()))+ " objects")
    pr_div()

    timestamp, resultsFile = create_folder(template="results_template.html")

    with tf.device(DEVICE):
        model = modellib.MaskRCNN(mode="inference", model_dir=PATH_IN_WEIGHTS, config=conf2)

    print("Loading network weights from the file: ", PATH_WEIGHTS)
    model.load_weights(PATH_WEIGHTS, by_name=True)

    def get_ax(rows=1, cols=1, size=8):
        _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
        return ax

    # Get all images from test folder

    all_prediction_images = glob.glob(PATH_IMAGES_TO_SEGMENT + "*.*")

    pr_div()
    processing_times = []

    # Process all images
    for item in all_prediction_images:
        pr_div()
        # Start of processing time
        start_time = timer()
        print("Processing image " + item + " ...")
        img = mpimg.imread(item)
        results = model.detect([img], verbose=1)
        ax = get_ax(1)
        r1 = results[0]
        # End of processing time
        end_time = timer()
        total_item_time = timedelta(seconds=end_time - start_time)
        print("Detection time of current image: ", str(total_item_time)[:-4])
        processing_times.append(total_item_time)

        # Get image name
        img_name = item.split('/')[-1]

        # Get YOLO coordinates for current image
        item_yolo_coordinates = yl.get_coordinates(item)

        item_warning = False

        if item_yolo_coordinates:

            print("Found YOLO hand coordinates for current image: "+ str(item_yolo_coordinates))

            item_warning = True

            for coor in item_yolo_coordinates:

                # Creates an empty mask for YOLO hands
                r_masks = np.concatenate((r1['masks'], np.full((*r1['masks'].shape[:-1], 1), False)), axis=-1)
                
                # Adds bounding box coordinates for YOLO hands
                r1['rois'] = np.append(r1['rois'], [coor], axis=0)

                # Adds next class ID for YOLO hands
                r1['class_ids'] = np.append(r1['class_ids'], [5], axis=0)

                # Adds next score for YOLO hands
                r1['scores'] = np.append(r1['scores'], [0.99], axis=0)

        else:
            r_masks = r1['masks']

        # Display results
        visualize.display_instances(img,
                                    r1['rois'],
                                    r_masks,
                                    r1['class_ids'],
                                    seg_dataset.class_names,
                                    r1['scores'],
                                    ax=ax,
                                    title="PCB Segmentation - Image " + img_name,
                                    path=PATH_RESULTS + "/" + timestamp + "/img/" + img_name)

        # Adds image result to results table
        html_img_path = str(PATH_RESULTS + "/" + timestamp + "/img/" + img_name).replace('/', '/')

        resultsFile.add_to_segmentation_table(path=html_img_path,
                                              name=img_name,
                                              segm_duration=str(total_item_time)[:-4],
                                              instances=r1,
                                              class_names=seg_dataset.class_names,
                                              trigger_warning=item_warning)

    # Additional information
    pr_div()
    print("Total number of processed images:\t", str(len(all_prediction_images)))
    resultsFile.replace('$SEGM_IMAGES_TOTAL$', str(len(all_prediction_images)))

    print("Total processing time:\t\t\t", np.sum(processing_times))
    resultsFile.replace('$SEGM_DURATION_TOTAL$', datetime_to_str(np.sum(processing_times)))

    print("Average processing time (one image):\t", np.mean(processing_times))
    resultsFile.replace('$AVG_SEGM_DURATION_TOTAL$', datetime_to_str(np.mean(processing_times)))
    pr_div()

    # Export the HTML file only if --r equals 'y'
    if arguments_list.r == 'y':
        resultsFile.export()

# Compute metrics on validation dataset -------------------------------------------------------------------------------#

def compute_metrics():

    timestamp, resultsFile = create_folder(template="metrics_template.html")

    with tf.device(DEVICE):
        model = modellib.MaskRCNN(mode="inference", model_dir=PATH_IN_WEIGHTS, config=conf2)

    print("Loading network weights from the file: ", PATH_WEIGHTS)
    model.load_weights(PATH_WEIGHTS, by_name=True)

    # New code

    m_func = PCB_metrics_func.MetFunc()

    mAP25, mAP50, mAP75, mF1, mF2 = m_func.compute_metrics_details(rfile=resultsFile,
                dataset=test_dataset,
                in_model=model,
                print_enabled=True,
                img_skip=0) # Compute every image  

    resultsFile.add_to_metrics_table_footer(total=str(len(test_dataset.image_ids)),
                                            MAP_25=round(mAP25, 2),
                                            MAP_50=round(mAP50, 2),
                                            MAP_75=round(mAP75, 2),
                                            MF1=round(mF1, 2),
                                            MF2=round(mF2, 2))

    # Additional information
    pr_div()
    print("Mean MAP (IoU 25):\t", round(mAP25, 2))
    print("Mean MAP (IoU 50):\t", round(mAP50, 2))
    print("Mean MAP (IoU 75):\t", round(mAP75, 2))
    print("Mean F1 score:\t", round(mF1, 2))
    print("Mean F2 score:\t", round(mF2, 2))
    pr_div()

    # Export the HTML file only if --r equals 'y'
    if arguments_list.r == 'y':
        resultsFile.export()


    '''
    
    # old code is here - moved to the PCB_metrics_func.py

    APs_IoU_25 = []
    APs_IoU_50 = []
    APs_IoU_75 = []
    F1s = []
    F2s = []
    
    # Loop through test dataset
    for image_id in test_dataset.image_ids:

        # Show divider, processed image and progress
        pr_div()
        print("Image ID: " + str(image_id+1))
        pr_prog(image_id + 1, len(test_dataset.image_ids))

        # Load image
        image, image_meta, gt_class_id, gt_bbox, gt_mask = \
            modellib.load_image_gt(test_dataset,
                                   conf2,
                                   image_id,
                                   use_mini_mask=False)
        #molded_images = np.expand_dims(modellib.mold_image(image, conf2), 0)

        # Run object detection
        results = model.detect([image], verbose=0)
        r = results[0]

        # Determine APs

        # IoU 0.25
        ap, precisions, recalls, overlaps = \
            utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                             r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.25)
        print("AP IoU20:\t" + str(round(ap, 2)))
        APs_IoU_25.append(ap)
        ap_25 = ap

        # IoU 0.75
        ap, precisions, recalls, overlaps = \
            utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                             r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.75)
        print("AP IoU75:\t"+str(round(ap, 2)))
        APs_IoU_75.append(ap)
        ap_75 = ap

        # IoU 0.50
        ap, precisions, recalls, overlaps = \
            utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                             r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.5)
        print("AP IoU50:\t" + str(round(ap, 2)))
        APs_IoU_50.append(ap)
        ap_50 = ap

        # Determine F1 and F2 score ( using IoU 0.5 )
        F1 = np.mean((2 * precisions * recalls) / (precisions + recalls))
        F2 = np.mean((5 * (precisions * recalls)) / (5 * precisions + recalls))
        
        # Fix nans
        if np.isnan(F1):
          F1 = 0
        if np.isnan(F2):
          F2 = 0

        print("Image F1 score:\t", round(F1, 2))
        print("Image F2 score:\t", round(F2, 2))
        F1s.append(F1)
        F2s.append(F2)

        resultsFile.add_to_metrics_table(path=str(image_id+1),
                                         AP_25=round(ap_25, 2),
                                         AP_50=round(ap_50, 2),
                                         AP_75=round(ap_75, 2),
                                         F1=round(F1, 2),
                                         F2=round(F2, 2))
        resultsFile.add_to_labels(str(image_id+1))

    mAP_IoU_25  = np.mean(APs_IoU_25)
    mAP_IoU_50  = np.mean(APs_IoU_50)
    mAP_IoU_75  = np.mean(APs_IoU_75)
    mF1         = np.mean(F1s)
    mF2         = np.mean(F2s)
    

    resultsFile.add_to_metrics_table_footer(total=str(len(test_dataset.image_ids)),
                                            MAP_25=round(mAP_IoU_25, 2),
                                            MAP_50=round(mAP_IoU_50, 2),
                                            MAP_75=round(mAP_IoU_75, 2),
                                            MF1=round(mF1, 2),
                                            MF2=round(mF2, 2))


    # Additional information
    pr_div()
    print("Mean MAP (IoU 25):\t", round(mAP_IoU_25, 2))
    print("Mean MAP (IoU 50):\t", round(mAP_IoU_50, 2))
    print("Mean MAP (IoU 75):\t", round(mAP_IoU_75, 2))
    print("Mean F1 score:\t", round(mF1, 2))
    print("Mean F2 score:\t", round(mF2, 2))
    pr_div()

    # Export the HTML file only if --r equals 'y'
    if arguments_list.r == 'y':
        resultsFile.export()

    '''

# Create timestamp & folder -------------------------------------------------------------------------------------------#
def create_folder(template):

    # Get date and time timestamp
    ts = str(datetime.datetime.now()).split('.')[0].replace(' ', '-').replace(':', '-')

    # Create results folder and subfolder
    os.mkdir(PATH_RESULTS + "/" + ts, mode=0o777)
    os.mkdir(PATH_RESULTS + "/" + ts + "/img")

    # Create HTML results file
    rf = PCB_results.HTMLresults(template_path=PATH_ROOT + "templates/",
                             export_path=PATH_RESULTS + "/" + ts,
                             template=template)

    # Adds general information
    rf.replace('$RUN_DATE$', ts[:10].replace('-', '/'))
    rf.replace('$RUN_TIME$', ts[11:].replace('-', ':'))

    return ts, rf

# Select action -------------------------------------------------------------------------------------------------------
if arguments_list.m == 'segmentation':
    do_segmentation()
elif arguments_list.m == 'metrics':
    compute_metrics()
else:
    # Train network ---------------------------------------------------------------------------------------------------#

    print("Training the network...")

    # Set up the model ------------------------------------------------------------------------------------------------#

    model = modellib.MaskRCNN(mode="training",
                              config=conf1,
                              model_dir=PATH_LOGS)
    model.load_weights(PATH_IN_WEIGHTS,
                       by_name=True)

    # Add augmentation to training ------------------------------------------------------------------------------------#


    # Heavy augmentation
    if arguments_list.a == "heavy":

      augmentation = iaa.Sequential([

          # Rotation augmentation
          iaa.OneOf([
              iaa.Affine(rotate=0),
              iaa.Affine(rotate=90),
              iaa.Affine(rotate=180),
              iaa.Affine(rotate=270),
          ]),

          # Cropping
          iaa.Crop(percent=(0, 0.1)),

          # Flip augmentation
          iaa.Fliplr(0.2),    # Horizontal
          iaa.Flipud(0.2),    # Vertical

          # Pixel dropout augmentation
          #iaa.OneOf([
          #    iaa.Cutout(fill_mode="constant", cval=255),
          #    iaa.CoarseDropout((0.0, 0.055), size_percent=(0.02, 0.25)),
          #]),

          # Brightness, contrast augmentation
          iaa.Multiply((0.7, 1.2), per_channel=0.2),
          iaa.contrast.LinearContrast((0.85, 1.1)),

          # Gaussian blue, sharpen augmentation
          iaa.OneOf([
              iaa.GaussianBlur(sigma=(0.0, 0.1)),
              iaa.Sharpen(alpha=(0.0, 0.1)),
          ])
      ],
          random_order=True
      )
    
    # Light augmentation
    if arguments_list.a == "light":

      augmentation = iaa.Sequential([

          # Rotation augmentation
          iaa.OneOf([
              iaa.Affine(rotate=0),
              iaa.Affine(rotate=90),
              iaa.Affine(rotate=180),
              iaa.Affine(rotate=270),
          ]),

          # Flip augmentation
          iaa.Fliplr(0.1),    # Horizontal
          iaa.Flipud(0.1),    # Vertical

      ],
          random_order=True
      )
    
    timestamp, resultsFile = create_folder(template="training_template.html")

    # Training start time timestamp
    train_start = timer()

    modeli = modellib.MaskRCNN(mode="inference", model_dir=PATH_IN_WEIGHTS, config=conf2)

    m_func = PCB_metrics_func.MetFunc()

    # Callback function to compute metrics on testing dataset during training of the network
    cb = modellib.Metrics(
      m_func,
      model,
      modeli,
      test_dataset,
      verbose=1)

    # Start training

    if arguments_list.a == "no":

      # Training without augmentation
      model.train(training_dataset,
                  validation_dataset,
                  learning_rate=LEARNING_RATE,
                  epochs=EPOCHS,
                  layers='heads',
                  custom_callbacks=[cb])
    else:

      # Training with augmentation
      model.train(training_dataset,
                  validation_dataset,
                  learning_rate=LEARNING_RATE,
                  epochs=EPOCHS,
                  layers='heads',
                  augmentation=augmentation,
                  custom_callbacks=[cb])
    
    # Training end time timestamp
    train_end = timer()

    # Total time of the training including metrics computation
    total_train_time = timedelta(seconds=train_end - train_start)

    resultsFile.create_train_metrics_table(cb.get_log())
    resultsFile.replace("$EPOCHS$", str(EPOCHS))
    resultsFile.replace("$STEPS$", str(EPOCH_STEPS))
    resultsFile.replace("$LEARNING_RATE$", str(LEARNING_RATE))
    if BB_NET == 'resnet50':
      resultsFile.replace("$BACKBONE$", "ResNet50")  
    else:
      resultsFile.replace("$BACKBONE$", "ResNet101")  
    resultsFile.replace("$TRAIN_TIME$", datetime_to_str(total_train_time))
    
    print("Total train time:", datetime_to_str(total_train_time))

    # Export the HTML file only if --r equals 'y'
    if arguments_list.r == 'y':
        resultsFile.export()