import logging
import numpy as np
from mpi4py import MPI
from .logger import Logger
import sympy
from .field import comp_laplacian


def comp_pressure(
    phi,
    phi_q,
    psi,
    hamiltonian,
    velocities,
    config,
    phi_fourier,
    phi_laplacian,
    phi_transfer,
    positions,
    bond_pr,
    angle_pr,
    comm=MPI.COMM_WORLD,
):
    """
    Computes total internal pressure of the system.
    Kinetic pressure is trivially calculated from the kinetic energy.
    Already computed bond and angle pressure terms are inserted into the
    total internal pressure.
    The field pressure equation is implemented.

    Parameters
    ----------
    phi : list[pmesh.pm.RealField], (M,)
        Pmesh :code:`RealField` objects containing discretized particle number
        density values on the computational grid; one for each particle type
        :code:`M`. Pre-allocated, but empty; any values in this field are discarded.
        Changed in-place. Local for each MPI rank--the full computaional grid
        is represented by the collective fields of all MPI ranks.
    hamiltonian : Hamiltonian
        Particle-field interaction energy handler object. Defines the
        grid-independent filtering function, :math:`H`.
    velocities : (N, D) numpy.ndarray
        Array of velocities of N particles in D dimensions.
    config : Config
        Configuration dataclass containing simulation metadata and parameters.
    phi_fourier : list[pmesh.pm.ComplexField], (M,)
        Pmesh :code:`ComplexField` objects containing discretized particle
        number density values in reciprocal space on the computational grid;
        one for each particle type. Pre-allocated, but empty; any values in
        this field are discarded Changed in-place. Local for each MPI rank--the
        full computaional grid is represented by the collective fields of all
        MPI ranks.
    phi_laplacian : list[pmesh.pm.RealField], (M, 3)
        Like phi, but containing the laplacian of particle number densities.
    phi_transfer : list[pmesh.pm.ComplexField], (3,)
        Like phi_fourier, used as an intermediary to perform FFT operations
        to obtain the gradient or laplacian of particle number densities.
    positions : (N,D) numpy.ndarray
        Array of positions for :code:`N` particles in :code:`D` dimensions.
        Local for each MPI rank.
    bond_pr : (3,) numpy.ndarray
        Total bond pressure due all two-particle bonds.
    angle_pr : (3,) numpy.ndarray
        Total angle pressure due all three-particle bonds
    comm : MPI.Intracomm, optional
        MPI communicator to use for rank commuication. Defaults to
        MPI.COMM_WORLD.

    Returns
    -------
    pressure : (18,) numpy.ndarray
        Pressure contributions from various energy terms.
        0: due to kinetic energy
        1-5: due to field energy
        6-8: due to two-particle bonded terms
        9-11: due to three-particle bonded terms (called angle terms)
        12-14: due to four-particle bonded terms (called dihedral terms)
        (defaults to 0 currently. Yet to be implemented)
        15-17: total pressure in x,y,z directions.
    """
    rank = comm.Get_rank()
    size = comm.Get_size()

    V = np.prod(config.box_size)
    n_mesh__cells = np.prod(np.full(3, config.mesh_size))
    volume_per_cell = V / n_mesh__cells

    w_0 = hamiltonian.w_0(phi) * volume_per_cell
    w_elec = 0.0
    if (config.coulombtype == "PIC_Spectral_GPE") or (
        config.coulombtype == "PIC_Spectral"
    ):
        w_elec = hamiltonian.w_elec([phi_q, psi]) * volume_per_cell
    w = w_0 + w_elec

    # Kinetic term
    kinetic_energy = 0.5 * config.mass * np.sum(velocities**2)
    p_kin = 2 / (3 * V) * kinetic_energy

    # Term 1
    p0 = -1 / V * np.sum(w)

    # Term 2
    V_bar = np.array([hamiltonian.V_bar[k]([phi, psi]) for k in range(config.n_types)])

    p1 = np.sum((volume_per_cell / V) * V_bar * phi)

    # Term 3
    comp_laplacian(
        phi_fourier,
        phi_transfer,
        phi_laplacian,
        hamiltonian,
        config,
    )

    p2 = np.sum(
        volume_per_cell
        / V
        * config.sigma**2
        * np.repeat(V_bar[:, np.newaxis, :, :, :], 3, axis=1)
        * phi_laplacian,
        axis=(0, 2, 3, 4),
    )

    # Bonded force term: linking 2 particles
    p_bond = bond_pr / V

    # Angle force term: linking 3 particles
    p_angle = angle_pr / V

    # TODO: Dihedral angle force term: linking 4 atoms
    p_dihedral = np.zeros(3)

    # Add formal parameter dihedral_forces as: comp_pressure(..., dihedral_forces)
    # Define dictionary:
    # forces = {
    #          'x': dihedral_forces[:,0],
    #          'y': dihedral_forces[:,1],
    #          'z': dihedral_forces[:,2]
    #           }
    # Compute the pressure due to dihedrals as:
    # p_dihedral = {
    #          'x': np.sum( np.multiply(forces['x'],positions[:,0]) )*(1/V),
    #          'y': np.sum( np.multiply(forces['y'],positions[:,1]) )*(1/V),
    #          'z': np.sum( np.multiply(forces['z'],positions[:,2]) )*(1/V)
    #              }

    # Total pressure in x, y, z
    p_tot = p_kin + p0 + p1 + p2 + p_bond + p_angle + p_dihedral

    return_value = np.array(
        [
            p_kin,
            p0,
            p1,
            p2[0],
            p2[1],
            p2[2],
            p_bond[0],
            p_bond[1],
            p_bond[2],
            p_angle[0],
            p_angle[1],
            p_angle[2],
            p_dihedral[0],
            p_dihedral[1],
            p_dihedral[2],
            p_tot[0],
            p_tot[1],
            p_tot[2],
        ]
    )
    return_value = comm.allreduce(return_value, MPI.SUM)

    return return_value
