import numpy as np
import inspect

class VTUclass:
  def _to_numpy(self, obj):
    if obj is None:
        return None
    if isinstance(obj, np.ndarray):
        return obj
    if hasattr(obj, "to_numpy"):
        try:
            return obj.to_numpy()
        except Exception:
            pass
    if hasattr(obj, "numpy"):
        try:
            return obj.numpy()
        except Exception:
            pass
    try:
        return np.asarray(obj)
    except Exception:
        return None

  def extract_pos_density_vel_pressure(self, ps, use_vis_buffer=True):
    """
    Return (pos, density, vel, pressure) as numpy arrays (or None).
    Prefers x_vis_buffer, density, v_buffer, pressure_buffer where available.
    """
    # ensure vis buffer updated if available
    if use_vis_buffer and hasattr(ps, "copy_to_vis_buffer"):
        try:
            ps.copy_to_vis_buffer()
        except Exception:
            pass

    pos_candidates = ["x_vis_buffer", "x_buffer", "x"]
    dens_candidates = ["density", "density_buffer", "density_adv", "density_adv_buffer"]
    vel_candidates = ["v_buffer", "v"]
    pres_candidates = ["pressure", "pressure_buffer"]

    def pick_one_from(obj, candidates):
        # use self._to_numpy to convert possible taichi fields/buffers to numpy
        for name in candidates:
            if hasattr(obj, name):
                raw = getattr(obj, name)
                arr = self._to_numpy(raw)
                if arr is not None:
                    # ensure non-empty
                    try:
                        if arr.size > 0:
                            return arr
                    except Exception:
                        # fallback: if arr is scalar-like treat as valid
                        return arr
        return None

    pos = pick_one_from(ps, pos_candidates)
    density = pick_one_from(ps, dens_candidates)
    vel = pick_one_from(ps, vel_candidates)
    pressure = pick_one_from(ps, pres_candidates)

    # reshape conveniences
    if pos is not None:
        pos = np.asarray(pos)
        if pos.ndim == 1 and pos.size % 3 == 0:
            pos = pos.reshape(-1, 3)
        elif pos.ndim == 2 and pos.shape[1] != 3:
            # try transposing if shape looks swapped
            if pos.shape[0] == 3 and pos.shape[1] != 3:
                pos = pos.T
    if vel is not None:
        vel = np.asarray(vel)
        if vel.ndim == 1 and vel.size % 3 == 0:
            vel = vel.reshape(-1, 3)
    if density is not None:
        density = np.asarray(density).reshape(-1,)
    if pressure is not None:
        pressure = np.asarray(pressure).reshape(-1,)

    return pos, density, vel, pressure

  def write_vtu_with_pressure(self, filename, pos, density=None, velocity=None,
                            pressure=None, extra_point_data=None):
    """
    Write VTU (ASCII) including optional pressure field.
    extra_point_data: dict name->1D array of length N
    """
    pos = np.asarray(pos, dtype=np.float32)
    n = pos.shape[0]
    assert pos.ndim == 2 and pos.shape[1] == 3, "pos must be (N,3)"
    if density is not None:
        density = np.asarray(density, dtype=np.float32).reshape(n,)
    if velocity is not None:
        velocity = np.asarray(velocity, dtype=np.float32).reshape(n,3)
    if pressure is not None:
        pressure = np.asarray(pressure, dtype=np.float32).reshape(n,)
    extra_point_data = extra_point_data or {}

    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: include density, pressure, velocity, extra
        if (density is not None) or (pressure is not None) or (velocity is not None) or extra_point_data:
            f.write('      <PointData Scalars="scalars">\n')
            if density is not None:
                f.write('        <DataArray type="Float32" Name="density" format="ascii">\n')
                for d in density:
                    f.write(f'          {float(d)}\n')
                f.write('        </DataArray>\n')
            if pressure is not None:
                f.write('        <DataArray type="Float32" Name="pressure" format="ascii">\n')
                for p in pressure:
                    f.write(f'          {float(p)}\n')
                f.write('        </DataArray>\n')
            if velocity is not None:
                f.write('        <DataArray type="Float32" NumberOfComponents="3" Name="velocity" format="ascii">\n')
                for v in velocity:
                    f.write(f'          {v[0]} {v[1]} {v[2]}\n')
                f.write('        </DataArray>\n')
            for name, arr in extra_point_data.items():
                arr = np.asarray(arr, dtype=np.float32).reshape(n,)
                f.write(f'        <DataArray type="Float32" Name="{name}" format="ascii">\n')
                for val in arr:
                    f.write(f'          {float(val)}\n')
                f.write('        </DataArray>\n')
            f.write('      </PointData>\n')

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

  # VTUclass のメソッド例（クラス内に入れる想定）
  def write_vtu_with_material(self, filename, pos, density=None, velocity=None,
                            pressure=None, extra_point_data=None, material=None, object_id=None):
    """
    pos: (N,3)
    material: length-N integer array (0: fluid, 1: solid など)
    object_id: length-N integer array （任意）
    extra_point_data: dict name->1D array
    """
    import numpy as np
    pos = np.asarray(pos, dtype=np.float32)
    n = pos.shape[0]
    assert pos.ndim == 2 and pos.shape[1] == 3

    if density is not None: density = np.asarray(density, dtype=np.float32).reshape(n,)
    if velocity is not None: velocity = np.asarray(velocity, dtype=np.float32).reshape(n,3)
    if pressure is not None: pressure = np.asarray(pressure, dtype=np.float32).reshape(n,)
    if material is not None: material = np.asarray(material, dtype=np.int32).reshape(n,)
    if object_id is not None: object_id = np.asarray(object_id, dtype=np.int32).reshape(n,)
    extra_point_data = extra_point_data or {}

    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
        if any(x is not None for x in (density, pressure, velocity, material, object_id)) or extra_point_data:
            f.write('      <PointData Scalars="scalars">\n')
            if density is not None:
                f.write('        <DataArray type="Float32" Name="density" format="ascii">\n')
                for d in density: f.write(f'          {float(d)}\n')
                f.write('        </DataArray>\n')
            if pressure is not None:
                f.write('        <DataArray type="Float32" Name="pressure" format="ascii">\n')
                for p in pressure: f.write(f'          {float(p)}\n')
                f.write('        </DataArray>\n')
            if velocity is not None:
                f.write('        <DataArray type="Float32" NumberOfComponents="3" Name="velocity" format="ascii">\n')
                for v in velocity: f.write(f'          {v[0]} {v[1]} {v[2]}\n')
                f.write('        </DataArray>\n')
            if material is not None:
                f.write('        <DataArray type="Int32" Name="material" format="ascii">\n')
                for m in material: f.write(f'          {int(m)}\n')
                f.write('        </DataArray>\n')
            if object_id is not None:
                f.write('        <DataArray type="Int32" Name="object_id" format="ascii">\n')
                for oid in object_id: f.write(f'          {int(oid)}\n')
                f.write('        </DataArray>\n')
            for name, arr in extra_point_data.items():
                arr = np.asarray(arr).reshape(n,)
                f.write(f'        <DataArray type="Float32" Name="{name}" format="ascii">\n')
                for val in arr: f.write(f'          {float(val)}\n')
                f.write('        </DataArray>\n')
            f.write('      </PointData>\n')

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

