import numpy as np

def solve_poisson_fft_tilde(f_tilde, Qx, Qy):
    """
    Solve the Poisson equation in Fourier space directly.
    """
    k2 = Qx**2 + Qy**2
    k2[0, 0] = 1  # avoid division by zero
    u_tilde = -f_tilde / k2  # Attention, here is for solving the Poisson equation with \Delta u + f = 0, so here the f_tilde could be right hand side of the Poisson equation
    u_tilde[0, 0] = 0  # set the zero-frequency component to zero
    return u_tilde

def fft_gradient(u_tilde, Qx, Qy):
    """
    For a given u_tilde in Fourier space, compute its gradient in real space.
    """
    ux_tilde = 1j * Qx * u_tilde
    uy_tilde = 1j * Qy * u_tilde
    ux = np.fft.ifft2(ux_tilde).real
    uy = np.fft.ifft2(uy_tilde).real
    return ux, uy

def divergence_to_tilde(fx, fy, Qx, Qy):
    """
    Compute the divergence of a vector field (fx, fy) in real space, then return its Fourier transform.
    """
    fx_tilde = np.fft.fft2(fx)
    fy_tilde = np.fft.fft2(fy)
    div_tilde = 1j * Qx * fx_tilde + 1j * Qy * fy_tilde
    return div_tilde

def solve_microfluctuation_poisson(G, G_0, f, Qx, Qy, epsilon=1e-8, max_iter=5000):
    """
    To solve the equation: ∇·(G∇N) + f = 0
    With fixed point iteration: G_0 ∇²N + ∇·((G-G_0)∇N) + f = 0
    Solve N^(0) with ∇²N = -f/G_0 as initial guess
    """
    f_tilde = np.fft.fft2(f)
    
    N_init_tilde = solve_poisson_fft_tilde(-f_tilde / G_0, Qx, Qy)
    N_tilde = N_init_tilde.copy()

    for _ in range(max_iter):                                   # The loop simply repeats the loop body max_iter times, without requiring a count value in each loop iteration.
        # Compute the gradient of N in real space
        Nx, Ny = fft_gradient(N_tilde, Qx, Qy)
        # Compute (G-G_0)*grad(u) in real space, then compute its divergence in Fourier space
        temp_x = (G - G_0)*Nx
        temp_y = (G - G_0)*Ny
        f_n_tilde = divergence_to_tilde(temp_x, temp_y, Qx, Qy)

        # Slove the Poisson equation in Fourier space
        # G_0 ΔN = -f - f_n
        # ΔN = -(f_n + f)/G_0
        N_next_tilde = solve_poisson_fft_tilde(-(f_n_tilde + f_tilde)/G_0, Qx, Qy)
        # Check the convergence
        if np.linalg.norm(np.fft.ifft2(N_next_tilde - N_tilde)) < epsilon:
            N_tilde = N_next_tilde
            break
        N_tilde = N_next_tilde

    N = np.fft.ifft2(N_tilde).real
    return N


def homogeneous_effective_modulus(G, G_0, Qx, Qy, n, m):
    # For N1_1:
    # f1 = -∇·(G(y)*e_1) = -∇·(G(y), 0)
    # ∇·(G(y),0) = dG/dx
    dG_dx = np.fft.ifft2(1j * Qx * np.fft.fft2(G)).real
    f1 = dG_dx

    N1_1 = solve_microfluctuation_poisson(G, G_0, f1, Qx, Qy)

    # For N1_2:
    # f2 = -∇·(G(y)*e_2) = -∇·(0, G(y))
    # ∇·(0,G(y)) = dG/dy
    dG_dy = np.fft.ifft2(1j * Qy * np.fft.fft2(G)).real
    f2 = dG_dy

    N1_2 = solve_microfluctuation_poisson(G, G_0, f2, Qx, Qy)

    # C_0 = ∫ G(y) [∇N1(y) + I] dy
    # N1 = [N1_1, N1_2]

    N1_1_tilde = np.fft.fft2(N1_1)
    N1_2_tilde = np.fft.fft2(N1_2)

    N1_1x, N1_1y = fft_gradient(N1_1_tilde, Qx, Qy)
    N1_2x, N1_2y = fft_gradient(N1_2_tilde, Qx, Qy)

    # ∇N1 is [[N1_1x, N1_1y],
    #         [N1_2x, N1_2y]]
    # ∇N1 + I is [[N1_1x+1, N1_1y   ],
    #            [N1_2x  , N1_2y+1]]
    N1_grad_plus_I = np.zeros((n,m,2,2))
    N1_grad_plus_I[:,:,0,0] = N1_1x + 1.0
    N1_grad_plus_I[:,:,0,1] = N1_1y
    N1_grad_plus_I[:,:,1,0] = N1_2x
    N1_grad_plus_I[:,:,1,1] = N1_2y + 1.0

    # Compute C_0 = ∫ G(y)*[∇N1+I] dy
    C0 = np.zeros((2,2))
    for i in range(2):
        for j in range(2):
            # Integral of unit cell, the area of the unit cell is 1
            C0[i,j] = np.mean(G * N1_grad_plus_I[:,:,i,j])

    return C0, N1_1, N1_2, N1_grad_plus_I

# Define gradient function  to compute ∇v0
def fft_gradient_2d(u_tilde, Qx, Qy):
    ux_tilde = (1j * Qx) * u_tilde
    uy_tilde = (1j * Qy) * u_tilde
    ux = np.fft.ifft2(ux_tilde).real
    uy = np.fft.ifft2(uy_tilde).real
    return ux, uy


def first_order_homogenization(F_macro, Qx_macro, Qy_macro, C0, N1_1, N1_2, eta):
    F_macro_tilde = np.fft.fft2(F_macro)
    denominator = (C0[0,0]*Qx_macro**2 + (C0[0,1]+C0[1,0])*Qx_macro*Qy_macro + C0[1,1]*Qy_macro**2) #Deducing process is in the appendix

    denominator[0,0] = 1.0  # avoid division by zero
    v0_tilde = F_macro_tilde / denominator                                                              # v0_tilde = -F_tilde / denominator     ！！！ there should not be a minus sign here
    v0_tilde[0,0] = 0.0
    v0 = np.fft.ifft2(v0_tilde).real

    v0_ux, v0_uy = fft_gradient_2d(np.fft.fft2(v0), Qx_macro, Qy_macro)
    # u(x) ≈ v_0(x) + η N_1(x/η) · ∇v_0(x)

    n = int(1/eta)
    N1_1_global = N1_1[::n,::n]
    N1_2_global = N1_2[::n,::n]
    # tile
    N1_1_global = np.tile(N1_1_global, (n,n))
    N1_2_global = np.tile(N1_2_global, (n,n))


    u = v0 + eta * (N1_1_global * v0_ux + N1_2_global * v0_uy)
    return u, v0, v0_ux, v0_uy