#!/usr/bin/env python3

import collections
import os
import pathlib
import subprocess
import sys

sys.path.append(str(pathlib.Path(__name__).resolve().parent))  # noqa


import click
import cv2 as cv
import numpy as np
import tqdm
import scipy.signal as signal

from metrikai import cleanup
from metrikai import show  # noqa
from metrikai import show_samples
from metrikai import scanxy  # noqa
from metrikai import SmartScan  # noqa
from metrikai import im_to_X
from metrikai import wait_for_key  # noqa
from metrikai import line_scanxy
from metrikai import rotate
from metrikai import find_lines
from metrikai import find_slope
from metrikai import move
from metrikai import italic


@click.group()
def main():
    pass


@main.command('download-samples')
def download_samples():
    Sample = collections.namedtuple('Sample', ('url', 'num'))
    languages = {
        'la': [
            Sample('http://www.epaveldas.lt/recordImageSmall/ARCH/1521/1/20?exId=24294&seqNr={num}', 39),
            Sample('http://www.epaveldas.lt/recordImageSmall/ARCH/1375/1/16?exId=230351&seqNr={num}', 46),
        ],
        'ru': [
            Sample('http://www.epaveldas.lt/recordImageSmall/ARCH/1338/1/10?exId=320185&seqNr={num}', 203),
            Sample('http://www.epaveldas.lt/recordImageSmall/ARCH/1360/1/19?exId=97388&seqNr={num}', 155),
            Sample('http://www.epaveldas.lt/recordImageSmall/ARCH/1375/1/8?exId=230381&seqNr={num}', 196),
        ],
        'lt': [
            Sample('http://www.epaveldas.lt/recordImageSmall/ARCH/1375/1/27?exId=230345&seqNr={num}', 60),
            Sample('http://www.epaveldas.lt/recordImageSmall/ARCH/1375/1/24?exId=230346&seqNr={num}', 164),
        ]
    }

    total = sum(sample.num for samples in languages.values() for sample in samples)

    with tqdm.tqdm(total=total) as pbar:
        for language, samples in languages.items():
            for i, sample in enumerate(samples, 1):
                os.makedirs(f'samples/epaveldas/{language}/{i:03d}', exist_ok=True)
                for num in range(1, sample.num + 1):
                    subprocess.run(['wget', '-q', sample.url.format(num=num),
                                    '-O', f'samples/epaveldas/{language}/{i:03d}/{num:03d}.jpg'])
                    pbar.update(1)


@main.command()
def lines():
    sample = 0
    samples = [
        'samples/epaveldas/ru/001/004.jpg',
        'samples/epaveldas/ru/002/004.jpg',
        'samples/epaveldas/ru/003/004.jpg',
        'samples/epaveldas/lt/001/004.jpg',
        'samples/epaveldas/lt/002/004.jpg',
        'samples/epaveldas/la/001/004.jpg',
        'samples/epaveldas/la/002/004.jpg',
    ]

    im = cleanup(cv.imread(samples[sample], cv.IMREAD_GRAYSCALE))
    ih, iw = im.shape

    winsize = 20
    windows = [
        ('blackman', signal.blackman(winsize)),
        ('blackmanharris', signal.blackmanharris(winsize)),
        ('hann', signal.hann(winsize)),
        ('barthann', signal.barthann(winsize)),
        ('bartlett', signal.bartlett(winsize)),
        ('boxcar', signal.boxcar(winsize)),
        ('cosine', signal.cosine(winsize)),
        ('exponential', signal.exponential(winsize)),
        ('flattop', signal.flattop(winsize)),
        ('gaussian(std=7)', signal.gaussian(winsize, std=7)),
        ('hamming', signal.hamming(winsize)),
        ('hanning', signal.hanning(winsize)),
        ('kaiser(beta=1)', signal.kaiser(winsize, beta=1)),
        ('nuttall', signal.nuttall(winsize)),
        ('parzen', signal.parzen(winsize)),
        ('slepian(width=10)', signal.slepian(winsize, width=10)),
        ('triang', signal.triang(winsize)),
        ('tukey', signal.tukey(winsize)),
    ]
    # currwin = 0
    gaussian_window_std = 7

    initialx = 284
    x = initialx
    while True:
        w = 200

        y = im[:, x:x + w].mean(axis=1)
        # wintitle, win = windows[currwin]
        win = signal.gaussian(winsize, std=gaussian_window_std)
        wintitle = f'gaussian(std={gaussian_window_std})'
        y = signal.convolve(y, win, mode='same') / sum(win)
        minimums = signal.argrelextrema(y, np.less_equal, order=10)[0]
        minimums = minimums[signal.convolve(np.diff(minimums), [.5, .5]) > 5]

        # w = int(np.ceil(np.diff(minimums).mean()))

        white = (255, 255, 255)
        green = (64, 255, 64)
        blue = (255, 64, 64)

        imx = cv.cvtColor(im[:, x:x + w].copy(), cv.COLOR_GRAY2BGR)
        imx[minimums, :] = green

        for ly, lh in zip(minimums, np.diff(minimums)):
            window = signal.blackman(5)
            slopes = []
            median_maximums = []
            for i in range(21):
                slope = i * -0.05
                medians = np.median(im[move(italic(slope, np.arange(w), w, lh), ly, x)], axis=1)
                convolved = signal.convolve(medians, window, mode='same') / sum(window)
                maximums = signal.argrelextrema(convolved, np.greater, order=6)[0]
                slopes.append(slope)
                median_maximums.append(convolved[maximums].mean() if maximums.size else 0)

            # slope = -0.7
            slope = slopes[median_maximums.index(max(median_maximums))]

            # means = im[ly:ly + lh, x:x + w].mean(axis=0)
            means = im[move(italic(slope, np.arange(w), w, lh), ly, x)].mean(axis=1)

            window = signal.gaussian(6, std=gaussian_window_std)
            convolved = signal.convolve(means, window, mode='same') / sum(window)
            line_minimums = signal.argrelextrema(convolved, np.less_equal, order=6)[0]
            if line_minimums.size > 2:
                line_minimums = line_minimums[signal.convolve(np.diff(line_minimums), [.5, .5], mode='full') > 5]
                # imx[ly:ly + lh, line_minimums] = green
                imx[move(italic(slope, line_minimums, w - 1, lh), ly, 0)] = green

        plot = np.zeros((imx.shape[0], 255, 3), dtype=np.uint8)
        cv.polylines(plot, [np.array([y, range(y.shape[0])]).astype(np.int32).T], False, blue, 1, cv.LINE_AA)
        cv.putText(plot, wintitle, (10, 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, white, 1, cv.LINE_AA)
        imx = np.append(imx, plot, axis=1)
        cv.imshow('im', imx)

        reload = False
        key = cv.waitKey(0)
        if key == 27:
            break
        elif key & 0xFF == 81:  # right
            x = x - 2 if x - 2 >= 0 else 0
        elif key & 0xFF == 83:  # right
            x = x + 2 if x + 2 < im.shape[1] - w else im.shape[1] - w
        elif key & 0xFF == 82:  # up
            # currwin = rotate(currwin, len(windows), 1)
            gaussian_window_std += 1
        elif key & 0xFF == 84:  # down
            # currwin = rotate(currwin, len(windows), -1)
            if gaussian_window_std > 1:
                gaussian_window_std -= 1
        elif key == ord('a'):
            sample = rotate(sample, len(samples), 1)
            reload = True
        elif key == ord('z'):
            sample = rotate(sample, len(samples), -1)
            reload = True

        if reload:
            im = cleanup(cv.imread(samples[sample], cv.IMREAD_GRAYSCALE))
            ih, iw = im.shape
            x = initialx

    cv.destroyAllWindows()


@main.command()
def train():
    sample = 0
    samples = [
        'samples/epaveldas/ru/001/004.jpg',
        'samples/epaveldas/ru/002/004.jpg',
        'samples/epaveldas/ru/003/004.jpg',
        'samples/epaveldas/lt/001/004.jpg',
        'samples/epaveldas/lt/002/004.jpg',
        'samples/epaveldas/la/001/004.jpg',
        'samples/epaveldas/la/002/004.jpg',
    ]

    w = 200
    reload = True
    initialx = x = 284
    while True:
        if reload:
            reload = False

            # im = cv.imread(samples[sample])
            # cv.imshow('orig', im)

            im = cv.imread(samples[sample], cv.IMREAD_GRAYSCALE)
            im = im[:400, :500]
            im = cleanup(im, remove_lines=True)
            ih, iw = im.shape
            x = initialx

            winsize = 20
            kernel = signal.gaussian(winsize, std=2)
            slope = find_slope(im, axis=0)
            lines = find_lines(im, kernel, slope, axis=0)
            im[italic(slope, lines, iw, ih)] = 255

            cv.imshow('im', im)

        key = cv.waitKey(0)
        if key == 27:
            break
        elif key & 0xFF == 81:  # right
            x = x - 2 if x - 2 >= 0 else 0
        elif key & 0xFF == 83:  # right
            x = x + 2 if x + 2 < im.shape[1] - w else im.shape[1] - w
        elif key & 0xFF == 82:  # up
            pass
            # currwin = rotate(currwin, len(windows), 1)
        elif key & 0xFF == 84:  # down
            pass
            # currwin = rotate(currwin, len(windows), -1)
        elif key == ord('a'):
            sample = rotate(sample, len(samples), 1)
            reload = True
        elif key == ord('z'):
            sample = rotate(sample, len(samples), -1)
            reload = True

    cv.destroyAllWindows()


@main.command()
def clusters():
    im = cleanup(cv.imread('samples/epaveldas/ru/001/004.jpg', cv.IMREAD_GRAYSCALE))
    ih, iw = im.shape
    s = 25
    xstep = 20

    print('scaning image...', flush=True)
    samples = np.empty((0, s * s), dtype=np.uint8)
    # for x, y in itertools.islice(line_scanxy(im, 50, 2), 8000):
    oldx = 0
    with tqdm.tqdm(total=iw - 50) as pbar:
        for x, y in line_scanxy(im, 50, xstep):
            if x != oldx:
                oldx = x
                pbar.update(xstep)
            samples = np.append(samples, im_to_X(im[y:y + s, x:x + s]), axis=0)

    print('samples after scan: %r', samples.shape, flush=True)

    n_clusters = 10

    def predict_kmeans(samples):
        from sklearn.cluster import KMeans
        model = KMeans(n_clusters=n_clusters, random_state=0)
        return model.fit_predict(samples)

    def predict_gausian_mixture(samples):
        from sklearn.mixture import GaussianMixture
        model = GaussianMixture(n_components=n_clusters, random_state=0)
        model.fit(samples)
        return model.predict(samples)

    def predict_agglomerative_clustering(samples):
        from sklearn import manifold
        from sklearn.cluster import AgglomerativeClustering
        X = manifold.SpectralEmbedding(n_components=n_clusters).fit_transform(samples)
        model = AgglomerativeClustering(linkage='ward', n_clusters=n_clusters)
        model.fit(X)
        return model.labels_

    def predict_dbscan(samples):
        from sklearn.cluster import DBSCAN
        model = DBSCAN(eps=0.3, min_samples=n_clusters).fit(samples)
        return model.labels_

    predict = predict_kmeans

    print('predicting clusters...', end=' ', flush=True)
    stack = [{
        'samples': samples,
        'clusters': predict(samples),
        'cluster': 0
    }]
    stack[-1]['labels'] = sorted(set(stack[-1]['clusters']))
    print('done.', flush=True)

    show_samples('im', stack[-1]['samples'][stack[-1]['clusters'] == stack[-1]['cluster']], s=s)

    while True:
        key = cv.waitKey(0)
        next_cluster = None

        if key == 27:
            break

        elif key & 0xFF == 81:  # left
            next_cluster = stack[-1]['cluster'] - 1 if stack[-1]['cluster'] > 0 else len(stack[-1]['labels']) - 1

        elif key & 0xFF == 83:  # right
            next_cluster = stack[-1]['cluster'] + 1 if stack[-1]['cluster'] < len(stack[-1]['labels']) - 1 else 0

        elif key & 0xFF == 82:  # up
            stack.pop()
            print('stack size:', len(stack), flush=True)
            show_samples('im', stack[-1]['samples'][stack[-1]['clusters'] == stack[-1]['cluster']], s=s)

        elif key & 0xFF == 84:  # down
            print('predicting clusters...', end=' ', flush=True)
            stack.append({
                'samples': stack[-1]['samples'][stack[-1]['clusters'] == stack[-1]['cluster']],
                'clusters': predict(stack[-1]['samples'][stack[-1]['clusters'] == stack[-1]['cluster']]),
                'cluster': 0,
            })
            stack[-1]['labels'] = sorted(set(stack[-1]['clusters']))
            print('done.', flush=True)
            show_samples('im', stack[-1]['samples'][stack[-1]['clusters'] == stack[-1]['cluster']], s=s)

        if next_cluster is not None:
            stack[-1]['cluster'] = next_cluster
            if stack[-1]['samples'][stack[-1]['clusters'] == stack[-1]['cluster']].size > 0:
                show_samples('im', stack[-1]['samples'][stack[-1]['clusters'] == stack[-1]['cluster']], s=s)

    cv.destroyAllWindows()


if __name__ == "__main__":
    main()
