// CED-Tiny (Consistent Ensemble Distillation ViT) AudioSet audio tagger. Self-contained Unity Sentis inference.
using System;
using System.Collections.Generic;
using System.IO;
using System.Threading.Tasks;
using Unity.InferenceEngine;
using UnityEngine;

namespace SentisModels
{
// Mono waveform -> 527 AudioSet class probabilities (multi-label: independent sigmoids, they do not sum to 1) and a
// 192-d clip embedding. The log-mel frontend (STFT as a fixed-basis Conv1d) is baked into the model, so the input is
// the raw 16000 Hz waveform in [-1, 1]; one inference covers a 3 s window. Other sample rates are resampled here.
public sealed class CedTiny : IDisposable
{
    public const int SampleRate = 16000;
    public const int WindowSamples = 48000;
    public const int ClassCount = 527;
    public const int EmbeddingSize = 192;
    const string DefaultModelFile = "ced_tiny_fp32.sentis";
    const string DefaultLabelFile = "audioset_labels.txt";

    public struct Prediction
    {
        public int Index;       // AudioSet class index (ontology order, class_labels_indices.csv)
        public string Label;
        public float Score;     // sigmoid probability
    }

    readonly BackendType m_BackendType;

    Worker m_Worker;
    Tensor<float> m_InputTensor;
    string[] m_Labels;
    readonly float[] m_Window = new float[WindowSamples];
    readonly float[] m_Stream = new float[WindowSamples];   // sliding buffer for Push / ClassifyLatest

    // Embedding of the most recent Classify call (penultimate features; use for your own classes / similarity).
    public float[] Embedding { get; } = new float[EmbeddingSize];
    public IReadOnlyList<string> Labels => m_Labels;

    public CedTiny(BackendType backendType = BackendType.GPUCompute)
    {
        m_BackendType = backendType == BackendType.GPUCompute && !SystemInfo.supportsComputeShaders
            ? BackendType.CPU
            : backendType;
    }

    public bool IsLoaded => m_Worker != null;

    // modelRoot: absolute path to the folder holding ced_tiny_fp32.sentis and audioset_labels.txt.
    public void Load(string modelRoot, string modelFile = DefaultModelFile)
    {
        Load(ModelLoader.Load(Path.Combine(modelRoot, modelFile)), File.ReadAllLines(Path.Combine(modelRoot, DefaultLabelFile)));
    }

    // For platforms where the files must be read as bytes / text (e.g. Android StreamingAssets).
    public void Load(byte[] modelBytes, string labelsText)
    {
        using var stream = new MemoryStream(modelBytes);
        Load(ModelLoader.Load(stream), labelsText.Split(new[] { "\r\n", "\n" }, StringSplitOptions.RemoveEmptyEntries));
    }

    void Load(Model model, string[] labels)
    {
        if (labels.Length != ClassCount)
            throw new ArgumentException($"expected {ClassCount} labels, got {labels.Length}");
        DisposeResources();
        m_Labels = labels;
        m_Worker = new Worker(model, m_BackendType);
        m_InputTensor = new Tensor<float>(new TensorShape(1, WindowSamples));
        Array.Clear(m_Stream, 0, m_Stream.Length);
    }

    // One silent pass so the first real call doesn't pay shader compilation / allocation.
    public void Warmup()
    {
        Classify(new float[WindowSamples]);
    }

    // samples: mono PCM at sampleRate. The LAST 3 s are classified (zero-padded in front when shorter).
    public float[] Classify(float[] samples, int sampleRate = SampleRate)
    {
        Schedule(samples, sampleRate);
        return Read();
    }

    // Async readback variant; prefer this for per-frame use on GPU backends (no main-thread stall).
    public async Task<float[]> ClassifyAsync(float[] samples, int sampleRate = SampleRate)
    {
        Schedule(samples, sampleRate);
        using var scores = await (m_Worker.PeekOutput("scores") as Tensor<float>).ReadbackAndCloneAsync();
        using var embedding = await (m_Worker.PeekOutput("embedding") as Tensor<float>).ReadbackAndCloneAsync();
        embedding.DownloadToArray().CopyTo(Embedding, 0);
        return scores.DownloadToArray();
    }

    // Streaming: append new microphone / clip samples (mono, any rate) to the sliding window.
    public void Push(float[] samples, int sampleRate = SampleRate)
    {
        var x = sampleRate == SampleRate ? samples : Resample(samples, sampleRate, SampleRate);
        if (x.Length >= WindowSamples)
        {
            Array.Copy(x, x.Length - WindowSamples, m_Stream, 0, WindowSamples);
            return;
        }
        Array.Copy(m_Stream, x.Length, m_Stream, 0, WindowSamples - x.Length);
        Array.Copy(x, 0, m_Stream, WindowSamples - x.Length, x.Length);
    }

    // Classifies the last 3 s pushed with Push (call e.g. every 0.25-0.5 s for a sliding window).
    public float[] ClassifyLatest() => Classify(m_Stream);

    public Task<float[]> ClassifyLatestAsync() => ClassifyAsync(m_Stream);

    // Highest-scoring classes, best first.
    public List<Prediction> Top(float[] scores, int k = 5, float minScore = 0f)
    {
        var order = new int[scores.Length];
        for (var i = 0; i < order.Length; i++)
            order[i] = i;
        Array.Sort(order, (a, b) => scores[b].CompareTo(scores[a]));
        var result = new List<Prediction>();
        for (var i = 0; i < order.Length && result.Count < k; i++)
        {
            if (scores[order[i]] < minScore)
                break;
            result.Add(new Prediction { Index = order[i], Label = m_Labels[order[i]], Score = scores[order[i]] });
        }
        return result;
    }

    // Average the channels of interleaved PCM (AudioClip.GetData / Microphone data).
    public static float[] ToMono(float[] interleaved, int channels)
    {
        if (channels <= 1)
            return interleaved;
        var mono = new float[interleaved.Length / channels];
        for (var i = 0; i < mono.Length; i++)
        {
            var sum = 0f;
            for (var c = 0; c < channels; c++)
                sum += interleaved[i * channels + c];
            mono[i] = sum / channels;
        }
        return mono;
    }

    // Linear-interpolation resampler (adequate for tagging; use a band-limited resampler for best fidelity).
    public static float[] Resample(float[] src, int srcRate, int dstRate)
    {
        if (srcRate == dstRate || src.Length == 0)
            return src;
        var n = (int)((long)src.Length * dstRate / srcRate);
        var dst = new float[n];
        var step = (double)srcRate / dstRate;
        for (var i = 0; i < n; i++)
        {
            var p = i * step;
            var j = (int)p;
            var t = (float)(p - j);
            dst[i] = j + 1 < src.Length ? src[j] + (src[j + 1] - src[j]) * t : src[src.Length - 1];
        }
        return dst;
    }

    void Schedule(float[] samples, int sampleRate)
    {
        if (m_Worker == null)
            throw new InvalidOperationException("CedTiny.Load must be called before Classify.");
        var x = sampleRate == SampleRate ? samples : Resample(samples, sampleRate, SampleRate);
        var n = Math.Min(x.Length, WindowSamples);
        Array.Clear(m_Window, 0, WindowSamples - n);
        Array.Copy(x, x.Length - n, m_Window, WindowSamples - n, n);
        m_InputTensor.Upload(m_Window);
        m_Worker.Schedule(m_InputTensor);
    }

    float[] Read()
    {
        using var scores = (m_Worker.PeekOutput("scores") as Tensor<float>).ReadbackAndClone();
        using var embedding = (m_Worker.PeekOutput("embedding") as Tensor<float>).ReadbackAndClone();
        embedding.DownloadToArray().CopyTo(Embedding, 0);
        return scores.DownloadToArray();
    }

    void DisposeResources()
    {
        m_Worker?.Dispose();
        m_Worker = null;
        m_InputTensor?.Dispose();
        m_InputTensor = null;
    }

    public void Dispose()
    {
        DisposeResources();
    }
}
}