from pystencils.simp.assignment_collection import AssignmentCollection
import sympy as sp
from pystencils import Assignment

from lbmpy.moments import MOMENT_SYMBOLS
from lbmpy.methods.centeredcumulant.centered_cumulants import statistical_quantity_symbol
from lbmpy.methods.centeredcumulant.cumulant_transform import PRE_COLLISION_CUMULANT

x, y, z = MOMENT_SYMBOLS
corrected_polynomials = [x**2 - y**2, x**2 - z**2, x**2 + y**2 + z**2]


def contains_corrected_polynomials(polynomials):
    return all(polynomials.containts(p) for p in corrected_polynomials)


def add_galilean_correction(poly_relaxation_eqs, polynomials, correction_terms):
    #   Call PC1 = (x^2 - y^2), PC2 = (x^2 - z^2), PC3 = (x^2 + y^2 + z^2)
    try:
        index_pc1 = polynomials.index(corrected_polynomials[0])
        index_pc2 = polynomials.index(corrected_polynomials[1])
        index_pc3 = polynomials.index(corrected_polynomials[2])
    except ValueError:
        raise ValueError("For the galilean correction, all three polynomial cumulants"
                         + "(x^2 - y^2), (x^2 - z^2) and (x^2 + y^2 + z^2) need to be present!")

    cor1 = correction_terms.main_assignments[0].lhs
    cor2 = correction_terms.main_assignments[1].lhs
    cor3 = correction_terms.main_assignments[2].lhs

    poly_relaxation_eqs[index_pc1] += cor1
    poly_relaxation_eqs[index_pc2] += cor2
    poly_relaxation_eqs[index_pc3] += cor3

    return poly_relaxation_eqs


def get_galilean_correction_terms(cumulant_to_relaxation_info_dict, rho, u_vector,
                                  pre_collision_cumulant_base=PRE_COLLISION_CUMULANT):

    pc1 = corrected_polynomials[0]
    pc2 = corrected_polynomials[1]
    pc3 = corrected_polynomials[2]

    try:
        omega_1 = cumulant_to_relaxation_info_dict[pc1].relaxation_rate
        assert omega_1 == cumulant_to_relaxation_info_dict[pc2].relaxation_rate, \
            "Cumulants (x^2 - y^2) and (x^2 - z^2) must have the same relaxation rate"
        omega_2 = cumulant_to_relaxation_info_dict[pc3].relaxation_rate
    except IndexError:
        raise ValueError("For the galilean correction, all three polynomial cumulants"
                         + "(x^2 - y^2), (x^2 - z^2) and (x^2 + y^2 + z^2) must be present!")

    dx, dy, dz = sp.symbols('Dx, Dy, Dz')
    c_xx = statistical_quantity_symbol(pre_collision_cumulant_base, (2, 0, 0))
    c_yy = statistical_quantity_symbol(pre_collision_cumulant_base, (0, 2, 0))
    c_zz = statistical_quantity_symbol(pre_collision_cumulant_base, (0, 0, 2))

    cor1, cor2, cor3 = sp.symbols("corr_:3")

    #   Derivative Approximations
    subexpressions = [
        Assignment(dx, - omega_1 / (2 * rho) * (2 * c_xx - c_yy - c_zz)
                   - omega_2 / (2 * rho) * (c_xx + c_yy + c_zz - rho)),
        Assignment(dy, dx + (3 * omega_1) / (2 * rho) * (c_xx - c_yy)),
        Assignment(dz, dx + (3 * omega_1) / (2 * rho) * (c_xx - c_zz))]

    #   Correction Terms
    main_assignments = [
        Assignment(cor1, - 3 * rho * (1 - omega_1 / 2) * (u_vector[0]**2 * dx - u_vector[1]**2 * dy)),
        Assignment(cor2, - 3 * rho * (1 - omega_1 / 2) * (u_vector[0]**2 * dx - u_vector[2]**2 * dz)),
        Assignment(cor3, - 3 * rho * (1 - omega_2 / 2)
                   * (u_vector[0]**2 * dx + u_vector[1]**2 * dy + u_vector[2]**2 * dz))]
    return AssignmentCollection(main_assignments=main_assignments, subexpressions=subexpressions)