#!/usr/bin/env python

""" Trains a binner color matcher for cone detection """

from __future__ import print_function
import cv2
import numpy as np
import sys
import time
import glob
import rospy
import json

class ColorTrainer:

    def __init__(self):
        self.numBins = 16

    def main(self):
        rospy.init_node('color_trainer_binned')
        self.numBins = rospy.get_param('~bins', 16)
        self.threshold = rospy.get_param('~threshold', 0.75)
        image_pattern = rospy.get_param('~images')
        mask_pattern = rospy.get_param('~masks')
        outputPath = rospy.get_param('~output')

        images = glob.glob(image_pattern)
        images.sort()
        masks = glob.glob(mask_pattern)
        masks.sort()
        if len(images) == 0:
            print('No images found using pattern:', image_pattern)
            sys.exit(1)

        if len(masks) != len(images):
            print('Images and masks must be 1-1', '#images:', len(images), '#masks:', len(masks))
            sys.exit(1)

        self.goodBins = np.zeros(self.numBins**3, np.int32)
        self.badBins = np.zeros(self.numBins**3, np.int32)

        for i in range(len(images)):
            self.processImage(images[i], masks[i])
            
        maskBins = []
        for i in range(len(self.goodBins)):
            n = self.goodBins[i] + self.badBins[i]
            if n >= 10:
                prob = self.goodBins[i] / float(n)
                if prob >= self.threshold:
                    maskBins.append(i)

        result = {
            'numBins': self.numBins,
            'maskBins': maskBins
        }

        with open(outputPath, 'w') as outfile:
            json.dump(result, outfile)
        
        sys.exit(0)

    def processImage(self, imgFile, maskFile):
        print('Processing', imgFile)
        im = cv2.imread(imgFile)
        mask = cv2.cvtColor(cv2.imread(maskFile), cv2.COLOR_BGR2GRAY)

        h, w = im.shape[0:2]
        im = im.reshape((h*w, im.shape[2]))
        mask = mask.reshape(h*w)

        scale = 256 / self.numBins
        toIndex = np.array([1, self.numBins, self.numBins**2], np.int32)
        imBinned = np.matmul(im / scale, toIndex)

        for i in range(len(imBinned)):
            if mask[i] > 0:
                self.goodBins[imBinned[i]] += 1
            else:
                self.badBins[imBinned[i]] += 1

if __name__ == "__main__":
    try:
        ColorTrainer().main()

    except rospy.ROSInterruptException:
        pass
