import numpy as np
import struct
import taichi as ti
import meshio

@ti.data_oriented 
class VTUclass:
  def write_vtuASCII( self, filename, ps: ti.template() ):
    pos = ps.x
    velocity = ps.v
    density = ps.density
    pressure = ps.pressure
    id = ps.object_id
    rigid_density = ps.staticRigidDensity

    n = pos.shape[0]
    with open(filename, "w") as f:
        f.write('<?xml version="1.0"?>\n')
        f.write('<VTKFile type="UnstructuredGrid" version="0.1" byte_order="LittleEndian">\n')
        f.write('  <UnstructuredGrid>\n')
        f.write(f'    <Piece NumberOfPoints="{n}" NumberOfCells="{n}">\n')
        # Points
        f.write('      <Points>\n')
        f.write('        <DataArray type="Float32" NumberOfComponents="3" format="ascii">\n')
        for i in range(n):
            f.write(f'          {pos[i][0]} {pos[i][1]} {pos[i][2]}\n')
        f.write('        </DataArray>\n')
        f.write('      </Points>\n')
        # Cells
        f.write('      <Cells>\n')
        f.write('        <DataArray type="Int32" Name="connectivity" format="ascii">\n')
        f.write('          ' + " ".join(map(str, range(n))) + '\n')
        f.write('        </DataArray>\n')
        f.write('        <DataArray type="Int32" Name="offsets" format="ascii">\n')
        f.write('          ' + " ".join(str(i+1) for i in range(n)) + '\n')
        f.write('        </DataArray>\n')
        f.write('        <DataArray type="UInt8" Name="types" format="ascii">\n')
        f.write('          ' + " ".join(["1"] * n) + '\n')
        f.write('        </DataArray>\n')
        f.write('      </Cells>\n')

        # PointData
        f.write('      <PointData Scalars="scalars">\n')
        f.write('        <DataArray type="Float32" Name="density" format="ascii">\n')
        for i in range(density.shape[0]):
            if id[i] == 0: 	# liquid
               f.write(f'          {float(density[i])}\n')
            else:  		# rigid body
               f.write(f'          {float(rigid_density)}\n')
        f.write('        </DataArray>\n')
        f.write('        <DataArray type="Float32" Name="pressure" format="ascii">\n')
        for i in range(pressure.shape[0]):
            f.write(f'          {float(pressure[i])}\n')
        f.write('        </DataArray>\n')
        f.write('        <DataArray type="Float32" NumberOfComponents="3" Name="velocity" format="ascii">\n')
        for i in range(velocity.shape[0]):
            f.write(f'          {velocity[i][0]} {velocity[i][1]} {velocity[i][2]}\n')
        f.write('        </DataArray>\n')
        f.write('        <DataArray type="Int32" Name="object_id" format="ascii">\n')
        for i in range(id.shape[0]): 
            f.write(f'          {int(id[i])}\n')
        f.write('        </DataArray>\n')
        f.write('      </PointData>\n')

        f.write('    </Piece>\n')
        f.write('  </UnstructuredGrid>\n')
        f.write('</VTKFile>\n')

  def write_vtuBINARY(self, filename, ps: ti.template()):
        pos = ps.x
        velocity = ps.v
        density = ps.density
        pressure = ps.pressure
        id = ps.object_id

        n = int(pos.shape[0])
        pos_np = np.empty((n, 3), dtype=np.float32)
        vel_np = np.empty((n, 3), dtype=np.float32)
        dens_np = np.empty((n,), dtype=np.float32)
        pres_np = np.empty((n,), dtype=np.float32)
        id_np  = np.empty((n,), dtype=np.int32)
        for i in range(n):				 
            pos_np[i, 0] = float(pos[i][0])
            pos_np[i, 1] = float(pos[i][1])
            pos_np[i, 2] = float(pos[i][2])

            vel_np[i, 0] = float(velocity[i][0])
            vel_np[i, 1] = float(velocity[i][1])
            vel_np[i, 2] = float(velocity[i][2])

            if id[i] == 0: dens_np[i] = float(density[i])		# for liquid
            else:          dens_np[i] = float(ps.staticRigidDensity)    # for rigid

            pres_np[i] = float(pressure[i])
            id_np[i] = int(id[i])

        # Build appended blocks
        blocks = []
        blocks.append(pos_np.astype("<f4").tobytes())
        blocks.append(np.arange(n, dtype="<i4").tobytes())
        blocks.append(np.arange(1, n+1, dtype="<i4").tobytes())
        blocks.append((np.ones(n, dtype=np.uint8)).tobytes())
        blocks.append(dens_np.astype("<f4").tobytes())
        blocks.append(pres_np.astype("<f4").tobytes())
        blocks.append(vel_np.astype("<f4").tobytes())
        blocks.append(id_np.astype("<i4").tobytes())

        appended_offsets = []
        cur = 0
        for b in blocks:
            appended_offsets.append(cur)
            cur += 4 + len(b)

        header_lines = []
        header_lines.append('<?xml version="1.0"?>\n')
        header_lines.append('<VTKFile type="UnstructuredGrid" version="0.1" byte_order="LittleEndian">\n')
        header_lines.append('  <UnstructuredGrid>\n')
        header_lines.append(f'    <Piece NumberOfPoints="{n}" NumberOfCells="{n}">\n')
        header_lines.append('      <Points>\n')
        header_lines.append(f'        <DataArray type="Float32" NumberOfComponents="3" format="appended" offset="{appended_offsets[0]}"/>\n')
        header_lines.append('      </Points>\n')
        header_lines.append('      <Cells>\n')
        header_lines.append(f'        <DataArray type="Int32" Name="connectivity" format="appended" offset="{appended_offsets[1]}"/>\n')
        header_lines.append(f'        <DataArray type="Int32" Name="offsets" format="appended" offset="{appended_offsets[2]}"/>\n')
        header_lines.append(f'        <DataArray type="UInt8" Name="types" format="appended" offset="{appended_offsets[3]}"/>\n')
        header_lines.append('      </Cells>\n')
        header_lines.append('      <PointData Scalars="scalars">\n')
        header_lines.append(f'        <DataArray type="Float32" Name="density" format="appended" offset="{appended_offsets[4]}"/>\n')
        header_lines.append(f'        <DataArray type="Float32" Name="pressure" format="appended" offset="{appended_offsets[5]}"/>\n')
        header_lines.append(f'        <DataArray type="Float32" NumberOfComponents="3" Name="velocity" format="appended" offset="{appended_offsets[6]}"/>\n')
        header_lines.append(f'        <DataArray type="Int32" Name="object_id" format="appended" offset="{appended_offsets[7]}"/>\n')
        header_lines.append('      </PointData>\n')
        header_lines.append('    </Piece>\n')
        header_lines.append('  </UnstructuredGrid>\n')
        header_lines.append('  <AppendedData encoding="raw">\n')
        header = ''.join(header_lines).encode('utf-8')

        with open(filename, "wb") as f:
            f.write(header)
            f.write(b'_')
            for b in blocks:
                f.write(struct.pack('<I', len(b)))
                f.write(b)
            f.write(b'\n  </AppendedData>\n')
            f.write(b'</VTKFile>\n')

  @ti.kernel
  def compute_stress_and_vm(self, F_field: ti.template(),   # pass the Taichi matrix field
                          mu: ti.f32,
                          lam: ti.f32,
                          stress_out: ti.types.ndarray(),   # (n,9) float32
                          vm_out: ti.types.ndarray(),       # (n,) float32
                          n: ti.i32):
    for i in range(n):
        F = F_field[i]                     	# assume F is 3x3 taichi matrix field
        J = ti.max( ti.abs(F.determinant()), 1e-8 )			# avoid negative jitter
        B = F @ F.transpose()
        I = ti.Matrix([[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]])
        tau = mu * (B - I) + lam * ti.log(J) * I
        sigma = tau / J

        stress_out[i, 0] = sigma[0,0]
        stress_out[i, 1] = sigma[0,1]
        stress_out[i, 2] = sigma[0,2]
        stress_out[i, 3] = sigma[1,0]
        stress_out[i, 4] = sigma[1,1]
        stress_out[i, 5] = sigma[1,2]
        stress_out[i, 6] = sigma[2,0]
        stress_out[i, 7] = sigma[2,1]
        stress_out[i, 8] = sigma[2,2]

        mean = (sigma[0,0] + sigma[1,1] + sigma[2,2]) / 3.0
        sxx = sigma[0,0] - mean
        syy = sigma[1,1] - mean
        szz = sigma[2,2] - mean
        sxy = 0.5 * (sigma[0,1] + sigma[1,0])
        sxz = 0.5 * (sigma[0,2] + sigma[2,0])
        syz = 0.5 * (sigma[1,2] + sigma[2,1])
        vm_out[i] = ti.sqrt(1.5 * (sxx*sxx + syy*syy + szz*szz + 2.0*(sxy*sxy + sxz*sxz + syz*syz)))

  def write_vtuMPMASCII(self, filename, mpm: ti.template()):
      n = int( mpm.n_particles[None] )
      stress_np = np.zeros((n, 9), dtype=np.float32)
      vm_np = np.zeros((n,), dtype=np.float32)

      # Taichi -> NumPy
      pos = mpm.x.to_numpy()[:n]           # (n, 3)
      vel = mpm.v.to_numpy()[:n]           # (n, 3)
      mat = mpm.material.to_numpy()[:n]    # (n,)

      mu_ = mpm.mu_0
      lambda_ = mpm.lambda_0

      self.compute_stress_and_vm(mpm.F, mu_, lambda_, stress_np, vm_np, n)

      # VTU は vertex cell を使う
      cells = [("vertex", np.arange(n).reshape(-1, 1))]

      point_data = {
          "velocity": vel,     # [vx, vy, vz]
          "material": mat,     # scalar
          "von_mises": vm_np,
          "stress": stress_np,   # (n,9)
      }

      mesh = meshio.Mesh(
          points=pos,
          cells=cells,
          point_data=point_data
      )

      mesh.write(filename)




