using System.IO;
using System.Collections.Generic;
using System.Reflection;
using UnityEditor;
using UnityEngine;


namespace FourierWater.Editor
{
    public class FourierWater : EditorWindow
    {
        [Tooltip("Size in meters to simulate")]
        public float sizeM = 256f;

        [SerializeField]
        private int bigN = 1024;

        [Tooltip("Gravitational acceleration (m/s^2)")]
        private float g = 9.81f;
        [Tooltip("Surface tension (N/m)")]
        private float surfaceTension = 0.076f;
        [Tooltip("Density (kg / m^3)")]
        private float density = 1000f;
        [Tooltip("Wind speed (m/s)")]
        private float U = 30f;
        [Tooltip("Fetch (m)")]
        private float F = 100f * 1000f;
        [Tooltip("Depth (m)")]
        private float h = 60f;
        [Tooltip("Peak enhancement (unitless)")]
        private float gamma = 4f;

        [SerializeField]
        private MeshFilter waterMesh;

        // Viscous deep water wave relation. See:
        //   https://yummers.dev/wave-dispersion-relations-with-derivatives.html
        private float Omega(float k)
        {
            float k3 = k * k * k;
            return Mathf.Sqrt(g * k + surfaceTension * k3 / density);
        }

        private float DOmegaDK(float k)
        {
            float k2 = k * k;
            float k3 = k2 * k;
            return (g + 3 * surfaceTension * k2 / density) / (2 * Mathf.Sqrt(g * k + surfaceTension * k3 / density));
        }

        // JONSWAP spectrum
        private float SJonswap(float omega) {
          // Correction from equation 8 of Barfuss et. al. "The Impact of
          // Offshore Wind Farms on Sea State Demonstrated by Airborne LiDAR
          // Measurements."
          float omega_p = 22.0f * Mathf.Pow((g * g / (U * F)), 0.33333f);
          float sigma = (omega <= omega_p) ? 0.07f : 0.09f;
          float alpha = 0.076f * Mathf.Pow(U * U / (F * g), 0.22f);
          float r = Mathf.Exp(-(omega - omega_p) * (omega - omega_p) / (2 * sigma * sigma * omega_p * omega_p));
          float omega2 = omega * omega;
          float omega4 = omega2 * omega2;
          float omega5 = omega4 * omega;
          float meow4 = (omega_p / omega);
          meow4 = meow4 * meow4;
          meow4 = meow4 * meow4;
          float peak_enhancement = Mathf.Pow(gamma, r);
          return (alpha * g * g / omega5) * Mathf.Exp(-1.25f * meow4) * peak_enhancement;
        }

        // TMA spectrum
        private float STma(float omega) {
          float omega_h = omega * Mathf.Sqrt(h / g);
          if (omega_h <= 1.0f ) {
            return SJonswap(omega) * 0.5f * omega_h * omega_h;
          } else if (omega_h < 2.0f) {
            return SJonswap(omega) * (1.0f - 0.5f * (2.0f - omega_h) * (2.0f - omega_h));
          } else {
            return SJonswap(omega);
          }
        }

        // Flat directional spreading function
        private float DFlat()
        {
          return 1.0f / (3.14159265f * 2.0f);
        }

        // https://stackoverflow.com/a/218600
        private float RandomGaussian(System.Random r)
        {
          float u1 = (float) (1.0f - r.NextDouble());
          float u2 = (float) (1.0f - r.NextDouble());
          float randStdNormal = Mathf.Sqrt(-2.0f * Mathf.Log(u1)) * Mathf.Sin(2.0f * Mathf.PI * u2);
          // Want mean=0 stdev=1 so just return randStdNormal
          return randStdNormal;
        }

        private Texture2D MakePowerSpectrum(string name, int sizePx, float sizeM)
        {
            Texture2D tex = new Texture2D(sizePx, sizePx, TextureFormat.RGBAHalf,
                mipChain: false, linear: true)
            {
                name = name,
                filterMode = FilterMode.Point,
                wrapMode = TextureWrapMode.Clamp,
            };

            System.Random rand = new System.Random();
            float tau = Mathf.PI * 2.0f;
            for (int y = 0; y < sizePx; y++)
            {
                // y-axis wave number
                float kY = (tau / sizeM) * (y < sizePx / 2 ? y : y - sizePx);
                for (int x = 0; x < sizePx; x++)
                {
                    // x-axis wave number
                    float kX = (tau / sizeM) * (x < sizePx / 2 ? x : x - sizePx);
                    float k = Mathf.Sqrt(kX * kX + kY * kY);

                    // Per Horvath 2015 eq. 61:
                    //   S(k_x, k_y) = S(omega, theta) * DOmegaDK(k) / k.
                    float omega = Omega(k);
                    float s = DFlat() * STma(omega) * DOmegaDK(k) / k;

                    float dk = tau / sizeM;
                    float nu = RandomGaussian(rand);
                    // Horvath eq. 47
                    float amp = nu * Mathf.Sqrt(2.0f * s * dk * dk);
                    float phase = (float) rand.NextDouble() * tau;

                    Color res = new Color(amp * Mathf.Cos(phase), amp * Mathf.Sin(phase), 0, 1);
                    if (k == 0) {
                      res = new Color(0, 0, 0, 1);
                    }
                    tex.SetPixel(x, y, res);
                }
            }

            tex.Apply(updateMipmaps: false);
            return tex;
        }

        readonly string pipelineName = "Fourier Water";

        [MenuItem("Tools/yum_food/Fourier Water")]
        public static void ShowWindow()
        {
            GetWindow<FourierWater>("Fourier Water");
        }

        private string GetGeneratedAssetsDir()
        {
            MonoScript script = MonoScript.FromScriptableObject(this);
            string top = Path.GetDirectoryName(Path.GetDirectoryName(
                AssetDatabase.GetAssetPath(script)));
            string genDir = Path.Combine(top, "Generated").Replace('\\', '/');
            if (AssetDatabase.IsValidFolder(genDir))
            {
                AssetDatabase.DeleteAsset(genDir);
            }
            Directory.CreateDirectory(genDir);
            AssetDatabase.Refresh();
            return genDir;
        }

        private CustomRenderTexture MakeCRT(string name, int size, Material mat)
        {
            CustomRenderTexture tex = new CustomRenderTexture(
                size, size,
                RenderTextureFormat.ARGBHalf,
                RenderTextureReadWrite.Linear)
            {
                name = name,
                material = mat,
                shaderPass = 0,
                filterMode = FilterMode.Bilinear,
                wrapMode = TextureWrapMode.Repeat,
                useMipMap = false,
                autoGenerateMips = false,
                doubleBuffered = false,
                initializationMode = CustomRenderTextureUpdateMode.OnDemand,
                initializationSource = CustomRenderTextureInitializationSource.TextureAndColor,
                initializationColor = Color.black,
                updateMode = CustomRenderTextureUpdateMode.OnDemand,
                updatePeriod = 0f
            };

            tex.Create();
            return tex;
        }

        private static int[] PlanRadices(int size)
        {
            if (size < 1 || (size & (size - 1)) != 0)
            {
                throw new System.ArgumentException("FFT size must be a power of two.");
            }
            var stages = new List<int>();
            while (size > 1) {
                int radix = size % 8 == 0 ? 8 : size % 4 == 0 ? 4 : 2;
                stages.Add(radix);
                size /= radix;
            }
            return stages.ToArray();
        }

        private void OnGUI()
        {
            bigN = EditorGUILayout.IntField("FFT resolution", bigN);
            sizeM = EditorGUILayout.FloatField("Patch size (meters)", sizeM);
            U = EditorGUILayout.FloatField("Wind speed (m/s)", U);
            F = EditorGUILayout.FloatField("Fetch (meters)", F);
            h = EditorGUILayout.FloatField("Depth (meters)", h);
            gamma = EditorGUILayout.FloatField("Peak enhancement", gamma);

            waterMesh = (MeshFilter) EditorGUILayout.ObjectField(
                "Water mesh", waterMesh, typeof(MeshFilter), true);

            if (!GUILayout.Button("Generate"))
            {
                return;
            }

            if (bigN < 1 || bigN > 1024 || (bigN & (bigN - 1)) != 0)
            {
                Debug.LogError("FFT resolution must be a power of two from 1 to 1024.");
                return;
            }

            if (waterMesh == null || waterMesh.sharedMesh == null)
            {
                Debug.LogError("Assign a water mesh before generating.");
                return;
            }
            Renderer waterRenderer = waterMesh.GetComponent<Renderer>();
            if (waterRenderer == null || waterRenderer.sharedMaterial == null)
            {
                Debug.LogError("The water mesh needs a renderer and material.");
                return;
            }
            Material waterMaterial = waterRenderer.sharedMaterial;
            float meshSize = waterMesh.sharedMesh.bounds.size.x;
            if (sizeM <= 0 || meshSize <= 0)
            {
                Debug.LogError("Patch size and mesh size must be positive.");
                return;
            }

            // Delete old shit
            GameObject oldPipeline = GameObject.Find(pipelineName);
            if (oldPipeline != null)
            {
                Undo.DestroyObjectImmediate(oldPipeline);
            }
            string genDir = GetGeneratedAssetsDir();

            GameObject parent = new GameObject(pipelineName);

            Shader shader = Shader.Find("yum_food/fourier_water_compute");
            Shader displayShader = Shader.Find("yum_food/fourier_water_display");

            Texture spectrum = MakePowerSpectrum("spectrum", bigN, sizeM);
            string spectrumPath = Path.Combine(genDir, "spectrum.asset");
            AssetDatabase.CreateAsset(spectrum, spectrumPath);

            string[] outputProperties = {
                "_Fourier_Water_Height",
                "_Fourier_Water_Normals",
                "_Fourier_Water_Displacement",
                "_Fourier_Water_Displacement_Derivatives"
            };
            string[] inputKeywords = {
                "_TIME_EVOLVE", "_DERIVATIVE", "_DISPLACEMENT", "_DISPLACEMENT_DERIVATIVE"
            };
            int n_pipelines = outputProperties.Length;
            var outputs = new CustomRenderTexture[n_pipelines];
            int[] radices = PlanRadices(bigN);
            for (int which_pipeline = 0; which_pipeline < n_pipelines; which_pipeline++) {
              Texture input = spectrum;
              int n_stages = radices.Length;
              for (int axis = 0; axis < 2; axis++)
              {
                string axisStr = axis == 0 ? "x" : "y";
                int stageN = 1;
                int stride = bigN;
                for (int i = (axis == 0 ? -1 : 0); i < n_stages; i++)
                {
                  if (i > -1) {
                    stride /= radices[i];
                    stageN *= radices[i];
                  }

                  string stageName = $"pipeline_{which_pipeline}_axis_{axisStr}_stage_{i}";
                  Material mat = new Material(shader);
                  mat.EnableKeyword($"_FFT_SIZE_{bigN}");

                  GameObject quad = GameObject.CreatePrimitive(PrimitiveType.Quad);
                  quad.transform.SetParent(parent.transform);
                  quad.transform.localPosition = new Vector3(i, 2 * (n_pipelines - which_pipeline) - axis, 0);

                  if (i == -1) {
                    mat.EnableKeyword(inputKeywords[which_pipeline]);
                  } else if (radices[i] == 8) {
                    mat.EnableKeyword("_RADIX_8");
                  } else if (radices[i] == 4) {
                    mat.EnableKeyword("_RADIX_4");
                  }
                  mat.SetTexture("_Signal", input);
                  mat.SetInt("_Stage", i);
                  mat.SetInt("_Big_N", stageN);
                  mat.SetFloat("_Size_M", sizeM);
                  mat.SetInt("_Stride", stride);
                  mat.SetInt("_Direction", axis);

                  string matPath = Path.Combine(genDir,
                      $"{stageName}.mat").Replace('\\', '/');
                  AssetDatabase.CreateAsset(mat, matPath);

                  CustomRenderTexture crt = MakeCRT(stageName, bigN, mat);
                  string crtPath = Path.Combine(genDir,
                      $"{stageName}.asset").Replace('\\', '/');
                  AssetDatabase.CreateAsset(crt, crtPath);

                  Material displayMat = new Material(displayShader);
                  displayMat.SetTexture("_MainTex", crt);
                  displayMat.SetInt("_Stage", i);
                  displayMat.SetInt("_Num_Stages", Mathf.Max(1, n_stages));
                  displayMat.SetInt("_Direction", axis);
                  string displayMatPath = Path.Combine(genDir,
                      $"{stageName}_display.mat").Replace('\\', '/');
                  AssetDatabase.CreateAsset(displayMat, displayMatPath);
                  quad.GetComponent<Renderer>().sharedMaterial = displayMat;

                  var updater = quad.AddComponent<ShaderFes.CustomRenderTextureUpdater>();
                  var serializedUpdater = new SerializedObject(updater);
                  serializedUpdater.FindProperty("_targetTexture").objectReferenceValue = crt;
                  serializedUpdater.FindProperty("_updateCount").intValue = 1;
                  serializedUpdater.FindProperty("_updatePeriod").floatValue = 0;
                  serializedUpdater.ApplyModifiedPropertiesWithoutUndo();

                  input = crt;
                }
              }
              outputs[which_pipeline] = (CustomRenderTexture)input;
            }

            for (int pipeline = 0; pipeline < n_pipelines; pipeline++) {
                waterMaterial.SetTexture(outputProperties[pipeline], outputs[pipeline]);
            }
            waterMaterial.SetFloat("_Fourier_Water_Patch_Size", sizeM);
            waterMaterial.SetFloat("_Fourier_Water_Mesh_Size", meshSize);
            EditorUtility.SetDirty(waterMaterial);
            AssetDatabase.SaveAssets();
        }
    }
}