
import avians.util.image as aui
import avians.nn.common as anc
from datetime import datetime as dt

import argparse
import cv2
import numpy as np

import logging
LOG = logging.getLogger(__name__)

def contour_region(img, contour): 
    LOG.debug("=== img.shape ===")
    LOG.debug(img.shape)
    
    mask = np.zeros_like(img)
    LOG.debug("=== mask.shape ===")
    LOG.debug(mask.shape)
    LOG.debug("=== contour ===")
    LOG.debug(contour)
    
    cv2.drawContours(mask, [contour], -1, 255, cv2.FILLED)
    masked = np.bitwise_and(mask, img)
    LOG.debug("=== masked ===")
    LOG.debug(masked)
    (y, x, h, w) = cv2.boundingRect(contour)
    LOG.debug("=== (x, y, w, h) ===")
    LOG.debug((x, y, w, h))
    if len(img.shape) == 2:
        return masked[x:(x+w), y:(y+h)]
    else: 
        return masked[x:(x+w), y:(y+h), :]

def resized_contour_region(img, contour, sizex, sizey): 
    # contour, holes, level = contour_structure
    contour_img = contour_region(img, contour)
    LOG.debug("=== img.shape ===")
    LOG.debug(img.shape)
    LOG.debug("=== contour_img.shape ===")
    LOG.debug(contour_img.shape)
    
    resized_contour_img = cv2.resize(contour_img, (sizex, sizey))
    return resized_contour_img

def text_components(image, 
                    model, 
                    prob_threshold=0.01, 
                    rows=32, 
                    cols=32): 
    if len(image.shape) > 2: 
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    large, contours = aui.enlarge_and_get_contours(image)
    zoom_ratio = large.shape[0] // image.shape[0]
    features = np.zeros(shape=(len(contours), 1, rows, cols), dtype=np.uint8)
    locations = np.zeros(shape=(len(contours), 4), dtype=np.int)

    for i, contour in enumerate(contours): 
        features[i, 0] = resized_contour_region(large, contour, rows, cols)
        cnt, holes, level = contour
        locations[i] = cv2.boundingRect(cnt)

    predictions = model.predict_proba(features, batch_size=128, verbose=1)
    indices = (predictions > 0.01).nonzero()[0]
    dezoomed = locations / zoom_ratio
    return features[indices], dezoomed[indices]
    
def text_detection_v3(image, 
                      model, 
                      prob_threshold=0.01, 
                      rows=32, 
                      cols=32): 

    features, dezoomed = text_components(image, 
                                         model,
                                         prob_threshold,
                                         rows,
                                         cols)

    for loc in dezoomed[indices]:
        r1 = (int(loc[0]), int(loc[1]))
        r2 = (int(loc[0]) + int(loc[2]), int(loc[1]) + int(loc[3]))
        cv2.rectangle(image, r1, r2, (0, 0, 255), 1)

    return image

def load_model(fn):
    m = anc.load_keras_model(fn,
                             loss="binary_crossentropy",
                             optimizer="adadelta")
    m.summary()
    return m

def main_caller(model_f,
                image,
                output): 
    
    img = cv2.imread(image, 0)
    model = load_model(model_f)
    result = text_detection_v3(img, model)
    cv2.imwrite(output, result)
    print("Result Saved To: {}".format(output))
    assert False

def main(): 
    parser = argparse.ArgumentParser(description="""Detects textual regions on
    the given image using the supplied model generated by train_rnn """,
                                    formatter_class=argparse.RawDescriptionHelpFormatter) # noqa
    
    parser.add_argument('model', 
                        help="Keras model name to detect text")
    parser.add_argument('image', 
                        help="Image that we are going to detect the text")
    parser.add_argument('--output', 
                        help="File to save the image", 
                        default="text-detection-result-{}.png".format(
                            dt.now().strftime("%F-%T")))

    args = vars(parser.parse_args())
    print(args)
    main_caller(**args)

if __name__ == '__main__': 
    main()
