# A component of the BASEmesh pre-processing toolkit.
# Copyright (C) 2020  ETH Zürich

# This program 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.

# This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.

"""Main wrapper class and utility code."""

import datetime
import os
import platform
import stat
import subprocess
from typing import (Any, Callable, Collection, Dict, Iterable, List, Optional,
                    Tuple)
from .log import logger
from .markers import HoleMarker, RegionMarker
from .pslg import PSLGNode, PSLGSegment
from .triangle_io import write_nodes, write_poly

# Path to the Triangle installation
_TRIANGLE_PATH = 'bin/'


class Triangle:
    """Wrapper class for Jonathan Shewchuk's "Triangle".

    This class is responsible for creating the input files, spawning
    the "Triangle" subprocesses, and eventually parsing the output.

    See the triangulate_nodes and triangulate_poly for usage
    instructions.

    Attributes
    ----------
    executable : str
        The executable chosen for the current platform, only provided
        for debug purposes
    """

    def __init__(self) -> None:
        """Initialize the Triangle namespace class.

        This mainly sets the internal executable for the current
        system platform.
        """
        logger.info('Initializing Triangle...')
        self.executable = self._find_executable()

    @staticmethod
    def _find_executable() -> str:
        """Retrieve the appropriate binary for the current system."""
        # Microsoft NT (Windows)
        if os.name == 'nt':
            filename = 'triangle_32.exe'
        # Unix
        elif os.name == 'posix':
            # Linux
            if '32bit' in platform.architecture():
                filename = 'triangle_linux_32'
            else:
                filename = 'triangle_linux_64'
        # Other (unsupported)
        else:
            raise RuntimeError(f'Unsupported operating system "{os.name}"')

        # Look for the appropriate executable
        folder = os.path.dirname(__file__)
        triangle_path = os.path.join(folder, _TRIANGLE_PATH, filename)
        if not os.path.isfile(triangle_path):
            raise RuntimeError(
                f'Unable to locate executable "{triangle_path}"')
        # Set executable flag for unix systems
        if os.name == 'posix':
            os.chmod(triangle_path, stat.S_IEXEC)
        logger.info('Compatible Triangle binary found: %s', filename)
        return triangle_path

    @staticmethod
    def _get_base_name() -> str:
        """Return a default triangle input filename.

        This name will generally not be visible to the user and is
        mainly intended to be unique for debug purposes.
        """
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
        return f'triangle_{timestamp}'

    @staticmethod
    def _get_output_names(
            filename: str, *, node_file: bool = True, ele_file: bool = True,
            poly_file: bool = True, edge_file: bool = False,
            neigh_file: bool = False) -> Tuple[str, ...]:
        """Return the output filenames for a given input filename.

        This function will take in a given input filename or path and
        return the expected output filename generated by Triangle.

        Parameters
        ----------
        filename : str
            The input file passed to Triangle
        node_file : bool, optional
            Whether to include NODE output, by default True
        ele_file : bool, optional
            Whether to include ELE output, by default True
        poly_file : bool, optional
            Whether to include POLY output, by default True
        edge_file : bool, optional
            Whether to include EDGE output, by default False
        neigh_file : bool, optional
            Whether to include NEIGH output, by default False

        Returns
        -------
        Tuple[str, ...]
            The output file names that will be generated by Triangle
        """
        # Strip off any path information
        base_name = os.path.basename(filename)
        # Right-split the base name up to two times, one for the iteration
        # number, the other for the file extension
        items = base_name.rsplit('.', 2)
        iter_num = 0
        if len(items) >= 3:
            base_name = items[0]
            # If the second word is an int, treat it as an iteration number,
            # otherwise assume it's part of the filename.
            try:
                iter_num = int(items[1])
            except ValueError:
                base_name += '.' + items[1]
        else:
            base_name = items[0]
        # Increment iteration number
        iter_num += 1
        # Determine the output extensions to include
        extensions = {'node': node_file, 'ele': ele_file, 'poly': poly_file,
                      'edge': edge_file, 'neigh': neigh_file}
        output_files: Tuple[str, ...] = tuple(
            f'{base_name}.{iter_num}.{k}' for k, v in extensions.items() if v)
        return output_files

    @staticmethod
    def _move_to_current(dir_: str, *args: str) -> None:
        """Move any files in *args to the current working directory.

        Parameters
        ----------
        dir_ : str
            The directory containing the input files
        *args: str
            The filenames to move to the CWD
        """
        for filename in args:
            temp_name = os.path.join(dir_, filename)
            os.replace(temp_name, filename)

    def triangulate_nodes(self, nodes: Collection[PSLGNode],
                          segments: Optional[Collection[PSLGSegment]] = None,
                          node_attr_keys: Optional[List[str]] = None,
                          *, dir_: str = '', **kwargs: Any) -> Tuple[str, str]:
        """Triangulate the given input nodes.

        This effectively creates an elevation mesh of the input points.

        Mesh boundary polygons must be split into lines and included in
        the segments argument.

        Any extra keyword arguments will be passed on to the run()
        method.

        Parameters
        ----------
        nodes : Collection[PSLGNode]
            The input nodes to process
        dir_ : str, optional
            The directory to run Triangle in, by default '.'
        segments : Collection[PSLGSegment], optional
            Any input segments to enforce during triangulation, by
            default None
        node_attr_keys : List[str], optional
            The list of node attributes to use during triangulation, by
            default None
        **kwargs: Any
            Additional keyword arguments to pass to the run() method

        Returns
        -------
        Tuple[str, str]
            The Triangle output files
        """
        base_name = os.path.join(dir_, self._get_base_name())

        ext = 'node' if segments is None else 'poly'
        input_name = f'{base_name}.{ext}'
        input_file = os.path.join(dir_, input_name)
        if segments is None:
            write_nodes(input_file, nodes, node_attr_keys=node_attr_keys)
        else:
            write_poly(input_file, nodes=nodes, segments=segments,
                       node_attr_keys=node_attr_keys)
        try:
            self.run(input_file, input_is_poly=segments is not None,
                     no_output_poly=True, redirect_stdout=_inline_print,
                     **kwargs)
        except subprocess.CalledProcessError as err:
            print('An error occurred during the execution of Triangle')
            raise err
        node_file, ele_file = self._get_output_names(
            base_name, poly_file=False)
        return os.path.join(dir_, node_file), os.path.join(dir_, ele_file)

    def triangulate_poly(self, nodes: Collection[PSLGNode],
                         segments: Optional[Collection[PSLGSegment]] = None,
                         holes: Optional[Iterable[HoleMarker]] = None,
                         regions: Optional[Iterable[RegionMarker]] = None,
                         *, dir_: str = '',
                         node_attr_keys: Optional[List[str]] = None,
                         use_matid: bool = True, **kwargs: Any
                         ) -> Tuple[str, str]:
        """Triangulate the given input data.

        This supports all the functionality of triangulate_nodes, but
        also provides access to segment definitions, holes and region
        markers.

        Parameters
        ----------
        nodes : Collection[PSLGNode]
            Input nodes to triangulate
        dir_ : str, optional
            The directory to run Triangle in, by default '.'
        segments : Collection[PSLGSegment], optional
            Input segments to enforce during triangulation, by default
            None
        holes : Iterable[HoleMarker], optional
            Optional hole markers, by default None
        regions : Iterable[RegionMarker], optional
            Optionaln region markers, by default None
        node_attr_keys : List[str], optional
            The node attribute keys to use during triangulation, by
            default None
        use_matid : bool, optional
            Whether to include material IDs, by default True

        Returns
        -------
        Tuple[str, str]
            Triangle output files
        """
        base_name = self._get_base_name()

        input_name = f'{base_name}.poly'
        input_file = os.path.join(dir_, input_name)
        extra_args: Dict[str, Any] = {}
        if segments is not None:
            extra_args['segments'] = segments
        if holes is not None:
            extra_args['holes'] = holes
        reg_area = False
        if regions:
            extra_args['regions'] = regions
            reg_area = True
        write_poly(input_file, nodes=nodes, node_attr_keys=node_attr_keys,
                   include_matid=use_matid, **extra_args)
        try:
            self.run(input_file, input_is_poly=True,
                     no_output_poly=True, redirect_stdout=_inline_print,
                     use_region_attributes=use_matid,
                     use_region_areas=reg_area,
                     **kwargs)
        except subprocess.CalledProcessError as err:
            print('An error occurred during the execution of Triangle.')
            raise err
        node_file, ele_file = self._get_output_names(
            base_name, poly_file=False)
        return os.path.join(dir_, node_file), os.path.join(dir_, ele_file)

    def run(self, input_file: str, *, suppress_output: bool = False,
            redirect_stdout: Optional[Callable[[str], None]] = None,
            flags: Optional[Dict[str, Any]] = None, **kwargs: Any) -> None:
        """Run Triangle using the provided input parameters.

        This is the most low-level interface for interfacing with
        "Triangle". Note that this file does not automatically generate
        its own input files.

        Parameters
        ----------
        input_file : str
            The input POLY or NODE file to run through Triangle
        suppress_output : bool, optional
            Whether to suppress all Triangle output during execution,
            by default False
        redirect_stdout : Callable[[str], None], optional
            A callable to pass any Triangle output to, rather than
            printing it, by default None
        flags : Dict[str, Any], optional
            Custom Triangle command line flags; flags defined here will
            overwrite those defined by other sources, by default None

        Raises
        ------
        ValueError
            Raised if the input file is neither a NODE nor POLY file
        subprocess.CalledProcessError
            Raised if Triangle returns a non-zero exit code
        """
        # Validate input filename
        if os.path.splitext(input_file)[1].lower() not in ('.poly', '.node'):
            raise ValueError('Input file must be a NODE or POLY file')
        # Extract custom triangle arguments
        custom_args = kwargs.pop('custom_args', '')
        # Generate the arguments and command string for Triangle execution
        flag_dict = dict(kwargs)
        if flags is not None:
            flag_dict.update(flags)
        flag_str = self._parse_flags(**flag_dict)
        command = f'{self.executable} -{flag_str}{custom_args} {input_file}'
        # Instantiate a new subprocess
        logger.info('Running Triangle: %s', command)
        kwargs = {}
        if os.name == 'nt':
            kwargs.update({'creationflags': subprocess.CREATE_NO_WINDOW})
        else:
            kwargs.update({'shell': True})
        process = subprocess.Popen(command, stdout=subprocess.PIPE,
                                   universal_newlines=True, **kwargs)
        logger.debug('Triangle subprocess spawned (PID %d)', process.pid)
        try:
            # Process subprocess output as it comes in
            for line in iter(process.stdout.readline, ''):  # type: ignore
                line = str(line)
                logger.debug('Subprocess: %s', line.rstrip())
                if suppress_output:
                    continue
                if redirect_stdout:
                    redirect_stdout(line)
        finally:
            process.stdout.close()  # type: ignore
        # Wait for subprocess completion and raise appropriate errors
        return_code = process.wait()
        logger.debug(
            'Triangle subprocess exited with exit status %d', return_code)
        if return_code:
            raise subprocess.CalledProcessError(return_code, command)

    @staticmethod
    def _parse_flags(**kwargs: str) -> str:
        """Convert input keyword arguments to "Triangle" flags.

        This method parses any number of input kwargs and returns the
        appropriate Triangle comand string.

        Returns
        -------
        str
            The Triangle command string

        Raises
        ------
        ValueError
            Raised if unexpected keys are found in the input
        """
        # The following dict maps the argument-less command line switches to
        # simplify the elif chain below.
        switches = {
            'input_is_poly': 'p',
            'conforming_delaunay': 'D',
            'jettison_nodes': 'j',
            'write_edges': 'e',
            'write_neigh': 'n',
            'no_boundary_in_output': 'B',
            'no_iteration_numbers': 'I',
            'ignore_holes': 'O',
            'no_steiner_on_boundary': 'Y',
            'no_steiner_on_segment': 'YY',
            'quiet': 'Q',
            'verbose': 'V',
            'keep_convex_hull': 'c',
            # 'refine': 'r',
            'no_output_node': 'N',
            'no_output_ele': 'E',
            'no_output_poly': 'P',
            'use_region_areas': 'a',
            'use_region_attributes': 'A'
        }
        flags = ''
        for key, value in kwargs.items():
            # Basic switches
            if key in switches:
                if value:
                    flags += switches[key]
                continue
            # Switches with arguments
            if key == 'min_angle':
                flags += f'q{value}'
            elif key == 'max_area':
                flags += f'a{value}'
            elif key == 'max_steiner':
                flags += f'S{value}'
            elif key == 'debug_level':
                flags += 'V' * int(value)
            # Fallback
            else:
                raise ValueError(f'Unknown triangle flag: {key}={value}')
        return flags


def _inline_print(string: str) -> None:
    """Like the default print, but without automatic newlines.

    Parameters
    ----------
    string : str
        The string to print
    """
    print(string, end='')
