import taichi as ti
from taichi.linalg import MatrixFreeCG, LinearOperator
ti.init(arch=ti.gpu)
dim = 2
q = 0.5
n_particles = int(8192 * q**2)
n_grid = int(128 * q)
dx, inv_dx = 1 / n_grid, float(n_grid)
dt = 2e-5

p_rho = 1
p_vol = (dx * 0.5) ** 2
p_mass = p_vol * p_rho
gravity = 98
bound = 3
M = 10
G = 1e-4
E, nu = 9e3, 0.3  # Young's modulus and Poisson's ratio
mu, la = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu))
kappa = .666666 * mu + la
x = ti.Vector.field(2, float, n_particles)
psi_max = ti.field(float, n_particles)
c = ti.field(float, n_particles)

alpha = ti.field(float, n_particles)
nacc = False
M = 2.36         # Slope of the critical state line in p-q space
alpha0 = -0.2
beta = 1     # Parameter related to hardening
xi = 2         # Parameter related to hardening
hardening_on = True  # Toggle hardening on or off

color = ti.Vector.field(3, float, n_particles)
v = ti.Vector.field(2, float, n_particles)
C = ti.Matrix.field(2, 2, float, n_particles)
F = ti.Matrix.field(2, 2, dtype=float, shape=n_particles)  # deformation gradient

grid_v = ti.Vector.field(2, float, (n_grid, n_grid))
grid_m = ti.field(float, (n_grid, n_grid))
rhs  = ti.field(float, (n_grid, n_grid))
b =   ti.field(float, (n_grid, n_grid))
grid_c = ti.field(float, (n_grid, n_grid))
gradc_scp = ti.Vector.field(2, float, n_particles)
grad = ti.field(float, (n_grid, n_grid))
diag =  ti.field(float, (n_grid, n_grid))
@ti.data_oriented
class ConjugateGradientSolver2D:
    def __init__(self, n, m):
        self.n = n
        self.m = m
        self.x = ti.field(dtype=ti.f32, shape=(n, m))   # Solution
        self.b = ti.field(dtype=ti.f32, shape=(n, m))   # Right-hand side
        self.r = ti.field(dtype=ti.f32, shape=(n, m))  # Residual
        self.p = ti.field(dtype=ti.f32, shape=(n, m))  # Search direction
        self.z = ti.field(dtype=ti.f32, shape=(n, m))  # Preconditioned residual
        self.q = ti.field(dtype=ti.f32, shape=(n, m))  # A * p
        self.diagonal = ti.field(dtype=ti.f32, shape=(n, m))   # Diagonal for Jacobi preconditioner
        self.user_matvec = None  # User-provided matrix-vector multiplication function

    def matvec(self, p, q):
        """Matrix-free multiplication: q = A * p (user-provided)"""
        self.user_matvec(p, q)

    @ti.kernel
    def apply_jacobi_preconditioner(self, r: ti.template(), z: ti.template()):
        """Apply Jacobi preconditioner: z = M^-1 * r in 2D"""
        for i, j in ti.ndrange(self.n, self.m):
            if self.diagonal[i, j] != 0:
                z[i, j] = r[i, j] / self.diagonal[i, j]

    @ti.kernel
    def dot(self, v1: ti.template(), v2: ti.template()) -> ti.f32:
        """Dot product of two vectors in 2D"""
        sum = 0.0
        for i, j in ti.ndrange(self.n, self.m):
            sum += v1[i, j] * v2[i, j]
        return sum

    @ti.kernel
    def vec_add(self, v1: ti.template(), v2: ti.template(), alpha: ti.f32):
        """Element-wise operation: v1 += alpha * v2 in 2D"""
        for i, j in ti.ndrange(self.n, self.m):
            v1[i, j] += alpha * v2[i, j]

    @ti.kernel
    def vec_sub(self, v1: ti.template(), v2: ti.template(), alpha: ti.f32):
        """Element-wise operation: v1 -= alpha * v2 in 2D"""
        for i, j in ti.ndrange(self.n, self.m):
            v1[i, j] -= alpha * v2[i, j]

    @ti.kernel
    def vec_mul(self, v1: ti.template(), alpha: ti.f32):
        """Element-wise operation: v1 *= alpha in 2D"""
        for i, j in ti.ndrange(self.n, self.m):
            v1[i, j] *= alpha

    @ti.kernel
    def vec_add_scaled(self, v1: ti.template(), v2: ti.template(), beta: ti.f32):
        """Element-wise operation: v1 = beta * v1 + v2 in 2D"""
        for i, j in ti.ndrange(self.n, self.m):
            v1[i, j] = beta * v1[i, j] + v2[i, j]

    def conjugate_gradient(self, tol=1e-6, max_iter=1000):
        """Conjugate Gradient method in 2D"""
        # Initial residual r = b - A * x
        self.matvec(self.x, self.q)
        self.r.copy_from(self.b)
        self.vec_sub(self.r, self.q, 1.0)

        # Preconditioner: z = M^-1 * r (Jacobi)
        self.apply_jacobi_preconditioner(self.r, self.z)

        # Initial direction p = z
        self.p.copy_from(self.z)

        rz_old = self.dot(self.r, self.z)

        for iter in range(max_iter):
            # q = A * p
            self.q.fill(0)
            self.matvec(self.p, self.q)

            # alpha = rz_old / (p^T * A * p)
            pq = self.dot(self.p, self.q)
            alpha = rz_old / pq

            # Update x = x + alpha * p
            self.vec_add(self.x, self.p, alpha)

            # Update residual r = r - alpha * q
            self.vec_sub(self.r, self.q, alpha)

            # Check convergence: ||r|| < tol
            rz_new = self.dot(self.r, self.r)
            print(rz_new ** 0.5)
            if rz_new < tol ** 2:
                print(f"Converged in {iter} iterations.")
                break

            # Preconditioner: z = M^-1 * r
            self.apply_jacobi_preconditioner(self.r, self.z)

            # Update rz_new = r^T * z
            rz_new = self.dot(self.r, self.z)

            # Beta = rz_new / rz_old
            beta = rz_new / rz_old

            # Update search direction: p = z + beta * p
            self.vec_mul(self.p, beta)
            self.vec_add(self.p, self.z, 1.0)

            rz_old = rz_new

        else:
            print(f"Did not converge in {max_iter} iterations.")
@ti.func
def sinh(x):
    return ti.exp(x) / 2 - ti.exp(-x) / 2

@ti.func
def project_strain(F_input, alpha_input):
    # Copy input variables to local variables
    F = F_input
    logJp = alpha_input

    # Small epsilon value to prevent division by zero and negative square roots
    eps = 1e-8

    # Perform Singular Value Decomposition (SVD) on F
    U, sig, V = ti.svd(F)
    sigma = ti.Vector([sig[0, 0], sig[1, 1]])  # Extract singular values into a vector

    # Compute determinant J
    J = sigma[0] * sigma[1]
    J = ti.max(J, eps)  # Ensure J is not too small to prevent division by zero

    # Calculate p0 based on hardening
    p0 = kappa * (1e-5 + sinh(xi * ti.max(-logJp, 0.0)))
    p0 = ti.max(p0, 0.0)  # Ensure p0 is non-negative

    # Compute trial deviatoric stress
    B_hat_trial = sigma * sigma
    B_hat_trial_avg = B_hat_trial.sum() / dim
    s_hat_trial = mu * J**(-2.0 / dim) * (B_hat_trial - B_hat_trial_avg)

    # Compute trial pressure
    prime = kappa / 2.0 * (J - 1.0 / J)
    p_trial = -prime * J

    # Compute yield function value y
    y_s_half_coeff = (6.0 - dim) / 2.0 * (1.0 + 2.0 * beta)
    y_p_half = M * M * (p_trial + beta * p0) * (p_trial - p0)
    y = y_s_half_coeff * s_hat_trial.norm_sqr() + y_p_half

    # Project sigma if p_trial exceeds bounds
    p_min = beta * p0
    updated = False  # Flag to check if we need to update F and alpha

    tolerance = 1e-6  # Define a small tolerance for numerical errors

    if p_trial > p0 + tolerance:
        # Overflow in compression
        Je_new_argument = -2.0 * p0 / kappa + 1.0
        Je_new_argument = ti.max(Je_new_argument, eps)  # Ensure non-negative argument
        Je_new = ti.sqrt(Je_new_argument)
        sigma_new = ti.Vector([Je_new ** (1.0 / dim)] * dim)
        F = U @ ti.Matrix([[sigma_new[0], 0.0], [0.0, sigma_new[1]]]) @ V.transpose()
        if hardening_on:
            logJp += ti.log(J / (Je_new + eps))  # Prevent division by zero
        updated = True

    elif p_trial < -p_min - tolerance:
        # Overflow in tension
        Je_new_argument = 2.0 * p_min / kappa + 1.0
        Je_new_argument = ti.max(Je_new_argument, eps)  # Ensure non-negative argument
        Je_new = ti.sqrt(Je_new_argument)
        sigma_new = ti.Vector([Je_new ** (1.0 / dim)] * dim)
        F = U @ ti.Matrix([[sigma_new[0], 0.0], [0.0, sigma_new[1]]]) @ V.transpose()
        if hardening_on:
            logJp += ti.log(J / (Je_new + eps))  # Prevent division by zero
        updated = True

    elif y < 1e-4:
        # Inside the yield surface; no update needed
        pass
    else:
        # Project onto yield surface
        s_hat_trial_norm = s_hat_trial.norm()
        denominator = s_hat_trial_norm + eps  # Prevent division by zero
        # Compute numerator
        sqrt_argument = -y_p_half / y_s_half_coeff
        sqrt_argument = ti.max(sqrt_argument, 0.0)  # Ensure non-negative for sqrt
        numerator = ti.sqrt(sqrt_argument)
        factor = numerator / denominator
        B_hat_new = J**(2.0 / dim) / mu * factor * s_hat_trial + B_hat_trial_avg

        # Ensure B_hat_new is valid
        B_hat_new = ti.max(B_hat_new, 0.0)

        # Update sigma and F
        sigma_new = ti.sqrt(B_hat_new + eps)  # Prevent sqrt of zero
        F = U @ ti.Matrix([[sigma_new[0], 0.0], [0.0, sigma_new[1]]]) @ V.transpose()

        # Hardening update if necessary
        if p0 > 1e-4 and p_trial < p0 - 1e-4 and p_trial > 1e-4 - p_min:
            p_center = (p0 - p_min) * 0.5
            q_trial = ti.sqrt((6.0 - dim) / 2.0) * s_hat_trial_norm

            dir_x = p_center - p_trial
            dir_y = -q_trial
            dir_norm = ti.sqrt(dir_x * dir_x + dir_y * dir_y) + eps  # Prevent division by zero
            direction = ti.Vector([dir_x / dir_norm, dir_y / dir_norm])

            # Solve quadratic for lambda
            C = M * M * (p_center + beta * p0) * (p_center - p0)
            B_quad = M * M * direction[0] * (2.0 * p_center - p0 + beta * p0)
            A_quad = M * M * direction[0] ** 2 + (1.0 + 2.0 * beta) * direction[1] ** 2 + eps  # Prevent division by zero

            discriminant = B_quad * B_quad - 4.0 * A_quad * C
            discriminant = ti.max(discriminant, 0.0)  # Ensure discriminant is non-negative
            if discriminant >= 0:
                sqrt_discriminant = ti.sqrt(discriminant)
                l1 = (-B_quad + sqrt_discriminant) / (2.0 * A_quad)
                l2 = (-B_quad - sqrt_discriminant) / (2.0 * A_quad)

                # Choose the correct lambda
                delta_p1 = p_center + l1 * direction[0] - p_center
                delta_p2 = p_center + l2 * direction[0] - p_center
                lambda_candidate = l1 if delta_p1 * (p_trial - p_center) > 0 else l2
                p_fake = p_center + lambda_candidate * direction[0]
                Je_new_fake_argument = abs(-2.0 * p_fake / kappa + 1.0)
                Je_new_fake_argument = ti.max(Je_new_fake_argument, eps)  # Ensure non-negative
                Je_new_fake = ti.sqrt(Je_new_fake_argument)
                if Je_new_fake > 1e-4 and hardening_on:
                    logJp += ti.log(J / (Je_new_fake + eps))  # Prevent division by zero
                updated = True

    # Return updated F and alpha
    alpha_updated = logJp if updated and hardening_on else alpha_input
    return F, alpha_updated
@ti.func
def compute_stress(Fp,g=1.):
    B = Fp@Fp.transpose()
    devB = B - ti.Matrix.identity(float, 2)  * 1. / 2.*B.trace()
    J = Fp.determinant()
    tau_dev = mu * 1./J * devB
    prime = kappa/2. * (J-1./J)
    tau_vol = J*prime*ti.Matrix.identity(float, 2)
    out = g * (tau_dev + tau_vol) if J >= 1. else g * tau_dev + tau_vol
    return out
@ti.func
def get_psi_pos(Fp):
    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))
    out = psi_dev + psi_vol if J >= 1. else  psi_dev
    return out
@ti.kernel
def reset_grid():
    for i, j in grid_m:
        grid_v[i, j] = [0, 0]
        grid_m[i, j] = 0
        rhs[i,j] = 0
        b[i,j] = 0
        grid_c[i,j] = 0
        diag[i,j] = 0
        grad[i,j] = 0
@ti.kernel
def phase_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]
        if get_psi_pos(F[p]) > psi_max[p]:
            psi_max[p] = get_psi_pos(F[p])
        tw = 0.
        for i, j in ti.static(ti.ndrange(3, 3)):
            weight = w[i].x * w[j].y
            tw += weight
        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
            rhs[grid_idx] += p_vol*F[p].determinant()*(M+c[p]/dt)*weight
            grid_c[grid_idx] += c[p]*weight
            diag[grid_idx] += p_vol*F[p].determinant()*weight*(4*dx*M*psi_max[p]*0.999/G + M + 1/dt)

@ti.kernel
def mvp(v:ti.template(), mv:ti.template()):
    for p in x:
        gradc_scp[p] = 0
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base

        # Weight functions
        w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1) ** 2, 0.5 * (fx - 0.5) ** 2]

        # Gradient of weight functions (necessary for Laplacian)
        dw = [-1.5 + fx, 1 - 2*(fx - 1), fx - 0.5]

        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            weight = w[i].x * w[j].y

            gradc_scp[p] += p_vol * F[p].determinant() *4 * dx**2 * M * ti.Vector([dw[i].x * w[j].y, w[i].x * dw[j].y])
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        fx = Xp - base

        # Weight functions
        w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1) ** 2, 0.5 * (fx - 0.5) ** 2]

        # Gradient of weight functions (necessary for Laplacian)
        dw = [-1.5 + fx, 1 - 2*(fx - 1), fx - 0.5]

        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            weight = w[i].x * w[j].y

            # Compute gradient (MLS shape function derivative)
            grad_weight = ti.Vector([dw[i].x * w[j].y, w[i].x * dw[j].y])

            # Update grid momentum (with Laplacian term)
            mv[base+offset] += p_vol * F[p].determinant() * weight * (4 * dx * M * psi_max[p] * 0.999 / G + M + 1/dt) * v[base+offset]
            mv[base+offset] += gradc_scp[p].dot(grad_weight)* v[base+offset]
@ti.kernel
def p2g():
    for p in x:
        Xp = x[p] / dx
        base = int(Xp - 0.5)
        color[p][0] = alpha[p]
        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])
        J = F[p].determinant()
        stress = compute_stress(F[p], c[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 grid_op():
    for i, j in grid_m:
        if grid_m[i, j] > 0:
            grid_v[i, j] /= grid_m[i, j]
        grid_v[i, j].y -= dt * gravity
        if i < bound and grid_v[i, j].x < 0:
            grid_v[i, j].x = 0
        if i > n_grid - bound and grid_v[i, j].x > 0:
            grid_v[i, j].x = 0
        if j < bound and grid_v[i, j].y < 0:
            grid_v[i, j].y = 0
        if j > n_grid - bound and grid_v[i, j].y > 0:
            grid_v[i, j].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, grid_v.shape[0] - 1),ti.math.clamp(grid_idx[1], 0, grid_v.shape[1] - 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 phase_g2p():
    for p in range(n_particles):
        color[p][0] = 1-c[p]
        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_c = 0.
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            grid_idx = base + offset
            weight = w[i][0] * w[j][1]
            new_c += (sol[grid_idx])*weight
        #c[p] = max(0,min(c[p],new_c))
@ti.kernel
def init():
    for i in range(n_particles):
        x[i] = [ti.random() * 0.1 + 0.25, ti.random() * 0.8 + 0.2]
        v[i] = [0, -1]
        F[i] = ti.Matrix([[1, 0], [0, 1]])
        alpha[i] = alpha0
        c[i] = 1
# Initialize the particles
init()

# Main GUI loop
gui = ti.ui.Window("MPM88", (1024, 1024))
canvas = gui.get_canvas()
solver = ConjugateGradientSolver2D(n_grid,n_grid)
while gui.running:
    canvas.set_background_color((1, 1, 1))
    for s in range(5):
        reset_grid()
        phase_p2g()
        solver.user_matvec = mvp
        solver.x = b
        solver.b = rhs
        solver.diagonal = diag
        solver.conjugate_gradient(tol=1e-6, max_iter=10)
        sol = solver.x
        print(sol.to_numpy().max())
        phase_g2p()
        p2g()
        grid_op()
        g2p()
    canvas.circles(x, radius=0.005, per_vertex_color=color)
    canvas.set_image(sol.to_numpy())
    gui.show()

Phase field isn't implemented correctly, fix that