using System;
using UnityEngine;

namespace Audio2Lipsync
{
    /// <summary>Post-processing knobs — same names/defaults as server.py LipSyncRequest.</summary>
    [Serializable]
    public struct LipsyncSettings
    {
        [Range(1, 15)] public int smoothWindow;           // moving-average window (1 = off)
        [Range(0.1f, 5f)] public float gain;              // centered amplification around the rest pose
        [Range(0f, 1f)] public float silenceGate;         // RMS mouth gate strength (0 = off)
        [Range(-80f, -10f)] public float silenceThresholdDb;
        [Range(0.5f, 5f)] public float cfgScale;          // classifier-free guidance (1 = off)

        public static LipsyncSettings Default => new LipsyncSettings
        {
            smoothWindow = 3, gain = 1.5f, silenceGate = 0f, silenceThresholdDb = -40f, cfgScale = 1f,
        };
    }

    /// <summary>
    /// Pure-C# port of server.py run_inference() post-processing. Frames are row-major [n, 52].
    /// </summary>
    public static class LipsyncPostProcess
    {
        const int N = A2LConstants.NumBlendshapes;
        const int FadeFrames = 6;

        /// <param name="cond">denormalized model output [n*52]</param>
        /// <param name="uncond">silent-audio output [n*52] (only read when cfgScale != 1)</param>
        public static float[] Apply(float[] cond, float[] uncond, int nFrames, float[] audio16k, LipsyncSettings s)
        {
            var bs = new float[nFrames * N];
            Array.Copy(cond, bs, bs.Length);

            if (s.cfgScale != 1f && uncond != null)
                for (int i = 0; i < bs.Length; i++)
                    bs[i] = uncond[i] + s.cfgScale * (bs[i] - uncond[i]);
            Clip01(bs);

            // Zero non-mouth channels (brows/eyes/nose belong to the upper-face layer).
            for (int f = 0; f < nFrames; f++)
                for (int c = 0; c < N; c++)
                    if (!A2LConstants.KeepChannel[c]) bs[f * N + c] = 0f;

            if (s.gain != 1f)
            {
                for (int f = 0; f < nFrames; f++)
                    for (int c = 0; c < N; c++)
                    {
                        if (!A2LConstants.KeepChannel[c]) continue;
                        float rest = A2LConstants.Mean[c];
                        bs[f * N + c] = rest + s.gain * (bs[f * N + c] - rest);
                    }
                Clip01(bs);
            }

            if (s.smoothWindow > 1 && nFrames > s.smoothWindow)
            {
                bs = UniformFilterFrames(bs, nFrames, s.smoothWindow);
                Clip01(bs);
            }

            if (s.silenceGate > 0f)
            {
                var gate = SilenceGate(audio16k, nFrames, s.silenceThresholdDb, s.silenceGate);
                for (int f = 0; f < nFrames; f++)
                    for (int c = 0; c < N; c++)
                        if (A2LConstants.KeepChannel[c]) bs[f * N + c] *= gate[f];
            }

            // ~0.1 s fade in/out so the mouth opens from and closes to neutral.
            for (int f = 0; f < Mathf.Min(FadeFrames, nFrames); f++)
            {
                float k = (f + 1) / (float)FadeFrames;
                for (int c = 0; c < N; c++)
                {
                    bs[f * N + c] *= k;
                    bs[(nFrames - 1 - f) * N + c] *= k;
                }
            }
            return bs;
        }

        static void Clip01(float[] a)
        {
            for (int i = 0; i < a.Length; i++) a[i] = Mathf.Clamp01(a[i]);
        }

        // scipy.ndimage.uniform_filter1d(axis=0, mode="nearest"): window [i - w/2, i - w/2 + w - 1].
        static float[] UniformFilterFrames(float[] src, int nFrames, int w)
        {
            var dst = new float[src.Length];
            int left = w / 2;
            for (int f = 0; f < nFrames; f++)
                for (int c = 0; c < N; c++)
                {
                    float sum = 0f;
                    for (int k = 0; k < w; k++)
                    {
                        int j = Mathf.Clamp(f - left + k, 0, nFrames - 1);
                        sum += src[j * N + c];
                    }
                    dst[f * N + c] = sum / w;
                }
            return dst;
        }

        // server.py compute_silence_gate()
        static float[] SilenceGate(float[] audio, int nFrames, float thresholdDb, float strength)
        {
            var gate = new float[nFrames];
            int spf = Mathf.Max(1, A2LConstants.AudioSampleRate / A2LConstants.FaceFps);
            int nAudio = audio.Length / spf;
            if (nAudio < 1) { for (int i = 0; i < nFrames; i++) gate[i] = 1f; return gate; }

            var rms = new float[nAudio];
            for (int i = 0; i < nAudio; i++)
            {
                double acc = 0;
                for (int k = 0; k < spf; k++) { float v = audio[i * spf + k]; acc += v * v; }
                rms[i] = (float)Math.Sqrt(acc / spf + 1e-12);
            }

            int smoothN = Mathf.Max(1, (int)(0.25f * A2LConstants.FaceFps));
            if (smoothN > 1 && nAudio > smoothN)
            {
                var sm = new float[nAudio];
                int left = smoothN / 2;
                for (int i = 0; i < nAudio; i++)
                {
                    float sum = 0f;
                    for (int k = 0; k < smoothN; k++) sum += rms[Mathf.Clamp(i - left + k, 0, nAudio - 1)];
                    sm[i] = sum / smoothN;
                }
                rms = sm;
            }

            var raw = new float[nAudio];
            for (int i = 0; i < nAudio; i++)
            {
                float db = 20f * Mathf.Log10(Mathf.Max(rms[i], 1e-8f));
                raw[i] = 1f / (1f + Mathf.Exp(-(db - thresholdDb) / 4f));
            }

            for (int f = 0; f < nFrames; f++)
            {
                float x = nFrames == 1 ? 0f : f / (float)(nFrames - 1) * (nAudio - 1);
                int i0 = Mathf.Clamp((int)x, 0, nAudio - 1), i1 = Mathf.Min(i0 + 1, nAudio - 1);
                float g = Mathf.Lerp(raw[i0], raw[i1], x - i0);
                gate[f] = (1f - strength) + strength * g;
            }
            return gate;
        }
    }
}