import flax
import deepxde as dde
import jax.numpy as np
import matplotlib.pyplot as plt

rho = 1  
mu = 1   
u_in = 1
D = 1
L = 2

geom = dde.geometry.Rectangle(xmin=[-L / 2, -D / 2], xmax=[L / 2, D / 2])
def boundary_wall(X, on_boundary):
    on_wall = np.logical_and(
        np.logical_or(
            np.isclose(X[1], -D / 2, rtol=1e-05, atol=1e-08),
            np.isclose(X[1], D / 2, rtol=1e-05, atol=1e-08),
        ),
        on_boundary,
    )
    return on_wall

def boundary_inlet(X, on_boundary):
    on_inlet = np.logical_and(
        np.isclose(X[0], -L / 2, rtol=1e-05, atol=1e-08), on_boundary
    )
    return on_inlet

def boundary_outlet(X, on_boundary):
    on_outlet = np.logical_and(
        np.isclose(X[0], L / 2, rtol=1e-05, atol=1e-08), on_boundary
    )
    return on_outlet

bc_wall_u = dde.DirichletBC(geom, lambda X: 0.0, boundary_wall, component=0)
bc_wall_v = dde.DirichletBC(geom, lambda X: 0.0, boundary_wall, component=1)

bc_inlet_u = dde.DirichletBC(geom, lambda X: u_in, boundary_inlet, component=0)
bc_inlet_v = dde.DirichletBC(geom, lambda X: 0.0, boundary_inlet, component=1)

bc_outlet_p = dde.DirichletBC(geom, lambda X: 0.0, boundary_outlet, component=2)
bc_outlet_v = dde.DirichletBC(geom, lambda X: 0.0, boundary_outlet, component=1)

def pde(X, Y):
    """Args:
    Y: network output [u, v, p],
    X: input coordinates [x, y]"""
    du_x,_ = dde.grad.jacobian(Y, X, i=0, j=0)
    du_y,_ = dde.grad.jacobian(Y, X, i=0, j=1)
    dv_x,_ = dde.grad.jacobian(Y, X, i=1, j=0)
    dv_y,_ = dde.grad.jacobian(Y, X, i=1, j=1)
    dp_x,_ = dde.grad.jacobian(Y, X, i=2, j=0)
    dp_y,_ = dde.grad.jacobian(Y, X, i=2, j=1)

    du_xx,_ = dde.grad.hessian(Y, X, component=0, i=0, j=0)
    du_yy,_ = dde.grad.hessian(Y, X, component=0, i=1, j=1)
    dv_xx,_ = dde.grad.hessian(Y, X, component=1, i=0, j=0)
    dv_yy,_ = dde.grad.hessian(Y, X, component=1, i=1, j=1)

    y_val, _ = Y  
    u_pred = y_val[:, 0:1]
    v_pred = y_val[:, 1:2]    
    pde_u = u_pred * du_x + v_pred * du_y + 1 / rho * dp_x - (mu / rho) * (du_xx + du_yy)
    pde_v = u_pred * dv_x + v_pred * dv_y + 1 / rho * dp_y - (mu / rho) * (dv_xx + dv_yy)
    pde_cont = du_x + dv_y

    return [pde_u, pde_v, pde_cont]

data = dde.data.PDE(
    geom,
    pde,
    [bc_wall_u, bc_wall_v, bc_inlet_u, bc_inlet_v, bc_outlet_p, bc_outlet_v],
    num_domain=2000,
    num_boundary=200,
    num_test=200,
)

net = dde.maps.FNN([2] + [64] * 5 + [3], "tanh", "Glorot uniform")
model = dde.Model(data, net)

model.compile("adam", lr=1e-3)
losshistory, train_state = model.train(epochs=10000)

plt.figure(figsize=(10, 8))
plt.scatter(data.train_x_all[:, 0], data.train_x_all[:, 1], s=0.5)
plt.xlabel("x")
plt.ylabel("y")
plt.show()


samples = geom.random_points(500000)
result = model.predict(samples)

color_legend = [[0, 1.5], [-0.3, 0.3], [0, 35]]
for idx in range(3):
    plt.figure(figsize=(20, 4))
    plt.scatter(samples[:, 0], samples[:, 1], c=result[:, idx], cmap="jet", s=2)
    plt.colorbar()
    plt.clim(color_legend[idx])
    plt.xlim((0 - L / 2, L - L / 2))
    plt.ylim((0 - D / 2, D - D / 2))
    plt.tight_layout()
    plt.show()