using System.Numerics;
using FftSharp;
using NAudio.Wave;

namespace Substation;

public record MusicPattern(List<WaveFrame> ChannelA, List<WaveFrame> ChannelB, TimeSpan Duration);

public static class MusicAnalyzer
{
    public const int WindowSize = 2048;
    public const int HopSize = 1024;
    public const double TickDuration = 0.1;

    const double RhythmLow = 20, RhythmHigh = 250;
    const double MelodyLow = 300, MelodyHigh = 4000;

    public static MusicPattern Analyze(string mp3Path, IProgress<int>? progress = null)
    {
        const int sampleRate = 44100;
        var samples = DecodeToMono(mp3Path, sampleRate);
        var duration = TimeSpan.FromSeconds((double)samples.Length / sampleRate);

        var tickCount = (int)Math.Ceiling((double)samples.Length / sampleRate / TickDuration);
        var channelA = new List<WaveFrame>(tickCount);
        var channelB = new List<WaveFrame>(tickCount);

        var window = new FftSharp.Windows.Hanning();
        var buffer = new double[WindowSize];
        int fftsPerTick = (int)Math.Round(TickDuration * sampleRate / HopSize);

        double[]? prevRhythmMag = null;
        double maxFlux = 0;
        double maxMelodyEnergy = 0;

        // First pass: collect per-tick features
        var tickFeatures = new List<TickFeature>(tickCount);
        int pos = 0;
        int fftIndex = 0;

        while (pos + WindowSize <= samples.Length)
        {
            for (int i = 0; i < WindowSize; i++)
                buffer[i] = samples[pos + i];

            window.ApplyInPlace(buffer);
            var spectrum = FFT.Forward(buffer);
            var mag = FFT.Magnitude(spectrum);

            var (rhythmEnergy, rhythmFlux, melodyEnergy, melodyFreq) =
                ExtractFeatures(mag, prevRhythmMag, sampleRate);

            if (rhythmFlux > maxFlux) maxFlux = rhythmFlux;
            if (melodyEnergy > maxMelodyEnergy) maxMelodyEnergy = melodyEnergy;

            prevRhythmMag = mag;

            int tickIdx = fftIndex / Math.Max(1, fftsPerTick);
            while (tickFeatures.Count <= tickIdx)
                tickFeatures.Add(new TickFeature());

            var tf = tickFeatures[tickIdx];
            tf.RhythmFlux = Math.Max(tf.RhythmFlux, rhythmFlux);
            tf.MelodyEnergy += melodyEnergy;
            tf.MelodyFreqSamples.Add(melodyFreq);
            tf.MelodyCount++;

            pos += HopSize;
            fftIndex++;

            if (progress != null && fftIndex % 50 == 0)
            {
                var pct = (int)((double)pos / samples.Length * 100);
                progress.Report(pct);
            }
        }

        progress?.Report(100);

        // Second pass: normalize and build WaveFrames
        foreach (var tf in tickFeatures)
        {
            var (frameA, frameB) = BuildWaveFrame(tf, maxFlux, maxMelodyEnergy);
            channelA.Add(frameA);
            channelB.Add(frameB);
        }

        return new MusicPattern(channelA, channelB, duration);
    }

    public static (double rhythmEnergy, double rhythmFlux, double melodyEnergy, double melodyFreq)
        ExtractFeatures(double[] magnitude, double[]? prevRhythmMag, int sampleRate)
    {
        double binWidth = (double)sampleRate / WindowSize;
        int rhythmLoBin = (int)(RhythmLow / binWidth);
        int rhythmHiBin = (int)(RhythmHigh / binWidth);
        int melodyLoBin = (int)(MelodyLow / binWidth);
        int melodyHiBin = (int)(MelodyHigh / binWidth);

        // Rhythm band energy
        double rhythmEnergy = 0;
        for (int i = rhythmLoBin; i <= rhythmHiBin && i < magnitude.Length; i++)
            rhythmEnergy += magnitude[i] * magnitude[i];
        rhythmEnergy = Math.Sqrt(rhythmEnergy / (rhythmHiBin - rhythmLoBin + 1));

        // Spectral flux (positive change in rhythm band)
        double rhythmFlux = 0;
        if (prevRhythmMag != null)
        {
            for (int i = rhythmLoBin; i <= rhythmHiBin && i < magnitude.Length; i++)
            {
                double diff = magnitude[i] - prevRhythmMag[i];
                if (diff > 0) rhythmFlux += diff;
            }
        }

        // Melody band: energy + dominant frequency (spectral peak)
        double melodyEnergy = 0;
        double peakMag = 0;
        int peakBin = melodyLoBin;
        for (int i = melodyLoBin; i <= melodyHiBin && i < magnitude.Length; i++)
        {
            double m = magnitude[i];
            melodyEnergy += m * m;
            if (m > peakMag)
            {
                peakMag = m;
                peakBin = i;
            }
        }
        melodyEnergy = Math.Sqrt(melodyEnergy / (melodyHiBin - melodyLoBin + 1));
        double melodyFreq = peakBin * binWidth;

        return (rhythmEnergy, rhythmFlux, melodyEnergy, melodyFreq);
    }

    public static int MapPitchToPeriod(double hz)
    {
        // Map melody frequency (300-4000Hz) to e-stim period (10-1000ms)
        // Logarithmic mapping: low notes → deep, high notes → buzzy
        double logFreq = Math.Log(Math.Clamp(hz, MelodyLow, MelodyHigh));
        double logMin = Math.Log(MelodyLow);
        double logMax = Math.Log(MelodyHigh);
        double t = (logFreq - logMin) / (logMax - logMin); // 0..1
        // Invert: high freq → short period (buzzy), low freq → long period (deep)
        int ms = (int)Math.Round(1000 - t * 990); // 1000ms..10ms
        return Math.Clamp(ms, 10, 1000);
    }

    static double[] DecodeToMono(string mp3Path, int sampleRate)
    {
        using var reader = new Mp3FileReader(mp3Path);
        var format = new WaveFormat(sampleRate, 16, 1);
        using var resampler = new MediaFoundationResampler(reader, format);
        resampler.ResamplerQuality = 60;

        var sampleList = new List<float>();
        var buffer = new byte[sampleRate * 2]; // 1s worth of 16-bit mono
        int read;
        while ((read = resampler.Read(buffer, 0, buffer.Length)) > 0)
        {
            for (int i = 0; i < read; i += 2)
            {
                short sample = (short)(buffer[i] | (buffer[i + 1] << 8));
                sampleList.Add(sample / 32768f);
            }
        }

        var result = new double[sampleList.Count];
        for (int i = 0; i < sampleList.Count; i++)
            result[i] = sampleList[i];
        return result;
    }

    public class TickFeature
    {
        public double RhythmFlux;
        public double MelodyEnergy;
        public double MelodyCount;
        public readonly List<double> MelodyFreqSamples = new();
    }

    public static (WaveFrame chA, WaveFrame chB) BuildWaveFrame(TickFeature tf, double maxFlux, double maxMelodyEnergy)
    {
        const int rhythmFreqMs = 150;

        double normalizedFlux = maxFlux > 0 ? tf.RhythmFlux / maxFlux : 0;
        int onsetIntensity = (int)Math.Round(normalizedFlux * 100);
        var intA = new[]
        {
            (byte)Math.Clamp(onsetIntensity, 0, 100),
            (byte)Math.Clamp(onsetIntensity * 6 / 10, 0, 100),
            (byte)Math.Clamp(onsetIntensity * 3 / 10, 0, 100),
            (byte)0
        };
        var freqA = Freq.Compress4(new[] { rhythmFreqMs, rhythmFreqMs, rhythmFreqMs, rhythmFreqMs });

        double avgEnergy = tf.MelodyCount > 0 ? tf.MelodyEnergy / tf.MelodyCount : 0;
        double normalizedEnergy = maxMelodyEnergy > 0 ? avgEnergy / maxMelodyEnergy : 0;
        int melodyIntensity = (int)Math.Round(normalizedEnergy * 80);
        melodyIntensity = Math.Clamp(melodyIntensity, 0, 100);

        double weightedFreq = 0;
        double totalWeight = 0;
        foreach (var f in tf.MelodyFreqSamples)
        {
            weightedFreq += f * f;
            totalWeight += f;
        }
        double avgMelodyHz = totalWeight > 0 ? weightedFreq / totalWeight : 500;
        int estimsMs = MapPitchToPeriod(avgMelodyHz);

        var intB = new[]
        {
            (byte)melodyIntensity, (byte)melodyIntensity,
            (byte)melodyIntensity, (byte)melodyIntensity
        };
        var freqB = Freq.Compress4(new[] { estimsMs, estimsMs, estimsMs, estimsMs });

        return (new WaveFrame(freqA, intA), new WaveFrame(freqB, intB));
    }
}