#!/usr/bin/python

import sys

import cv
import cv2

from pllm.vision.algo import erode, dilate, contours
from pllm.vision.util import draw_segments


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("Usage: segs_tune IMG_FILE")
        sys.exit(1)

    cv.NamedWindow("preview", cv.CV_WINDOW_AUTOSIZE)
    cv.NamedWindow("controls", cv.CV_WINDOW_AUTOSIZE)
    cv.CreateTrackbar("threshold", "controls", 230, 255, lambda x: x)
    cv.CreateTrackbar("erode_iter", "controls", 1, 20, lambda x: x)
    cv.CreateTrackbar("dilate_iter", "controls", 5, 20, lambda x: x)
    cv.CreateTrackbar("eksize1", "controls", 2, 20, lambda x: x)
    cv.CreateTrackbar("eksize2", "controls", 2, 20, lambda x: x)

    cv.CreateTrackbar("dksize1", "controls", 5, 40, lambda x: x)
    cv.CreateTrackbar("dksize2", "controls", 2, 20, lambda x: x)
    cv.CreateTrackbar("invert", "controls", 1, 1, lambda x: x)

    img = cv2.imread(sys.argv[1])

    while True:
        threshold = cv.GetTrackbarPos("threshold", "controls")
        iters = cv.GetTrackbarPos("erode_iter", "controls")
        iters = max(1, iters)

        diters = cv.GetTrackbarPos("dilate_iter", "controls")
        diters = max(1, diters)

        erode_k1_size = cv.GetTrackbarPos("eksize1", "controls")
        erode_k2_size = cv.GetTrackbarPos("eksize2", "controls")

        dilate_k1_size = cv.GetTrackbarPos("dksize1", "controls")
        dilate_k2_size = cv.GetTrackbarPos("dksize2", "controls")

        invert = cv.GetTrackbarPos("invert", "controls")

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        ret, im = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY)

        if invert:
            im = cv2.bitwise_not(im)

        im = erode(im, iters, erode_k1_size, erode_k2_size)
        cv2.imshow('er', im)

        im = dilate(im, diters, dilate_k1_size, dilate_k2_size)
        cv2.imshow('di', im)

        segs = contours(im)

        res = img.copy()
        res = draw_segments(res, segs)

        cv2.imshow('orig', img)
        cv2.imshow('res', res)

        if 0xFF & cv2.waitKey(5) == 27:
            break

    cv2.destroyAllWindows()
