import cv2
import numpy as np
import glob
import math

import StringIO

from scipy.signal import argrelmax, find_peaks

# ============================================================================

zones_a = [(50, 72), (144, 166), (243, 265), (328, 350), (436, 458)]
zones_b = [(42,  64), (122, 144), (207, 229), (276, 298), (369, 391), (496, 518)]

# ============================================================================

def plot_report(filename, report):
    from matplotlib import pyplot as plt
    from matplotlib.gridspec import GridSpec

    IMAGE_KEYS = ['img', 'thresh', 'thresh_1', 'thresh_2', 'canny', 'canny_1', 'canny_2']
    PLOT_SPAN = 5
    TEXT_SPAN = 2
    ROW_COUNT = (len(IMAGE_KEYS) + 1) + 3 * (PLOT_SPAN + 1) + 2 * (TEXT_SPAN)

    fig = plt.figure()
    plt.suptitle(filename)
    gs = GridSpec(ROW_COUNT, 2)

    row = 0

    for key in IMAGE_KEYS:
        plt.subplot(gs[row,:])
        plt.gca().set_title(key)
        plt.imshow(report[key], cmap='gray', aspect='equal')
        plt.axis('off')
        row += 1

    proj_width = len(report['projection'])
    proj_x = np.arange(proj_width)

    plt.subplot(gs[row+1:row+1+PLOT_SPAN,:])
    plt.gca().set_title('Vertical Projections (Raw and Smoothed)')
    plt.plot(proj_x, report['projection'], 'y-')
    plt.plot(proj_x, report['projection_1'], 'r-')
    plt.plot(proj_x, report['projection_2'], 'b-')
    plt.xlim((0, proj_width - 1))
    plt.ylim((0, 255))

    row += PLOT_SPAN + 1

    #bins_ab, weight_ab, details_ab, checkbox_count

    plt.subplot(gs[row+1:row+1+PLOT_SPAN,:])
    plt.gca().set_title('Smoothed Projection with Peaks and Zones')
    plt.plot(proj_x, report['projection_2'])

    for zone in zones_a:
        plt.axvspan(zone[0], zone[1], facecolor='y', alpha=0.1)
    for zone in zones_b:
        plt.axvspan(zone[0], zone[1], facecolor='r', alpha=0.1)
    for x in report['peaks']:
        plt.axvline(x=x, color='m')

    plt.xlim((0, proj_width - 1))
    plt.ylim((0, report['projection_2'].max()))

    row += PLOT_SPAN + 1

    plt.subplot(gs[row+1:row+1+TEXT_SPAN,0], frameon=False)
    plt.gca().set_title('Details - 5 boxes')
    plt.axis([0, 1, 0, 1])
    plt.gca().axes.get_yaxis().set_visible(False)
    plt.gca().axes.get_xaxis().set_visible(False)
    plt.text(0, 1, report['details_a'], family='monospace', fontsize=8, ha='left', va='top')

    plt.subplot(gs[row+1:row+1+TEXT_SPAN,1], frameon=False)
    plt.gca().set_title('Details - 6 boxes')
    plt.axis([0, 1, 0, 1])
    plt.gca().axes.get_yaxis().set_visible(False)
    plt.gca().axes.get_xaxis().set_visible(False)
    plt.text(0, 1, report['details_b'], family='monospace', fontsize=8, ha='left', va='top')

    row += TEXT_SPAN

    plt.subplot(gs[row+1:row+1+PLOT_SPAN,:])
    plt.gca().set_title('Weights')
    plt.barh([2, 1]
        , [report['weight_a'], report['weight_b']]
        , align='center'
        , color=['y', 'r']
        , tick_label=['5 boxes', '6 boxes'])
    plt.ylim((0.5, 2.5))

    row += PLOT_SPAN + 1
    row += 1

    plt.subplot(gs[row,:])
    plt.gca().set_title('Input Image')
    plt.imshow(report['img'], cmap='gray', aspect='equal')
    plt.axis('off')

    row += 1

    plt.subplot(gs[row:row+TEXT_SPAN,:], frameon=False)
    plt.axis([0, 1, 0, 1])
    plt.gca().axes.get_yaxis().set_visible(False)
    plt.gca().axes.get_xaxis().set_visible(False)
    result_text = "The image contains %d boxes." % report['checkbox_count']
    plt.text(0.5, 1, result_text, family='monospace', weight='semibold', fontsize=24, ha='center', va='top')

    fig.set_size_inches(12, ROW_COUNT * 0.8)
    plt.savefig('plot_%s.png' % filename[:2], bbox_inches="tight")
    plt.close(fig)


def bin_peaks(peaks, values, zones):
    bins = [[] for x in xrange(len(zones))]

    for peak, value in zip(peaks, values):
        for i, zone in enumerate(zones):
            if (peak >= zone[0]) and (peak <= zone[1]):
                bins[i].append((peak, value))

    return bins

def analyze_bins(bins):
    buf = StringIO.StringIO()

    total_weight = 0.0
    for i, bin in enumerate(bins):
        buf.write("Position %d: " % i)
        weight = 0.0
        if len(bin) == 0:
            buf.write("no peaks")
        else:
            best_bin = sorted(bin, key=lambda x: x[1], reverse=True)[0]
            weight = best_bin[1]
            if len(bin) == 1:
                buf.write("single peak @ %d (value=%0.3f)" % best_bin)
            else:
                buf.write("%d peaks, best @ %d (value=%0.3f)" % (len(bin), best_bin[0], best_bin[1]))

        buf.write(" | weight=%0.3f\n" % weight)

        total_weight += weight

    buf.write("Total weight = %0.3f" % total_weight)
    return total_weight, buf.getvalue()


def proc(filename):
    report = {}

    img = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
    report['img'] = img.copy()

    thresh = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 15, 2)
    report['thresh'] = thresh.copy()

    thresh = cv2.morphologyEx(thresh, cv2.MORPH_ERODE, np.ones((1,3),np.uint8))
    report['thresh_1'] = thresh.copy()

    thresh = cv2.morphologyEx(thresh, cv2.MORPH_DILATE, np.ones((3,1),np.uint8))
    report['thresh_2'] = thresh.copy()

    edges = cv2.Canny(thresh, 40, 120, apertureSize=5)
    report['canny'] = edges.copy()

    edges = cv2.morphologyEx(edges, cv2.MORPH_OPEN, np.ones((5,1),np.uint8))
    report['canny_1'] = edges.copy()

    edges = cv2.morphologyEx(edges, cv2.MORPH_DILATE, np.ones((1,3),np.uint8))
    report['canny_2'] = edges.copy()

    projection = np.mean(edges, 0).flatten()
    report['projection'] = projection.copy()

    projection = cv2.blur(projection, (1, 21)).flatten()
    report['projection_1'] = projection.copy()

    projection = cv2.blur(projection, (1, 11)).flatten()
    report['projection_2'] = projection.copy()

    peaks = find_peaks(projection)[0]
    report['peaks'] = peaks.copy()

    peak_values = projection[peaks]
    report['peak_values'] = peak_values.copy()

    bins_a = bin_peaks(peaks, peak_values, zones_a)
    report['bins_a'] = list(bins_a)

    bins_b = bin_peaks(peaks, peak_values, zones_b)
    report['bins_b'] = list(bins_b)

    weight_a, details_a = analyze_bins(bins_a)
    report['weight_a'] = weight_a
    report['details_a'] = details_a
    weight_b, details_b = analyze_bins(bins_b)
    report['weight_b'] = weight_b
    report['details_b'] = details_b

    checkbox_count = 5 if (weight_a > weight_b) else 6
    report['checkbox_count'] = checkbox_count

    return checkbox_count, report


for filename in glob.glob('*-*.png'):
    box_count, report = proc(filename)
    plot_report(filename, report)