# Grammar-Guided Feature Extraction (GGFE)
#
# Author:   Damian Eads
#
# File:     haar.py
#
# Purpose:  This module contains an implementation of generating
#           Viola and Jones Haar-like features.
#
# Date:     August 8, 2009

import numpy as np

def rand1n(n):
    """
    Generates a random number between 1 and n (exclusive)
    """
    return np.random.randint(n)+1

def rand0n(n):
    """
    Generates a random number between 0 and n (exclusive)
    """
    if n == 0:
        return 0
    else:
        return np.random.randint(n)

def get_kernel_string(rects):
    """
    Returns a string from a rectangle specification list that can be
    quickly parsed by the C++ routine that evaluates Viola and Jones
    features on an image.
    """
    first = []
    second = []
    for (val, L) in rects:
        if val == -1:
            first.append(L)
        else:
            second.append(L)
    return '{(%s)(%s)}' % (get_block_string(first), get_block_string(second))

def get_block_string(rects):
    """
    Returns a string encoding a rectangle that can be quickly parsed
    by the C++ routine.
    """
    L = []
    for rect in rects:
        L.append(get_rect_string(*rect))
    return str.join('',L)

def get_rect_string(X, Y, W, H):
    """
    Returns a string encoding a rectangle that can be quickly parsed
    by the C++ routine.
    """
    (x1, y1, x2, y2) = (X, Y, X+W-1, Y+H-1)
    return "[%d,%d,%d,%d]" % (x1-W/2, y1-H/2, x2-W/2, y2-W/2)

class Haar:
    """
    A base class for defining common operations for Viola and Jones
    feature banks.
    """
        
    def __init__(self, W, H):
        """
        Initializes a new object for generating random Haar-like features.

        Arguments:

           - W : int
             The width of the detection window and thus the maximum width of
             the Viola and Jones feature.

           - H : int
             The height of the detection window and thus the maximum height of
             the Viola and Jones feature.
        """
        self.W = W
        self.H = H

    def get_number(self):
        raise ValueError("Abstract")

class Haar2Horizontal(Haar):

    """
    A class for generating horizontal 2-rectangle Viola and Jones features.
    """

    def __init__(self, W, H):
        """
        Initializes a new object for generating random Haar-like features.

        Arguments:

           - W : int
             The width of the detection window and thus the maximum width of
             a horizontal pair of rectangles.

           - H : int
             The height of the detection window and thus the maximum height of
             a horizontal pair of rectangles.
        """
        Haar.__init__(self, W, H)
        self.hW = np.floor_divide(W, 2.0)
        self.ws = np.arange(1, (self.hW+1)) * 2.0
        self.hs = np.arange(1, (self.H+1))

    def get_number(self):
        """
        Returns the total number of Haar-like features that can be
        generated by this bank.
        """
        return (self.W-self.ws+1).sum()*(self.H-self.hs+1).sum()

    def get_random_condensed(self):
        """
        Returns a random horizontal Haar rectangle pair in condensed form,
        (rX, rY, rW, rH). The coordinates (rX, rY) represent the location
        of the feature in the detection window and (rW, rH), its
        dimensions.
        """
        rW = rand1n(self.hW)*2.0
        rH = rand1n(self.H)
        rX = rand0n(self.W-rW)
        rY = rand0n(self.H-rH)
        return (rX, rY, rW, rH)

    def get_random_rectangles(self):
        """
        Returns a random horizontal Haar rectangle pair as a list of
        rectangle specifications

            [(r1X, r1Y, r1W, r1H), (r2X, r2Y, r2W, r2H), ...].
        """
        (rX, rY, rW, rH) = self.get_random_condensed()        
        return [(-1, (rX,      rY, rW/2, rH)),
                (1, (rX+rW/2, rY, rW/2, rH))]

class Haar2Vertical(Haar):
    """
    A class for generating vertical 2-rectangle Viola and Jones features.
    """

    def __init__(self, W, H):
        """
        Initializes a new object for generating random Haar-like
        (Viola and Jones) features.

        Arguments:

           - W : int
             The width of the detection window and thus the maximum width of
             a vertical pair of rectangles.

           - H : int
             The height of the detection window and thus the maximum height of
             a vertical pair of rectangles.
        """
        Haar.__init__(self, W, H)
        self.hH = np.floor(H/2.0)
        self.hs = np.arange(1, (self.hH+1)) * 2.0
        self.ws = np.arange(1, (self.W+1))

    def get_number(self):
        """
        Returns the total number of Haar-like features that can be
        generated by this bank.
        """
        return (self.W-self.ws+1).sum()*(self.H-self.hs+1).sum()

    def get_random_condensed(self):
        """
        Returns a random vertical Haar rectangle pair in condensed form,
        (rX, rY, rW, rH). The coordinates (rX, rY) represent the location
        of the feature in the detection window and (rW, rH), its
        dimensions.
        """
        rW = rand1n(self.W)
        rH = rand1n(self.hH)*2.0
        rX = rand0n(self.W-rW)
        rY = rand0n(self.H-rH)
        return (rX, rY, rW, rH)

    def get_random_rectangles(self):
        (rX, rY, rW, rH) = self.get_random_condensed()        
        return [(-1, (rX, rY,      rW, rH/2)),
                (1,  (rX, rY+rH/2, rW, rH/2))]

class Haar3Horizontal(Haar):
    """
    A class for generating horizontal triplet Viola and Jones features.
    """

    def __init__(self, W, H):
        """
        Initializes a new object for generating random Haar-like
        (Viola and Jones) features.

        Arguments:

           - W : int
             The width of the detection window and thus the maximum width of
             a horizontal triplet of rectangles.

           - H : int
             The height of the detection window and thus the maximum height of
             a horizontal triplet of rectangles.
        """
        Haar.__init__(self, W, H)
        self.hW = np.floor(W/3.0)
        self.ws = np.arange(1, (self.hW+1)) * 3.0
        self.hs = np.arange(1, (self.H+1))

    def get_number(self):
        """
        Returns the total number of Haar-like features that can be
        generated by this bank.
        """
        return (self.W-self.ws+1).sum()*(self.H-self.hs+1).sum()

    def get_random_condensed(self):
        """
        Returns a random horizontal Haar rectangle triplet in condensed form,
        (rX, rY, rW, rH). The coordinates (rX, rY) represent the location
        of the feature in the detection window and (rW, rH), its
        dimensions.
        """
        rW = rand1n(self.hW)*3.0
        rH = rand1n(self.H)
        rX = rand0n(self.W-rW)
        rY = rand0n(self.H-rH)
        return (rX, rY, rW, rH)

    def get_random_rectangles(self):
        """
        Returns a random horizontal Haar rectangle triplet as a list of
        rectangle specifications

            [(r1X, r1Y, r1W, r1H), (r2X, r2Y, r2W, r2H), ...].
        """
        (rX, rY, rW, rH) = self.get_random_condensed()        
        return [(-1, (rX,             rY, rW/3.0, rH)),
                (1,  (rX+rW/3.0,      rY, rW/3.0, rH)),
                (-1, (rX+(2.0/3.0)*rW,rY, rW/3.0, rH))]

class Haar3Vertical(Haar):
    """
    A class for generating vertical triplet Viola and Jones features.
    """

    def __init__(self, W, H):
        """
        Initializes a new object for generating random Haar-like
        (Viola and Jones) features.

        Arguments:

           - W : int
             The width of the detection window and thus the maximum width of
             a vertical triplet of rectangles.

           - H : int
             The height of the detection window and thus the maximum height of
             a vertical triplet of rectangles.
        """
        Haar.__init__(self, W, H)
        self.hH = np.floor(H/3.0)
        self.hs = np.arange(1, (self.hH+1)) * 3.0
        self.ws = np.arange(1, (self.W+1))

    def get_number(self):
        """
        Returns the total number of Haar-like features that can be
        generated by this bank.
        """
        return (self.W-self.ws+1).sum()*(self.H-self.hs+1).sum()

    def get_random_condensed(self):
        """
        Returns a random vertical Haar rectangle triplet in condensed form,
        (rX, rY, rW, rH). The coordinates (rX, rY) represent the location
        of the feature in the detection window and (rW, rH), its
        dimensions.
        """
        rW = rand1n(self.W)
        rH = rand1n(self.hH)*3.0
        rX = rand0n(self.W-rW)
        rY = rand0n(self.H-rH)
        return (rX, rY, rW, rH)

    def get_random_rectangles(self):
        """
        Returns a random vertical Haar rectangle triplet as a list of
        rectangle specifications

            [(r1X, r1Y, r1W, r1H), (r2X, r2Y, r2W, r2H), ...].
        """
        (rX, rY, rW, rH) = self.get_random_condensed()        
        return [(-1, (rX,rY,              rW, rH/3.0)),
                (1,  (rX,rY+rH/3.0,       rW, rH/3.0)),
                (-1, (rX,rY+(2.0/3.0)*rH, rW, rH/3.0))]

class Haar4(Haar):
    """
    A class for generating 2x2 quad Viola and Jones features.
    """

    def __init__(self, W, H):
        """
        Initializes a new object for generating random Haar-like
        (Viola and Jones) features.

        Arguments:

           - W : int
             The width of the detection window and thus the maximum width of
             a horizontal pair of rectangles.

           - H : int
             The height of the detection window and thus the maximum height of
             a horizontal pair of rectangles.
        """
        Haar.__init__(self, W, H)
        self.hH = np.floor(H/2.0)
        self.hW = np.floor(W/2.0)
        self.hs = np.arange(1, (self.hH+1)) * 2.0
        self.ws = np.arange(1, (self.hW+1)) * 2.0

    def get_number(self):
        """
        Returns the total number of Haar-like features that can be
        generated by this bank.
        """
        return (self.W-self.ws+1).sum()*(self.H-self.hs+1).sum()

    def get_random_condensed(self):
        """
        Returns a random Viola and Jones 2x2 quad feature in condensed form,
        (rX, rY, rW, rH). The coordinates (rX, rY) represent the location
        of the feature in the detection window and (rW, rH), its
        dimensions.
        """
        rW = rand1n(self.hW)*2.0
        rH = rand1n(self.hH)*2.0
        rX = rand0n(self.W-rW)
        rY = rand0n(self.H-rH)
        return (rX, rY, rW, rH)

    def get_random_rectangles(self):
        (rX, rY, rW, rH) = self.get_random_condensed()        
        return [(-1, (rX,        rY,        rW/2.0, rH/2.0)),
                (1,  (rX+rW/2.0, rY,        rW/2.0, rH/2.0)),
                (1,  (rX,        rY+rH/2.0, rW/2.0, rH/2.0)),
                (-1, (rX+rW/2.0, rY+rH/2.0, rW/2.0, rH/2.0))]

class HaarBank:

    def __init__(self, W, H):
        """
        Initializes a new object for generating random Haar-like
        (Viola and Jones) features.

        Arguments:

           - W : int
             The width of the detection window and thus the maximum width of
             a horizontal pair of rectangles.

           - H : int
             The height of the detection window and thus the maximum height of
             a horizontal pair of rectangles.
        """
        self.haar2h = Haar2Horizontal(W, H)
        self.haar2v = Haar2Vertical(W, H)
        self.haar3h = Haar3Horizontal(W, H)
        self.haar3v = Haar3Vertical(W, H)
        self.haar4 = Haar4(W, H)
        self.haars = [self.haar2h, self.haar2v, self.haar3h, self.haar3v, self.haar4]
        self.W = W
        self.H = H

    def get_number(self):
        """
        Returns the total number of Viola and Jones features that
        can be generated given the detection window passed.
        """
        return (self.haar2h.get_number()
                +self.haar2v.get_number()
                +self.haar3h.get_number()
                +self.haar3v.get_number()
                +self.haar4.get_number())

    def get_random_rectangles(self):
        """
        Returns a random 1x2, 2x1, 3x1, 1x3, or 2x2 Viola and Jones
        feature as a random rectangle.
        """
        return self.haars[np.random.randint(5)].get_random_rectangles()

    def random_plot(self):
        """
        Plots random Viola and Jones feature onto the screen to examine
        what the classes in this module are doing.
        """
        import matplotlib.pylab as mpl
        I = np.zeros((self.W, self.H), dtype='uint8')
        rects = self.get_random_rectangles()
        for (val, (X, Y, W, H)) in rects:
            if val==-1:
                f = 128
            elif val==1:
                f = 255
            else:
                f = 0
            I[X:(X+W),Y:(Y+H)] = f
        mpl.imshow(I, interpolation='nearest')

