#! /usr/bin/env python3
#-------------------------------*- python -*-----------------------------------
# ==  == ====== ====   ====    |
#                   \\     ||  | Multiphase Code Repository by HZDR
# ======   //   ||  || ===//   | Website: https://doi.org/10.14278/rodare.767
# ||  ||  //    ||  // || \\   | License: GPL-3.0-or-later
# ==  == ====== ====   ==  ==  |
#------------------------------------------------------------------------------
# License
#     This file is part of the Multiphase Code Repository by HZDR.
#
#     Copyright (C) 2025 by Helmholtz-Zentrum Dresden-Rossendorf e.V. (HZDR),
#     Website: https://hzdr.de
#
#     Multiphase Code Repository by HZDR is based on the free software for
#     computational fluid dynamics (CFD) from the OpenFOAM Foundation.
#     Copyright (C) 2025 by OpenFOAM Foundation, Website: https://openfoam.org
#
#     If you are interested in which files are original OpenFOAM Foundation
#     files, which OpenFOAM Foundation files were modified, and which files were
#     newly created, see FILES.md.
#
#     Multiphase Code Repository by HZDR 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.
#
#     Multiphase Code Repository by HZDR 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 the Multiphase Code Repository by HZDR. If not, see
#     <http://www.gnu.org/licenses/>.
#
#------------------------------------------------------------------------------

"""Sort iso-surface point data in 2D plane by their closest neighbourhood, e.g.
to plot 3D interface position with a closed curve. At first step, 3D data
is converted to 2D by cutting out a slice of finite thickness."""

import sys
import argparse
import numpy as np

#==============================================================================
#                                    MAIN
#==============================================================================

def main():
    """Control, which data to process and in what order the individual steps are
    done.
    """
    for fname in args.files:
        print(f'Processing file {fname}')

        # Read
        if args.two_dimensional:
            raw_data = read_data(fname, args.x_column, args.y_column)
        else:
            raw_data = read_data(fname, args.x_column, args.y_column,
                                 args.z_column)

        # Cut out slices
        if args.bounds_x is not None:
            raw_data = cut_slice(raw_data, args.x_column, args.bounds_x)

        if args.bounds_y is not None:
            raw_data = cut_slice(raw_data, args.y_column, args.bounds_y)

        if (not args.two_dimensional) and (args.bounds_z is not None):
            raw_data = cut_slice(raw_data, args.z_column, args.bounds_z)

        raw_data = list(raw_data[:, select_plane()])

        # Open output file
        oname = fname + '.sorted'

        first_time = True
        while np.shape(raw_data)[0] > 0:

            # Sort remaining data
            sorted_data, raw_data = sort_data(raw_data)

            # Write
            if args.verbose:
                print(f'  Saving sorted data to {oname}')

            if first_time:

                np.savetxt(oname, sorted_data)
                first_time = False

            else:

                with open(oname, "a", encoding='utf-8') as ofile:
                    ofile.write("\n")

                    for row in sorted_data:
                        ofile.write(f'{row[0]:.18e} {row[1]:.18e}\n')

    return 0


#==============================================================================
#                                  READ DATA
#==============================================================================

def read_data(fname, x_col, y_col, z_col=None):
    """Read data from file name and columns.
    """
    if args.verbose:
        print(f'  Reading unsorted data from {fname}')

    if z_col is None:

        if args.verbose:
            print(f'  Try to read 2D data from columns {x_col} {y_col}')

        return np.loadtxt(fname)[:, [x_col, y_col]]

    if args.verbose:
        print(f'  Try to read 3D data from columns {x_col} {y_col} {z_col}')

    try:
        return np.loadtxt(fname)[:, [x_col, y_col, z_col]]
    except IOError:
        sys.exit('Error: Cannot read 3D data. Wrong columns?')


#==============================================================================
#                                  CUT SLICE
#==============================================================================

def cut_slice(data, direction, bounds):
    """Cut out slice data with minimum and maximum coordinate value for one
    direction.
    """
    if args.verbose:
        print(f'  Cutting out slice for direction {direction}')

    indices = np.where((bounds[0] < data[:, direction])
                       & (data[:, direction] < bounds[1]))

    return data[indices]


#==============================================================================
#                                SELECT PLANE
#==============================================================================

def select_plane():
    """Set column indices for selected 2D plane.
    """
    if (args.two_dimensional) or (args.plane == 'xy'):
        plane_directions = [0, 1]
    elif args.plane == 'xz':
        plane_directions = [0, 2]
    else:
        plane_directions = [1, 2]

    return plane_directions


#==============================================================================
#                                  SORT DATA
#==============================================================================

def sort_data(raw_data):
    """Create new array with sorted data. Break, if smallest distance is lower
    than minimal distance.
    """
    if args.verbose:
        print(f'  Start sorting {len(raw_data)} points')

    # Get first point
    sorted_data = [raw_data[0]]
    raw_data.remove(sorted_data[0])

    # Save closest point to sorted_data and remove from raw_data
    while raw_data:

        if args.verbose:
            print(f'  {len(raw_data)} point(s) left')

        point1 = sorted_data[-1]
        min_distance = 0

        for index, point2 in enumerate(raw_data):
            distance = get_distance(point1, point2)

            if (min_distance == 0) or (distance < min_distance):
                min_distance = distance
                closest_index = index

        if args.maximum_distance < min_distance:
            print('  Multiple phase structures detected')
            break

        sorted_data.append(raw_data.pop(closest_index))

    # Closed loop condition
    if get_distance(sorted_data[0], sorted_data[-1]) < args.maximum_distance:

        if args.verbose:
            print('  Applying closed loop condition')

        sorted_data.append(sorted_data[0])

    elif args.verbose:
        print("  Closed loop condition not applied - distance to long")

    if args.verbose:
        print(f'  Sorted data contains {len(sorted_data)} points')

    return np.array(sorted_data), raw_data


#==============================================================================
#                                  DISTANCE
#==============================================================================

def get_distance(point_a, point_b):
    """Calculate euclidean distance.
    """
    return float(np.linalg.norm(point_b - point_a))


#==============================================================================
#                                  CALL MAIN
#==============================================================================

if __name__ == '__main__':

    # Parse arguments
    parser = argparse.ArgumentParser(description='Sort 2D point data by closest\
    neighbourhood, e.g. to plot 2D interface position with a closed curve.')

    # Add positional arguments
    parser.add_argument('files', nargs='+', type=str,
                        help='File(s), containing 3D interface position data')

    parser.add_argument('-2D', '--two-dimensional', action='store_true',
                        help='Process data, which is already 2D')

    parser.add_argument('-p', '--plane', type=str, default='xy',
                        choices=['xy', 'xz', 'yz'],
                        help='Plane to project data into; default: xy')

    parser.add_argument('-x', '--x-column', type=int, default=0,
                        help='Index of x data column; default: 0')
    parser.add_argument('-y', '--y-column', type=int, default=1,
                        help='Index of y data column; default: 1')
    parser.add_argument('-z', '--z-column', type=int, default=2,
                        help='Index of z data column; default: 2')

    parser.add_argument('-bx', '--bounds-x', type=float, nargs=2, default=None,
                        help='Boundary values in x direction')
    parser.add_argument('-by', '--bounds-y', type=float, nargs=2, default=None,
                        help='Boundary values in y direction')
    parser.add_argument('-bz', '--bounds-z', type=float, nargs=2, default=None,
                        help='Boundary values in z direction')

    parser.add_argument('-maxd', '--maximum-distance', type=float, default=1.0,
                        help='Maximum distance between to points to connect; \
                        default: 1')

    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Increase output verbosity')

    # Save arguments
    args = parser.parse_args()
    if args.verbose:
        print('Arguments:\n', args)

    # Run main
    main()

#------------------------------------------------------------------------------
