path: root/Scripts/Impostors.cs

#if UNITY_EDITOR
using UnityEngine;
using UnityEditor;
using System.IO;

[ExecuteInEditMode]
public class Impostors : MonoBehaviour
{
  [Header("Bounding Sphere")]
  public float radiusScale = 1f;

  [Header("Grid Settings")]
  [Range(2, 20)] public int gridResolution = 5;

  [Header("Camera Settings")]
  [Range(1, 4096)] public int cameraResolution = 256;
  public float nearClippingDistance = 0.01f;
  public LayerMask cullingMask = -1;
  public bool renderSkybox = false;

  [Header("Original Mesh")]
  public GameObject originalMesh;

  [HideInInspector] public Camera[] cameras;

  private GameObject impostorObject;
  private Material impostorMaterial;

  public bool HasImpostor => impostorObject != null;
  private string OutputFolder => GetOutputFolder();

  Bounds GetCombinedBounds()
  {
    Renderer[] allRenderers = originalMesh.GetComponentsInChildren<Renderer>(true);
    Renderer[] renderers = System.Array.FindAll(allRenderers, r => r.gameObject != impostorObject);
    Bounds bounds = renderers[0].bounds;
    for (int i = 1; i < renderers.Length; i++)
      bounds.Encapsulate(renderers[i].bounds);
    return bounds;
  }

  float GetBoundingSphereRadius()
  {
    return GetCombinedBounds().extents.magnitude * radiusScale;
  }

  void OnEnable() => Camera.onPreRender += UpdateMainCameraPos;
  void OnDisable() => Camera.onPreRender -= UpdateMainCameraPos;

  static void UpdateMainCameraPos(Camera cam)
  {
    if (cam.cameraType == CameraType.Game || cam.cameraType == CameraType.SceneView)
      Shader.SetGlobalVector("_Impostors_Main_Camera_Pos", cam.transform.position);
  }

  void OnDrawGizmos()
  {
    if (originalMesh == null) return;

    Bounds bounds = GetCombinedBounds();
    Vector3 center = bounds.center;
    float radius = GetBoundingSphereRadius();

    Gizmos.color = Color.cyan;
    Gizmos.DrawWireSphere(center, radius);

    if (Application.isEditor && gridResolution > 0)
    {
      Gizmos.color = Color.yellow;
      for (int y = 0; y < gridResolution; y++)
      {
        for (int x = 0; x < gridResolution; x++)
        {
          Vector3 worldPos = center + PlaneToHemiOctahedron(x, y) * (radius + nearClippingDistance);
          Gizmos.DrawSphere(worldPos, radius * 0.05f);
          Gizmos.DrawLine(worldPos, center);
        }
      }
    }
  }

  string GetOutputFolder()
  {
    string hierarchyPath = GetHierarchyPath();
    return $"Assets/yum_food/3ner/Impostor_Generated/{hierarchyPath}";
  }

  string GetHierarchyPath()
  {
    string path = gameObject.name;
    Transform parent = transform.parent;
    while (parent != null)
    {
      path = parent.name + "/" + path;
      parent = parent.parent;
    }
    return path;
  }

  Vector3 PlaneToHemiOctahedron(int gridX, int gridY)
  {
    float x = (gridX / (float)(gridResolution - 1)) * 2f - 1f;
    float z = (gridY / (float)(gridResolution - 1)) * 2f - 1f;

    // Rotate 45° to fit square into diamond
    float x_rot = (x + z) * 0.5f;
    float z_rot = (z - x) * 0.5f;

    // Octahedral decode
    float y = Mathf.Max(0f, 1f - Mathf.Abs(x_rot) - Mathf.Abs(z_rot));

    // Normalize
    Vector3 oct_pos = new Vector3(x_rot, y, z_rot).normalized;

    // Rotate back by -45° around y
    float rcp_sqrt2 = 0.70710678f;
    float x_unrot = (oct_pos.x - oct_pos.z) * rcp_sqrt2;
    float z_unrot = (oct_pos.x + oct_pos.z) * rcp_sqrt2;

    return new Vector3(x_unrot, oct_pos.y, z_unrot);
  }

  public void CreateCameras()
  {
    DestroyExistingCameras();

    GameObject parent = new GameObject("Cameras");
    parent.transform.SetParent(transform, false);

    cameras = new Camera[gridResolution * gridResolution];
    int idx = 0;

    Bounds bounds = GetCombinedBounds();
    Vector3 center = bounds.center;
    float radius = GetBoundingSphereRadius();

    for (int y = 0; y < gridResolution; y++)
    {
      for (int x = 0; x < gridResolution; x++)
      {
        Vector3 localDir = PlaneToHemiOctahedron(x, y);
        Vector3 worldPos = center + (transform.rotation * localDir) * (radius + nearClippingDistance);

        GameObject camObj = new GameObject($"Camera_{x}_{y}");
        camObj.transform.SetParent(parent.transform, false);
        camObj.transform.position = worldPos;
        camObj.transform.LookAt(center, transform.up);

        Camera cam = camObj.AddComponent<Camera>();
        cam.orthographic = true;
        cam.orthographicSize = radius;
        cam.nearClipPlane = nearClippingDistance;
        cam.farClipPlane = radius * 2f + nearClippingDistance;
        cam.cullingMask = cullingMask;
        cam.clearFlags = renderSkybox ? CameraClearFlags.Skybox : CameraClearFlags.SolidColor;
        cam.backgroundColor = Color.clear;
        cam.enabled = false;

        cameras[idx++] = cam;
      }
    }
  }

  public void DestroyExistingCameras()
  {
    Transform camsTransform = transform.Find("Cameras");
    if (camsTransform != null) DestroyImmediate(camsTransform.gameObject);
    cameras = null;
  }

  struct BakePass
  {
    public string name;
    public string keyword;
    public Texture2D atlas;
    public bool isDepth;
    public bool linear;

    public BakePass(string name, string keyword, Texture2D atlas, bool isDepth = false, bool linear = false)
    {
      this.name = name;
      this.keyword = keyword;
      this.atlas = atlas;
      this.isDepth = isDepth;
      this.linear = linear;
    }
  }

  void RenderAtlasPass(Texture2D atlas, RenderTexture colorRT, RenderTexture depthOnlyRT, Material depthBlitMat = null)
  {
    bool isDepth = depthBlitMat != null;
    RenderTexture linearDepthRT = isDepth ? RenderTexture.GetTemporary(cameraResolution, cameraResolution, 0, RenderTextureFormat.RFloat) : null;

    int idx = 0;
    for (int y = 0; y < gridResolution; y++)
    {
      for (int x = 0; x < gridResolution; x++)
      {
        Camera cam = cameras[idx++];
        cam.SetTargetBuffers(colorRT.colorBuffer, depthOnlyRT.depthBuffer);
        cam.Render();

        if (isDepth)
        {
          depthBlitMat.SetTexture("_DepthTex", depthOnlyRT);
          depthBlitMat.SetFloat("_Near", cam.nearClipPlane);
          depthBlitMat.SetFloat("_Far", cam.farClipPlane);
          Graphics.Blit(null, linearDepthRT, depthBlitMat);
          RenderTexture.active = linearDepthRT;
        }
        else
        {
          RenderTexture.active = colorRT;
        }

        Texture2D temp = new Texture2D(cameraResolution, cameraResolution, isDepth ? TextureFormat.RFloat : TextureFormat.RGBA32, false);
        temp.ReadPixels(new Rect(0, 0, cameraResolution, cameraResolution), 0, 0);
        temp.Apply();
        atlas.SetPixels(x * cameraResolution, y * cameraResolution, cameraResolution, cameraResolution, temp.GetPixels());
        DestroyImmediate(temp);
      }
    }
    atlas.Apply();

    if (linearDepthRT != null) RenderTexture.ReleaseTemporary(linearDepthRT);
  }

  void SetMaterialDebugKeyword(string keyword, bool enabled)
  {
    if (originalMesh == null || string.IsNullOrEmpty(keyword)) return;
    foreach (Renderer r in originalMesh.GetComponentsInChildren<Renderer>(true))
    {
      foreach (Material mat in r.sharedMaterials)
      {
        if (mat != null)
        {
          if (enabled)
            mat.EnableKeyword(keyword);
          else
            mat.DisableKeyword(keyword);
        }
      }
    }
  }

  void ExecutePassesSequentially(BakePass[] passes, RenderTexture srgbRT, RenderTexture linearRT, RenderTexture depthOnlyRT, Material depthBlitMat, int passIndex = 0)
  {
    if (passIndex >= passes.Length)
    {
      // All passes complete
      RenderTexture.active = null;
      RenderTexture.ReleaseTemporary(srgbRT);
      RenderTexture.ReleaseTemporary(linearRT);
      RenderTexture.ReleaseTemporary(depthOnlyRT);
      DestroyImmediate(depthBlitMat);

      SaveAndConfigureTextures(passes[0].atlas, passes[1].atlas, passes[2].atlas, passes[3].atlas);
      return;
    }

    BakePass pass = passes[passIndex];
    Debug.Log($"Baking {pass.name} pass...");

    SetMaterialDebugKeyword(pass.keyword, true);

    UnityEditor.EditorApplication.delayCall += () => {
      RenderTexture colorRT = pass.linear ? linearRT : srgbRT;
      RenderAtlasPass(pass.atlas, colorRT, depthOnlyRT, pass.isDepth ? depthBlitMat : null);
      SetMaterialDebugKeyword(pass.keyword, false);
      ExecutePassesSequentially(passes, srgbRT, linearRT, depthOnlyRT, depthBlitMat, passIndex + 1);
    };
  }

  public void BakeTexture()
  {
    SetRenderersEnabled(true);
    CreateCameras();

    Shader depthBlitShader = Shader.Find("Hidden/yum_food/DepthBlit");
    if (depthBlitShader == null) { Debug.LogError("DepthBlit shader not found"); return; }
    Material depthBlitMat = new Material(depthBlitShader);

    int size = cameraResolution * gridResolution;
    BakePass[] passes = new BakePass[]
    {
      new BakePass("albedo", "_DEBUG_VIEW_UNLIT", new Texture2D(size, size, TextureFormat.RGBA32, false)),
      new BakePass("normals", "_DEBUG_VIEW_WORLD_SPACE_NORMALS", new Texture2D(size, size, TextureFormat.RGBA32, false), linear: true),
      new BakePass("metallic/gloss", "_DEBUG_VIEW_METALLIC_GLOSS", new Texture2D(size, size, TextureFormat.RGBA32, false), linear: true),
      new BakePass("depth", "", new Texture2D(size, size, TextureFormat.RFloat, false), isDepth: true)
    };

    RenderTextureDescriptor srgbDesc = new RenderTextureDescriptor(cameraResolution, cameraResolution, RenderTextureFormat.ARGB32, 24);
    srgbDesc.sRGB = true;
    RenderTexture srgbRT = RenderTexture.GetTemporary(srgbDesc);

    RenderTextureDescriptor linearDesc = new RenderTextureDescriptor(cameraResolution, cameraResolution, RenderTextureFormat.ARGB32, 24);
    linearDesc.sRGB = false;
    RenderTexture linearRT = RenderTexture.GetTemporary(linearDesc);

    RenderTextureDescriptor depthDesc = new RenderTextureDescriptor(cameraResolution, cameraResolution, RenderTextureFormat.Depth, 24);
    RenderTexture depthOnlyRT = RenderTexture.GetTemporary(depthDesc);

    // Ensure all debug keywords start disabled
    foreach (var pass in passes) SetMaterialDebugKeyword(pass.keyword, false);

    ExecutePassesSequentially(passes, srgbRT, linearRT, depthOnlyRT, depthBlitMat);
  }

  struct TextureExportSettings
  {
    public string suffix;
    public bool isEXR;
    public bool mipmaps;
    public bool sRGB;
    public FilterMode filter;
    public bool alphaTransparency;
    public bool uncompressed;
    public bool isNormalMap;

    public TextureExportSettings(string suffix, bool isEXR = false, bool mipmaps = true, bool sRGB = true,
                                  FilterMode filter = FilterMode.Trilinear, bool alphaTransparency = false,
                                  bool uncompressed = false, bool isNormalMap = false)
    {
      this.suffix = suffix;
      this.isEXR = isEXR;
      this.mipmaps = mipmaps;
      this.sRGB = sRGB;
      this.filter = filter;
      this.alphaTransparency = alphaTransparency;
      this.uncompressed = uncompressed;
      this.isNormalMap = isNormalMap;
    }
  }

  void DilateTexture(Texture2D tex, Texture2D alphaSource, bool preserveAlpha = false, int iterations = 8)
  {
    int w = tex.width, h = tex.height;
    Color[] pixels = tex.GetPixels();
    Color[] alpha = alphaSource.GetPixels();
    Color[] origAlpha = preserveAlpha ? (Color[])alpha.Clone() : null;
    Color[] buffer = (Color[])pixels.Clone();

    bool[] filled = new bool[pixels.Length];
    for (int i = 0; i < pixels.Length; i++)
      filled[i] = alpha[i].a > 0.01f;
    bool[] filledNext = new bool[filled.Length];

    int[] dx = { -1, 1, 0, 0 };
    int[] dy = { 0, 0, -1, 1 };

    for (int iter = 0; iter < iterations; iter++)
    {
      System.Array.Copy(pixels, buffer, pixels.Length);
      System.Array.Copy(filled, filledNext, filled.Length);

      for (int y = 0; y < h; y++)
      {
        for (int x = 0; x < w; x++)
        {
          int i = y * w + x;
          if (filled[i]) continue;

          Color sum = Color.clear;
          int count = 0;
          for (int d = 0; d < 4; d++)
          {
            int nx = x + dx[d], ny = y + dy[d];
            if (nx < 0 || nx >= w || ny < 0 || ny >= h) continue;
            int ni = ny * w + nx;
            if (filled[ni])
            {
              sum += pixels[ni];
              count++;
            }
          }
          if (count > 0)
          {
            buffer[i] = sum / count;
            buffer[i].a = 1;
            filledNext[i] = true;
          }
        }
      }
      var tmpFilled = filled; filled = filledNext; filledNext = tmpFilled;
      var tmpPixels = pixels; pixels = buffer; buffer = tmpPixels;
    }

    if (preserveAlpha)
      for (int i = 0; i < pixels.Length; i++) pixels[i].a = origAlpha[i].a;

    tex.SetPixels(pixels);
    tex.Apply();
  }

  void DilateDepthTexture(Texture2D depthTex, Texture2D alphaSource, int iterations = 8)
  {
    int w = depthTex.width, h = depthTex.height;
    Color[] depth = depthTex.GetPixels();
    Color[] mask = alphaSource.GetPixels();

    bool[] filled = new bool[depth.Length];
    for (int i = 0; i < depth.Length; i++)
      filled[i] = mask[i].a > 0.01f;

    Color[] buffer = (Color[])depth.Clone();
    bool[] filledNext = new bool[filled.Length];

    int[] dx = { -1, 1, 0, 0 };
    int[] dy = { 0, 0, -1, 1 };

    for (int iter = 0; iter < iterations; iter++)
    {
      // Preserve all previously-filled pixels; only write new pixels outside the alpha mask.
      System.Array.Copy(depth, buffer, depth.Length);
      System.Array.Copy(filled, filledNext, filled.Length);

      for (int y = 0; y < h; y++)
      {
        for (int x = 0; x < w; x++)
        {
          int i = y * w + x;
          if (filled[i]) continue;

          float sum = 0f;
          int count = 0;
          for (int d = 0; d < 4; d++)
          {
            int nx = x + dx[d], ny = y + dy[d];
            if (nx < 0 || nx >= w || ny < 0 || ny >= h) continue;
            int ni = ny * w + nx;
            if (!filled[ni]) continue;
            sum += depth[ni].r;
            count++;
          }

          if (count > 0)
          {
            buffer[i].r = sum / count;
            filledNext[i] = true;
          }
        }
      }

      var tmpFilled = filled; filled = filledNext; filledNext = tmpFilled;
      var tmpDepth = depth; depth = buffer; buffer = tmpDepth;
    }

    depthTex.SetPixels(depth);
    depthTex.Apply();
  }

  void SaveAndConfigureTexture(Texture2D atlas, TextureExportSettings settings, string baseName, out string path)
  {
    path = Path.Combine(OutputFolder, $"{baseName}_{settings.suffix}.{(settings.isEXR ? "exr" : "png")}");
    byte[] data = settings.isEXR ? atlas.EncodeToEXR(Texture2D.EXRFlags.OutputAsFloat) : atlas.EncodeToPNG();
    File.WriteAllBytes(Path.Combine(Application.dataPath, "..", path), data);
    DestroyImmediate(atlas);
  }

  void ConfigureTextureImporter(string path, TextureExportSettings settings)
  {
    TextureImporter importer = AssetImporter.GetAtPath(path) as TextureImporter;
    if (importer != null)
    {
      if (settings.isNormalMap) importer.textureType = TextureImporterType.NormalMap;
      importer.mipmapEnabled = settings.mipmaps;
      importer.sRGBTexture = settings.sRGB;
      importer.wrapMode = TextureWrapMode.Clamp;
      importer.filterMode = settings.filter;
      if (settings.alphaTransparency) importer.alphaIsTransparency = true;
      if (settings.uncompressed) importer.textureCompression = TextureImporterCompression.Uncompressed;
      EditorUtility.SetDirty(importer);
      importer.SaveAndReimport();
      AssetDatabase.ImportAsset(path, ImportAssetOptions.ForceSynchronousImport);
    }
  }

  void SaveAndConfigureTextures(Texture2D albedoAtlas, Texture2D normalAtlas, Texture2D metallicGlossAtlas, Texture2D depthAtlas)
  {
    if (!AssetDatabase.IsValidFolder(OutputFolder))
    {
      string[] pathParts = OutputFolder.Split('/');
      string currentPath = "";
      for (int i = 0; i < pathParts.Length; i++)
      {
        string parentPath = i == 0 ? "" : string.Join("/", pathParts, 0, i);
        currentPath = i == 0 ? pathParts[0] : $"{parentPath}/{pathParts[i]}";
        if (!AssetDatabase.IsValidFolder(currentPath))
        {
          string folderName = pathParts[i];
          AssetDatabase.CreateFolder(parentPath, folderName);
        }
      }
    }

    string baseName = gameObject.name.Replace(" ", "_");

    // Dilate before combining
    DilateTexture(albedoAtlas, albedoAtlas, preserveAlpha: true);
    DilateTexture(normalAtlas, albedoAtlas);
    DilateTexture(metallicGlossAtlas, albedoAtlas);
    DilateDepthTexture(depthAtlas, albedoAtlas);

    // Combine metallic_gloss (RG) with depth (B) into a single texture
    Texture2D combinedAtlas = new Texture2D(metallicGlossAtlas.width, metallicGlossAtlas.height, TextureFormat.RGBA32, false);
    Color[] metallicGlossPixels = metallicGlossAtlas.GetPixels();
    Color[] depthPixels = depthAtlas.GetPixels();
    Color[] combinedPixels = new Color[metallicGlossPixels.Length];

    for (int i = 0; i < combinedPixels.Length; i++)
    {
      combinedPixels[i] = new Color(
        metallicGlossPixels[i].r,  // Metallic
        metallicGlossPixels[i].g,  // Gloss
        depthPixels[i].r,          // Depth
        1.0f                       // Alpha (unused)
      );
    }
    combinedAtlas.SetPixels(combinedPixels);
    combinedAtlas.Apply();

    // Clean up the separate textures
    DestroyImmediate(metallicGlossAtlas);
    DestroyImmediate(depthAtlas);

    var exportSettings = new (TextureExportSettings settings, string materialProp)[]
    {
      (new TextureExportSettings("albedo", mipmaps: true, sRGB: true, alphaTransparency: true), "_Impostors_Atlas"),
      (new TextureExportSettings("normal", mipmaps: true, sRGB: false), "_Impostors_Normal_Atlas"),
      (new TextureExportSettings("metallic_gloss_depth", mipmaps: true, sRGB: false, uncompressed: true), "_Impostors_Metallic_Gloss_Depth_Atlas")
    };

    Texture2D[] atlases = { albedoAtlas, normalAtlas, combinedAtlas };
    string[] paths = new string[exportSettings.Length];
    for (int i = 0; i < exportSettings.Length; i++)
      SaveAndConfigureTexture(atlases[i], exportSettings[i].settings, baseName, out paths[i]);

    AssetDatabase.Refresh();

    for (int i = 0; i < paths.Length; i++)
      ConfigureTextureImporter(paths[i], exportSettings[i].settings);

    // Wait for shader compiler then create impostor
    WaitForCompilerThenCreateImpostor(paths, exportSettings, baseName);
  }

  void WaitForCompilerThenCreateImpostor(string[] paths, (TextureExportSettings settings, string materialProp)[] exportSettings, string baseName)
  {
    if (EditorApplication.isCompiling)
    {
      EditorApplication.delayCall += () => WaitForCompilerThenCreateImpostor(paths, exportSettings, baseName);
      return;
    }

    Texture2D[] textures = new Texture2D[paths.Length];
    for (int i = 0; i < paths.Length; i++)
      textures[i] = AssetDatabase.LoadAssetAtPath<Texture2D>(paths[i]);

    if (textures[0] != null)
    {
      DestroyExistingImpostor();

      Shader shader = Shader.Find("yum_food/3ner");
      if (shader == null) { Debug.LogError("Shader not found"); return; }

      Bounds bounds = GetCombinedBounds();
      float radius = GetBoundingSphereRadius();

      impostorMaterial = new Material(shader);
      impostorMaterial.enableInstancing = true;
      for (int i = 0; i < textures.Length; i++)
        impostorMaterial.SetTexture(exportSettings[i].materialProp, textures[i]);
      impostorMaterial.SetInt("_Impostors_Grid_Resolution", gridResolution);
      impostorMaterial.SetFloat("_Impostors_Sphere_Radius", radius);
      impostorMaterial.SetFloat("_Impostors_Near_Clip", nearClippingDistance);
      impostorMaterial.SetFloat("_Impostors_Far_Clip", radius * 2f + nearClippingDistance);
      impostorMaterial.SetFloat("_ZTest", (float)UnityEngine.Rendering.CompareFunction.LessEqual);
      impostorMaterial.SetFloat("_Impostors_Enabled", 1);
      impostorMaterial.SetFloat("_Impostors_Parallax", 1);
      impostorMaterial.SetFloat("_Cull", (float)UnityEngine.Rendering.CullMode.Front);
      AssetDatabase.CreateAsset(impostorMaterial, Path.Combine(OutputFolder, $"{baseName}_mat.mat"));

      impostorObject = GameObject.CreatePrimitive(PrimitiveType.Quad);
      impostorObject.name = "Impostor";
      impostorObject.transform.position = bounds.center;
      impostorObject.transform.localScale = Vector3.one * radius * 2.0f;
      impostorObject.transform.SetParent(transform, true);
      DestroyImmediate(impostorObject.GetComponent<Collider>());
      impostorObject.GetComponent<MeshRenderer>().sharedMaterial = impostorMaterial;

      SetRenderersEnabled(false);
      Debug.Log("Impostor baking complete!");
    }
  }

  public void DestroyExistingImpostor()
  {
    if (impostorObject != null) DestroyImmediate(impostorObject);
    impostorObject = null;
    impostorMaterial = null;
  }

  public void ToggleRenderers()
  {
    if (originalMesh == null) return;
    bool showing = originalMesh.GetComponentInChildren<Renderer>()?.enabled ?? false;
    SetRenderersEnabled(!showing);
  }

  void SetRenderersEnabled(bool enabled)
  {
    if (originalMesh == null) return;
    foreach (Renderer r in originalMesh.GetComponentsInChildren<Renderer>(true))
      r.enabled = enabled;
    if (impostorObject != null)
    {
      Renderer impostorRenderer = impostorObject.GetComponent<Renderer>();
      if (impostorRenderer != null) impostorRenderer.enabled = !enabled;
    }
  }
}

[CustomEditor(typeof(Impostors))]
public class ImpostorsEditor : Editor
{
  public override void OnInspectorGUI()
  {
    DrawDefaultInspector();
    Impostors s = (Impostors)target;

    GUILayout.Space(10);
    GUILayout.Label("Impostor Management", EditorStyles.boldLabel);

    EditorGUILayout.BeginHorizontal();
    if (GUILayout.Button("Create Impostor", GUILayout.Height(40))) s.BakeTexture();
    GUI.enabled = s.HasImpostor;
    if (GUILayout.Button("Destroy Impostor", GUILayout.Height(40)))
    {
      s.DestroyExistingImpostor();
      s.DestroyExistingCameras();
    }
    GUI.enabled = true;
    EditorGUILayout.EndHorizontal();

    if (s.originalMesh != null && GUILayout.Button("Toggle Visibility", GUILayout.Height(30)))
      s.ToggleRenderers();
  }
}
#endif  // UNITY_EDITOR