#! /usr/bin/env python3
"""------------------------------*- python -*----------------------------------
  ==  == ====== ====   ====    |
                    \\     ||  | HZDR Multiphase Case Collection for OpenFOAM
  ======   //   ||  || ===//   | Website: https://doi.org/10.14278/rodare.811
  ||  ||  //    ||  // || \\   | License: Creative Commons Attribution 4.0
  ==  == ====== ====   ==  ==  |
-------------------------------------------------------------------------------
License
    This file is part of the HZDR Multiphase Case Collection for OpenFOAM.

    Copyright (C) 2022 Helmholtz-Zentrum Dresden-Rossendorf

    The HZDR Multiphase Case Collection for OpenFOAM is licensed under the
    Creative Commons Attribution 4.0 International License. To view a copy of
    this license, visit http://creativecommons.org/licenses/by/4.0/ or send a
    letter to Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.

Description
    Process data from trackPhaseStructures functionObject, which describe a
    single structure, which is spatially split by periodic boundary conditions
    and reconstruct the quantities for one single structure. By default, the
    trajectory stays folded, which is necessary for computing radial profiles.
    Alternatively, the trajectory may be unfolded to compute a continuous
    trajectory, e.g., for purpose of plotting.

----------------------------------------------------------------------------"""

import sys
import argparse
import numpy as np

#==============================================================================
#                                  PARAMETERS
#==============================================================================

# Geometry
X_MIN = [-2., -2., 0.]
plength = [4., 4.]
pdirs = [0, 1]

# Data file
fname = 'postProcessing/trackBubble/0/trackBubble.dat'

# Data format after deleting first column (structure index)
t_col = 0
X_COLS = [1, 2, 3]
U_COLS = [4, 5, 6]
vol_col = 7


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

def main(args=None):
    """Main function"""

    # Read data and remove index column
    arr = np.loadtxt(fname)[:, 1:9]

    # Detect sequences
    seqs = sequences(arr, col=t_col)

    # Create results data, remove sequences, keep single structure times only
    arr_out = remove_sequences(arr, seqs)

    # Average quantities of multiple structures for each sequence
    arr_out = process_sequences(arr, arr_out, seqs)

    # Sort by time column
    arr_out = sort_by_column(arr_out, t_col)

    # Remove periodic jumps
    if args.unfold:
        for periodic_dir in pdirs:
            arr_out = remove_periodic_jumps(arr_out, periodic_dir)

    # Velocity magnitude
    arr_out = add_mag(arr_out, U_COLS)

    # Save data to file
    if args.unfold:
        suffix = '_processed'
    else:
        suffix = '_4radialProfiles'

    np.savetxt(fname+suffix, arr_out,
               header='t\t\tx\t\ty\t\tz\t\tu\t\tv\t\tw\t\tV\t\tU')


#==============================================================================
#                              REMOVE SEQUENCES
#==============================================================================

def remove_sequences(array, seqs):
    """Remove sequence lines from array.
    """
    for seq in reversed(seqs):
        beg = seq[0]
        end = seq[1] + 1
        array = np.delete(array, (range(beg, end)), axis=0)

    return array


#==============================================================================
#                               DETECT SEQUENCES
#==============================================================================

def sequences(array, col=0):
    """Return list of lists of start and end index of sequences of elements in
    given column of given array.
    """
    match_row_pairs = np.equal(array[:-1, col], array[1:, col])

    # Detect begin of sequences
    seq_start = np.logical_and(np.logical_not(match_row_pairs[:-1]),
                               match_row_pairs[1:])
    seq_first = np.where(seq_start)
    seq_first = [i+1 for i in seq_first][0]

    if match_row_pairs[0]:
        # Sequence starts at first element
        seq_first = np.insert(seq_first, 0, 0)

    # Detect end of sequences
    seq_end = np.logical_and(match_row_pairs[:-1],
                             np.logical_not(match_row_pairs[1:]))
    seq_last = np.where(seq_end)
    seq_last = [i+1 for i in seq_last][0]

    if match_row_pairs[-1]:
        # Sequence ends at last element
        seq_last = np.append(seq_last, np.shape(array)[0]-1)

    if len(seq_first) != len(seq_last):
        sys.exit('ERROR: Length of lists for sequence start and end indices \
                 not matching')

    return np.hstack((seq_first[:, np.newaxis], seq_last[:, np.newaxis]))


#==============================================================================
#                              PROCESS SEQUENCES
#==============================================================================

def process_sequences(array, result_array, seqs):
    """Average centre of gravity and velocity and sum up volume for all
    structures for every sequence.
    """
    for seq in seqs:

        # Get sub array for sequence only
        beg = seq[0]
        end = seq[1] + 1
        seq_arr = array[beg:end, :]

        # Get sequence time and total structure volume
        time = seq_arr[0, t_col]
        volume = np.sum(seq_arr[:, vol_col])

        # Periodic boundaries
        for i in pdirs:
            x_vals = seq_arr[:, X_COLS[i]]
            x_diff_max = np.max(x_vals) - np.min(x_vals)
            if x_diff_max > plength[i]/2:
                seq_arr[:, X_COLS[i]][x_vals > X_MIN[i] + plength[i]/2] \
                -= plength[i]

        # Calculate centre of gravity and velocity
        seq_x = np.average(seq_arr[:, X_COLS], axis=0,
                           weights=seq_arr[:, vol_col])
        for i in pdirs:
            if seq_x[i] < X_MIN[i]:
                seq_x[i] += plength[i]

        seq_u = np.average(seq_arr[:, U_COLS], axis=0,
                           weights=seq_arr[:, vol_col])

        # Format data
        seq_data = np.hstack((np.array([time]), seq_x, seq_u,
                              np.array([volume])))
        result_array = np.append(result_array, seq_data[np.newaxis, :], axis=0)

    return result_array


#==============================================================================
#                               SORT BY COLUMN
#==============================================================================

def sort_by_column(array, column):
    """Sort given array by given column.
    """
    return array[np.argsort(array[:, column]), :]


#==============================================================================
 #                                  ADD MAG
#==============================================================================

def add_mag(array, columns):
    """Add a column with magnitude of components in given columns.
    """
    t_mag = array[:, columns[0]]**2

    for col in columns[1:]:
        t_mag += array[:, col]**2

    t_mag = t_mag**0.5

    return np.hstack((array, t_mag[:, np.newaxis]))


#==============================================================================
#                            REMOVE PERIODIC JUMPS
#==============================================================================

def remove_periodic_jumps(array, directory):
    """Detect jumps in position due to periodic boundaries in given directory
    and remove them.
    """
    length = plength[directory]
    column = X_COLS[directory]
    x_vals = array[:, column]
    delta_xs = x_vals[1:] - x_vals[:-1]

    # Detect jumps
    jump_poss = np.where(abs(delta_xs) > length/2.)

    for jump_pos in jump_poss[0]:
        delta_x = delta_xs[jump_pos]

        # Calculate correction
        jump_factor, jump_residual = divmod(delta_x, length)
        if jump_residual > length/2.:
            jump_factor += 1
            jump_residual -= length

        # Remove jump
        array[jump_pos+1:, column] = array[jump_pos+1:, column] - \
                                     jump_factor*length

    return array


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

if __name__ == '__main__':

    # Parse arguments
    PARSER = argparse.ArgumentParser(description='Reconstruct single structure \
    from structures split by periodic boundary conditions.')

    # Add positional arguments
    PARSER.add_argument('-u', '--unfold', action='store_true', \
        help='Unfold trajectory for continuous line')

    # Save arguments
    ARGS = PARSER.parse_args()

    main(ARGS)
