# -*- coding: utf-8 -*-
import numpy as np
from .general import blkdiag, transform_unit

'''
FE-related tools.
'''

def intpoints_from_elements(nodes, elements, sort_axis=0):
    """
    Calculates the integration points (midpoints) for each element based on node coordinates.

    Parameters
    ----------
    nodes : np.ndarray
        Array of node coordinates with shape (n_nodes, 4), where columns represent node index and x, y, z coordinates.
    elements : np.ndarray
        Array of element connectivity with shape (n_elements, 2), where each row contains indices of the two nodes forming an element.
    sort_axis : int, optional
        Axis (0 for x, 1 for y, 2 for z) to sort the integration points by. Default is 0 (x-axis).

    Returns
    -------
    x : np.ndarray
        Array of x-coordinates of the integration points, sorted by the specified axis.
    y : np.ndarray
        Array of y-coordinates of the integration points, sorted by the specified axis.
    z : np.ndarray
        Array of z-coordinates of the integration points, sorted by the specified axis.

    Notes
    -----
    Assumes that `nodeix_from_elements` is a function that returns the indices of the nodes for each element. Docstring is generated by Github Copilot.
    """
   
    nodeix = nodeix_from_elements(elements, nodes).astype('int')

    intpoints = (nodes[nodeix[:,0], 1:4]+nodes[nodeix[:,1], 1:4])/2
    sortix = np.argsort(intpoints[:,sort_axis])
    intpoints = intpoints[sortix, :]

    x = intpoints[:, 0]
    y = intpoints[:, 1]
    z = intpoints[:, 2]

    return x, y, z


def nodeix_from_elements(element_matrix, node_matrix, return_as_flat=False):
    """
    Map element node labels to their corresponding indices in the node matrix.

    Parameters
    ----------
    element_matrix : np.ndarray
        Array of elements with shape (n_elements, m), where columns 1 and 2 contain node IDs for each element.
    node_matrix : np.ndarray
        Array of nodes with shape (n_nodes, k), where column 0 contains node IDs.
    return_as_flat : bool, optional
        If True, returns a 1D array of unique node indices used by all elements.
        If False, returns a 2D array of shape (n_elements, 2) with node indices for each element.
        Default is False.

    Returns
    -------
    np.ndarray
        If return_as_flat is False, returns a 2D array of node indices for each element (shape: n_elements, 2).
        If return_as_flat is True, returns a 1D array of unique node indices.

    Notes
    -----
    Each element is assumed to be defined by two node labels in columns 1 and 2 of element_matrix. Docstring is generated by GitHub Copilot.
    """
    nodeix = [None] * len(element_matrix[:, 0])
    for element_ix, __ in enumerate(element_matrix[:, 0]):
        node1 = element_matrix[element_ix, 1]
        node2 = element_matrix[element_ix, 2]

        nodeix1 = np.where(node_matrix[:, 0] == node1)[0][0]
        nodeix2 = np.where(node_matrix[:, 0] == node2)[0][0]
        nodeix[element_ix] = [nodeix1, nodeix2]

    nodeix = np.array(nodeix)

    if return_as_flat:
        nodeix = np.unique(nodeix.flatten())

    return nodeix


def create_node_dict(element_dict, node_labels, x_nodes, y_nodes, z_nodes):
    """
    Creates a dictionary mapping element keys to their corresponding node data.

    Parameters
    ----------
    element_dict : dict
        A dictionary where each key corresponds to an element and each value contains information
        about the nodes associated with that element.
    node_labels : array-like
        An array of node labels/IDs.
    x_nodes : array-like
        An array of x-coordinates for each node.
    y_nodes : array-like
        An array of y-coordinates for each node.
    z_nodes : array-like
        An array of z-coordinates for each node.

    Returns
    -------
    node_dict : dict
        A dictionary where each key matches an element key from `element_dict`, and each value is
        an array containing the node label and coordinates (label, x, y, z) for the nodes
        associated with that element.

    Notes
    -----
    This function relies on the helper function `nodeix_from_elements` to determine the indices
    of nodes associated with each element. Docstring is generated by GitHub Copilot.
    """

    node_dict = dict()
    node_matrix = np.vstack([node_labels, x_nodes, y_nodes, z_nodes]).T

    for key in element_dict.keys():
        node_ix = nodeix_from_elements(element_dict[key], node_matrix, return_as_flat=True)
        node_dict[key] = node_matrix[node_ix, :]

    return node_dict


def elements_with_node(element_matrix, node_label):
    """
    Finds elements containing a specific node and returns their labels, indices, and local node indices.
    
    Parameters
    ----------
    element_matrix : np.ndarray
        A 2D array where each row represents an element. The first column contains element labels,
        and the second and third columns contain node labels associated with each element.
    node_label : int or float
        The node label to search for within the element matrix.

    Returns
    -------
    element_labels : np.ndarray
        Array of element labels that contain the specified node.
    element_ixs : np.ndarray
        Array of indices in `element_matrix` where the specified node is found.
    local_node_ix : np.ndarray
        Array indicating the local node index (0 or 1) within each element where the node is found.

    Examples
    --------
    >>> element_matrix = np.array([[1, 10, 20],
    ...                            [2, 20, 30],
    ...                            [3, 10, 30]])
    >>> elements_with_node(element_matrix, 10)
    (array([1, 3]), array([0, 2]), array([0., 0.]))

    Notes
    ---------
    Docstring is generated by GitHub Copilot.
    """
    element_ixs1 = np.array(np.where(element_matrix[:,1]==node_label)).flatten()
    element_ixs2 = np.array(np.where(element_matrix[:,2]==node_label)).flatten()

    element_ixs = np.hstack([element_ixs1, element_ixs2])
    element_labels = element_matrix[element_ixs, 0]

    local_node_ix = np.zeros(np.shape(element_ixs))
    local_node_ix[np.arange(0,len(element_ixs1))] = 0
    local_node_ix[np.arange(len(element_ixs1), len(element_ixs1) + len(element_ixs2))] = 1

    return element_labels, element_ixs, local_node_ix


def nodes_to_beam_element_matrix(node_labels, first_element_label=1):
    """
    Generates an element connectivity matrix for beam elements from a sequence of node labels.

    Parameters
    ----------
    node_labels : array_like
        Sequence of node labels (integers or floats) defining the order of nodes along the beam.
    first_element_label : int, optional
        The label to assign to the first element. Default is 1.

    Returns
    -------
    element_matrix : ndarray of shape (n_elements, 3)
        Array where each row represents a beam element. The columns are:
        [element_label, start_node_label, end_node_label].

    Examples
    --------
    >>> nodes = [10, 20, 30, 40]
    >>> nodes_to_beam_element_matrix(nodes)
    array([[ 1., 10., 20.],
           [ 2., 20., 30.],
           [ 3., 30., 40.]])

    Notes
    ---------
    Docstring is generated by GitHub Copilot.
    """

    n_nodes = len(node_labels)
    n_elements = n_nodes-1
    
    element_matrix = np.empty([n_elements, 3])
    element_matrix[:, 0] = np.arange(first_element_label,first_element_label+n_elements)
    element_matrix[:, 1] = node_labels[0:-1]
    element_matrix[:, 2] = node_labels[1:]

    return element_matrix


def node_ix_to_dof_ix(node_ix, n_dofs=6):
    """
    Converts a node index to a degree of freedom (DOF) index.
    Each node has n_dofs degrees of freedom, and this function returns the corresponding DOF indices.

    Parameters
    ----------
    node_ix : int
        Index of the node for which to find the DOF indices.
    n_dofs : int, optional
        Number of degrees of freedom per node. Default is 6.
    
    Returns
    -------
    dof_ix : np.ndarray
        Array of DOF indices corresponding to the given node index.

    Examples
    --------
    >>> node_ix = 2
    >>> n_dofs = 6
    >>> dof_ix = node_ix_to_dof_ix(node_ix, n_dofs)
    >>> print(dof_ix)
    [12 13 14 15 16 17]

    Notes
    --------
    Docstring is generated by GitHub Copilot.
    
    """
    start = node_ix*n_dofs
    stop = node_ix*n_dofs+n_dofs
    dof_ix = []
    for (i,j) in zip(start,stop):
        dof_ix.append(np.arange(i,j))

    dof_ix = np.array(dof_ix).flatten()

    return dof_ix


def dof_sel(arr, dof_sel, n_dofs=6, axis=0):
    """
    Selects degrees of freedom (DOFs) from an array along a specified axis.

    Parameters
    ----------
    arr : np.ndarray
        Input array from which to select DOFs.
    dof_sel : array-like
        Indices of the DOFs to select (relative to each block of n_dofs).
    n_dofs : int, optional
        Number of DOFs per node or block. Default is 6.
    axis : int, optional
        Axis along which to select DOFs. Default is 0.

    Returns
    -------
    arr_sel : np.ndarray
        Array containing only the selected DOFs along the specified axis.

    Examples
    --------
    >>> arr = np.arange(18).reshape(3, 6)
    >>> dof_sel(arr, [0, 2], n_dofs=6, axis=1)
    array([[ 0,  2],
           [ 6,  8],
           [12, 14]])

    Notes
    -------
    Docstring is generated by GitHub Copilot.
    """

    N = np.shape(arr)[axis]
    all_ix = [range(dof_sel_i, N, n_dofs) for dof_sel_i in dof_sel]
    sel_ix = np.array(all_ix).T.flatten()
    
    # Select the elements
    arr_sel = np.take(arr, sel_ix, axis=axis)
    
    return arr_sel
    


def elements_from_common_nodes(element_matrix, selected_nodes):
    """
    Find elements that have both nodes in `selected_nodes`.

    Parameters
    ----------
    element_matrix : np.ndarray
        Array of elements with shape (n_elements, 3), where columns are [element_id, node1, node2].
    selected_nodes : array-like
        List or array of node labels to search for.

    Returns
    -------
    selected_element_matrix : np.ndarray
        Subset of `element_matrix` where both nodes are in `selected_nodes`.
    sel_ix : np.ndarray
        Indices of the selected elements in the original `element_matrix`.

    Notes
    -----
    Only elements where both node1 and node2 are in `selected_nodes` are selected. Docstring is generated by GitHub Copilot.
    """
    
    mask = np.isin(element_matrix[:, 1], selected_nodes) & np.isin(element_matrix[:, 2], selected_nodes)
    sel_ix = np.where(mask)[0]
    selected_element_matrix = element_matrix[sel_ix, :]

    return selected_element_matrix, sel_ix


def transform_elements(node_matrix, element_matrix, e2p, repeats=1):
    """
    Transforms elements from global to local coordinates.
    Given a node matrix and an element matrix, this function computes the transformation matrices
    from global to local coordinates for each element. The transformation is based on the direction
    vector between the nodes of each element. The transformation matrices are constructed using
    the `transform_unit` function and are repeated as specified.

    Parameters
    ----------
    node_matrix : np.ndarray
        Array of node data. The first column should contain node IDs, and the remaining columns
        should contain node coordinates.
    element_matrix : np.ndarray
        Array of element data. Each row corresponds to an element, with the first column as the
        element ID and the next columns as node IDs defining the element.
    e2p : np.ndarray or similar
        Additional parameter passed to `transform_unit` for transformation construction.
    repeats : int, optional
        Number of times to repeat the transformation block (default is 1).

    Returns
    -------
    list of np.ndarray
        List of transformation matrices, one for each element, mapping global to local coordinates.

    Notes
    -----
    - Assumes that `transform_unit` and `blkdiag` functions are defined elsewhere.
    - The function matches node IDs between `element_matrix` and `node_matrix` to determine node positions.
    Docstring is generated by GitHub Copilot.
    """

    n_elements = np.shape(element_matrix)[0]
    T_g2el = [None]*n_elements

    for el in range(0, n_elements):
        n1_ix = np.where(node_matrix[:,0]==element_matrix[el, 1])
        n2_ix = np.where(node_matrix[:,0]==element_matrix[el, 2])

        X1 = node_matrix[n1_ix, 1:]
        X2 = node_matrix[n2_ix, 1:]
        dx = X2-X1
        e1 = dx/np.linalg.norm(dx)

        T_g2el[el] = blkdiag(transform_unit(e1, e2p), repeats)   # Transform from global to local coordinates VLocal = T*VGlobal

    return T_g2el