"""
Functions:
loading dataset instances
aplying detection on picture
training
"""
from skimage import data, io
from matplotlib import pyplot as plt
import cv2
import csv
import os
import sys
import json
import datetime
import numpy as np
import skimage.draw
from config import Config
import utils
import model as modellib

ROOT_DIR = os.getcwd()
sys.path.append(ROOT_DIR)


class BillboardConfig(Config):
    """Configfor training on billboard dataset using Nvidia 
    titan XP with 12 GB of memory
    Based on config class with some vales overriden
    """

    NAME = "Billboards"

    # Two images fit in 12 GB of memory
    IMAGES_PER_GPU = 2

    # Number of classes
    NUM_CLASSES = 1 + 1  # Background + billboard

    # 1000 is default, could be set to number of instances in training dataset
    STEPS_PER_EPOCH = 1000
    
    IMAGE_MIN_DIM = 1024
    IMAGE_MAX_DIM = 1024
    
    LEARNING_RATE = 0.001

    # Use smaller masks, saves memeory and makes training faster
    USE_MINI_MASK=True

class Billboard(object):
    coord_list = []
    avg_color = np.empty(shape=(3))
    pixel_count = 0

    def __init__(self):
        self.coord_list = []
        self.avg_color = np.empty(shape=(3))
        self.avg_color[0]=0
        self.avg_color[1]=0
        self.avg_color[2]=0
        self.pixel_count = 0
        
    def addPixel(self,x,y,r,g,b):
        self.coord_list.append((x,y))
        self.avg_color[0]= (self.avg_color[0]*self.pixel_count+r)/(self.pixel_count+1)
        self.avg_color[1]= (self.avg_color[1]*self.pixel_count+g)/(self.pixel_count+1)
        self.avg_color[2]= (self.avg_color[2]*self.pixel_count+b)/(self.pixel_count+1)
        self.pixel_count=self.pixel_count+1
    
    def addMultiPixels(self,new_coord_list,r,g,b,count):
        self.coord_list=self.coord_list+new_coord_list
        self.avg_color[0]= ((self.avg_color[0]*self.pixel_count)+(r*count))/(self.pixel_count+count)
        self.avg_color[1]= ((self.avg_color[1]*self.pixel_count)+(r*count))/(self.pixel_count+count)
        self.avg_color[2]= ((self.avg_color[2]*self.pixel_count)+(r*count))/(self.pixel_count+count)
        self.pixel_count=self.pixel_count+count

class BillboardDataset(utils.Dataset):

    def load_billboards(self, subset):
        """Load a subset of the Balloon dataset.
        subset: train or val
        """
        # define single class - billboard
        self.add_class("billboard", 1, "billboard")

        # load data depending on what set we want
        assert subset in ["train", "val"]
        if(subset == "train"):
            subset=""
        if(subset == "val"):
            subset="val_"
            
        dataset_dir = os.path.join(os.getcwd(), 'billboards')
        
        # add all images to dataset
        for i in range(8000):            
            name = subset+"real0"+str(i)+".jpeg"
            if(os.path.isfile(os.path.join(dataset_dir,name))):
                image_path = os.path.join(dataset_dir, name)
                image = skimage.io.imread(image_path)
                height = image.shape[0]
                width = image.shape[1]            
                self.add_image(
                    "billboard",
                    image_id=i,
                    path=image_path,
                    width=width, height=height,
                    image_name=subset+"real0"+str(i))

    def load_mask(self, image_id):
        """
        Creats insatnce mask based on *seg.png file and csv file containing
        instance color information
        Returns an array of masks for each of the billboard instances and
        class ids for each instance - 1D array of ones in this case
        """
        
        image_info = self.image_info[image_id]
        
        #path to *seg.png files
        dataset_dir = os.path.join(os.getcwd(), 'billboards')
        img_name = str(image_info["image_name"])+"_seg.png"
        #path to seg.png image based on the non-seg image from dataset
        image_path = os.path.join(dataset_dir, img_name)
        #load seg image
        image = cv2.imread(image_path)
        
        #image = skimage.io.imread(image_path)
        #io.imshow(image)
        #plt.show()
        #cv2.imshow("Image" , image)
        #print(img_name)
        
        #load csv file with seg billboard colors
        csv_name = str(image_info["image_name"])+"_seg.csv"
        csv_path = os.path.join(dataset_dir, csv_name)
        csv_file = open(csv_path, "r")
        reader = csv.reader(csv_file)
        
        #load colors from csv file
        colors_list = []
        for line in reader:
            if line!="":
                colors_list.append((int(line[0]),int(line[1]),int(line[2])))            
        
        #print("pocet bb: "+str(len(colors_list)))
        
        #defines how much different can a color be in each of RGB channels
        #and still be recognzed as a part of an instance
        color_diff = 15
        color_num = 0

        #create empty mask array with dimensions mask_height x mask_width x instance count
        mask = np.zeros([image.shape[0], image.shape[1], len(colors_list)],dtype=np.uint8)
        
        #For each color, applt an inRange filter for segImage
        #every pixel of the seg picture that has values between 
        #pixel_color+color_diff and pixel_color-color_diff will be
        #1 after applying filter, rest will be 0
        for color in colors_list:
            #cv2 works with BRG color space instead of RGB
            lower = np.array([max((color[2]-color_diff),0), max((color[1]-color_diff),0), max((color[0]-color_diff),0)])
            upper = np.array([min((color[2]+color_diff),255), min((color[1]+color_diff),255), min((color[0]+color_diff),255)])
            tmp_mask = cv2.inRange(image, lower, upper)
            
            #res = cv2.bitwise_and(image, image, mask= tmp_mask)
            #cv2.imshow('res',res)
            #cv2.waitKey(0)
            
            mask[:,:,color_num]=tmp_mask
            color_num=color_num+1

        return mask, np.ones([mask.shape[-1]], dtype=np.int32)

    #method that returns path for a specified image_id
    #used mainly for debuging
    def image_reference(self, image_id):
        #info = self.image_info[image_id]
        return image_id


def train(model,config,epoch_num):
    """Train the model."""
    # Load training dataset
    dataset_train = BillboardDataset()
    dataset_train.load_billboards("train")
    dataset_train.prepare()

    # Load validation dataset
    dataset_val = BillboardDataset()
    dataset_val.load_billboards("val")
    dataset_val.prepare()

    # Training itself. We are training all layers,
    # not just head layers with other layers "freezed",
    # hyperparameters based on config
    print("Training all layers")
    model.train(dataset_train, dataset_val,
                learning_rate=config.LEARNING_RATE,
                epochs=epoch_num,
                layers='all')


def apply_mask(image, mask):
    """
    Combine RGB image with instance segmentation mask
    Image - [height, width, 3]
    Mask - [height, width, instance_count]
    """
    # Since we have mask for each of the instance, we need to combine the masks
    # so that we have one mask for all instances
    # [height, width, instance_count] -> [height, width]
    mask = (np.sum(mask, -1, keepdims=True) == 0)
    
    #Check if there are valid masks - if thhere has not been anything detected,
    #the mask will be empty
    if mask.shape[0] > 0:
        mask = mask[:,:,0]
        #apply gaussian blur so that edges of the mask are less visible
        mask = skimage.filters.gaussian(mask, (image.shape[0]/100)/3)
    
        #since we need to apply the mask for each of the RGB channels
        #we need to repeat the mask 3x
        mask2 = np.zeros([image.shape[0],image.shape[1],3])
        mask2[:,:,0] = mask
        mask2[:,:,1] = mask
        mask2[:,:,2] = mask
        
        #divied pixel value for each channel by 255 which gives us 0-1 range
        #for each of the channels of a pixel and then multiply it by a 0-1 
        #value of a mask
        return np.multiply(image/255,mask2)
    else:
        return image