#!/usr/bin/env python
# *-* coding: utf-8 *-*
"""Generate a decouplings.dat for CRTomo

From a given extra_lines.dat file (see triangular grid generation in manual),
generate a decouplings.dat for the decoupling of cell interfaces from the
regularisation.

As only the start and end points are given, the algorithm must find all nodes
generated by GMSH along those lines. This is done by comparing the orthogonal
distance of each node to the line. All nodes within a certain threshold
distance are assumed to lie on the line. The threshold is computed by taking
the minimal node distance of the grid (pdist function) and then multiplying
with 0.75. This probably is a little bit overkill and we could also use some
measure of machine precision.

For convenience this threshold can be overwritten using the --eps option.

Some of the for-loops can probably be optimized or replaced by some numpy
broadcasting rules.

Examples
--------

    # enter the output directory of cr_trig_create.py
    grid_extralines_gen_decouplings.py

    # this creates the file decouplings.dat
    grid_extralines_gen_decouplings.py --plot_node_nrs --plot_elm_nrs

"""
import sys
import os
# from optparse import OptionParser
from argparse import ArgumentParser
import numpy as np
import crtomo.mpl
import crtomo.grid as CRGrid

import matplotlib.pylab as plt
import matplotlib as mpl

import IPython
IPython

crtomo.mpl.setup()


def handle_cmd_options_v2():
    parser = ArgumentParser()
    parser.add_argument(
        '-e',
        "--elem",
        dest="elem_file",
        type=str,
        help="elem.dat file (default: elem.dat)",
        default="elem.dat",
    )
    parser.add_argument(
        "--eps",
        dest="eps",
        type=float,
        help="User override for distance eps",
        default=None,
    )
    parser.add_argument(
        '-w',
        "--eta",
        action="append",
        dest="eta",
        type=float,
        help="User override for coupling coefficient (default: 0)",
        # default=0.0,
    )
    parser.add_argument(
        '-l',
        "--linefile",
        # dest="linefile",
        action='append',
        # type=str,
        # help="Line file (default: extra_lines.dat)",
        # default='extra_lines.dat',
    )
    parser.add_argument(
        "-o",
        "--output",
        dest="output",
        help="Output file (default: decouplings.dat)",
        metavar="FILE", default="decouplings.dat",
    )

    parser.add_argument(
        '-linenr',
        "--ln",
        action="store",
        dest="line_nr",
        type=int,
        help="Process only one line with index N (zero indexed)",
        default=None,
    )

    parser.add_argument(
        '--debug_plot',
        action='store_true',
        dest='debug_plot',
        default=True,
        help='plot a debug plot',
    )

    parser.add_argument(
        '--plot_node_nrs',
        action='store_true',
        dest='plot_node_nrs',
        default=False,
        help='plot node numbers in the debug plot. default: False',
    )

    parser.add_argument(
        '--plot_elm_nrs',
        action='store_true',
        dest='plot_elm_nrs',
        default=False,
        help='plot elements numbers in the debug plot. default: False',
    )

    args = parser.parse_args()
    return args


# def handle_cmd_options():
#     parser = OptionParser()
#     parser.add_option(
#         '-e',
#         "--elem",
#         dest="elem_file",
#         type="string",
#         help="elem.dat file (default: elem.dat)",
#         default="elem.dat",
#     )
#     parser.add_option(
#         "--eps",
#         dest="eps",
#         type="float",
#         help="User override for distance eps",
#         default=None,
#     )
#     parser.add_option(
#         "--eta",
#         action="store",
#         dest="eta",
#         type="float",
#         help="User override for coupling coefficient (default: 0)",
#         default=0,
#     )
#     parser.add_option(
#         '-l',
#         "--linefile",
#         dest="linefile",
#         help="Line file (default: extra_lines.dat)",
#         default='extra_lines.dat',
#     )
#     parser.add_option(
#         "-o",
#         "--output",
#         dest="output",
#         help="Output file (default: decouplings.dat)",
#         metavar="FILE", default="decouplings.dat",
#     )

#     parser.add_option(
#         "--ln",
#         action="store",
#         dest="line_nr",
#         type="int",
#         help="Process only one line with index N (zero indexed)",
#         default=None,
#     )

#     parser.add_option(
#         '--debug_plot',
#         action='store_true',
#         dest='debug_plot',
#         default=True,
#         help='plot a debug plot',
#     )

#     parser.add_option(
#         '--plot_node_nrs',
#         action='store_true',
#         dest='plot_node_nrs',
#         default=False,
#         help='plot node numbers in the debug plot. default: False',
#     )

#     parser.add_option(
#         '--plot_elm_nrs',
#         action='store_true',
#         dest='plot_elm_nrs',
#         default=False,
#         help='plot elements numbers in the debug plot. default: False',
#     )

#     (options, args) = parser.parse_args()
#     return options


def line_line_intersect(x, y):
    """Compute the intersection point of two lines

    Parameters
    ----------
    x = x4 array: x1, x2, x3, x4
    y = x4 array: y1, y2, y3, y4
    line 1 is defined by p1,p2
    line 2 is defined by p3,p4

    Returns
    -------
    Ix: x-coordinate of intersection
    Iy: y-coordinate of intersection
    """
    A = x[0] * y[1] - y[0] * x[1]
    B = x[2] * y[3] - y[2] * x[4]
    C = (x[0] - x[1]) * (y[2] - y[3]) - (y[0] - y[1]) * (x[2] - x[3])

    Ix = (A * (x[2] - x[3]) - (x[0] - x[1]) * B) / C
    Iy = (A * (y[2] - y[3]) - (y[0] - y[1]) * B) / C
    return Ix, Iy


def distances(x, y, px, py, plot=False):
    """
    Compute shortest distance of points (px, py) to the line described by the
    end points in (x, y).

    Parameters
    ----------
    x,y: x and y coordinates of line
    px: x coordinates of data points
    py: y coordiante of data points

    Returns
    -------
    distances: list of distances
    """

    length_line = np.sqrt((x[1] - x[0]) ** 2 + (y[1] - y[0]) ** 2)

    horizontal = (np.diff(y) == 0)
    vertical = (np.diff(x) == 0)

    dist = []
    dpl = []
    counter = 0
    for xp, yp in zip(px, py):
        # compute distances to the end points
        dp = np.sqrt((xp - x) ** 2 + (yp - y) ** 2)

        dpl.append(dp)
        if np.any(dp > length_line):
            point_dist = np.min(dp)
        elif horizontal:
            point_dist = np.abs(yp - y[0])
        elif vertical:
            point_dist = np.abs(xp - x[0])
        else:
            nominator = np.abs(
                (y[1] - y[0]) * xp -
                (x[1] - x[0]) * yp +
                x[1] * y[0] -
                y[1] * x[0])
            denominator = np.sqrt((y[1] - y[0]) ** 2 + (x[1] - x[0]) ** 2)
            point_dist = nominator / denominator

        dist.append(point_dist)
        counter += 1

    # debug plot
    if plot:
        fig, ax = plt.subplots(1, 1, figsize=(10, 8))
        ax.scatter(x, y)
        ax.scatter(px, py, color='r')

        index = 0
        for xp, yp in zip(px, py):
            ax.annotate('{0:.5}, {1:.5}, {2:.5}'.format(
                dist[index], dpl[index][0], dpl[index][1]),
                xy=(xp, yp))
            index += 1
        ax.plot(x, y)
        ax.set_aspect('equal')
    return dist


def get_decouplings_for_line(grid, line, settings, fids=None):
    """Compute the element pairs for regularisation decoupling, given a grid
    object and a line, denoted by start and end coordinates, in the grid.

    Parameters
    ----------
    grid: a crtomo.grid.crt_grid object
    line: np.array/list with 4 entries: x1, y1, x2, x2, denoting start and end
          point
    fids: None, or: list with file ids for
        debug_nodes



    Returns
    -------
    neighbors: Mx2 array with the numbers of adjoining elements along the given
               line
    """
    key = 'presort'
    # x/z coordinates of 'macro'-line
    x = [line[0], line[2]]
    y = [line[1], line[3]]

    # x/z coordinates of all nodes in the grid
    nx = grid.nodes[key][:, 1]
    ny = grid.nodes[key][:, 2]
    nxy = np.vstack((nx, ny)).T

    # shortest distance of all nodes to this line
    dist = np.array(distances(x, y, nx, ny))

    # set the epsilon environment
    if settings.get('eps', None) is not None:
        eps = settings['eps']
        print('Taking user supplied eps value: {0}'.format(eps))
    else:
        import scipy.spatial.distance as spdist
        pair_distances = spdist.pdist(nxy)
        print('automatic eps determination')
        print('minimal distance: {0}'.format(pair_distances.min()))
        eps = np.sqrt(
            pair_distances.min() ** 2 - (pair_distances.min() ** 2) / 4
        )
        print('estimated smallest distance to next node: {0}'.format(eps))
        eps *= 0.5
        print('final eps (0.5 * estimate): {0}'.format(eps))
        if eps == 0:
            raise Exception('eps == 0')

    # only consider nodes that lie in the eps environment
    indices = np.where(dist < eps)[0]

    # extract these nodes
    nodes = grid.nodes[key][indices, 0]
    nodes = indices

    nodes_full = np.hstack((
        grid.nodes[key][indices, :].squeeze(),
        np.atleast_2d(dist[indices]).T,
    ))
    if fids is not None:
        np.savetxt(
            fids[0],
            nodes_full.squeeze()
        )
        np.savetxt(
            fids[1],
            np.atleast_2d(dist).T,
        )

    elm_indices = []
    elm_nodes = []
    # find elements with two nodes in it
    for elmnr, element in enumerate(grid.elements):
        if len(np.intersect1d(element, nodes)) == 2:
            elm_indices.append(elmnr)
            elm_nodes.append(element)

    elms = np.array(elm_nodes)

    # find neighboring elements
    neighbors = []
    eta = settings['eta']
    # for each element on the line
    for index, elm in enumerate(elms):
        found_it = False
        for index1, elm1 in enumerate(elms):
            ints = np.intersect1d(elm, elm1)
            # only two nodes can be on a line
            if len(ints) == 2:
                # this check ensures that we do not identify the
                # boundary between two adjacent cells on the line
                # as a decoupling line.
                if np.all(dist[ints] < eps):
                    found_it = True
                    break
                else:
                    pass
                    # print(
                    #     'distances of neighbor not correct:', dist[ints], eps
                    # )
            # probably the element itself
            elif len(ints) > 2:
                # print('found more than two common nodes!')
                # print('element nodes:', elm)
                # print('element1 nodes:', elm1)
                pass
        if found_it:
            nb = (elm_indices[index] + 1, elm_indices[index1] + 1, eta)
            # reversed neighbor
            # nb_rev = (nb[1], nb[0], nb[2])
            # if nb_rev not in neighbors:
            if True:
                neighbors.append(
                    nb
                )
        else:
            pass
            print('No neighbors found. Strange...')

    return np.array(neighbors)


def check_options(options):
    if options.linefile is None:
        options.linefile = ['extra_lines.dat', ]
    if options.eta is None:
        options.eta = [0.0, ]

    assert len(options.linefile) == len(options.eta), \
        "require same number of --linefile and --eta parameters"

    for filename in [options.elem_file, ] + options.linefile:
        if not os.path.isfile(filename):
            raise IOError('File not found: {0}'.format(filename))


def debug_plot(grid, extra_lines, decoupling_elements, options,):
    print('Creating debug plot...')
    fig, ax = plt.subplots(1, 1, figsize=(10, 10))

    all_xz = []
    for x, z in zip(grid.grid['x'], grid.grid['z']):
        tmp = np.vstack((x, z)).T
        all_xz.append(tmp)

    collection = mpl.collections.PolyCollection(
        all_xz,
        edgecolor='k',
        facecolor='none',
        linewidth=0.4,
    )
    ax.add_collection(collection)
    if grid.electrodes is not None:
        ax.scatter(
            grid.electrodes[:, 1],
            grid.electrodes[:, 2],
            color='k',
            clip_on=False,
            label='electrodes',
        )

    # plot extra lines
    for line in np.atleast_2d(extra_lines):
        ax.plot(
            [line[0], line[2]],
            [line[1], line[3]],
            '.-',
            color='c',
            label='extra line',
        )
        # only line 0
        break

    # plot decouplings
    label = 'decoupling line'
    for (el1, el2, coef) in decoupling_elements:
        n1 = grid.elements[int(el1) - 1]
        n2 = grid.elements[int(el2) - 1]

        ints = np.intersect1d(n1, n2)

        x = grid.nodes['presort'][ints, 1]
        z = grid.nodes['presort'][ints, 2]

        ax.plot(
            x,
            z,
            '.-',
            color='r',
            # linestyle='dashed',
            label=label,
        )
        label = ''

    # plot search nodes
    nodes = np.loadtxt('debug_search_nodes.dat')
    ax.scatter(
        nodes[:, 1],
        nodes[:, 2],
        color='g',
        s=40,
        clip_on=False,
        label='search nodes for decoupling',
    )
    for node in nodes:
        circle = plt.Circle(
            (node[1], node[2]),
            radius=node[3],
            fc='none',
        )
        ax.add_patch(circle)

    # plot grid numbers
    if options.plot_elm_nrs:
        for nr, quads in enumerate(grid.elements):
            # nr = grid.nodes['presort'][quads, 0],
            x = grid.nodes['presort'][quads, 1],
            z = grid.nodes['presort'][quads, 2],

            # print('NXZ', nr, x, z)
            ax.annotate(
                '{0}'.format(nr + 1),
                xy=(
                    np.mean(x),
                    np.mean(z),
                ),
                xycoords='data',
            )

    # # plot node numbers
    if options.plot_node_nrs:
        for nr, node in enumerate(grid.nodes['sorted']):
            ax.annotate(
                '{0}({1})'.format(nr, node[0]),
                xy=(
                    node[1],
                    node[2],
                ),
                xycoords='data',
                color='y',
            )

    ax.set_xlim(grid.grid['x'].min(), grid.grid['x'].max())
    ax.set_ylim(grid.grid['z'].min(), grid.grid['z'].max())
    # ax.autoscale_view()
    ax.set_aspect('equal')
    ax.set_xlabel('x')
    ax.set_ylabel('z')
    ax.set_title('Debug plot decoupling lines')
    fig.tight_layout()
    ax.legend(
        loc="lower center",
        ncol=4,
        bbox_to_anchor=(0, 0, 1, 1),
        bbox_transform=fig.transFigure,
    )
    fig.subplots_adjust(bottom=0.1)

    fig.savefig('debug_plot_decoupling.png', dpi=300, bbox_inches='tight')


def main():
    # options = handle_cmd_options()
    options = handle_cmd_options_v2()

    check_options(options)
    # import IPython
    # IPython.embed()
    # exit()
    grid = CRGrid.crt_grid()
    grid.load_elem_file(options.elem_file)

    fids = []
    fids.append(
        open('debug_search_nodes.dat', 'wb')
    )
    fids.append(
        open('debug_dist.dat', 'wb')
    )
    neighbors = None

    for linefile, eta in zip(options.linefile, options.eta):
        extra_lines = np.atleast_2d(np.loadtxt(linefile))

        if options.line_nr is None:
            lines = extra_lines
        else:
            print('processing only line: {0}'.format(options.line_nr))
            lines = (extra_lines[options.line_nr], )
        for line in lines:
            data = get_decouplings_for_line(
                grid,
                line,
                {
                    'eps': options.eps,
                    'eta': eta,
                },
                fids
            )
            if neighbors is None:
                neighbors = data
            else:
                #     neighbors = np.vstack(
                #         (neighbors, data)
                #     )

                try:
                    if data.size > 0:
                        neighbors = np.vstack(
                            (neighbors, data)
                        )
                except Exception as e:
                    print('Exception')
                    print(e)
                    print(data)
                    import IPython
                    IPython.embed()

    for fid in fids:
        fid.close()

    with open(options.output, 'w') as fid:
        fid.write('{0}\n'.format(neighbors.shape[0]))
    with open(options.output, 'ab') as fid:
        np.savetxt(fid, np.array(neighbors), fmt='%i %i %.3f')

    if options.debug_plot:
        debug_plot(
            grid=grid,
            extra_lines=extra_lines,
            decoupling_elements=np.array(neighbors),
            options=options,
        )


if __name__ == '__main__':
    main()
