# This file is part of eddylicious
# (c) Timofey Mukha
# The code is released under the GNU GPL Version 3 licence.
# See LICENCE.txt and the Legal section in the User Guide for more information

"""Functions for generating inflow velocity fields by
interpolation of data (from a precursor).

"""
from __future__ import print_function
from __future__ import division
import numpy as np
from mpi4py import MPI
from scipy.interpolate import NearestNDInterpolator 
from .helper_functions import chunks_and_offsets
from ..writers.ofnative_writers import write_velocity_to_ofnative
from ..writers.hdf5_writers import write_velocity_to_hdf5

__all__ = ["interpolation_generate"]

def interpolation_generate(readerFunction,
                           writer, writePath,
                           dt, t0, tEnd, timePrecision,
                           points,
                           pointsInfl,
                           idxPrec,
                           times):
    """Generate the the inflow velocity interpolation.

    This function will take some precursor data and interpolate it
    on the grid at the inflow patch.

    Parameters
    ----------
    readerFunction : function
        The function to use for reading in data, generated by the
        reader. Should contain the reader's name in the attribute
        "reader".
    writer: str
        The writer that will be used to save the values of the velocity
        field.
    writePath : str
        The path for the writer.
    dt : float
        The time-step to be used in the simulation. This will be used to
        associate a time-value with the produced velocity fields.
    t0 : float
        The starting time to be used in the simulation. This will be
        used to associate a time-value with the produced velocity.
    timePrecision : int
        Number of points after the decimal to keep for the time value.
    tEnd : float
        The ending time for the simulation.
    points : ndarray
        A 2d array containing the values the points of the
        source geometry or their Delaunay triangulation.
    pointsInfl : ndarray
        A 2d array containing the values the points of the
        inlet geometry.
    idxPrec : ndarray
        Indices for filtering the read-in velocity
    times : list of floats or strings
        The times for which the velocity field was sampled in the
        precursor simulation.

    """
    # Grab info regarding parallelization
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    nProcs = comm.Get_size()

    # Get the total amount of rescalings to be done
    size = int((tEnd-t0)/dt+1)
    
    # Calculate the amount of rescalings each processor is responsible for
    [chunks, offsets] = chunks_and_offsets(nProcs, size)

    # Perform the rescaling
    for i in range(chunks[rank]):
        t = t0 + dt*i + dt*int(offsets[rank])
        t = float(("{0:."+str(timePrecision)+"f}").format(t))
        position = int(offsets[rank]) + i

        if (rank == 0) and (np.mod(i, int(chunks[rank]/10)) == 0):
            print("     Interpolated about "+str(int(i/chunks[rank]*100))+"%")

        # Read U data
        if readerFunction.reader == "foamFile":
            assert position < len(times)
            [uX, uY, uZ] = readerFunction(times[position])
        elif readerFunction.reader == "hdf5":
            assert position < len(times)
            [uX, uY, uZ] = readerFunction(position)
        else:
            raise ValueError("Unknown reader")

        uXInterp = NearestNDInterpolator(points.points, uX[idxPrec])
        uYInterp = NearestNDInterpolator(points.points, uY[idxPrec])
        uZInterp = NearestNDInterpolator(points.points, uZ[idxPrec])

        uXInfl = uXInterp(pointsInfl)
        uYInfl = uYInterp(pointsInfl)
        uZInfl = uZInterp(pointsInfl)

        # Write
        if writer == "ofnative":
            write_velocity_to_ofnative(writePath, t, uXInfl, uYInfl, uZInfl)
        elif writer == "hdf5":
            write_velocity_to_hdf5(writePath, t, uXInfl, uYInfl, uZInfl,
                                   position)
        else:
            raise ValueError("Unknown writer")
