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

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

E, nu = 3e2, 0.2  # Young's modulus and Poisson's ratio
mu, la = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu))
kappa = .666666 * mu + la
M, G = 1000, 1e-5
x = ti.Vector.field(2, float, n_particles)
c = ti.field(float, n_particles)
psi_max = ti.field(float, n_particles)
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))
grid_c = ti.field(float, (n_grid, n_grid))
rhs = ti.field(float, (n_grid, n_grid))
zeros  = ti.field(float, (n_grid, n_grid))

b = ti.field(float, (n_grid, n_grid))
gradc_scp = ti.Vector.field(2, float, n_particles)
attractor_strength = ti.field(dtype=float, shape=())
attractor_pos = ti.Vector.field(2, dtype=float, shape=())
diag = ti.field(float, (n_grid, n_grid))
@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_c[i, j] = 0
        grid_m[i, j] = 0
        rhs[i,j] = 0
        b[i,j] = 0
        diag[i,j] = 0
    for p in x:
        gradc_scp[p] = 0
@ti.kernel
def mvp(v: ti.template(), mv: ti.template()):
    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]
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            dpos = (offset - fx) * dx
            weight = w[i][0] * w[j][1]
            gradc_scp[p] += F[p].determinant()*p_vol*16*dpos*weight*v[base + offset]
    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]
        for i, j in ti.static(ti.ndrange(3, 3)):
            offset = ti.Vector([i, j])
            dpos = (offset - fx) * dx
            weight = w[i][0] * w[j][1]
            mv[base + offset] += F[p].determinant()*p_vol* (4*dx*0.999*psi_max[p]/G+M + 1 / dt) * weight * v[base + offset]
            mv[base + offset] += gradc_scp[p].dot(dpos*weight) * M
@ti.kernel
def get_preconditioner():
    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]
        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][0] * w[j][1]
            diag[grid_idx[0],grid_idx[1]] += (F[p].determinant()*p_vol* (4*dx*0.999*psi_max[p]/G+M + 1 / dt) * weight)
    for i,j in diag:
        if diag[i,j] > 0:
            diag[i,j] = 1/diag[i,j]
@ti.data_oriented
class GMRES_Solver:
    def __init__(self, rows, cols, n_iter=1, tol=1e-6):
        self.rows = rows
        self.cols = cols
        self.n_iter = n_iter
        self.tol = tol  # Convergence tolerance

        # Taichi fields for 2D vectors and matrices
        self.V = ti.field(dtype=ti.f32, shape=(n_iter + 1, rows, cols))
        self.H = ti.field(dtype=ti.f32, shape=(n_iter + 1, n_iter))
        self.w = ti.field(dtype=ti.f32, shape=(rows, cols))
        self.r0 = ti.field(dtype=ti.f32, shape=(rows, cols))
        self.diagonal_inv = ti.field(dtype=ti.f32, shape=(rows, cols))

        # Fields for Givens rotations
        self.cs = ti.field(dtype=ti.f32, shape=self.n_iter)
        self.sn = ti.field(dtype=ti.f32, shape=self.n_iter)
        self.s = ti.field(dtype=ti.f32, shape=(self.n_iter + 1))

        # Field to hold V[j] for use in Python scope
        self.Vj = ti.field(dtype=ti.f32, shape=(self.rows, self.cols))

    # Kernel for vector subtraction: result = v1 - v2
    @ti.kernel
    def subtract_vectors(self, v1: ti.template(), v2: ti.template(), result: ti.template()):
        for i, j in v1:
            result[i, j] = v1[i, j] - v2[i, j]

    # Kernel for vector normalization and computing the norm
    @ti.kernel
    def normalize_and_compute_norm(self, v: ti.template(), norm: ti.template()):
        norm_val = 0.0
        for i, j in v:
            norm_val += v[i, j] * v[i, j]
        norm[None] = ti.sqrt(norm_val)
        for i, j in v:
            v[i, j] /= norm[None]

    # Kernel to copy one vector to another
    @ti.kernel
    def copy_vector(self, src: ti.template(), dst: ti.template(), idx: ti.i32):
        for i, j in src:
            dst[idx, i, j] = src[i, j]

    # Kernel for preconditioned matrix-vector multiplication: result = D^{-1} * v
    @ti.kernel
    def apply_preconditioner(self, v: ti.template(), result: ti.template()):
        for i, j in v:
            result[i, j] = self.diagonal_inv[i, j] * v[i, j]  # Applying diagonal inverse

    # Kernel to extract V[j] into self.Vj
    @ti.kernel
    def extract_Vj(self, j: ti.i32):
        for i, k in ti.ndrange(self.rows, self.cols):
            self.Vj[i, k] = self.V[j, i, k]

    # Taichi kernel for orthogonalization
    @ti.kernel
    def orthogonalize(self, j: ti.i32):
        for i in range(j + 1):
            # Compute H[i, j] = dot_product(w, V[i])
            h_ij = 0.0
            for m, n in ti.ndrange(self.rows, self.cols):
                h_ij += self.w[m, n] * self.V[i, m, n]
            self.H[i, j] = h_ij

            # w = w - H[i, j] * V[i]
            for m, n in ti.ndrange(self.rows, self.cols):
                self.w[m, n] -= h_ij * self.V[i, m, n]

        # H[j + 1, j] = norm(w)
        norm_w = 0.0
        for m, n in ti.ndrange(self.rows, self.cols):
            norm_w += self.w[m, n] ** 2
        norm_w = ti.sqrt(norm_w)
        self.H[j + 1, j] = norm_w

        # V[j + 1] = w / H[j + 1, j]
        for m, n in ti.ndrange(self.rows, self.cols):
            self.V[j + 1, m, n] = self.w[m, n] / norm_w

    # Function for applying Givens rotations
    @ti.kernel
    def apply_givens_rotations(self, j: ti.i32):
        # Apply Givens rotations to H and s
        for i in range(j):
            temp = self.cs[i] * self.H[i, j] + self.sn[i] * self.H[i + 1, j]
            self.H[i + 1, j] = -self.sn[i] * self.H[i, j] + self.cs[i] * self.H[i + 1, j]
            self.H[i, j] = temp

        # Compute the next rotation parameters
        delta = ti.sqrt(self.H[j, j] ** 2 + self.H[j + 1, j] ** 2)
        self.cs[j] = self.H[j, j] / delta
        self.sn[j] = self.H[j + 1, j] / delta

        # Eliminate H[j + 1, j]
        self.H[j, j] = self.cs[j] * self.H[j, j] + self.sn[j] * self.H[j + 1, j]
        self.H[j + 1, j] = 0.0

        # Update the residual vector s
        temp = self.cs[j] * self.s[j] + self.sn[j] * self.s[j + 1]
        self.s[j + 1] = -self.sn[j] * self.s[j] + self.cs[j] * self.s[j + 1]
        self.s[j] = temp

    # Kernel for back substitution to solve the upper triangular system
    @ti.kernel
    def back_substitution(self, y: ti.template(), n_iter: ti.i32):
        for idx in range(n_iter):
            i = n_iter - 1 - idx  # Reverse index
            sum_val = 0.0
            for j in range(i + 1, n_iter):
                sum_val += self.H[i, j] * y[j]
            y[i] = (self.s[i] - sum_val) / self.H[i, i]


    # Kernel to update the solution x_new
    @ti.kernel
    def update_solution(self, x0: ti.template(), y: ti.template(), x_new: ti.template(), n_iter: ti.i32):
        for i, j in ti.ndrange(self.rows, self.cols):
            sum_val = 0.0
            for k in range(n_iter):
                sum_val += self.V[k, i, j] * y[k]
            x_new[i, j] = x0[i, j] + sum_val

    # GMRES solver function with user-provided linear_map (as a @ti.kernel)
    def gmres_solve(self, b, x0, linear_map, diagonal):
        # Set up the preconditioner (diagonal inverse)
        @ti.kernel
        def setup_diagonal_inv():
            for i, j in diagonal:
                if diagonal[i, j] != 0:
                    self.diagonal_inv[i, j] = 1.0 / diagonal[i, j]
                else:
                    self.diagonal_inv[i, j] = 0.0  # Handle zeros in the diagonal

        setup_diagonal_inv()

        # Step 1: r0 = b - linear_map(x0)
        linear_map(x0, self.r0)
        self.subtract_vectors(b, self.r0, self.r0)

        # Apply Jacobi preconditioner to r0: r0 = D^{-1} * r0
        self.apply_preconditioner(self.r0, self.r0)

        # Step 2: beta = norm(r0)
        beta = ti.field(dtype=ti.f32, shape=())
        self.normalize_and_compute_norm(self.r0, beta)

        # Step 3: V[0] = r0
        self.copy_vector(self.r0, self.V, 0)

        # Initialize s vector
        self.s.fill(0)
        self.s[0] = beta[None]

        # GMRES iteration loop
        for j in range(self.n_iter):
            # Extract Vj from V into self.Vj
            self.extract_Vj(j)

            # Compute w = linear_map(Vj)
            linear_map(self.Vj, self.w)

            # Apply preconditioner to w
            self.apply_preconditioner(self.w, self.w)

            # Orthogonalization process
            self.orthogonalize(j)

            # Apply Givens rotations
            self.apply_givens_rotations(j)

            # Check for convergence
            resid = abs(self.s[j + 1]) / beta[None]
            if resid < self.tol:
                self.n_iter = j + 1
                break

        # Solve the upper triangular system
        y = ti.field(dtype=ti.f32, shape=(self.n_iter))
        self.back_substitution(y, self.n_iter)

        # Update the solution x_new
        x_new = ti.field(dtype=ti.f32, shape=(self.rows, self.cols))
        self.update_solution(x0, y, x_new, self.n_iter)

        return x_new


@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])
        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 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 psi_max[p] < get_psi_pos(F[p]):
            psi_max[p] = get_psi_pos(F[p])
        color[p][2] = 1-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
            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
            grid_c[grid_idx] += weight * c[p] -0.13
            rhs[grid_idx] += F[p].determinant()*p_vol*(M+c[p]/dt)*weight


@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_c[i, j] /= grid_m[i, j]
        grid_v[i, j].y -= dt * gravity
        dist = attractor_pos[None] - dx * ti.Vector([i, j])
        grid_v[i, j] += dist / (0.01 + dist.norm()) * attractor_strength[None] * dt * 100
        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):
        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
            if 0 <= grid_idx[0] < n_grid and 0 <= grid_idx[1] < n_grid:
                weight = w[i][0] * w[j][1]
                new_c += (sol[grid_idx]-grid_c[grid_idx])*weight
        c[p] = max(0, min(c[p],c[p]+new_c))
@ti.kernel
def init():
    for i in range(n_particles):
        x[i] = [ti.random() * 0.05 + 0.2, ti.random() * 0.7 + 0.0]
        v[i] = [0, 0]
        c[i] = 1
        F[i] = ti.Matrix([[1, 0], [0, 1]])
# Initialize the particles
init()

# Main GUI loop
gui = ti.ui.Window("MPM88", (1024, 1024))
canvas = gui.get_canvas()
i = 0
while gui.running:
    canvas.set_background_color((1, 1, 1))
    for s in range(2):
        reset_grid()
        phase_p2g()
        get_preconditioner()
        solver = GMRES_Solver(n_grid, n_grid, n_iter=5)
        sol =  solver.gmres_solve(rhs,zeros,mvp,diag)
        phase_g2p()
        p2g()
        grid_op()
        g2p()
        print(c.to_numpy().mean())
    canvas.circles(x, radius=0.0025, per_vertex_color=color)
    #canvas.set_image(sol.to_numpy())
    gui.show()
    i += 1
    gui.save_image("./out/"+str(i)+".png")
Field c quickly falls to zero, fix that