"""Monolayer epithelium objects
"""

import logging
import numpy as np
import pandas as pd

from scipy.spatial import cKDTree
from .objects import Epithelium
from .sheet import Sheet
from ..generation import extrude, subdivide_faces
from ..geometry.bulk_geometry import BulkGeometry

logger = logging.getLogger(name=__name__)


class Monolayer(Epithelium):
    """
    3D monolayer epithelium
    """

    def __init__(self, name, datasets, specs=None, coords=None):

        super().__init__(name, datasets, specs, coords)
        self.vert_df["is_active"] = 1
        self.cell_df["is_alive"] = 1
        self.face_df["is_alive"] = 1
        self.reset_topo()
        BulkGeometry.update_all(self)

    @classmethod
    def from_flat_sheet(cls, name, apical_sheet, specs, thickness=1):
        datasets = extrude(
            apical_sheet.datasets, method="translation", vector=[0, 0, -thickness]
        )

        mono = cls(name, datasets, specs)
        mono.reset_topo()
        return mono

    def segment_index(self, segment, element):
        df = getattr(self, "{}_df".format(element))
        return df[df["segment"] == segment].index

    @property
    def lateral_faces(self):
        return self.segment_index("lateral", "face")

    @property
    def apical_faces(self):
        return self.segment_index("apical", "face")

    @property
    def basal_faces(self):
        return self.segment_index("basal", "face")

    @property
    def lateral_edges(self):
        return self.segment_index("lateral", "edge")

    @property
    def apical_edges(self):
        return self.segment_index("apical", "edge")

    @property
    def basal_edges(self):
        return self.segment_index("basal", "edge")

    @property
    def apical_verts(self):
        return self.segment_index("apical", "vert")

    @property
    def basal_verts(self):
        return self.segment_index("basal", "vert")

    @property
    def lateral_verts(self):
        return self.segment_index("lateral", "vert")

    def get_sub_sheet(self, segment):
        """ Returns a :class:`Sheet` object of the corresponding
        segment

        Parameters
        ----------
        segment: str, the corresponding segment, wether 'apical' or 'basal'

        """
        datasets = {
            element: self.datasets[element].loc[self.segment_index(segment, element)]
            for element in ["edge", "face", "vert"]
        }
        specs = {k: self.specs[k] for k in ["face", "edge", "vert", "settings"]}
        return Sheet(self.identifier + segment, datasets, specs)

    def guess_vert_segment(self, vert):
        """Infers the vertex segment from its surrounding edges.

        """
        v_edges = self.edge_df[self.edge_df["srce"] == vert]
        if v_edges.shape[0] == 0:
            logger.info("Vertex %d not found", vert)
            return
        if v_edges.shape[0] == 12:
            self.vert_df.loc[vert, ["segment"]] = "lateral"
            return
        intersect = {"apical", "basal"}.intersection(v_edges["segment"])
        if len(intersect) == 2:
            logger.info("Segment of vertex %d could not be determined", vert)
            self.vert_df.loc[vert, ["segment"]] = "unknown"
        elif not intersect:
            self.vert_df.loc[vert, ["segment"]] = "lateral"
        else:  # intersect is {"apical"} or {"basal"}
            self.vert_df.loc[vert, ["segment"]], = intersect

    def guess_face_segment(self, face):
        """Infers the face segment from its surrounding edges.

        """
        face_edges = self.edge_df[self.edge_df.face == face]
        if face_edges.shape[0] == 0:
            logger.info("face %d not found", face)
        v_segments = set(self.vert_df.loc[face_edges["srce"], "segment"])
        if len(v_segments) == 2:
            self.face_df.loc[face, "segment"] = "lateral"
        elif len(v_segments) == 1:
            new_segment, = v_segments
            self.face_df.loc[face, "segment"] = new_segment


class MonolayerWithLamina(Monolayer):
    """
    3D monolayer epithelium with a lamina meshing
    """

    def __init__(self, name, datasets, specs=None, coords=None):

        super().__init__(name, datasets, specs, coords)

        BulkGeometry.update_all(self)
        self.reset_index()

        subdivided = subdivide_faces(self, self.basal_faces)
        for name, df in subdivided.items():
            setattr(self, "{}_df".format(name), df)
        self.reset_index()
        self.reset_topo()

        subdiv_edges = self.edge_df[self.edge_df["subdiv"] == 1].index
        self.edge_df.loc[subdiv_edges, "segment"] = "basal"

        subdiv_verts = self.vert_df[self.vert_df["subdiv"] == 1].index
        self.vert_df.loc[subdiv_verts, "segment"] = "basal"
        self.vert_df.loc[subdiv_verts, "basal_shift"] = 0.0
        self.vert_df.loc[subdiv_verts, "is_active"] = 1.0

        subdiv_verts = self.vert_df[self.vert_df["subdiv"] == 1].index
        focal_adhesions = self.vert_df.loc[subdiv_verts]

        max_dist = self.edge_df.length.dropna().median() * 1.74
        lamina_tree = cKDTree(focal_adhesions[self.coords].values)
        lamina_edges = pd.DataFrame(
            [[i, j] for i, j in lamina_tree.query_pairs(max_dist, eps=1e-3)],
            columns=["srce", "trgt"],
        )
        lamina_edges.index.name = "edge"
        lamina_edges["srce"] = focal_adhesions.index[lamina_edges["srce"]]
        lamina_edges["trgt"] = focal_adhesions.index[lamina_edges["trgt"]]
        # place holder face and cell
        lamina_face = self.face_df.index.max() + 1
        lamina_edges["face"] = lamina_face
        self.face_df.append(self.face_df.iloc[0].copy())
        self.face_df.loc[lamina_face, "is_alive"] = 0

        lamina_cell = self.cell_df.index.max() + 1
        lamina_edges["cell"] = lamina_cell
        self.cell_df.append(self.cell_df.iloc[0].copy())
        self.cell_df.loc[lamina_cell, "is_alive"] = 0

        lamina_edges.index += self.edge_df.index.max() + 1
        lamina_edges["segment"] = "lamina"
        lamina_edges["subdiv"] = 0
        self.edge_df = pd.concat([self.edge_df, lamina_edges], sort=True)
        self.reset_topo()
        BulkGeometry.update_all(self)

    @property
    def lamina_edges(self):
        return self.segment_index("lamina", "edge")
