"""Define a graph of nodes-links-patches.

Examples
--------

>>> from landlab.graph import Graph

>>> node_x, node_y = [0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]
>>> graph = Graph((node_y, node_x))
>>> graph.x_of_node
array([ 0.,  1.,  2.,  0.,  1.,  2.,  0.,  1.,  2.])
>>> graph.y_of_node
array([ 0.,  0.,  0.,  1.,  1.,  1.,  2.,  2.,  2.])

>>> links = ((0, 1), (1, 2),
...          (0, 3), (1, 4), (2, 5),
...          (3, 4), (4, 5),
...          (3, 6), (4, 7), (5, 8),
...          (6, 7), (7, 8))
>>> graph = Graph((node_y, node_x), links=links)
>>> graph.nodes_at_link # doctest: +NORMALIZE_WHITESPACE
array([[0, 1], [1, 2],
       [0, 3], [1, 4], [2, 5],
       [3, 4], [4, 5],
       [3, 6], [4, 7], [5, 8],
       [6, 7], [7, 8]])
>>> graph.node_at_link_head
array([1, 2, 3, 4, 5, 4, 5, 6, 7, 8, 7, 8])
>>> graph.node_at_link_tail
array([0, 1, 0, 1, 2, 3, 4, 3, 4, 5, 6, 7])

>>> graph.links_at_node # doctest: +NORMALIZE_WHITESPACE
array([[ 0,  2, -1, -1], [ 1,  3,  0, -1], [ 4,  1, -1, -1],
       [ 5,  7,  2, -1], [ 6,  8,  5,  3], [ 9,  6,  4, -1],
       [10,  7, -1, -1], [11, 10,  8, -1], [11,  9, -1, -1]])

>>> graph.link_dirs_at_node # doctest: +NORMALIZE_WHITESPACE
array([[-1, -1,  0,  0], [-1, -1,  1,  0], [-1,  1,  0,  0],
       [-1, -1,  1,  0], [-1, -1,  1,  1], [-1,  1,  1,  0],
       [-1,  1,  0,  0], [-1,  1,  1,  0], [ 1,  1,  0,  0]])

>>> patches = ((5, 3, 0, 2), (6, 4, 1, 3), (10, 8, 5, 7), (11, 9, 6, 8))
>>> graph = Graph((node_y, node_x), links=links, patches=patches)
>>> graph.links_at_patch
array([[ 3,  5,  2,  0],
       [ 4,  6,  3,  1],
       [ 8, 10,  7,  5],
       [ 9, 11,  8,  6]])
>>> graph.nodes_at_patch
array([[4, 3, 0, 1],
       [5, 4, 1, 2],
       [7, 6, 3, 4],
       [8, 7, 4, 5]])
"""
from six.moves import range

import numpy as np

from ..core.utils import as_id_array, argsort_points_by_x_then_y
from ..utils.jaggedarray import flatten_jagged_array
from ..utils.decorators import store_result_in_grid, read_only_array
from .sort import sort_graph, reindex_by_xy, reorder_links_at_patch
from .object.at_node import get_links_at_node
from .object.at_patch import get_nodes_at_patch
from .quantity.of_link import (get_angle_of_link, get_length_of_link,
                               get_midpoint_of_link)
from .quantity.of_patch import get_centroid_of_patch, get_area_of_patch

from .sort.sort import reverse_one_to_many, reorient_link_dirs

def _parse_sorting_opt(sorting):
    SORTING_OPTS = ('xy', 'ccw', 'ne')

    as_dict = None

    if isinstance(sorting, bool):
        as_dict = dict([(opt, format) for opt in SORTING_OPTS])
    elif isinstance(sorting, dict):
        as_dict = dict(sorting.items())
        for opt in SORTING_OPTS:
            sorting.setdefault(opt, True)

    return as_dict

def find_perimeter_nodes(graph):
    """Find nodes on the perimeter of a graph.

    Uses a convex hull to locate the perimeter nodes of a graph.

    Parameters
    ----------
    graph : graph_like
        A Graph of nodes (just requires *xy_of_node*).

    Returns
    -------
    ndarray of int
        Identifiers of the perimeter nodes.
    """
    from scipy.spatial import ConvexHull

    hull = ConvexHull(graph.xy_of_node, qhull_options='Qt')
    return as_id_array(hull.vertices)

class Graph(object):

    """Define the connectivity of a graph of nodes, links, and patches."""

    def __init__(self, nodes, links=None, patches=None, sorting=True):
        """Define a graph of connected nodes.

        Parameters
        ----------
        nodes : tuple of array_like
            Coordinates of nodes as (*y*, *x*).
        links : array_like of tuple
            Tail node and head node for each link in the graph.
        patches : array_like of tuple
            Links that define each patch.
        sort : bool, optional
            Sort elements by their *x* and then *y* coordinate and use
            counter-clockwise element ordering when ordering one set
            of elements around another.
        """
        sorting = _parse_sorting_opt(sorting)
        if sorting is None:
            raise ValueError('bad argument for sorting keyword')
        if len(nodes[0]) != len(nodes[1]):
            raise ValueError('length mismatch in node coordinates')

        self._sorting = sorting

        nodes = [np.asarray(coord, dtype=float) for coord in nodes]

        if patches is not None:
            if len(patches) > 0:
                patches = flatten_jagged_array(patches, dtype=int)
            else:
                patches = None

        self._xy_of_node = np.stack((nodes[1], nodes[0])).T.copy()
        # self._y_of_node, self._x_of_node = nodes[0], nodes[1]
        self._nodes = np.arange(len(nodes[0]), dtype=int)

        self._create_nodes_at_link(links)
        self._create_links_at_patch(patches)

        not sorting['ne'] or reorient_link_dirs(self)
        not sorting['xy'] or reindex_by_xy(self)
        not sorting['ccw'] or reorder_links_at_patch(self)

        self._origin = (0., 0.)

    def _create_nodes_at_link(self, links):
        """Set up node-link data structures."""
        if links is not None:
            self._nodes_at_link = np.asarray(links, dtype=np.int)
            return self._nodes_at_link

    def _create_links_at_patch(self, patches):
        """Set up patch data structures."""
        from .matrix.at_patch import links_at_patch

        if patches is not None:
            self._links_at_patch = links_at_patch(
                patches, nodes_at_link=self.nodes_at_link)
            return self._links_at_patch

    @property
    def ndim(self):
        return 2

    @property
    def xy_of_node(self):
        """Get x and y-coordinates of node.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1]
        >>> graph = Graph((node_y, node_x))
        >>> graph.xy_of_node[:, 0]
        array([ 0.,  1.,  2.,  0.,  1.,  2.])
        >>> graph.xy_of_node[:, 1]
        array([ 0.,  0.,  0.,  1.,  1.,  1.])
        """
        return self._xy_of_node

    @property
    def x_of_node(self):
        """Get x-coordinate of node.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1]
        >>> graph = Graph((node_y, node_x))
        >>> graph.x_of_node
        array([ 0.,  1.,  2.,  0.,  1.,  2.])
        """
        return self._xy_of_node[:, 0]
        # return self._x_of_node

    @property
    def y_of_node(self):
        """Get y-coordinate of node.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1]
        >>> graph = Graph((node_y, node_x))
        >>> graph.y_of_node
        array([ 0.,  0.,  0.,  1.,  1.,  1.])
        """
        return self._xy_of_node[:, 1]
        # return self._y_of_node

    @property
    def nodes(self):
        """Get identifier for each node.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1]
        >>> graph = Graph((node_y, node_x))
        >>> graph.nodes
        array([0, 1, 2, 3, 4, 5])
        """
        return self._nodes

    @property
    @store_result_in_grid()
    def perimeter_nodes(self):
        return find_perimeter_nodes(self)

    @property
    def number_of_nodes(self):
        """Get total number of nodes.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1]
        >>> graph = Graph((node_y, node_x))
        >>> graph.number_of_nodes
        6
        """
        return self._nodes.size

    @property
    def nodes_at_link(self):
        """Get nodes at either end of links.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> graph = Graph((node_y, node_x), links=links)
        >>> graph.nodes_at_link # doctest: +NORMALIZE_WHITESPACE
        array([[0, 1], [1, 2],
               [0, 3], [1, 4], [2, 5],
               [3, 4], [4, 5],
               [3, 6], [4, 7], [5, 8],
               [6, 7], [7, 8]])
        """
        return self._nodes_at_link

    @property
    def node_at_link_tail(self):
        """Get nodes at link tail.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> graph = Graph((node_y, node_x), links=links)
        >>> graph.node_at_link_tail
        array([0, 1, 0, 1, 2, 3, 4, 3, 4, 5, 6, 7])
        """
        return self._nodes_at_link[:, 0]

    @property
    def node_at_link_head(self):
        """Get nodes at link head.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> graph = Graph((node_y, node_x), links=links)
        >>> graph.node_at_link_head
        array([1, 2, 3, 4, 5, 4, 5, 6, 7, 8, 7, 8])
        """
        return self._nodes_at_link[:, 1]

    @property
    def number_of_links(self):
        """Get nodes at link head.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> graph = Graph((node_y, node_x), links=links)
        >>> graph.number_of_links
        12
        """
        try:
            return len(self._nodes_at_link)
        except AttributeError:
            return 0

    @property
    def links_at_patch(self):
        """Get the links that define a patch.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> patches = ((0, 3, 5, 2), (1, 4, 6, 3))
        >>> graph = Graph((node_y, node_x), links=links, patches=patches)
        >>> graph.links_at_patch
        array([[3, 5, 2, 0],
               [4, 6, 3, 1]])
        """
        return self._links_at_patch

    @property
    def nodes_at_patch(self):
        """Get the nodes that define a patch.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = ([0, 1, 2, 0, 1, 2, 0, 1, 2],
        ...                   [0, 0, 0, 1, 1, 1, 2, 2, 2])
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> patches = ((0, 3, 5, 2), (1, 4, 6, 3))
        >>> graph = Graph((node_y, node_x), links=links, patches=patches)
        >>> graph.nodes_at_patch
        array([[4, 3, 0, 1],
               [5, 4, 1, 2]])
        """
        try:
            return self._nodes_at_patch
        except AttributeError:
            self._nodes_at_patch = get_nodes_at_patch(self)
            return self._nodes_at_patch

    @property
    def patches_at_node(self):
        """Get the patches that touch each node.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = ([0, 1, 2, 0, 1, 2],
        ...                   [0, 0, 0, 1, 1, 1])
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5))
        >>> patches = ((0, 3, 5, 2), (1, 4, 6, 3))
        >>> graph = Graph((node_y, node_x), links=links, patches=patches)
        >>> graph.patches_at_node # doctest: +NORMALIZE_WHITESPACE
        array([[ 0, -1], [ 0,  1], [ 1, -1],
               [ 0, -1], [ 0,  1], [ 1, -1]])
        """
        try:
            return self._patches_at_node
        except AttributeError:
            self._patches_at_node = reverse_one_to_many(self.nodes_at_patch)
            return self._patches_at_node

    @property
    @store_result_in_grid()
    @read_only_array
    def patches_at_link(self):
        """Get the patches on either side of each link.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = ([0, 1, 2, 0, 1, 2],
        ...                   [0, 0, 0, 1, 1, 1])
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5))
        >>> patches = ((0, 3, 5, 2), (1, 4, 6, 3))
        >>> graph = Graph((node_y, node_x), links=links, patches=patches)
        >>> graph.patches_at_link # doctest: +NORMALIZE_WHITESPACE
        array([[ 0, -1], [ 1, -1],
               [ 0, -1], [ 0,  1], [ 1, -1],
               [ 0, -1], [ 1, -1]])
        """
        return reverse_one_to_many(self._links_at_patch)

    @property
    def number_of_patches(self):
        """Get the number of patches.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> patches = ((0, 3, 5, 2), (1, 4, 6, 3))
        >>> graph = Graph((node_y, node_x), links=links, patches=patches)
        >>> graph.number_of_patches
        2
        """
        try:
            return len(self._links_at_patch)
        except AttributeError:
            return 0

    @property
    def links_at_node(self):
        """Get links touching a node.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x = [0, 1, 2, 0, 1, 2, 0, 1, 2]
        >>> node_y = [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> graph = Graph((node_y, node_x), links=links)
        >>> graph.links_at_node # doctest: +NORMALIZE_WHITESPACE
        array([[ 0,  2, -1, -1], [ 1,  3,  0, -1], [ 4,  1, -1, -1],
               [ 5,  7,  2, -1], [ 6,  8,  5,  3], [ 9,  6,  4, -1],
               [10,  7, -1, -1], [11, 10,  8, -1], [11,  9, -1, -1]])
        """
        try:
            return self._links_at_node
        except AttributeError:
            (self._links_at_node,
             self._link_dirs_at_node) = self._create_links_and_dirs_at_node()
            return self._links_at_node

    def _create_links_and_dirs_at_node(self):
        return get_links_at_node(self, sort=self._sorting['ccw'])

    @property
    def link_dirs_at_node(self):
        """Get directions of links touching a node.

        Examples
        --------
        >>> from landlab.graph import Graph
        >>> node_x, node_y = [0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2]
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5),
        ...          (3, 6), (4, 7), (5, 8),
        ...          (6, 7), (7, 8))
        >>> graph = Graph((node_y, node_x), links=links)
        >>> graph.link_dirs_at_node # doctest: +NORMALIZE_WHITESPACE
        array([[-1, -1,  0,  0], [-1, -1,  1,  0], [-1,  1,  0,  0],
               [-1, -1,  1,  0], [-1, -1,  1,  1], [-1,  1,  1,  0],
               [-1,  1,  0,  0], [-1,  1,  1,  0], [ 1,  1,  0,  0]])
        """
        try:
            return self._link_dirs_at_node
        except AttributeError:
            (self._links_at_node,
             self._link_dirs_at_node) = self._create_links_and_dirs_at_node()
            return self._link_dirs_at_node

    @property
    @store_result_in_grid()
    @read_only_array
    def angle_of_link(self):
        """Get the angle of each link.

        Examples
        --------
        >>> import numpy as np
        >>> from landlab.graph import Graph

        >>> node_x, node_y = ([0, 1, 2, 0, 1, 2],
        ...                   [0, 0, 0, 1, 1, 1])
        >>> links = ((0, 1), (1, 2),
        ...          (0, 3), (1, 4), (2, 5),
        ...          (3, 4), (4, 5))
        >>> graph = Graph((node_y, node_x), links=links)
        >>> graph.angle_of_link * 180. / np.pi
        array([  0.,   0.,  90.,  90.,  90.,   0.,   0.])
        """
        return get_angle_of_link(self)

    @property
    @store_result_in_grid()
    @read_only_array
    def length_of_link(self):
        """Get the length of links.

        Examples
        --------
        >>> import numpy as np
        >>> from landlab.graph import UniformRectilinearGraph

        >>> graph = UniformRectilinearGraph((2, 3), spacing=(1, 2))
        >>> graph.length_of_link
        array([ 2.,  2.,  1.,  1.,  1.,  2.,  2.])
        """
        return get_length_of_link(self)

    @property
    @store_result_in_grid()
    @read_only_array
    def midpoint_of_link(self):
        """Get the middle of links.

        Examples
        --------
        >>> import numpy as np
        >>> from landlab.graph import UniformRectilinearGraph

        >>> graph = UniformRectilinearGraph((2, 3), spacing=(1, 2))
        >>> graph.midpoint_of_link # doctest: +NORMALIZE_WHITESPACE
        array([[ 1. ,  0. ], [ 3. ,  0. ],
               [ 0. ,  0.5], [ 2. ,  0.5], [ 4. ,  0.5],
               [ 1. ,  1. ], [ 3. ,  1. ]])
        """
        return get_midpoint_of_link(self)

    @property
    @store_result_in_grid()
    @read_only_array
    def xy_of_link(self):
        return get_midpoint_of_link(self)

    @property
    @store_result_in_grid()
    @read_only_array
    def xy_of_patch(self):
        return get_centroid_of_patch(self)

    @property
    @store_result_in_grid()
    @read_only_array
    def area_of_patch(self):
        return get_area_of_patch(self)


#!/usr/bin/env python

from pathlib import Path
from sys import path as pythonpath
try:
    from commandlink import command_link as discard; discard
except:
    def _fix_pythonpath():
        pythonpath[:0] = str(Path(__file__).resolve().parent),
    _fix_pythonpath()
from commandlink import command_link
from commandsearchpath import which
from errors import WaitError

from argparse import ArgumentParser
from errno import ENOENT
from itertools import chain, count
from os import chdir
from subprocess import Popen
from sys import argv
from tarfile import TarFile

def main(args):
    args = _ArgumentParser().parse_args(args)
    this = command_link(__file__)
    chdir(str(this.parent))
    vm, misc = _provider_specific(this)
    vm = _just_one(vm)
    tar_cmd = which(r'tar')
    if args.force_python_tar or tar_cmd is None:
        tar = _tar_py
    else:
        tar = _tar_cmd(tar_cmd)
    vmdks = tuple(Path().files(r'*.vmdk'))
    box = Path(args.box)
    _vmdk_must_not_be_older(vmdks, box)
    tar(box, chain((r'metadata.json', str(vm),), misc, (str(y) for y in vmdks)))

class _ArgumentParser(ArgumentParser):
    def __init__(self):
        super(_ArgumentParser, self).__init__()
        self.add_argument(r'-b', r'--box', default=r'/var/large/tmp.box')
        self.add_argument(r'-p', r'--force-python-tar', action=r'store_true')

def _provider_specific(this):
    # no better way than `if else` seems to be
    this = str(this)
    if r'/vmware/' in this:
        rv, provider = (r'*.vmx', ()), r'vmware_desktop'
    elif r'/virtualbox/' in this:
        rv, provider = (r'box.ovf', (r'Vagrantfile',)), r'vmware'
    else:
        raise ProviderNotAvailableFor(this)
    with Path(r'metadata.json').open(r'wb') as ostream:
        ostream.write(_metadata_json.format(provider))
    return rv

_metadata_json = r'''
{{
  "provider": "{}"
}}
'''[1:]

def _just_one(pattern):
    i = 0
    for i, name in zip(count(1), Path().files(pattern)):
        pass
    if i != 1:
        raise TooFewOrTooMany(i, pattern)
    return name

def _vmdk_must_not_be_older(vmdks, box):
    if not box.exists():
        return
    time = box.stat().st_mtime
    for vmdk in vmdks:
        if vmdk.stat().st_mtime < time:
            raise IsOlder(vmdk)

def _tar_cmd(tar_cmd):
    def rv(box, names):
        str_new_box = str(box) + r'.new'
        new_box = Path(str_new_box)
        p = None
        try:
            p = Popen(chain((str(tar_cmd), r'-zcf', str_new_box), names))
        except OSError:
            pass
        finally:
            done = False
            try:
                if p is None:
                    raise
                w = p.wait()
                if w != 0:
                    raise WaitError(w)
                done = True
            finally:
                _finish_tar(done, new_box, box)
    return rv

def _tar_py(box, names):
    str_new_box = str(box) + r'.new'
    new_box = Path(str_new_box)
    done = False
    try:
        with new_box.open(r'wb') as ostream:
            with TarFile.gzopen(None, 'w', ostream) as tar:
                for name in names:
                    tar.add(name)
        done = True
    finally:
        _finish_tar(done, new_box, box)

def _finish_tar(done, new_box, box):
    if done:
        new_box.rename(box)
    else:
        try:
            print('***UNLINK***')
            new_box.unlink()
        except OSError as e:
            if e.args[0] != ENOENT:
                raise

class TooFewOrTooMany(Exception):
    pass

class ProviderNotAvailableFor(Exception):
    pass

class IsOlder(Exception):
    pass

if __name__ == r'__main__':
    main(argv[1:])


""" 项目运行常量参数 """

import os
import yaml

# 当前文件路径
_CURRENT_DIR = os.path.abspath(os.path.dirname(__file__))
# 项目根路径
PROJECT_ROOT = os.path.normpath(os.path.join(_CURRENT_DIR, '..'))
# 配置文件路径
CONFIG_DIR = os.path.join(PROJECT_ROOT, 'config')
# 日志文件路径
LOGGING_DIR = os.path.join(PROJECT_ROOT, 'logs')
# 资源目录
RESOURCE_DIR = os.path.join(PROJECT_ROOT, 'resource')

# 配置文件
CONFIG_PATH = os.path.join(CONFIG_DIR, 'config.yml')
YAML_CONFIG = yaml.load(open(CONFIG_PATH, encoding='utf-8'), Loader=yaml.FullLoader)

# ElasticSearch config
ES_CONFIG = YAML_CONFIG['Elastic']

# 加密值
SHA_SALT = YAML_CONFIG['API']['salt']
# json_web_token
JWT_SECRET = YAML_CONFIG['API']['jwtSecret']
# 每页数据量
PAGE_SIZE = YAML_CONFIG['API']['pageSize']


import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import initializers
# np.seterr(all='raise')
from scipy.linalg import solve_discrete_are
from scipy.linalg import sqrtm
import matplotlib
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import datetime
import os
import ipdb
from scipy.stats import expon
from ipdb import launch_ipdb_on_exception
from tqdm import trange

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.experimental.set_visible_devices([], 'GPU')
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=5120)])
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
  except RuntimeError as e:
    # Visible devices must be set before GPUs have been initialized
    print(e)

# # Set CPU as available physical device
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
tf.config.experimental.set_visible_devices(devices= cpus, device_type='CPU')

def feature_vectors(traj, n_fv):
    # Delx as the fv
    fv = np.zeros((n_fv, traj.shape[-1]))
    # fv[0,:] = ((traj[0,:] - traj[2,:])**2)

    # x and xobj as the fv
    # fv[0,:] = traj[0,:]
    # fv[1,:] = traj[2,:]
    # fv[2,:] = ((traj[0,:] - traj[2,:])**2)

    fv[0,:] = np.sum((traj[:2,:] - traj[4:6,:])**2, axis=0)
    fv[1,:]  = np.sum((traj[4:6,:] - traj[8:10,:])**2, axis=0)
    fv[2,:]  = np.sum((traj[:2,:] - traj[8:10,:])**2, axis=0)

    # fv = np.zeros((n_fv, traj.shape[-1]-1))
    # fv[0,:] = ((traj[0,:-1] - traj[0,1:])**2)
    # fv[0,:] = ((traj[2,:-1] - traj[2,1:])**2)
    return fv

def main():
    n_modes = 3
    n = 4
    m = 2
    n_fv = 3
    n_iter = 100
    lr = 0.0001
    alpha = 0.0

    n_trajs = 280
    T = 200

    # Import data
    data = np.load('../data/train_expert_iLQR_N280_T250_dt0_05_self_xz_nmodes_exp1.npz')
    # train_xinit = data['train_xinit']
    train_expert_trajs = data['train_expert_trajs']

    # massage data to usable format [xt, at, xt+1]
    # train_traj = np.zeros((n_trajs, n+m+n, T))
    # train_traj[:, :n, :] = train_expert_trajs[:, :n, :T]
    # train_traj[:, n:n+m, :] = train_expert_trajs[:, n:n+m, :T]
    # train_traj[:, n+m:, :] = train_expert_trajs[:, :n, 1:T+1]

    train_traj = np.concatenate(train_expert_trajs, axis = 1)

    # subset of data for classification
    # train_traj = train_traj[:, :(T+1)*n_trajs]

    N = train_traj.shape[-1]
    fv = feature_vectors(train_traj, n_fv)

    means = np.random.uniform(-5, 5, (n_fv, n_modes))
    # means = np.array([[0, 0.2]])
    means = np.array([[4, 0.0064, 0.0064],[4, 2, 0.0064], [5, 3, 0.0064]])
    converged = 0
    count = -1
    # for i in range(n_iter):
    while converged==0 and count < n_iter:
        #Classify
        diff = np.zeros((n_modes, N))
        for i in range(n_modes):
            diff[i,:] = np.linalg.norm( (fv - means[:,i].reshape(n_fv,1)), axis = 0)
        mode = np.argmin(diff, axis = 0)
        mode_onehot = np.zeros((n_modes, N))
        mode_onehot[mode, np.arange(N)] = 1
        #Recalculate cluster centers
        means_old = np.copy(means)
        for i in range(n_modes):
            means[:, i] = np.sum(mode_onehot[i,:]*fv, axis=1)/np.sum(mode_onehot[i,:])
        if ((means - means_old)**2).sum() < 1e-4:
            converged = 1
        count = count + 1
    print('count:', count)

    # train_traj = train_traj[:,:-1]
    mode0 = train_traj[:, mode==0]
    mode1 = train_traj[:, mode==1]
    mode2 = train_traj[:, mode==2]
    fig, axes = plt.subplots(2, 2, figsize=(10, 10))
    axes[0, 0].scatter(mode0[0,:],mode0[4,:], cmap= [1,0,0], label = 'mode1')
    axes[0, 0].scatter(mode1[0,:],mode1[4,:], cmap= [0,1,0], label = 'mode2')
    axes[0, 0].scatter(mode2[0,:],mode2[4,:], cmap= [0,0,1], label = 'mode3')
    axes[0, 0].set_xlabel('x gripper')
    axes[0, 0].set_ylabel('x obj1')
    axes[0, 0].legend()

    axes[0, 1].scatter(mode0[0,:],mode0[8,:], cmap= [1,0,0], label = 'mode1')
    axes[0, 1].scatter(mode1[0,:],mode1[8,:], cmap= [0,1,0], label = 'mode2')
    axes[0, 1].scatter(mode2[0,:],mode2[8,:], cmap= [0,0,1], label = 'mode3')
    axes[0, 1].set_xlabel('x gripper')
    axes[0, 1].set_ylabel('x obj2')
    axes[0, 1].legend()

    plt.savefig('kmeans_clustering_rel.png')
    ipdb.set_trace()

    # fig, axes = plt.subplots(2, 2, figsize=(10, 10))
    # axes[0, 0].scatter(fv[0,:],mode)
    # axes[0, 0].set_xlabel('x-xobj')
    # axes[0, 0].set_ylabel('Mode')

    # axes[0, 1].scatter(train_traj[0,:],mode)
    # axes[0, 1].set_xlabel('x')
    # axes[0, 1].set_ylabel('Mode')

    # axes[1, 0].scatter(train_traj[1,:],mode)
    # axes[1, 0].set_xlabel('xdot')
    # axes[1, 0].set_ylabel('Mode')

    # axes[1, 1].scatter(train_traj[2,:],mode)
    # axes[1, 1].set_xlabel('xobj')
    # axes[1, 1].set_ylabel('Mode')
    # plt.savefig('kmeans_clustering_rel2.png')
    print('CLASSIFICATION DONE')
    # plt.show()

    ## Supervised learning using NN classifier
    classifier = Sequential()
    classifier.add(Dense(8, activation='relu',input_shape=(3,), dtype = tf.float64))#, kernel_initializer='glorot_uniform', bias_initializer='zeros'))
    # classifier.add(Dense(8, activation='relu', dtype = tf.float64))#, kernel_initializer='glorot_uniform', bias_initializer='zeros'))
    classifier.add(Dense(n_modes, activation='softmax', dtype = tf.float64))# kernel_initializer='glorot_uniform', bias_initializer='zeros'))
    classifier.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
              metrics=['accuracy'])

    # train_traj_class = np.concatenate((train_traj[:n, :-1], train_traj[n*2:, :-1], train_traj[:n, 1:]), axis=0)
    classifier.fit(fv.T, mode, epochs = 1)
    classifier.save("../data/Classifier_model_nmodes")
    print(classifier.variables)

    modes_predicted = classifier.predict(fv.T)
    modes_predicted_argmax = np.argmax(modes_predicted, axis = 1)
    # fig, axes = plt.subplots(2, 2, figsize=(10, 10))
    # axes[0, 0].scatter(train_traj[0,:],modes_predicted[:,0])
    # axes[0, 0].set_xlabel('x')
    # axes[0, 0].set_ylabel('Mode')

    # axes[0, 1].scatter(train_traj[2,:],modes_predicted[:,0])
    # axes[0, 1].set_xlabel('xobj')
    # axes[0, 1].set_ylabel('Mode')

    # axes[1, 0].scatter(train_traj[1,:],modes_predicted[:,0])
    # axes[1, 0].set_xlabel('xdot')
    # axes[1, 0].set_ylabel('Mode')

    # axes[1, 1].scatter(train_traj[3,:],modes_predicted[:,0])
    # axes[1, 1].set_xlabel('xobjdot')
    # axes[1, 1].set_ylabel('Mode')
    # plt.savefig('NNclassification.png')

    mode0 = train_traj[:, modes_predicted_argmax==0]
    mode1 = train_traj[:, modes_predicted_argmax==1]
    mode2 = train_traj[:, modes_predicted_argmax==2]

    fig, axes = plt.subplots(2, 2, figsize=(10, 10))
    axes[0, 0].scatter(mode0[0,:],mode0[4,:], cmap= [1,0,0])
    axes[0, 0].scatter(mode1[0,:],mode1[4,:], cmap= [0,1,0])
    axes[0, 0].scatter(mode2[0,:],mode2[4,:], cmap= [0,0,1])
    axes[0, 0].set_xlabel('x gripper')
    axes[0, 0].set_ylabel('x obj1')

    axes[0, 1].scatter(mode0[0,:],mode0[8,:], cmap= [1,0,0])
    axes[0, 1].scatter(mode1[0,:],mode1[8,:], cmap= [0,1,0])
    axes[0, 1].scatter(mode2[0,:],mode2[8,:], cmap= [0,0,1])
    axes[0, 1].set_xlabel('x gripper')
    axes[0, 1].set_ylabel('x obj2')

    plt.savefig('NNclassification_rel.png')
    # reconstructed_model = tf.keras.models.load_model("../data/Classifier_model")
    # modes_predicted = reconstructed_model.predict(train_traj[:2*n,:].T)
    # modes_predicted_argmax = np.argmax(modes_predicted, axis = 1)
    # fig, axes = plt.subplots(2, 2, figsize=(10, 10))
    # axes[0, 0].scatter(train_traj[0,:],modes_predicted_argmax)
    # axes[0, 0].set_xlabel('x')
    # axes[0, 0].set_ylabel('Mode')

    # axes[0, 1].scatter(train_traj[1,:],modes_predicted_argmax)
    # axes[0, 1].set_xlabel('x')
    # axes[0, 1].set_ylabel('Mode')

    # axes[1, 0].scatter(train_traj[2,:],modes_predicted_argmax)
    # axes[1, 0].set_xlabel('x')
    # axes[1, 0].set_ylabel('Mode')

    # axes[1, 1].scatter(train_traj[3,:],modes_predicted_argmax)
    # axes[1, 1].set_xlabel('x')
    # axes[1, 1].set_ylabel('Mode')
    # plt.savefig('NN2classification.png')
if __name__ == '__main__':
    # with launch_ipdb_on_exception():
    main()


from itertools import combinations, permutations
import numpy as np
from sklearn import svm, linear_model, cross_validation
import pandas as pd

class RankSVM(svm.LinearSVC):

    def __init__(self, C=1.0, w=None):
        super().__init__(C=C)
        self.w = w

    # X,y为原始缺陷训练集中软件模块的特征和缺陷个数
    # 输出为模块对
    def transform_pairwise(self, X, y):
        X_new = []
        y_new = []
        y = np.asarray(y)
        if y.ndim == 1:
            y = np.c_[y, np.ones(y.shape[0])]
        comb = combinations(range(X.shape[0]), 2)
        for k, (i, j) in enumerate(comb):
            # if y[i, 0] == y[j, 0] or y[i, 1] != y[j, 1]:
            if y[i, 0] == y[j, 0]:
                # skip if same target or different group
                continue
            X_new.append(X[i] - X[j])
            y_new.append(np.sign(y[i, 0] - y[j, 0]))  # -1/1

        X_pairs = np.asarray(X_new, dtype=np.float)
        y_pairs = np.asarray(y_new, dtype=np.float)
        return X_pairs, y_pairs

    # X为测试集中软件模块的特征
    def tst_transform_pairwise(self, X):
        X_new = []
        perm = permutations(range(X.shape[0]), 2)
        # print(X.shape[0])
        for k, (i, j) in enumerate(perm):
            X_new.append(X[i] - X[j])
        return np.asarray(X_new)

    def fit(self, X, y):
        X_pairs, y_pairs = self.transform_pairwise(X, y)
        # print(X_trans,y_trans)
        super(RankSVM, self).fit(X_pairs, y_pairs)
        return self

    def predict2(self, X):
        X_tests = self.tst_transform_pairwise(X)
        y_pred = super(RankSVM, self).predict(X_tests)

        length = X.shape[0]

        count = self.rank_list(y_pred, length)
        pred_bug_rank = self.trans(count=count)

        return pred_bug_rank

    def predict3(self, X):
        """
        由genetic ALgorithm求出来的w进行预测
        """
        X_tests = self.tst_transform_pairwise(X)
        y_pred = self.predict_with_w(X_tests)

        length = X.shape[0]

        count = self.rank_list(y_pred, length)
        pred_bug_rank = self.trans(count=count)

        return pred_bug_rank

    def predict_with_w(self, X):
        """
        Mi > Mk <==> <w, xj-xk> > 0
        return: [1,-1,1,1,1-,1,...,1,-1,1]
        """
        y_pred = []
        for xi in X:
            if np.dot(self.w, xi) > 0:
                y_pred.append(1)
            else:
                y_pred.append(-1)
        return y_pred

    def rank_list(self, y, length):
        '''
        比如测试集中软件模块为x0,x1,x2,x3，y为[1,-1,1,   -1,-1,1,   1,1,1,   -1,-1,-1]
        则表明x0>x1,x0<x2,x0>x3,   x1<x0,x1<x2,x1<x3  x2>x0,x2>x1,x2>x3   x3<x0,x3<x1,x3<x2
        最终得出x3<x1<x0<x2
        length: 测试集中软件模块个数 这里length即为4
        :return: 排序后从小到大的下标，这里为3,1,0,2
        '''
        # 计数数组
        count_list = []
        # 计算xi -1 的个数
        for i in range(length):
            count = 0
            for j in range(length - 1):
                n = i * (length - 1) + j
                if(y[n] == -1):
                    count = count + 1
            count_list.append(count)

        k = 0
        max_list = []
        # 选出最小的，处理相同个数的 -1 ，并且把已处理的位置标记为 -1
        while(k < len(count_list)):
            k = k + 1
            large = max(count_list)
            max_index = [m for m in range(length) if count_list[m] == large]
            if len(max_index) > 1:
                for i in range(len(max_index) - 1):
                    max_i = max_index[i]
                    n = max_i * (length - 1) + max_index[i + 1] - 1
                    if(y[n] == 1):
                        max_i = max_index[i + 1]
                max_list.append(max_i)
                count_list[max_i] = -1
            else:
                max_list.append(max_index[0])
                count_list[max_index[0]] = -1
        # 返回的就是一个从小到大的下标
        return max_list

    # 对得到的从小到大的模块排序，赋予每个模块不同的缺陷数目
    # 比如预测模块的缺陷个数排序为x3<x1<x0<x2，则y0=2，y1=1，y2=3，y3=0，返回值为[2,1,3,0]
    def trans(self, count):
        t = []
        for i, j in enumerate(count):
            t.append((i, j))

        t = sorted(t, key=lambda k: k[1])
        res = []

        for i in t:
            res.append(i[0])

        return res

'''
if __name__ == '__main__':
    X = np.array([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6],[7, 7], [8, 8], [9, 9], [10, 10], [11, 11], [12, 12]])
    y = np.array([5, 4, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0])

    X_trans = RankSVM().rank_list([1,-1,1,-1,-1,1,1,1,1,-1,-1,-1],4)
    x_count=RankSVM().trans(X_trans)
    print(X_trans)
    print(x_count)
'''