import numpy as np
import cmath

from ..mesh.mesh_attributes import ArrayAttribute, Attribute
from ..mesh.datatypes import *
from .. import geometry as geom
from ..geometry import Vec

@allowed_mesh_types(SurfaceMesh, VolumeMesh)
def face_area(mesh : Mesh, name="area", persistent:bool=True, dense:bool=True) -> Attribute:
    """
    Computes the barycenter point of each face.

    Parameters:
        mesh (SurfaceMesh): the input mesh
        name (str, optional): Name given to the attribute. Defaults to "area".
        persistent (bool, optional): If the attribute is persistent (stored in the mesh object) or not. Defaults to True.
        dense (bool, optional): Is the attribute dense (numpy array) or not (dict). Defaults to True

    Returns:
        Attribute: a float per faces
    """
    if persistent:
        area = mesh.faces.create_attribute(name, float, dense=dense)
    else:
        area = ArrayAttribute(float, len(mesh.faces)) if dense else Attribute(float)
    for T in mesh.id_faces:
        pts = [mesh.vertices[u] for u in mesh.faces[T]]
        npt = len(pts)
        if npt==3:
            area[T] = geom.triangle_area(*pts)
        elif npt==4:
            area[T] = geom.quad_area(*pts)
        else:
            bary = sum(pts)/npt
            for i in range(npt):
                A = pts[i]
                B = pts[(i+1)%npt]
                area[T] += geom.triangle_area(A,B,bary)            
    return area

@allowed_mesh_types(SurfaceMesh)
def face_normals(mesh : SurfaceMesh, name="normals", persistent:bool=True, dense:bool=True) -> Attribute:
    """
    Computes the barycenter point of each face.

    Parameters:
        mesh (SurfaceMesh): the input mesh
        name (str, optional): Name given to the attribute. Defaults to "normals".
        persistent (bool, optional): If the attribute is persistent (stored in the mesh object) or not. Defaults to True.
        dense (bool, optional): Is the attribute dense (numpy array) or not (dict). Defaults to True

    Returns:
        Attribute: 3 floats per faces
    """
    if persistent :
        normals = mesh.faces.create_attribute(name, float, 3, dense=dense)
    else:
        normals = ArrayAttribute(float, len(mesh.faces), 3) if dense else Attribute(float, 3)
    for iT,T in enumerate(mesh.faces):
        pA,pB,pC = (mesh.vertices[u] for u in T[:3])
        normals[iT] = Vec.normalized(geom.cross(pB-pA, pC-pA))
    return normals

@allowed_mesh_types(SurfaceMesh, VolumeMesh)
def face_barycenter(mesh : SurfaceMesh, name="barycenter", persistent:bool = True, dense:bool = True) -> Attribute:
    """
    Computes the barycenter point of each face.

    Parameters:
        mesh (SurfaceMesh): the input mesh
        name (str, optional): Name given to the attribute. Defaults to "barycenter".
        persistent (bool, optional): If the attribute is persistent (stored in the mesh object) or not. Defaults to True.
        dense (bool, optional): Is the attribute dense (numpy array) or not (dict). Defaults to True

    Returns:
        Attribute: one 3D vector per face
    """
    if persistent :
        bary = mesh.faces.create_attribute(name, float, 3, dense=dense)
    else:
        bary = ArrayAttribute(float, len(mesh.faces), 3) if dense else Attribute(float, 3)
    for iT,T in enumerate(mesh.faces):
        bary[iT] = sum(mesh.vertices[u] for u in T) / len(T)
    return bary

@allowed_mesh_types(SurfaceMesh, VolumeMesh)
def face_circumcenter(mesh : SurfaceMesh, name="circumcenter", persistent:bool=True, dense:bool=True) -> Attribute:
    """
    Computes the circumcenter point of each triangular face.

    Parameters:
        mesh (SurfaceMesh): the input mesh
        name (str, optional): Name given to the attribute. Defaults to "barycenter".
        persistent (bool, optional): If the attribute is persistent (stored in the mesh object) or not. Defaults to True.
        dense (bool, optional): Is the attribute dense (numpy array) or not (dict). Defaults to True

    Returns:
        Attribute: one float per face

    Raises:
        Exception: fails if a face of the mesh is not a triangle
    """
    if persistent :
        circum = mesh.faces.create_attribute(name, float, 3, dense=dense)
    else:
        circum = ArrayAttribute(float, len(mesh.faces), 3) if dense else Attribute(float, 3)
    for iF,F in enumerate(mesh.faces):
        if len(F)!=3:
            raise Exception("circumcenters can be computed only for triangular faces. Received face n°{} containing {}!=3 vertices".format(iF,len(F)))
        circum[iF] = geom.circumcenter(*(mesh.vertices[u] for u in F))
    return circum

@allowed_mesh_types(SurfaceMesh)
def face_near_border(mesh : SurfaceMesh, dist:int = 2, name = "near_border", persistent:bool = True, dense:bool = False) -> Attribute:
    """Returns a boolean attribute that is True if the face is at most at 'dist' neighbours from the boundary.
    Proceeds by region growing, starting from all faces touching the boundary and going inwards.

    Parameters:
        mesh (SurfaceMesh): the input mesh
        dist (int, optional): Extend to which we flag the faces. All faces with a path of length < dist will be set to True. Defaults to 2.
        name (str, optional): Name given to the attribute. Defaults to "near_border".
        persistent (bool, optional): If the attribute is persistent (stored in the mesh object) or not. Defaults to True.
        dense (bool, optional): Is the attribute dense (numpy array) or not (dict). Defaults to False

    Returns:
        Attribute: one bool per face.
    """
    
    near = set()
    for e in mesh.boundary_edges:
        T1,T2 = mesh.connectivity.edge_to_faces(*mesh.edges[e])
        if T1 is None: near.add(T2)
        else: near.add(T1)
    for _ in range(dist-1):
        new_near = set()
        for T in mesh.id_faces:
            if T in near : continue
            for T2 in mesh.connectivity.face_to_faces(T):
                if T2 in near: new_near.add(T)
        near = near | new_near
    
    if persistent:
        near_attr = mesh.faces.create_attribute(name, bool, dense=dense)
    else:
        near_attr = ArrayAttribute(bool, len(mesh.faces)) if dense else Attribute(bool)
    for f in near: near_attr[f] = True
    return near_attr

@allowed_mesh_types(SurfaceMesh)
def triangle_aspect_ratio(mesh : SurfaceMesh, name : str="aspect_ratio", persistent : bool = True, dense : bool =True) -> Attribute:
    """Computes the aspect ratio of every triangular faces. Sets aspect ratio to -1 for every other faces

    Parameters:
        mesh (SurfaceMesh): the input mesh
        name (str, optional): Name given to the attribute. Defaults to "aspect_ratio".
        persistent (bool, optional): If the attribute is persistent (stored in the mesh object) or not. Defaults to True.
        dense (bool, optional): Is the attribute dense (numpy array) or not (dict). Defaults to False.

    Returns:
        Attribute: One float per face.
    """
    if persistent :
        ratio = mesh.faces.create_attribute(name, float, dense=dense)
    else:
        ratio = ArrayAttribute(float, len(mesh.faces)) if dense else Attribute(float)
    for iF,F in enumerate(mesh.faces):
        if len(F)!=3: ratio[iF] = -1 
        ratio[iF] = geom.aspect_ratio(*(mesh.vertices[u] for u in F))

@allowed_mesh_types(SurfaceMesh)
def parallel_transport_curvature(mesh : SurfaceMesh, PT:"SurfaceConnectionVertices", name : str="curvature", persistent : bool=True, dense : bool = True):
    """
    Compute the curvature of each face associated to a given parallel transport pT

    Args:
        mesh (SurfaceMesh): input mesh
        PT (SurfaceConnectionVertices): local bases and parallel transport for vertices.
        name (str, optional): Name given to the attribute.. Defaults to "curvature".
        persistent (bool, optional): If the attribute is persistent (stored in the mesh object) or not. Defaults to True.
        dense (bool, optional): Is the attribute dense (numpy array) or not (dict). Defaults to False.
    
    Returns:
        Attribute: One float per face.
    """
    if persistent :
        curv = mesh.faces.create_attribute(name, float, dense=dense)
    else:
        curv = ArrayAttribute(float, len(mesh.faces)) if dense else Attribute(float)
    for iF, (A,B,C) in enumerate(mesh.faces):
        v = 1+0j
        for a,b in [(A,B) ,(B,C), (C,A)]:
            v *= cmath.rect(1., PT.transport(b,a) - PT.transport(a,b) - np.pi)
        curv[iF] = cmath.phase(v)
    return curv