#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
# This file is part of PyHOPE
#
# Copyright (c) 2024 Numerics Research Group, University of Stuttgart, Prof. Andrea Beck
#
# PyHOPE is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# PyHOPE is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# PyHOPE. If not, see <http://www.gnu.org/licenses/>.

# ==================================================================================================================================
# Mesh generation library
# ==================================================================================================================================
# ----------------------------------------------------------------------------------------------------------------------------------
# Standard libraries
# ----------------------------------------------------------------------------------------------------------------------------------
import gc
import re
from typing import Final, Optional, cast
from collections.abc import Iterable
# ----------------------------------------------------------------------------------------------------------------------------------
# Third-party libraries
# ----------------------------------------------------------------------------------------------------------------------------------
import numpy as np
# ----------------------------------------------------------------------------------------------------------------------------------
# Local imports
# ----------------------------------------------------------------------------------------------------------------------------------
import pyhope.mesh.mesh_vars as mesh_vars
from pyhope.common.common_numba import jit, types
from pyhope.basis.basis_basis import change_basis_2D
from pyhope.mesh.mesh_common import face_to_nodes
# ==================================================================================================================================


# Use Pool initializer to attach process-local data to the worker function
def init_worker(function, VdmEqToGP, DGP, weights) -> None:
    """Initializer to set process-local attributes on the worker function
    """
    function.VdmEqToGP = VdmEqToGP
    function.DGP       = DGP
    function.weights   = weights


@jit(types.Tuple((types.float64, types.float64, types.float64))(
    types.float64[:, :, ::1],
    types.float64[:, :, ::1],
    types.float64[:,    ::1]
), nopython=True, cache=True, nogil=True)
def eval_dotprod(dXdetaGP: np.ndarray, dXdxiGP: np.ndarray, weights: np.ndarray) -> tuple[np.ndarray, ...]:
    NSurf0 = -np.sum(weights * (dXdxiGP[1] * dXdetaGP[2] - dXdxiGP[2] * dXdetaGP[1]))
    NSurf1 = -np.sum(weights * (dXdxiGP[2] * dXdetaGP[0] - dXdxiGP[0] * dXdetaGP[2]))
    NSurf2 = -np.sum(weights * (dXdxiGP[0] * dXdetaGP[1] - dXdxiGP[1] * dXdetaGP[0]))
    return NSurf0, NSurf1, NSurf2


def eval_nsurf(XGeo: np.ndarray, Vdm: np.ndarray, DGP: np.ndarray, weights: np.ndarray) -> np.ndarray:
    """ Evaluate the surface integral for normals over a side of an element
    """
    # Change basis to Gauss points
    xGP      = change_basis_2D(Vdm, XGeo)

    # Compute derivatives at all Gauss points
    # dXdxiGP  = np.tensordot(DGP, xGP, axes=(1, 1)).transpose(1, 0, 2)  # Shape: (3, N_GP+1, N_GP+1)
    # # dXdxiGP  = np.moveaxis(dXdxiGP , 0, 1).reshape(3, -1)              # Flatten for cross computation (slower)
    # dXdxiGP  = dXdxiGP .reshape(3, -1)                                 # Flatten for cross computation
    #
    # dXdetaGP = np.tensordot(DGP, xGP, axes=(1, 2)).transpose(1, 0, 2)  # Shape: (3, N_GP+1, N_GP+1)
    # # dXdetaGP = np.moveaxis(dXdetaGP, 0, 1).reshape(3, -1)              # Flatten for cross computation (slower)
    # dXdetaGP = dXdetaGP.reshape(3, -1)                                 # Flatten for cross computation

    # Compute derivatives at all Gauss points using matrix multiplications to avoid extra copies
    dXdxiGP  = np.empty_like(xGP)
    dXdetaGP = np.empty_like(xGP)
    DT       = DGP.T
    for k in range(3):
        dXdxiGP[ k] = DGP @ xGP[k]
        dXdetaGP[k] = xGP[k] @ DT

    # Compute the cross product at each Gauss point
    # VDMSize  = Vdm.shape[-1]
    # # nVec     = np.cross(dXdxiGP, dXdetaGP, axis=0)  # Shape: (3, N_GP*N_GP)
    # # > Manually compute cross product
    # nVec = np.empty_like(dXdxiGP)
    # nVec[0] = dXdxiGP[1] * dXdetaGP[2] - dXdxiGP[2] * dXdetaGP[1]
    # nVec[1] = dXdxiGP[2] * dXdetaGP[0] - dXdxiGP[0] * dXdetaGP[2]
    # nVec[2] = dXdxiGP[0] * dXdetaGP[1] - dXdxiGP[1] * dXdetaGP[0]
    # nVec     = nVec.reshape(3, VDMSize, VDMSize)    # Reshape to (3, N_GP+1, N_GP+1)

    # Compute the weighted normals
    # nVecW    = nVec * weights                       # Broadcast weights to shape (3, N_GP+1, N_GP+1)

    # Integrate over the Gauss points
    # return -np.sum(nVecW, axis=(1, 2))              # Sum over the last two axes

    # Compute the weighted cross product integral directly
    NSurf0, NSurf1, NSurf2 = eval_dotprod(dXdetaGP, dXdxiGP, weights)

    return np.array((NSurf0, NSurf1, NSurf2), dtype=xGP.dtype)


def check_sides(elem,
                # points   : np.ndarray,
                VdmEqToGP: np.ndarray,
                DGP      : np.ndarray,
                weights  : np.ndarray,
                # sides    : list
                failed_only: bool = False,
               ) -> Optional[list[tuple]]:
    """ Check if connected sides are watertight under high-order curving
    """
    results = None
    points  = mesh_vars.mesh.points
    elems   = mesh_vars.elems
    sides   = mesh_vars.sides
    nGeo    = mesh_vars.nGeo

    # Calculate the cube root of the element volume
    elemTol  = np.cbrt(np.prod(np.ptp(points[elem.nodes], axis=0)))
    elemType = elem.type

    for SideID in elem.sides:
        side   = sides[SideID]

        # Only connected sides and not small mortar sides
        # > Small mortar sides connect to big mortar side, so we will never match
        if side.connection is None or side.sideType < 0:
            continue

        # Big mortar side
        elif side.connection < 0:
            mortarType = abs(side.connection)
            # INFO: This should be faster but I could not confirm the speedup in practice
            # nSurf   = eval_nsurf(np.moveaxis( points[  nodes], 2, 0), VdmEqToGP, DGP, weights)
            # nSurf   = eval_nsurf(np.transpose(np.take(points,   nodes, axis=0), axes=(2, 0, 1)), VdmEqToGP, DGP, weights)
            idx     = elem.nodes[face_to_nodes(side.face, elemType, nGeo)]
            nSurf   = eval_nsurf(points[idx].transpose(2, 0, 1), VdmEqToGP, DGP, weights)

            # Calculate the L2 norm of the side and take the maximum
            sideTol = np.linalg.norm(nSurf, ord=2)
            tol     = np.maximum(elemTol, sideTol) * mesh_vars.tolInternal

            # Mortar sides are the following virtual sides
            nMortar = 4 if mortarType == 1 else 2
            nnbSurf = np.zeros((3,), dtype=float)
            for mortarSide in range(nMortar):
                # Get the matching side
                # INFO: This should be faster but I could not confirm the speedup in practice
                # nnbSurf += eval_nsurf(np.moveaxis(points[nbnodes], 2, 0), VdmEqToGP, DGP, weights)
                nbside   = sides[sides[SideID + mortarSide + 1].connection]
                nbelem   = elems[nbside.elemID]
                idx      = nbelem.nodes[face_to_nodes(nbside.face, nbelem.type, nGeo)]
                nnbSurf += eval_nsurf(np.transpose(points[idx], axes=(2, 0, 1)), VdmEqToGP, DGP, weights)

            # Check if side normals are within tolerance
            nSurfErr = np.sum(np.abs(nnbSurf + nSurf))
            success  = nSurfErr < tol

        # Internal side
        elif side.connection >= 0:
            # Only process the side with the smaller ID
            if SideID > side.connection:
                continue

            # Ignore the virtual mortar sides
            if side.locMortar is not None:
                continue

            # INFO: This should be faster but I could not confirm the speedup in practice
            # nSurf   = eval_nsurf(np.moveaxis( points[  nodes], 2, 0), VdmEqToGP, DGP, weights)
            idx     = elem.nodes[face_to_nodes(side.face, elemType, nGeo)]
            nSurf   = eval_nsurf(np.transpose(points[idx]), VdmEqToGP, DGP, weights)

            # Calculate the L2 norm of the side and take the maximum
            sideTol = np.linalg.norm(nSurf, ord=2)
            tol     = np.max((elemTol, sideTol)) * mesh_vars.tolInternal

            # Connected side
            nbside  = sides[side.connection]
            nbelem  = elems[nbside.elemID]
            # INFO: This should be faster but I could not confirm the speedup in practice
            # nnbSurf = eval_nsurf(np.moveaxis(points[nbnodes], 2, 0), VdmEqToGP, DGP, weights)
            idx     = nbelem.nodes[face_to_nodes(nbside.face, nbelem.type, nGeo)]
            nnbSurf = eval_nsurf(np.transpose(points[idx]), VdmEqToGP, DGP, weights)

            # Check if side normals are within tolerance
            nSurfErr = np.sum(np.abs(nnbSurf + nSurf))
            success  = nSurfErr < tol

        else:
            continue

        # If requested, only return errors
        if failed_only and success:
            continue

        # Lazily initialize results on first failure
        if results is None:
            results = []
        results.append((success, SideID, nSurf, nnbSurf, nSurfErr, tol))

    # Avoid creating empty lists on elem_results
    if results is None:
        return None if failed_only else []

    return results


def process_chunk(chunk) -> list:
    """Process a chunk of elements by checking surface normal orientation
    """
    # Only keep failures to reduce memory and avoid building large arrays of successes
    chunk_results = []
    for elem in chunk:
        elem_result = check_sides(elem,
                                   process_chunk.VdmEqToGP,  # pyright: ignore[reportFunctionMemberAccess] # ty: ignore[unresolved-attribute]
                                   process_chunk.DGP,        # pyright: ignore[reportFunctionMemberAccess] # ty: ignore[unresolved-attribute]
                                   process_chunk.weights,    # pyright: ignore[reportFunctionMemberAccess] # ty: ignore[unresolved-attribute]
                                   failed_only=True)
        # Append a lightweight sentinel (None) for successes, actual failure list otherwise
        chunk_results.append(elem_result)
    return chunk_results


def CheckWatertight() -> None:
    """ Check if the mesh is watertight
    """
    # Local imports ----------------------------------------
    import pyhope.output.output as hopout
    import pyhope.mesh.mesh_vars as mesh_vars
    from pyhope.basis.basis_basis import barycentric_weights, legendre_gauss_nodes
    from pyhope.basis.basis_basis import calc_vandermonde, polynomial_derivative_matrix
    from pyhope.common.common_parallel import run_in_parallel
    from pyhope.common.common_vars import np_mtp
    from pyhope.readintools.readintools import GetLogical
    # ------------------------------------------------------

    hopout.separator()
    hopout.info('CHECK WATERTIGHTNESS...')
    hopout.sep()

    checkWatertightness = GetLogical('CheckWatertightness')
    if not checkWatertightness:
        return None

    nGeo:      Final[int] = mesh_vars.nGeo

    # Compute the equidistant point set used by meshIO
    xEq:       Final[np.ndarray] = np.linspace(-1., 1., nGeo+1)
    wBaryEq:   Final[np.ndarray] = barycentric_weights(nGeo+1, xEq)

    xGP, wGP  = legendre_gauss_nodes(nGeo+1)
    DGP:       Final[np.ndarray] = polynomial_derivative_matrix(nGeo+1, xGP)
    VdmEqToGP: Final[np.ndarray] = calc_vandermonde(nGeo+1, nGeo+1, wBaryEq, xEq, xGP)

    # Compute the weights
    weights:   Final[np.ndarray] = np.outer(wGP, wGP)                   # Shape: (N_GP+1, N_GP+1)

    # Check all sides
    elems:     Final[list] = mesh_vars.elems
    sides:     Final[list] = mesh_vars.sides
    mesh:      Final             = mesh_vars.mesh
    points:    Final[np.ndarray] = mesh.points
    # points    = mesh_vars.mesh.points
    # checked   = np.zeros((len(sides)), dtype=bool)

    # Only consider hexahedrons
    if any(e.type % 100 != 8 for e in elems):
        elemTypes = list(set([e.type for e in elems if e.type % 100 != 8]))
        print(hopout.warn('Ignored element type: {}'.format(
            [re.sub(r"\d+$", "", mesh_vars.ELEMTYPE.inam[e][0]) for e in elemTypes]
        )))
        return

    # Prepare elements for parallel processing
    if np_mtp > 0:
        # Run in parallel with a chunk size
        # > Dispatch the tasks to the workers, minimum 10 tasks per worker, maximum 1000 tasks per worker
        res     = run_in_parallel(process_chunk,
                                  elems,
                                  chunk_size  = max(1, min(1000, max(10, int(len(elems)/(40.*np_mtp))))),  # noqa: E251
                                  initializer = init_worker,                                               # noqa: E251
                                  init_args   = (process_chunk, VdmEqToGP, DGP, weights))                  # noqa: E251
    else:
        res     = [elem for elem in elems if check_sides(elem,
                                                         VdmEqToGP, DGP, weights,
                                                         failed_only=True)]

    if len(res) > 0:
        # Flatten per-element results (skip None placeholders)
        results = tuple(result for elem_results in res if isinstance(elem_results, Iterable) and elem_results is not None
                               for result       in elem_results)  # noqa: E272

        # Compute total number of checked connections without materializing all results
        nconn = 0
        for SideID, side in enumerate(sides):
            # Only connected sides and not small mortar sides
            if side.connection is None or side.sideType < 0:
                continue
            # Big mortar side is counted once
            elif side.connection < 0:
                nconn += 1
            # Internal side: only count the canonical representative and ignore virtual mortar sides
            elif side.connection >= 0:
                if SideID > side.connection:
                    continue
                if side.locMortar is not None:
                    continue
                nconn += 1

        for result in cast(tuple[tuple], results):
            # Unpack the results
            side    = sides[result[1]]
            elem    = elems[side.elemID]
            nbside  = sides[side.connection]
            nbelem  = elems[nbside.elemID]

            nSurf, nbnSurf, nSurfErr, tol = result[2], result[3], result[4], result[5]

            nodes   =   elem.nodes[face_to_nodes(  side.face,   elem.type, nGeo)]
            nbnodes = nbelem.nodes[face_to_nodes(nbside.face, nbelem.type, nGeo)]

            print()
            # Check if side is oriented inwards
            errStr  =      'Side is oriented inwards!' if nSurfErr < 0 \
                  else 'Surface normals are not within tolerance {:9.6e} > {:9.6e}'.format(nSurfErr, tol)
            print(hopout.warn(errStr, length=len(errStr)+16))

            # Print the information
            strLen  = max(len(str(side.sideID+1)), len(str(nbside.sideID+1)))
            print(hopout.warn(f'> Element {  elem.elemID+1:>{strLen}}, Side {  side.face}, Side {  side.sideID+1:>{strLen}}'))  # noqa: E501
            print(hopout.warn('> Normal vector: [' + ' '.join('{:12.3f}'.format(s) for s in   nSurf) + ']'))                    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[  nodes[ 0,  0]]) + ']'))    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[  nodes[ 0, -1]]) + ']'))    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[  nodes[-1,  0]]) + ']'))    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[  nodes[-1, -1]]) + ']'))    # noqa: E271
            # print()
            print(hopout.warn(f'> Element {nbelem.elemID+1:>{strLen}}, Side {nbside.face}, Side {nbside.sideID+1:>{strLen}}'))  # noqa: E501
            print(hopout.warn('> Normal vector: [' + ' '.join('{:12.3f}'.format(s) for s in nbnSurf) + ']'))                    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[nbnodes[ 0,  0]]) + ']'))    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[nbnodes[ 0, -1]]) + ']'))    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[nbnodes[-1,  0]]) + ']'))    # noqa: E271
            print(hopout.warn('- Coordinates  : [' + ' '.join('{:12.3f}'.format(s) for s in points[nbnodes[-1, -1]]) + ']'))    # noqa: E271

        hopout.error(f'Watertightness check failed for {len(results)} / {nconn} connections!')

    # Run garbage collector to release memory
    gc.collect()
