#! /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
    From data in a plane, obtained with surfaces functionObject, and the centre
    point, obtained with trackPhaseStructures functionObject, compute radial
    profiles, averaged over time and angle phi.

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

#pylint: disable=invalid-name,too-many-locals,too-many-statements

import sys
import os
import math
import argparse
import numpy as np

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

NX = 60.        # Number of grid points
LX = 4.         # Length of domain
OFFSET = [2, 2] # Offset of origin
TOL = 5.e-6     # Tolerance to match cell centre positions

DX = LX/NX

R_MAX = 2.
R_RES = 0.4*DX
PHI_RES = 2.5*math.pi/180.

CWD = os.getcwd()

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

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

    #---Get line coordinates
    radii = np.arange(0, R_MAX, R_RES)
    phis = (np.arange(0, 2*math.pi, PHI_RES)).tolist()

    #---Load bubble position data
    filename='trackBubble.dat_4radialProfiles'
    filedir=f'{CWD}/postProcessing/trackBubble/{args.zero_time:g}'
    cogs = np.loadtxt(f'{filedir}/{filename}')

    #---Get times
    path = f'postProcessing/{args.surface}'
    ofile = f'radialVelocityProfile_{args.surface}.dat'
    times = sorted(os.listdir(os.path.join(CWD,path)), key=float)

    #---Shorten times if cogs contains less data
    if cogs[-1, 0] < float(times[-1]):
        float_times = np.array(times).astype(np.float)
        t_max = find_index(float_times, cogs[-1, 0])
        times = times[:t_max+1]

    #---Start from a later time step
    i_0 = find_index(np.asarray([float(i) for i in times]), args.start_time)
    times = times[i_0:]

    #---Average
    u_sum = np.zeros((len(radii), 3))
    number_sum = np.zeros((len(radii)))
    for time in times:

        fname = f'{CWD}/{path}/{time}/cutPlane.xy'

        print(f'  Processing mean velocity for time {time}')

        #---Load fluid data
        data = np.loadtxt(fname)[:,(0,1,2,4,5,6)]
        data = data[np.lexsort((data[:, 1], data[:, 0]))]
        _, vel = rearrange_data(data, OFFSET)

        #---Get bubble position
        t_index = find_index(cogs[:, 0], time)
        cog = cogs[t_index, :][1:3]

        for phi in phis:
            for i in range(np.shape(radii)[0]):
                radius = radii[i]

                point = periodic(cog + cartes_coordinates(radius, phi), -2, 2)
                index = [c_to_idx(point[0], OFFSET[0]), \
                         c_to_idx(point[1], OFFSET[1])]

                vel_loc = vel[index[0], index[1]]
                vel_loc_rot = rotate(vel_loc, phi)
                u_sum[i] += vel_loc_rot

                number_sum[i] += 1

    u_sum /= number_sum[:, np.newaxis]
    print('\n')

    #---Fluctuations
    vel_fluc_sum = np.zeros((len(radii), 3))
    cross_corr_sum = np.zeros((len(radii), 3))
    k_sum = np.zeros((len(radii)))
    number_sum = np.zeros((len(radii)))
    for time in times:
        print(f'  Processing higher order statistics for time {time}')

        #---Load fluid data
        fname = f'{CWD}/{path}/{time}/cutPlane.xy'

        data = np.loadtxt(fname)[:,(0,1,2,4,5,6)]
        data = data[np.lexsort((data[:, 1], data[:, 0]))]
        _, vel = rearrange_data(data, OFFSET)

        #---Get bubble position
        t_index = find_index(cogs[:, 0], time)
        cog = cogs[t_index, :][1:3]

        for phi in phis:
            for i in range(np.shape(radii)[0]):
                radius = radii[i]

                point = periodic(cog + cartes_coordinates(radius, phi), -2, 2)
                index = [c_to_idx(point[0], OFFSET[0]), \
                         c_to_idx(point[1], OFFSET[1])]

                vel_loc = vel[index[0], index[1]]
                vel_loc_rot = rotate(vel_loc, phi)

                vel_fluc = vel_loc_rot - u_sum[i]

                vel_fluc_sum[i] += vel_fluc**2
                cross_corr_sum[i, 0] += (vel_fluc[0])*(vel_fluc[1])
                cross_corr_sum[i, 1] += (vel_fluc[0])*(vel_fluc[2])
                cross_corr_sum[i, 2] += (vel_fluc[1])*(vel_fluc[2])
                k_sum[i] += 0.5*(vel_fluc[0]**2 + vel_fluc[1]**2 + \
                            vel_fluc[2]**2)
                number_sum[i] += 1

    vel_fluc_sum /= number_sum[:, np.newaxis]
    cross_corr_sum /= number_sum[:, np.newaxis]
    k_sum /= number_sum[:]
    print('\n')

    #---Write-results
    print('Writing data\n')
    radii_vel_vel_flux = np.hstack((radii[:, np.newaxis], \
                          u_sum, \
                          vel_fluc_sum, \
                          cross_corr_sum, \
                          k_sum[:, np.newaxis]))
    np.savetxt(f'{CWD}/postProcessing/{ofile}', radii_vel_vel_flux, header="""
    r\tUr\tUphi\tUz\tUrPrime\tUphiPrime\tUzPrime\ttaurphi\ttaurz\ttauphiz\tk""")


#==============================================================================
#                                  FIND INDEX
#==============================================================================

def find_index(array, value):
    """Return the index of an array, where its element, closest to value, is
    located.
    """
    return np.argmin(np.abs(array - float(value)))


#==============================================================================
#                               ARRANGE MATRICES
#==============================================================================

def rearrange_data(data, offset=None, delta_x=DX):
    """Select data at cell centre positions (*.5) and save them into an array
    of the shape (xIndex, yIndex) to be able to access the data quickly by
    its positional coordinate.
    """
    idx = (data[:, :2]+offset)//delta_x

    if offset is None:
        offset = [0, 0]

    #---Extract Cell Values
    indices_pos = []

    for i in range(np.shape(idx)[0]):
        # Distinct between point and cell centre locations

        if ((abs(idx[i, 0] - round(idx[i, 0])) < TOL) and
            (abs(idx[i, 1] - round(idx[i, 1])) < TOL)):
            # modulo (%) gives one wrong match for pos
            indices_pos.append(i)

    num_cells = int(NX) + 1

    pos = np.zeros((num_cells, num_cells, 2))
    vel = np.zeros((num_cells, num_cells, 3))
    for index_pos in indices_pos:
        # Save coordinates (pos) and velocity (vel) in array in the
        # structure x,y
        idloc = [c_to_idx(data[index_pos, 0], offset=offset[0]), \
                 c_to_idx(data[index_pos, 1], offset=offset[1])]

        pos[idloc[0], idloc[1], 0] = data[index_pos, 0]
        pos[idloc[0], idloc[1], 1] = data[index_pos, 1]

        vel[idloc[0], idloc[1], 0] = data[index_pos, 3]
        vel[idloc[0], idloc[1], 1] = data[index_pos, 4]
        vel[idloc[0], idloc[1], 2] = data[index_pos, 5]

    return pos, vel


#==============================================================================
#                       CONVERT COORDINATES AND INDICES
#==============================================================================

def c_to_idx(pos, offset=0, delta_x=DX):
    """Convert coordinate to index.
    """
    # format fixes precision problems
    return int((float(format(pos, '.8f')) + offset)/delta_x + 0.5)


#==============================================================================
#                                 PERIODIC BC
#==============================================================================

def periodic(actual_pos, lower_bound=0, upper_bound=1):
    """Return a value or its periodic twin inside the boundaries, if it is
    outside the given boundaries.
    """
    length = upper_bound - lower_bound

    new_pow = np.zeros((2))

    for i in [0, 1]:

        if actual_pos[i] < lower_bound:
            new_pow[i] = actual_pos[i] + length

        elif actual_pos[i] >= upper_bound:
            new_pow[i] = actual_pos[i] - length

        else:
            new_pow[i] = actual_pos[i]

    return new_pow


#==============================================================================
#                        POLAR TO CARTESIAN COORDINATES
#==============================================================================

def cartes_coordinates(radius, phi):
    """Convert polar coordinates to cartesian ones.
    """
    return np.array([radius*math.cos(phi), radius*math.sin(phi)])


def normal_vector(phi):
    """Vector of magnitude unity in the direction of phi.
    """
    return np.array([math.cos(phi), math.sin(phi)])


#==============================================================================
#                          COORDINATE TRANSFORMATION
#==============================================================================

def rotate(actual_direction, phi):
    """Transform coordinates into a coordinate system, which is rotated by angle
    phi in the direction of z.
    """
    transformation_matrix = np.array([[math.cos(phi), -math.sin(phi), 0], \
                  [math.sin(phi), math.cos(phi), 0], \
                  [0, 0, 1]])

    return np.dot(actual_direction, transformation_matrix)


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

if __name__ == '__main__':

    # Parse arguments
    PARSER = argparse.ArgumentParser(description='Computes radial profiles \
    from surfaces and phaseStructureProperties data, averaged over all times  \
    and angles phi.')

    # Add positional arguments
    PARSER.add_argument('-s', '--surface', type=str, default='surfaces', \
        help='Name of surface functionObject')

    PARSER.add_argument('-st', '--start-time', type=float, default=0, \
        help='First time step to consider')

    PARSER.add_argument('-zt', '--zero-time', type=float, default='0', \
        help='Start time of simulation - name of postProcessing directory')

    # Save arguments
    ARGS = PARSER.parse_args()

    # Check data source path
    PATH = f'/postProcessing/{ARGS.surface}'
    try:
        os.listdir(CWD + PATH)
    except IOError:
        print(f'Directory not found: "{PATH}"')
        sys.exit()

    main(ARGS)
