from os import listdir
from os.path import isdir
from warnings import filterwarnings

from matplotlib.pyplot import imread, plot, show, xlim, yscale
from matplotlib.pyplot import ylabel, legend, gca, xticks
from numpy import abs, sum, argsort, min
from numpy import interp, linspace, zeros, power, cumsum, concatenate, all, logical_or, ones, where, array, logical_xor, ceil, round
from numpy.linalg import norm
from scipy.ndimage import binary_erosion
from scipy.signal import savgol_filter
from skimage.color import rgb2gray
from skimage.measure import find_contours

SETS_PATH = './images/ss10'
FILTER_WINDOW = 21
ANGLE_REL_IMP = 2
LEN_REL_IMP = 2


def load_image(path):
    # Load image and convert to grayscale
    return rgb2gray(imread(path))


def cut_polylines(image, polylines):
    # Obtain eroded version of foreground
    foreground = binary_erosion(input=logical_or(image < .25, image > .75), structure=ones((3, 3)))
    # Cut polylines
    new_polylines = []
    for polyline in polylines:
        # Find polyline vertices contained in foreground
        contained = foreground[round(polyline[:, 0]).astype(int), round(polyline[:, 1]).astype(int)]
        idx_cuts = where(logical_xor(contained[1:], contained[:-1]))[0] + 1
        # Cut polyline is original polyline
        if idx_cuts.size == 0:
            new_polylines.append(polyline)
            continue
        # Find non-wrapping polylines
        for idx_start, idx_end in zip(idx_cuts[:-1], idx_cuts[1:]):
            if contained[idx_start]:
                new_polylines.append(polyline[idx_start:idx_end, :])
        # Find potentially wrapping polylines
        if contained[idx_cuts[-1]]:
            # Closed polyline
            if all(polyline[0] == polyline[-1]):
                new_polylines.append(concatenate((polyline[idx_cuts[-1]:, :], polyline[:idx_cuts[0], :])))
            # Open polyline
            else:
                new_polylines.append(polyline[idx_cuts[-1]:, :])
                new_polylines.append(polyline[:idx_cuts[0], :])
    return array(new_polylines)


def extract_polylines(image, cut=False, min_len=0, resample=True):
    polylines = find_contours(array=image, level=.75)
    # Optionally cut polylines
    if cut:
        polylines = cut_polylines(image=image, polylines=polylines)
    # Optionally resample polylines
    if resample:
        res_polylines = []
        for polyline in polylines:
            # Find t in x=f(t), y=g(t), where t is the cumulative polyline length
            ts = concatenate(([0], cumsum(norm(polyline[1:] - polyline[:-1], axis=1))))
            if ts[-1] < min_len:
                continue
            # Resample
            new_ts = linspace(0, ts[-1], int(ceil(ts[-1])))
            # Interpolate
            new_polyline = zeros((len(new_ts), 2))
            new_polyline[:, 0] = interp(x=new_ts, xp=ts, fp=polyline[:, 0])
            new_polyline[:, 1] = interp(x=new_ts, xp=ts, fp=polyline[:, 1])
            res_polylines.append(new_polyline)
        polylines = res_polylines
    return polylines


def extract_der_angles(polyline, filter_window):
    filterwarnings("ignore", message="internal gelsd driver lwork query error, required iwork dimension not returned")
    # Differentiate between closed and open polylines
    if all(polyline[0] == polyline[-1]):
        r = polyline[:-1, 0]
        c = polyline[:-1, 1]
        mode = 'wrap'
    else:
        r = polyline[:, 0]
        c = polyline[:, 1]
        mode = 'nearest'
    # Obtain first and second derivates
    dr = savgol_filter(x=r, window_length=filter_window, polyorder=1, deriv=1, mode=mode)
    dc = savgol_filter(x=c, window_length=filter_window, polyorder=1, deriv=1, mode=mode)
    ddr = savgol_filter(x=dr, window_length=filter_window, polyorder=1, deriv=1, mode=mode)
    ddc = savgol_filter(x=dc, window_length=filter_window, polyorder=1, deriv=1, mode=mode)
    # Obtain angular derivates (derivative of arctan2)
    der_angles = -dr / (power(dc, 2) + power(dr, 2)) * ddc + dc / (power(dc, 2) + power(dr, 2)) * ddr
    return der_angles


def compute_score(image):
    polylines = extract_polylines(image=image, cut=True, min_len=FILTER_WINDOW)
    # Compute polyline scores and accumulate
    score = 0
    for polyline in polylines:
        der_angles = extract_der_angles(polyline=polyline, filter_window=FILTER_WINDOW)
        # Relative scoring according to angular derivative and polyline length
        score += sum(abs(der_angles) ** ANGLE_REL_IMP) * len(der_angles) ** LEN_REL_IMP

    return score


def evaluate_sets():
    # Iterate over test sets
    n_sets = len([dir for dir in listdir(path=SETS_PATH) if isdir('%s/%s' % (SETS_PATH, dir))])
    for s_idx in range(n_sets):
        # Obtain scores for images in a test set
        n_files = len([file for file in listdir(path="%s/s%d/" % (SETS_PATH, s_idx)) if file.endswith(".png")])
        images = [load_image(path='%s/s%d/f%d.png' % (SETS_PATH, s_idx, f_idx)) for f_idx in range(n_files)]
        scores = array([compute_score(image=image) for image in images])
        # Order and print relative scores
        print('s%-2s' % s_idx, end=' - ')
        for plot_idx, img_idx in enumerate(argsort(scores)):
            rel_score = scores[img_idx] / min(scores)
            # rel_score = scores[img_idx] / sorted(scores)[max(0, plot_idx - 1)] # Uncomment for relative scores w.r.t. previous image
            print('f%d: %-5.2f' % (img_idx, rel_score), end=' ')
        print('')
        # Plot absolute scores
        plot(scores, 'o:', label='%s' % s_idx)

    # Format y-axis (score)
    ylabel('score')
    yscale('log')

    # Format x-axis (images)
    _, max_x = xlim()
    xticks(range(int(max_x + 1)), ['f%d' % i for i in range(int(max_x + 1))])

    # Add legend (test set)
    ymin, ymax = gca().get_ylim()
    gca().set_ylim((ymin, ymax + .3 * (ymax - ymin)))
    legend(loc='upper center', mode='expand', ncol=6, title='', labels=['s%d' % i for i in range(n_sets)])

    show()


evaluate_sets()
