#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import arrayfire as af
import numpy as np

def compute_moments(self, moment_name, f=None):
    """
    Used in computing the moments of the distribution function.
    The moment definitions which are passed to physical system
    are used in computing these moment quantities.

    Parameters
    ----------

    moments_name : str
                   Pass the moment name which needs to be computed.
                   It must be noted that this needs to be defined by the
                   user under moments under src and passed to the 
                   physical_system object.
    
    f: af.Array
       Pass this argument as well when you want to compute the 
       moments of the input array and not the one stored by the state vector
       of the object.

    Examples
    --------
    
    >> solver.compute_moments('density')

    The above line will lookup the definition for 'density' and calculate the same
    accordingly
    """
    N_g_p = self.N_ghost_p
    
    # Considering p-arrays non-inclusive of the ghost zones:
    if(N_g_p != 0):
    
        p1 = af.flat(af.moddims(self.p1_center,
                                self.N_p1 + 2 * N_g_p,
                                self.N_p2 + 2 * N_g_p,
                                self.N_p3 + 2 * N_g_p,
                               )[N_g_p:-N_g_p, N_g_p:-N_g_p, N_g_p:-N_g_p]
                    ) 
        p2 = af.flat(af.moddims(self.p2_center,
                                self.N_p1 + 2 * N_g_p,
                                self.N_p2 + 2 * N_g_p,
                                self.N_p3 + 2 * N_g_p,
                               )[N_g_p:-N_g_p, N_g_p:-N_g_p, N_g_p:-N_g_p]
                    ) 
        p3 = af.flat(af.moddims(self.p3_center,
                                self.N_p1 + 2 * N_g_p,
                                self.N_p2 + 2 * N_g_p,
                                self.N_p3 + 2 * N_g_p,
                               )[N_g_p:-N_g_p, N_g_p:-N_g_p, N_g_p:-N_g_p]
                    ) 

    else:
        
        p1 = self.p1_center
        p2 = self.p2_center
        p3 = self.p3_center

    if(f is None):
        try:
            N_q1 = self.f.shape[1]
            N_q2 = self.f.shape[2]

        except:
            N_q1 = N_q2 = 1

        if(N_g_p != 0):

            f = af.moddims(self.\
                           _convert_to_p_expanded(self.f)[N_g_p:-N_g_p, 
                                                          N_g_p:-N_g_p,
                                                          N_g_p:-N_g_p
                                                         ],
                           self.N_p1 * self.N_p2 * self.N_p3, 
                           N_q1, N_q2
                          )

        else:
            f = self.f
    
    else:

        try:
            N_q1 = f.shape[1]
            N_q2 = f.shape[2]
        except:
            N_q1 = N_q2 = 1

        if(N_g_p != 0):
            f = af.moddims(self.\
                           _convert_to_p_expanded(f)[N_g_p:-N_g_p, 
                                                     N_g_p:-N_g_p,
                                                     N_g_p:-N_g_p
                                                    ],
                           self.N_p1 * self.N_p2 * self.N_p3, 
                           N_q1, N_q2
                          )
    
    moment = af.broadcast(getattr(self.physical_system.moments, 
                                  moment_name
                                 ), f, p1, p2, p3, self.dp3 * self.dp2 * self.dp1
                         )

    af.eval(moment)
    return (moment)
