path: root/Scripts/Impostors.cs

using UnityEngine;
using UnityEditor;
using System.IO;

[ExecuteInEditMode]
public class Impostors : MonoBehaviour
{
  [Header("Bounding Sphere")]
  public float sphere_radius_ = 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 float Radius => sphere_radius_ * transform.lossyScale.x;
  private const string OutputFolder = "Assets/yum_food/3ner/Impostor_Generated";

  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("_ImpostorMainCameraPos", cam.transform.position);
  }

  void OnDrawGizmos()
  {
    Gizmos.color = Color.cyan;
    Gizmos.DrawWireSphere(transform.position, 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 = transform.position + PlaneToHemiOctahedron(x, y) * (Radius + nearClippingDistance);
          Gizmos.DrawSphere(worldPos, Radius * 0.05f);
          Gizmos.DrawLine(worldPos, transform.position);
        }
      }
    }
  }

  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;

    for (int y = 0; y < gridResolution; y++)
    {
      for (int x = 0; x < gridResolution; x++)
      {
        Vector3 localDir = PlaneToHemiOctahedron(x, y);
        Vector3 worldPos = transform.position + (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(transform.position, transform.up);

        Camera cam = camObj.AddComponent<Camera>();
        cam.orthographic = true;
        cam.orthographicSize = sphere_radius_;
        cam.nearClipPlane = nearClippingDistance;
        cam.farClipPlane = sphere_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);
          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);
    if (cameras == null || cameras.Length != gridResolution * gridResolution || cameras[0] == null) 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 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))
    {
      if (!AssetDatabase.IsValidFolder("Assets/yum_food/3ner"))
        AssetDatabase.CreateFolder("Assets/yum_food", "3ner");
      AssetDatabase.CreateFolder("Assets/yum_food/3ner", "Impostor_Generated");
    }

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

    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", mipmaps: true, sRGB: false), "_Impostors_Metallic_Gloss_Atlas"),
      (new TextureExportSettings("depth", isEXR: true, mipmaps: false, sRGB: false, filter: FilterMode.Bilinear, uncompressed: true), "_Impostors_Depth_Atlas")
    };

    Texture2D[] atlases = { albedoAtlas, normalAtlas, metallicGlossAtlas, depthAtlas };
    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; }

      impostorMaterial = new Material(shader);
      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", sphere_radius_);
      impostorMaterial.SetFloat("_Impostors_Enabled", 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.SetParent(transform, false);
      impostorObject.transform.localScale = Vector3.one * sphere_radius_ * 2f;
      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) impostorObject.SetActive(!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();
  }
}