import copy
from typing import Any, Dict, Union

import numpy as np
from skfem import MeshTri1


class Domain:
    """Square domain [0,1]x[0,1] with initial coarse mesh and fine integration mesh."""

    def __init__(
        self,
        *,
        domain_config: Dict[Union[str, int], Any],
        random_state: np.random.RandomState,
    ):
        xmin, ymin, xmax, ymax = domain_config.get("boundary", [0.0, 0.0, 1.0, 1.0])
        self._boundary = np.array([xmin, ymin, xmax, ymax])
        self._random_state = random_state

        num_elements = domain_config.get("initial_num_elements", None)
        if num_elements is not None:
            domain_area = (xmax - xmin) * (ymax - ymin)
            self._max_volume = 2.0 * domain_area / float(num_elements)
        else:
            self._max_volume = domain_config.get("max_initial_element_volume", 0.05)

        self._initial_mesh = self._create_initial_mesh()

    @property
    def initial_mesh(self) -> MeshTri1:
        return copy.deepcopy(self._initial_mesh)

    def replace_initial_mesh(self, mesh: MeshTri1) -> None:
        """Replace the stored initial mesh (e.g. after Nyquist enforcement)."""
        self._initial_mesh = mesh

    def get_integration_mesh(self) -> MeshTri1:
        return self._initial_mesh.refined(4)

    @property
    def boundary_line_segments(self) -> np.ndarray:
        boundary_edges = self._initial_mesh.boundary_facets()
        boundary_node_indices = self._initial_mesh.facets[:, boundary_edges]
        return self._initial_mesh.p[:, boundary_node_indices].T.reshape(-1, 4)

    def _create_initial_mesh(self) -> MeshTri1:
        return self._meshpy_square()

    def _meshpy_square(self) -> MeshTri1:
        import meshpy.triangle as triangle

        xmin, ymin, xmax, ymax = self._boundary
        points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
        facets = [(0, 1), (1, 2), (2, 3), (3, 0)]

        info = triangle.MeshInfo()
        info.set_points(points)
        info.set_facets(facets)

        mesh = triangle.build(info, max_volume=self._max_volume)
        vertices = np.array(mesh.points).T
        triangles = np.array(mesh.elements).T
        return MeshTri1(vertices, triangles)


def create_domain(
    *, domain_config: Dict[Union[str, int], Any], random_state: np.random.RandomState
) -> Domain:
    return Domain(domain_config=domain_config, random_state=random_state)