import taichi as ti
import numpy as np
from taichi.linalg import make_cg, LinearOperator

ti.init(arch=ti.cpu, debug=False)

# Parameters (same as before)
q = 0.5
n_particles = int(8192 * q**2)
n_grid = int(128 * q)
dx, inv_dx = 1 / n_grid, float(n_grid)
dt = 1e-4

p_rho = 1
p_vol = (dx * 0.5) ** 2
p_mass = p_vol * p_rho
gravity = 98
bound = 3

E, nu = 3e3, 0.0  # Young's modulus and Poisson's ratio
mu, la = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu))
kappa = 2 * mu  # Bulk modulus approximation for plane stress condition

# Fields
x = ti.Vector.field(2, dtype=float, shape=n_particles)
v = ti.Vector.field(2, dtype=float, shape=n_particles)
C = ti.Matrix.field(2, 2, dtype=float, shape=n_particles)
F = ti.Matrix.field(2, 2, dtype=float, shape=n_particles)
phi = ti.field(dtype=float, shape=n_particles)  # Phase field variable
color = ti.Vector.field(3, dtype=float, shape=n_particles)

grid_v = ti.Vector.field(2, dtype=float, shape=(n_grid, n_grid))
grid_m = ti.field(dtype=float, shape=(n_grid, n_grid))

# Phase field grids
grid_phi = ti.field(dtype=float, shape=(n_grid, n_grid))
grid_phi_new = ti.field(dtype=float, shape=(n_grid, n_grid))
grid_phi_m = ti.field(dtype=float, shape=(n_grid, n_grid))
grid_psi = ti.field(dtype=float, shape=(n_grid, n_grid))
grid_psi_m = ti.field(dtype=float, shape=(n_grid, n_grid))

@ti.func
def degradation(phi):
    return (1 - phi)**2  # Quadratic degradation

@ti.func
def compute_strain_energy(Fp, phi_p):
    J = Fp.determinant()
    JaF = 1 / ti.sqrt(J) * Fp
    psi_dev = mu / 2 * ((JaF.transpose() @ JaF).trace() - 2)
    psi_vol = kappa / 2 * ((J**2 - 1) / 2 - ti.log(J))
    g_phi = degradation(phi_p)
    psi = g_phi * (psi_dev + psi_vol)
    return psi

@ti.func
def compute_stress(Fp, phi_p):
    B = Fp @ Fp.transpose()
    devB = B - ti.Matrix.identity(float, 2) * (B.trace() / 2)
    J = Fp.determinant()
    tau_dev = mu / J * devB
    prime = kappa / 2.0 * (J - 1.0 / J)
    tau_vol = J * prime * ti.Matrix.identity(float, 2)
    g_phi = degradation(phi_p)
    stress = g_phi * (tau_dev + tau_vol)
    return stress

@ti.kernel
def reset_grid():
    for I in ti.grouped(grid_m):
        grid_v[I] = [0, 0]
        grid_m[I] = 0
        grid_phi[I] = 0
        grid_phi_m[I] = 0
        grid_psi[I] = 0
        grid_psi_m[I] = 0

@ti.kernel
def p2g():
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base
        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
        F[p] = (ti.Matrix.identity(float, 2) + dt * C[p]) @ F[p]
        U, sig, V = ti.svd(F[p], ti.f32)
        J = F[p].determinant()
        F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J)
        stress = compute_stress(F[p], phi[p])
        stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
        affine = stress + p_mass * C[p]
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            dpos = (offset - fx) * dx
            grid_idx = base + offset
            weight = w[i].x * w[j].y
            grid_v[grid_idx] += weight * (p_mass * v[p] + affine @ dpos)
            grid_m[grid_idx] += weight * p_mass

@ti.kernel
def p2g_phi():
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base
        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
        phi_p = phi[p]
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            grid_idx = base + offset
            weight = w[i].x * w[j].y
            grid_phi[grid_idx] += weight * phi_p
            grid_phi_m[grid_idx] += weight

@ti.kernel
def normalize_grid_phi():
    for I in ti.grouped(grid_phi):
        if grid_phi_m[I] > 0:
            grid_phi[I] /= grid_phi_m[I]
        else:
            grid_phi[I] = 1.0

@ti.kernel
def p2g_psi():
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base
        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
        psi_p = compute_strain_energy(F[p], phi[p])
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            grid_idx = base + offset
            weight = w[i].x * w[j].y
            grid_psi[grid_idx] += weight * psi_p
            grid_psi_m[grid_idx] += weight

@ti.kernel
def normalize_grid_psi():
    for I in ti.grouped(grid_psi):
        if grid_psi_m[I] > 0:
            grid_psi[I] /= grid_psi_m[I]
        else:
            grid_psi[I] = 0.0

@ti.kernel
def grid_op():
    for I in ti.grouped(grid_m):
        if grid_m[I] > 0:
            grid_v[I] /= grid_m[I]
        grid_v[I].y -= dt * gravity
        if I[0] < bound and grid_v[I].x < 0:
            grid_v[I].x = 0
        if I[0] > n_grid - bound and grid_v[I].x > 0:
            grid_v[I].x = 0
        if I[1] < bound and grid_v[I].y < 0:
            grid_v[I].y = 0
        if I[1] > n_grid - bound and grid_v[I].y > 0:
            grid_v[I].y = 0

@ti.kernel
def g2p():
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base
        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
        new_v = ti.Vector.zero(float, 2)
        new_C = ti.Matrix.zero(float, 2, 2)
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            dpos = (offset - fx) * dx
            grid_idx = base + offset
            grid_idx = ti.Vector([ti.math.clamp(grid_idx[0], 0, n_grid - 1), ti.math.clamp(grid_idx[1], 0, n_grid - 1)])
            weight = w[i].x * w[j].y
            g_v = grid_v[grid_idx]
            new_v += weight * g_v
            new_C += 4 * weight * g_v.outer_product(dpos) / dx**2
        v[p] = new_v
        x[p] += dt * v[p]
        C[p] = new_C

@ti.kernel
def g2p_phi():
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base
        w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
        phi_new = 0.0
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            grid_idx = base + offset
            weight = w[i].x * w[j].y
            phi_new += weight * grid_phi[grid_idx]
        phi[p] = phi_new

@ti.kernel
def initialize_b(b: ti.types.ndarray()):
    for i, j in grid_phi:
        b[i, j] = grid_phi[i, j]

class PhaseFieldOperator(LinearOperator):
    def __init__(self, n):
        self.n = n  # Grid dimension

    @ti.kernel
    def matvec(self, y: ti.types.ndarray(), x: ti.types.ndarray()):
        for i, j in ti.ndrange(self.n, self.n):
            laplacian = (
                -4 * x[i, j]
                + x[min(i + 1, self.n - 1), j]
                + x[max(i - 1, 0), j]
                + x[i, min(j + 1, self.n - 1)]
                + x[i, max(j - 1, 0)]
            ) / dx**2
            psi = grid_psi[i, j]
            g_derivative = -2 * (1 - x[i, j])
            derivative = g_derivative * psi
            y[i, j] = x[i, j] - dt / tau * (Gc * laplacian - derivative)

class JacobiPreconditioner(LinearOperator):
    def __init__(self, diag):
        self.diag = diag

    def apply_precond(self, z, r):
        z[:] = r / self.diag

def solve_phase_field():
    # Transfer phi and psi to grid
    p2g_phi()
    normalize_grid_phi()
    p2g_psi()
    normalize_grid_psi()
    initialize_b(b_np)

    # Perform the CG solve
    x_np[:] = grid_phi.to_numpy()
    residual = cg_solver.solve(x_np, b_np)
    grid_phi.from_numpy(x_np)

@ti.kernel
def update_particle_color():
    for p in x:
        phi_p = phi[p]
        color[p] = ti.Vector([phi_p, 0.0, 1.0 - phi_p])  # Blue to red transition

@ti.kernel
def init():
    for i in range(n_particles):
        x[i] = [ti.random() * 0.5 + 0.25, ti.random() * 0.5 + 0.50]
        v[i] = [0, -1]
        F[i] = ti.Matrix([[1, 0], [0, 1]])
        phi[i] = 1.0  # Intact material

# Initialize the particles
init()

# Prepare the solver
n = n_grid
A = PhaseFieldOperator(n)
x_np = np.zeros((n, n), dtype=np.float32)
b_np = np.zeros((n, n), dtype=np.float32)
diag_np = np.ones((n, n), dtype=np.float32) * (1 + 4 * dt / tau * Gc / dx**2)
preconditioner = JacobiPreconditioner(diag_np)
cg_solver = make_cg(A, preconditioner=preconditioner, maxiter=1000, tol=1e-6)

# Main GUI loop
gui = ti.ui.Window("MPM88 with Phase Field Fracture", (1024, 1024))
canvas = gui.get_canvas()
while gui.running:
    canvas.set_background_color((1, 1, 1))
    for s in range(2):
        reset_grid()
        p2g()
        grid_op()
        g2p()
        # Phase field update
        solve_phase_field()
        g2p_phi()
    # Update color based on phi
    update_particle_color()
    canvas.circles(x, radius=0.005, per_vertex_color=color)
    gui.show()
Implement CG with preconditioner yourself