import itertools

import numpy as np

from scipy.spatial import Voronoi
from shapely import Polygon

from reyna.polymesher.two_dimensional._auxilliaries.abstraction import PolyMesh, Domain


def poly_mesher(domain: Domain, max_iterations: int = 100, **kwargs) -> PolyMesh:
    """
    The PolyMesher function. This function generates a bounded Voronoi tesselation of a given domain. The user must
    initialise the function with either the number of elements they require by submitting the 'n_points' keyword
    argument or initialise with a grid of points by inputting both the 'n_x' and 'n_y' keyword arguments. There is a
    further verbose argument also.

    Args:
        domain (Domain): This is a 'Domain' object, be that custom or a pre-defined one.
        max_iterations (int): This is the number of iterations required by the user. Generally, the more iterations, the
        more uniform the elements are in shape and number of edges etc.

    Returns:
        (PolyMesh): A PolyMesh object containing all the relevant information to be used in a geometry function.

    Raises:
        AttributeError: If the keyword 'n_points' or keywords 'n_x' and 'n_y' are missed.

    """
    if "n_points" in kwargs:
        n_points = kwargs.pop("n_points")
        points = _poly_mesher_init_point_set(domain, n_points=n_points)

    elif "n_x" in kwargs and "n_y" in kwargs:
        n_x = kwargs.pop("n_x")
        n_y = kwargs.pop("n_y")
        points = _poly_mesher_init_point_set(domain, n_x=n_x, n_y=n_y)

    else:
        raise AttributeError("key word error: must be just `n_points` or both `n_x` and `n_y`.")

    verbose: bool = False
    if 'verbose' in kwargs:
        verbose = kwargs.pop('verbose')

    fixed_points = domain.pFix()  # from here can call domain.fixed_points -- this initialises the property simult
    if fixed_points is not None:
        points = np.concatenate((fixed_points, points), axis=0)
        n_fixed = fixed_points.shape[0]
    else:
        n_fixed = 0

    iteration, error, tolerance = 0, 1.0, 1e-4
    bounding_box = domain.bounding_box
    area = (bounding_box[0, 1] - bounding_box[0, 0]) * (bounding_box[1, 1] - bounding_box[1, 0])
    n_points = points.shape[0]

    voronoi, nodes, elements = None, None, None

    while iteration <= max_iterations and error > tolerance:
        reflected_points = _poly_mesher_reflect(points, domain, area)

        voronoi = Voronoi(np.concatenate((points, reflected_points), axis=0), qhull_options='Qbb Qz')
        cond_len = len(voronoi.regions[0]) == 0

        nodes = voronoi.vertices
        if cond_len:
            sorting = [np.where(voronoi.point_region == x)[0][0] for x in range(1, len(voronoi.regions))]
            elements = [x for _, x in sorted(zip(sorting, voronoi.regions[1:]))]
        else:
            empty_ind = voronoi.regions.index([])
            sorting_1 = [np.where(voronoi.point_region == x)[0][0] for x in range(0, empty_ind)]
            sorting_2 = [np.where(voronoi.point_region == x)[0][0] for x in range(empty_ind+1, len(voronoi.regions))]
            sorting = sorting_1 + sorting_2
            regions = voronoi.regions
            del regions[empty_ind]
            elements = [x for _, x in sorted(zip(sorting, regions))]

        points, area, error = _poly_mesher_vorocentroid(points, nodes, elements)

        if fixed_points is not None:
            points[:n_fixed, :] = fixed_points

        iteration += 1

        if iteration % 10 == 0 and verbose:
            print(f"Iteration: {iteration}. Error: {error}")

    return PolyMesh(voronoi.vertices, elements[: n_points], points, voronoi.ridge_points, domain)


def poly_mesher_cleaner(poly_mesh: PolyMesh, tolerence: float = 0.1) -> PolyMesh:
    """
    This function is may be applied to a PolyMesh object to refine the outputs in any given manner. This collapses
    particularly small edges (which may cause errors in the following numerical methods) and extracts unnessary data and
    duplicates from the poly_mesher algorithm.

    Args:
        poly_mesh (PolyMesh): A PolyMesh object to be refined.
        tolerence (float): The tolerence to which the small edges are collapsed. This is a multiplier not a distance.

    Returns:
        (PolyMesh): Another PolyMesh object but (hopefully) more refined and manageable.
    """

    vertices, regions = _poly_mesher_extract_nodes(poly_mesh.vertices, poly_mesh.filtered_regions)
    vertices, regions = _poly_mesher_collapse_small_edges(vertices, regions, tolerence)

    return PolyMesh(vertices, regions, poly_mesh.filtered_points, poly_mesh.ridge_points, poly_mesh.domain)


def _poly_mesher_init_point_set(domain: Domain, **kwargs) -> np.ndarray:

    bounding_box = domain.bounding_box

    if "n_points" in kwargs:
        # This generates a random point set
        n_points = kwargs.get("n_points")
        points = np.full((n_points, 2), -np.inf)
        s = 0
        np.random.seed(1337)
        while s < n_points:
            p_1 = (bounding_box[0, 1] - bounding_box[0, 0]) * \
                np.random.uniform(size=(1, n_points)).T + bounding_box[0, 0]
            p_2 = (bounding_box[1, 1] - bounding_box[1, 0]) * \
                np.random.uniform(size=(1, n_points)).T + bounding_box[1, 0]

            p = np.concatenate((p_1, p_2), axis=1)
            d = domain.distances(p)
            last_index_negative = np.argwhere(d[:, -1] < 0.0)
            number_added = min(n_points - s, last_index_negative.shape[0])
            points[s:s + number_added, :] = p[last_index_negative[:number_added].T.flatten(), :]
            s += number_added

    elif "n_x" in kwargs and "n_y" in kwargs:
        # This generates a uniformly spread point set
        n_x = kwargs.get("n_x")
        n_y = kwargs.get("n_y")

        x = np.linspace(bounding_box[0, 0], bounding_box[0, 1], n_x + 1)
        y = np.linspace(bounding_box[1, 0], bounding_box[1, 1], n_y + 1)
        x_c = 0.5 * (x[1:] + x[:-1])
        y_c = 0.5 * (y[1:] + y[:-1])
        [X, Y] = np.meshgrid(x_c, y_c)

        X, Y = X.T, Y.T
        points = np.concatenate((np.reshape(X, (-1, 1), order='F'), np.reshape(Y, (-1, 1), order="F")), axis=1)
        d = domain.distances(points)
        log_ind = d[:, -1] < 0.0
        points = points[log_ind, :]

    else:
        raise AttributeError("key word error: must be just `n_points` or both `n_x` and `n_y`.")

    return points


def _poly_mesher_reflect(points: np.ndarray, domain: Domain, area: float) -> (np.ndarray, np.ndarray):

    epsilon = 1.0e-8
    n_points = points.shape[0]
    alpha = 1.5 * np.sqrt(area / float(n_points))

    d = domain.distances(points)
    n_boundary_segments = d.shape[1] - 1

    n_1 = 1.0 / epsilon * (domain.distances(points + np.tile(np.array([epsilon, 0.0]), (n_points, 1))) - d)
    n_2 = 1.0 / epsilon * (domain.distances(points + np.tile(np.array([0.0, epsilon]), (n_points, 1))) - d)

    log_ind = np.abs(d[:, :n_boundary_segments]) < alpha

    p_1 = np.tile(points[:, 0][:, np.newaxis], (1, n_boundary_segments))
    p_2 = np.tile(points[:, 1][:, np.newaxis], (1, n_boundary_segments))

    p_1 = np.concatenate([p_1[log_ind[:, i], i] for i in range(n_boundary_segments)], axis=0)[:, np.newaxis]
    p_2 = np.concatenate([p_2[log_ind[:, i], i] for i in range(n_boundary_segments)], axis=0)[:, np.newaxis]

    n_1 = np.concatenate([n_1[log_ind[:, i], i] for i in range(n_boundary_segments)], axis=0)[:, np.newaxis]
    n_2 = np.concatenate([n_2[log_ind[:, i], i] for i in range(n_boundary_segments)], axis=0)[:, np.newaxis]

    d = np.concatenate([d[log_ind[:, i], i] for i in range(n_boundary_segments)], axis=0)[:, np.newaxis]

    r_ps = np.concatenate((p_1, p_2), axis=1) - 2.0 * np.concatenate((n_1, n_2), axis=1) * np.tile(d, (1, 2))

    r_p_ds = domain.distances(r_ps)

    logical_rp = np.logical_and(r_p_ds[:, -1] > 0, np.abs(r_p_ds[:, -1]) >= 0.9 * np.abs(d).flatten())

    r_ps = r_ps[logical_rp, :]

    if not r_ps.size == 0:
        r_ps = np.unique(r_ps, axis=0)

    return r_ps


def _poly_mesher_vorocentroid(points: np.ndarray, vertices, elements) -> (np.ndarray, float, float):

    n_points = points.shape[0]
    center_points = np.full((n_points, 2), -np.inf)
    areas = np.full((n_points,), -np.inf)

    for i in range(n_points):
        region = Polygon(vertices[elements[i]])
        areas[i] = region.area
        center_points[i, :] = np.array(region.centroid.coords.xy).T.flatten()

    total_area = areas.sum()
    error = np.sqrt(np.sum((areas ** 2) * np.sum((center_points - points) ** 2, 1), 0)) * n_points / total_area ** 1.5

    return center_points, total_area, error


def _poly_mesher_extract_nodes(nodes: np.ndarray, filtered_elements: list) -> (np.ndarray, list):

    linked_elements = np.array(list((itertools.chain.from_iterable(filtered_elements))))
    unique_ = np.unique(linked_elements)  # these are the points solely used by the elements
    c_nodes = np.arange(nodes.shape[0])
    c_nodes[list(set(c_nodes).difference(set(unique_)))] = max(unique_)
    nodes, elements = _poly_mesher_rebuild_lists(nodes, filtered_elements, c_nodes)

    return nodes, elements


def _poly_mesher_collapse_small_edges(nodes: np.ndarray, elements: list, tolerance: float) -> (np.ndarray, list):

    while True:

        c_edge = np.array([[]])
        for element in elements:
            n_v = len(element)

            if n_v <= 3:
                continue  # this is the case when the element is a triangle -- can't collapse

            v_ = nodes[element]
            beta = np.arctan2(v_[:, 1] - np.sum(v_[:, 1])/n_v, v_[:, 0] - np.sum(v_[:, 0])/n_v)
            beta = np.mod(np.roll(beta, -1) - beta, 2 * np.pi)
            beta_ideal = 2 * np.pi / n_v
            edges = np.concatenate((np.array(element)[:, np.newaxis], np.roll(element, -1)[:, np.newaxis]), axis=1)
            idx = beta < tolerance * beta_ideal
            if (e_ := edges[idx, :]).size != 0:
                if c_edge.size == 0:
                    c_edge = np.atleast_2d(e_)
                else:
                    c_edge = np.concatenate((c_edge, np.atleast_2d(e_)), axis=0)

        if c_edge.size == 0:
            break

        c_edge = np.unique(np.sort(c_edge), axis=0)
        c_nodes = np.arange(nodes.shape[0])
        c_nodes[c_edge[:, 1]] = c_nodes[c_edge[:, 0]]

        nodes, elements = _poly_mesher_rebuild_lists(nodes, elements, c_nodes)

    return nodes, elements


def _poly_mesher_rebuild_lists(nodes: np.ndarray, elements: list, c_nodes: np.ndarray) -> (np.ndarray, list):

    _elems = []
    _, idx, r_idx = np.unique(c_nodes, return_index=True, return_inverse=True)

    # idx are the indices of the unique elements -- ie the unique nodes

    if nodes.shape[0] > len(idx):
        idx[-1] = np.max(c_nodes)

    nodes = nodes[idx, :]

    for elem in elements:
        temp_elem = np.unique(r_idx[elem])
        v_ = nodes[temp_elem]
        n_v = len(temp_elem)

        perm = np.argsort(np.arctan2(v_[:, 1] - np.sum(v_[:, 1])/n_v,
                                     v_[:, 0] - np.sum(v_[:, 0])/n_v))

        _elems.append(temp_elem[perm])

    return nodes, _elems
