#! /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/>.
#
#------------------------------------------------------------------------------

"""Search for files matching a pattern in a given directory. If data structure
and first column data are identical all other columns are averaged over all
files. Useful to average over multiple lines data or over time data sets."""

import os
import sys
import argparse
import numpy as np
from matplotlib import pyplot as plt

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

def main():
    """Average data from all files in a given directory. Only account if shape
    of data and first column are identical.
    """

    # Remove output file to exclude it from averaging
    outfile = os.path.join(args.directory, f'average.{args.pattern}.dat')
    try:
        os.remove(outfile)
    except OSError:
        pass

    # Create a list of files to process
    files = list_files(args.directory, args.pattern, dirpattern)

    if not files:
        sys.exit('Error: No matching files found.')

    # Initiate the average data set
    num = 1
    if args.verbose:
        print(f'\nProcessing first file: {files[0]}')
    data = read_data(files[0])

    # Add remaining data sets to the average data set
    for file in files[1:]:
        if args.verbose:
            print(f'Processing file: {file}')
        idata = read_data(file)

        # Exclude different matrix shapes
        if np.shape(idata) != np.shape(data):
            if args.verbose:
                print(f'Data shape of file: {file}')
                print(f'Does not match with file: {files[0]}')
                print('Skipping file')
            continue

        # Exclude different first columns
        if (idata[:, 0] != data[:, 0]).any():
            if args.verbose:
                print(f'First column of file: {file}')
                print(f'Does not match with file: {files[0]}')
                print('Skipping file')
            continue

        # Sum up data
        data[:, 1:] += idata[:, 1:]
        num += 1

    # Normalise data
    data[:, 1:] /= num
    print(f'\nNumber of averaged files: {num}\n')

    # Write output file
    if args.skip_zero:
        write_file(outfile, data[1:, 0], data[1:, 1:])
    else:
        write_file(outfile, data[:, 0], data[:, 1:])

    # Directly show data
    if args.plot:
        show_data(data[:, 0], data[:, 1:], \
                  title=f'Average {args.pattern}', log=True)


#==============================================================================
#                                 LIST FILES
#==============================================================================

def list_files(path, pattern, dpattern):
    """Return a list of all files, located in the given path and matching the
    given pattern and dpattern.
    """
    if args.verbose:
        print(f'Creating file list for: {path}')
        print(f'File name matching pattern: {pattern}')
        print(f'Directory matching pattern: {dpattern}')

    files = []

    # Search for all files in path
    for root, dummy, allfiles in os.walk(path):
        if dpattern in root+'/':            # Directory matches dpattern
            if allfiles:          # Actually found one or more files
                for fname in allfiles:
                    if pattern in fname:    # File matches pattern
                        files.append(root+'/'+fname)

    return files


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

def read_data(fname):
    """Try to read the file fName and, if valid, return content as numpy array.
    Special care is taken, if parentheses occur in the data, which is the case
    for vector probes from the OpenFOAM 'probes' functionObject.
    """
    if args.verbose:
        print(f'Reading file: {fname}')

    # Get number of columns
    ncols = np.loadtxt(fname, dtype='str').shape[1]

    # Remove parentheses, apply to all available columns
    convs = {k: lambda s: np.float64(s.strip('()')) for k in range(0, ncols)}

    # Save data to numpy array
    data = np.loadtxt(fname, converters=convs)

    return data


#==============================================================================
#                                  WRITE FILE
#==============================================================================

def write_file(fname, xdata, ydata, header=''):
    """Save header and all data to file of given name.
    """
    if args.verbose:
        print(f'Writing file: {fname}')

    if len(ydata.shape) == 1:
        alldata = np.vstack((xdata, ydata)).transpose()

    else:
        ycols = ydata.shape[1]
        for dummy in range(ycols):
            alldata = np.vstack((xdata, ydata.transpose())).transpose()

    np.savetxt(fname, alldata, header=header, fmt='%.10e')


#==============================================================================
#                                  SHOW DATA
#==============================================================================

def show_data(xdata, ydata, xlabel='', ylabel='', title='', show=True,
              log=False):
    """Directly shows a plot of given data. Intended for development purposes.
    If you want to plot multiple data sets, call the function with show=False
    every time except for the last call. In the latter also give the labels if
    required.
    """

    if args.verbose:
        print('Plotting data')

    # Only one single column for ydata
    if len(ydata.shape) == 1:
        labeli = ''
        if log:
            plt.loglog(xdata[1:], ydata[1:], label=labeli)

        else:
            plt.plot(xdata, ydata, label=labeli)

    # Multiple columns for ydata
    else:
        ycols = ydata.shape[1]

        for ycol in range(ycols):
            labeli = f'Column {str(ycol+1)}'

            if log:
                plt.loglog(xdata[1:], ydata[1:, ycol], label=labeli)

            else:
                plt.plot(xdata, ydata[:, ycol], label=labeli)

    if show:
        if title:
            plt.title(title)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.legend(loc=3)
        plt.grid()
        plt.show()


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

if __name__ == '__main__':

    # Parse arguments
    parser = argparse.ArgumentParser(description='Average data from all files \
    in a given directory. Only account if data structure and first column are \
    identical.')

    # Add positional arguments
    parser.add_argument('-directory', default=os.getcwd(), \
        help='Directory (default is the current work directory)')
    parser.add_argument('-pattern', default='', \
        help='Pattern to filter file names with (take all files by default)')
    parser.add_argument('-directory-pattern', default=argparse.SUPPRESS, \
        help='Directory, which has to be part of the file path, e.g. \
        time directory (take all files by default)')
    parser.add_argument('-skip-zero', action='store_true', \
        help='Skip first data point (zero wave number / frequency)')
    parser.add_argument('-plot', action='store_true', \
        help='Directly plot result')
    parser.add_argument('-v', '--verbose', action='store_true', \
        help='Increase output verbosity')

    # Save arguments
    args = parser.parse_args()

    # Process dir pattern
    dirpattern = ''
    try:
        # Argument directory-pattern is provided by the user
        dirpattern = args.directory_pattern
    except AttributeError:
        # No directory-pattern argument is provided by the user
        dirpattern = ''

    # Run main
    main()

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