#!/usr/bin/env python3

import aphros
try:
    import matplotlib.pyplot as plt
    aphros.plot.ApplyParams(plt)
except Exception:
    pass
import numpy as np
import argparse
import os
import sys

def printerr(m):
    sys.stderr.write('{:}\n'.format(m))
    sys.stderr.flush()

def read_lines_vtk(path, mirror=True):
    '''
    Returns array of lines from `sm_*.vtk` generated with `spacedim=2`
    [line0_x, line0_y, line1_x, line1_y, ...]
    where line0_x and line1_y are lists of coordinates
    '''
    points, poly, fields = aphros.ReadVtkPoly(path)

    # [[[x0a, y0b], [x0a, y0b]], ...]
    lines = points[poly][:, :, :2]
    # [[[x0a, x0b], [y0a, y0b]], ...]
    lines = np.transpose(lines, (0, 2, 1))
    # [[x0a, x0b], [y0a, y0b], ...]
    lines = lines.reshape((-1, 2))
    if mirror:
        lines_mirror = lines.copy()
        lines_mirror[1::2, ] *= -1
        lines = np.vstack((lines, lines_mirror))
    return lines

def read_lines_moonmd(path):
    '''
    Returns array of lines from `c1g3l4s.txt` generated by MooNMD
    [line0_x, line0_y, line1_x, line1_y, ...]
    where line0_x and line1_y are lists of coordinates
    '''
    lines = np.genfromtxt(path).T
    yy, xx = lines
    yy -= 0.5
    lines = (xx, yy)
    return lines

def read_lines(path):
    ext = os.path.splitext(path)[1]
    if ext == ".vtk":
        return read_lines_vtk(path)
    if ext == ".txt":
        return read_lines_moonmd(path)
    raise RuntimeError("Unknown extention: '{}' in path '{}'".format(ext, path))

def lines_to_points(lines):
    '''
    lines: list of lines [line0_x, line0_y, line1_x, line1_y, ...]
    where line0_x and line1_y are lists of coordinates

    Returns:
    list of points [[x0, y0], [x1, y1], ...] from line segments.
    '''
    xx = lines[::2]
    yy = lines[1::2]
    xx = np.hstack(xx)
    yy = np.hstack(yy)
    return np.vstack((xx, yy)).T

def directed_hausdorff(a, b):
    '''
    Returns the directed Hausdorff distance between sets of points.
    a, b: lists of points [[x0, y0], [x1, y1], ...]
    '''
    best_d = 0
    best_ia = None
    best_ib = None
    a = np.array(a)
    b = np.array(b)
    for ia in range(len(a)):
        dd = np.linalg.norm(a[ia] - b, axis=1)
        ib = np.argmin(dd)
        d = dd[ib]
        if d > best_d:
            best_d = d
            best_ia = ia
            best_ib = ib

    return best_d, best_ia, best_ib

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('plots',
                        type=str,
                        nargs='*',
                        help="List of plots as 'PATH;LABEL;COLOR;STYLE'. "
                        "PATH points to sm_*.vtk or ref/*s.txt")
    parser.add_argument('--output',
                        type=str,
                        default="a.pdf",
                        help="Path to output image")
    parser.add_argument('--dist',
                        type=int,
                        nargs=2,
                        help="Print directed Hausdorff distance between"
                        " lines of given indices")
    parser.add_argument('--figsize',
                        type=float,
                        nargs=2,
                        default=[3, 3],
                        help="Figure size")
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    fig = plt.figure(figsize=args.figsize)
    ax = plt.Axes(fig, [0, 0, 0.7, 1])
    fig.add_axes(ax)

    lines_all = []

    for plotarg in args.plots:
        plot = ['', '', '', '']
        for i, e in enumerate(plotarg.split(';')):
            plot[i] = e
        path, label, c, ls = plot
        lines = read_lines(path)
        lines_all.append(lines)
        d = dict()
        if c: d['c'] = c
        if ls: d['ls'] = ls
        l, = ax.plot([], [], label=label, **d)
        d['c'] = l.get_color()
        ax.plot(*lines, **d)

    if args.dist is not None:
        la, lb = args.dist
        pa = lines_to_points(lines_all[la])
        pb = lines_to_points(lines_all[lb])
        d, ia, ib = directed_hausdorff(pa, pb)
        ax.scatter(*pa[ia], c='k', s=0.5, zorder=100)
        ax.scatter(*pb[ib], c='k', s=0.5, zorder=100)
        print(d)


    ax.set_axis_off()
    ax.set_aspect(1)
    ax.legend(loc='upper left',
              fontsize=7,
              bbox_to_anchor=(1, 1),
              frameon=False)

    printerr(args.output)
    fig.savefig(args.output, dpi=300)
