import os
import typing
from dataclasses import dataclass

import dolfin

from . import Mesh
from . import numpy_mpi
from .dolfin_utils import compute_meshvolume
from .dolfin_utils import get_cavity_volume
from .dolfin_utils import map_vector_field
from .geometry_utils import load_geometry_from_h5
from .geometry_utils import move
from .geometry_utils import save_geometry_to_h5
from .utils import getLogger
from .utils import mpi_comm_world

logger = getLogger(__name__)


@dataclass
class Microstructure:
    """f0 - fibers, s0 - sheets, n0 - cross-sheets"""

    f0: typing.Optional[dolfin.Function] = None
    s0: typing.Optional[dolfin.Function] = None
    n0: typing.Optional[dolfin.Function] = None


@dataclass
class CRLBasis:
    """l0 - longitudinal, c0 - circumferential, r0 - radial"""

    c0: typing.Optional[dolfin.Function] = None
    r0: typing.Optional[dolfin.Function] = None
    l0: typing.Optional[dolfin.Function] = None


@dataclass
class MarkerFunctions:
    """
    # vfun - vertex function
    # efun - edge function
    # ffun - facet function
    # cfun - cell function
    """

    vfun: typing.Optional[dolfin.MeshFunction] = None
    efun: typing.Optional[dolfin.MeshFunction] = None
    ffun: typing.Optional[dolfin.MeshFunction] = None
    cfun: typing.Optional[dolfin.MeshFunction] = None


def get_attribute(obj, key1, key2, default=None):
    f = getattr(obj, key1, None)
    if f is None:
        f = getattr(obj, key2, default)
    return f


class Geometry(object):
    """
    Base class for geometry

    Arguments
    ---------
    mesh : :class:`dolfin.mesh`
        The mesh
    markers : dict
        A dictionary with markers for the mesh (optional)
    marker_functions : :class:`pulse.geometry.MarkerFunctions`
        A Markerfunction object with Meshfunctions for
        the mesh
    microstructure : :class:`pulse.geometry.Microstructure`
        A Markerfunction object with functions for the
        fiber, sheet and sheet normal (optional)
    crl_basis : :class:`pulse.geometry.CRLBasis`
        A CRLBasis objedt with funcions for the circumferential,
        radial and longitudinal vectors. (optional)


    Geometry can be intanciated directly

    Example
    -------

        .. code-block:: python

            import dolfin
            mesh = dolfin.UnitCubeMesh(3,3,3)
            geometry = Geometry(mesh)

    You can load create and instace be loading a geometry from a file

    Example
    -------

        .. code-block:: python

            # Geometry is stored in a file "geometry.h5"
            geo = Geometry.from_file("geometry.h5")

    """

    def __init__(
        self,
        mesh,
        markers=None,
        marker_functions=None,
        microstructure=None,
        crl_basis=None,
    ):
        self.mesh = mesh

        self.markers = markers or {}
        self.marker_functions = marker_functions or MarkerFunctions()
        self.microstructure = microstructure or Microstructure()
        self.crl_basis = crl_basis or CRLBasis()
        self._post_init()

    def _post_init(self):
        pass

    def __repr__(self):
        args = []
        for attr in ("f0", "s0", "n0", "vfun", "ffun", "cfun", "c0", "r0", "l0"):
            if getattr(self, attr) is not None:
                args.append(attr)

        return ("{self.__class__.__name__}" "({self.mesh}), " "{args})").format(
            self=self,
            args=", ".join(args),
        )

    @classmethod
    def from_file(cls, h5name, h5group="", comm=None):
        comm = comm if comm is not None else mpi_comm_world()

        return cls(**cls.load_from_file(h5name, h5group, comm))

    def _copy(self, u=None, factor=1.0):
        """
        Return a copy of this geometry and attach a moved
        mesh according to the provided displacement.
        If no displacement is provided this will just
        return a copy of the original geometry.
        """
        new_mesh = Mesh(self.mesh)

        if u is not None:
            U = move(new_mesh, u, factor)
        else:
            U = u

        marker_functions_ = {}
        for dim, fun in ((0, "vfun"), (1, "efun"), (2, "ffun"), (3, "cfun")):
            f_old = getattr(self, fun)
            if f_old is None:
                continue
            f = dolfin.MeshFunction("size_t", new_mesh, dim, new_mesh.domains())
            f.set_values(f_old.array())
            marker_functions_[fun] = f
        marker_functions = MarkerFunctions(**marker_functions_)

        microstructure_ = {}
        for field in ("f0", "s0", "n0"):
            v0 = getattr(self, field)
            if v0 is None:
                continue
            v = map_vector_field(v0, new_mesh, U)

            microstructure_[field] = v
        microstructure = Microstructure(**microstructure_)

        crl_basis_ = {}
        for basis in ("c0", "r0", "l0"):
            v0 = getattr(self, basis)
            if v0 is None:
                continue
            v = map_vector_field(v0, new_mesh, U)
            crl_basis_[basis] = v
        crl_basis = CRLBasis(**crl_basis_)

        return dict(
            mesh=new_mesh,
            markers=self.markers,
            marker_functions=marker_functions,
            microstructure=microstructure,
            crl_basis=crl_basis,
        )

    def copy(self, u=None, factor=1.0):
        """Return as copy of the geometry with possiblity to
        move the new geometry according to some factor times
        some displacement (u)
        """
        return self.__class__(**self._copy(u=u, factor=factor))

    @staticmethod
    def load_from_file(h5name, h5group, comm):
        logger.debug(f"Load geometry from file {h5name}")

        geo = load_geometry_from_h5(h5name, h5group, include_sheets=True, comm=comm)

        f0 = get_attribute(geo, "f0", "fiber", None)
        s0 = get_attribute(geo, "s0", "sheet", None)
        n0 = get_attribute(geo, "n0", "sheet_normal", None)

        c0 = get_attribute(geo, "c0", "circumferential", None)
        r0 = get_attribute(geo, "r0", "radial", None)
        l0 = get_attribute(geo, "l0", "longitudinal", None)

        vfun = getattr(geo, "vfun", None)
        ffun = getattr(geo, "ffun", None)
        cfun = get_attribute(geo, "cfun", "sfun", None)

        kwargs = {
            "mesh": geo.mesh,
            "markers": geo.markers,
            "marker_functions": MarkerFunctions(vfun=vfun, ffun=ffun, cfun=cfun),
            "microstructure": Microstructure(f0=f0, s0=s0, n0=n0),
            "crl_basis": CRLBasis(c0=c0, r0=r0, l0=l0),
        }

        return kwargs

    def save(
        self,
        h5name,
        h5group="",
        other_functions=None,
        other_attributes=None,
        overwrite_file=False,
        overwrite_group=True,
    ):
        h5name = os.path.splitext(h5name)[0] + ".h5"
        logger.debug(f"Save to {h5name}...")
        save_geometry_to_h5(
            self.mesh,
            h5name=h5name,
            h5group=h5group,
            markers=self.markers or None,
            fields=self.microstructure_list() or None,
            local_basis=self.crl_basis_list() or None,
            meshfunctions=self.meshfunction_list(),
            overwrite_file=overwrite_file,
            overwrite_group=overwrite_group,
        )
        logger.debug("Saved")

    @property
    def top_dim(self):
        """Topological Dimension"""
        return self.mesh.topology().dim()

    @property
    def geo_dim(self):
        """Geometric Dimension"""
        return self.mesh.geometry().dim()

    def crl_basis_list(self):
        """Return a list of the CRL basis in the order
        c0, r0, l0. Basis elements that are none will not
        be included
        """
        if not hasattr(self, "_crl_basis_list"):
            self._crl_basis_list = []
            for basis in ["c0", "r0", "l0"]:
                e = getattr(self.crl_basis, basis)
                if e is not None:
                    self._crl_basis_list.append(e)

        return self._crl_basis_list

    @property
    def nbasis(self):
        return len(self.crl_basis_list())

    @property
    def dx(self):
        """Return the volume measure using self.mesh as domain and
        self.cfun as subdomain_data
        """
        return dolfin.dx(domain=self.mesh, subdomain_data=self.cfun)

    @property
    def ds(self):
        """Return the surface measure of exterior facets using
        self.mesh as domain and self.ffun as subdomain_data
        """
        return dolfin.ds(domain=self.mesh, subdomain_data=self.ffun)

    @property
    def facet_normal(self):
        return dolfin.FacetNormal(self.mesh)

    @property
    def nregions(self):
        return len(self.regions)

    @property
    def regions(self):
        """Return a set of the unique values of the cell function (cfun)"""
        if not hasattr(self, "_regions"):
            try:
                # FIXME
                regions = set(
                    numpy_mpi.gather_broadcast(self.dx.subdomain_data().array()).astype(
                        int,
                    ),
                )
                self._regions = {int(i) for i in regions}
            except AttributeError:
                # self.dx.subdomain_data() is None
                logger.warning("No regions found")
                self._regions = {0}

        return self._regions

    @property
    def meshvols(self):
        """Return a list of the volume of each subdomain defined
        by the cell function (cfun)
        """
        if not hasattr(self, "_meshvols"):
            self._meshvols = [
                compute_meshvolume(dx=self.dx, subdomain_id=int(i))
                for i in self.regions
            ]
        return self._meshvols

    @property
    def meshvol(self):
        """Return the volume of the whole mesh"""
        if not hasattr(self, "_meshvol"):
            self._meshvol = compute_meshvolume(domain=self.mesh)

        return self._meshvol

    def microstructure_list(self):
        """Fibers, sheet and sheet-normals in a list"""
        fields = []
        for field in ["f0", "s0", "n0"]:
            e = getattr(self.microstructure, field)
            if e is not None:
                fields.append(e)
        return fields

    def meshfunction_list(self):
        meshfunctions = {}
        for dim, func in enumerate(["vfun", "efun", "ffun", "cfun"]):
            mf = getattr(self.marker_functions, func)
            if mf is None:
                mf = dolfin.MeshFunction("size_t", self.mesh, dim, self.mesh.domains())
            meshfunctions[dim] = mf
        return meshfunctions

    @property
    def vfun(self):
        """Vertex Function"""
        return self.marker_functions.vfun

    @property
    def efun(self):
        """Edge Function"""
        return self.marker_functions.efun

    @property
    def ffun(self):
        """Facet Function"""
        return self.marker_functions.ffun

    @property
    def cfun(self):
        """Cell Function"""
        return self.marker_functions.cfun

    @property
    def f0(self):
        """Fibers"""
        return self.microstructure.f0

    @property
    def s0(self):
        """Sheets"""
        return self.microstructure.s0

    @property
    def n0(self):
        """Cross-Sheets"""
        return self.microstructure.n0

    @property
    def l0(self):
        """Longitudinal"""
        return self.crl_basis.l0

    @property
    def c0(self):
        """Circumferential"""
        return self.crl_basis.c0

    @property
    def r0(self):
        """Radial"""
        return self.crl_basis.r0


class HeartGeometry(Geometry):
    def _post_init(self):
        self.xshift = 0.0

    @property
    def is_biv(self):
        if "ENDO_RV" in self.markers:
            # FIXME
            return self.markers["ENDO_RV"][0] in set(
                numpy_mpi.gather_broadcast(self.ffun.array()),
            )

        return False

    def cavity_volume(self, chamber="lv", u=None):
        return get_cavity_volume(self, chamber=chamber, u=u, xshift=self.xshift)

    @property
    def base_mean_position(self):
        """
        Return mean coordinates of the base
        (serial only?)
        """
        import numpy as np

        (facet_indices,) = np.where(self.ffun.array() == self.markers["BASE"][0])
        point_inidces = []
        for facet in dolfin.facets(self.mesh):
            if facet.index() in facet_indices:
                point_inidces.extend(facet.entities(0).tolist())

        return np.mean(self.mesh.coordinates()[point_inidces, :], 0)

    def update_xshift(self):
        self.xshift = self.base_mean_position[0]
