using System;
using System.Collections.Generic;
using Sandbox;

namespace RedSnail.EngineBankSlicer;

/// <summary>A mono, 16-bit view of a loaded wav.</summary>
public sealed class WavData
{
	public float[] Samples;
	public int SampleRate;

	public float DurationSeconds => SampleRate > 0 ? (float)Samples.Length / SampleRate : 0.0f;
}

/// <summary>One extracted loop, named for the engine speed it was found at.</summary>
public sealed class SlicedLayer
{
	public int Rpm;
	public short[] Samples;
	public int SampleRate;

	/// <summary>Detected firing period in samples, kept for the report.</summary>
	public float Period;

	/// <summary>Whole engine cycles the loop contains — never fractional, see the slicer remarks.</summary>
	public int Cycles;
}

public sealed class SliceOptions
{
	public int Cylinders = 6;
	public bool FourStroke = true;

	/// <summary>How many loops to cut. More gives smoother crossfades and costs more voices at runtime.</summary>
	public int LayerCount = 15;

	/// <summary>Roughly how long each loop should be. Rounded to whole cycles, so it is a target, not a promise.</summary>
	public float LoopSeconds = 0.35f;

	/// <summary>Detection bounds. Anything outside is treated as a failed read rather than a real engine speed.</summary>
	public float MinRpm = 500.0f;
	public float MaxRpm = 9000.0f;

	/// <summary>
	/// Label each clip with the engine speed MEASURED from the audio, rather than interpolated between
	/// <see cref="StartRpm"/> and <see cref="EndRpm"/>.
	/// </summary>
	/// <remarks>
	/// Off by default, because measurement turned out not to be trustworthy enough to label with — and a label is
	/// the one thing that must be right, since the synthesiser derives playback pitch from it.
	///
	/// The problem is not the algorithm. Engine audio is strongly periodic at THREE different rates at once: the
	/// firing rate, the crank rate below it, and whatever the exhaust rings at, which is fixed and frequently sits
	/// inside the firing range. Correlation finds all of them and nothing in the signal says which is which.
	/// Tested against a sweep of known speed, waveform correlation locked onto the exhaust resonance and reported
	/// the same figure from idle to redline; the envelope approach that should have stripped the carrier failed
	/// too, because separating carrier from rhythm needs them not to overlap, and they do.
	///
	/// Detection is still used for LOOP LENGTH, always, and is reliable there — that only needs some true period
	/// of the signal, and the resonance period cuts just as seamless a loop as the firing period does. It is
	/// labelling that needs the one specific period, and that is the part that cannot be resolved from audio
	/// alone.
	///
	/// Turn it on if you want to try it on a particular recording; check the reported figures before trusting them.
	/// </remarks>
	public bool UseDetectedRpm = false;

	/// <summary>
	/// Engine speed at the START of the usable audio, and at the END. Normally idle and redline.
	/// </summary>
	/// <remarks>
	/// THESE SET THE LABELS. Each probe is labelled by where it sits between the two, which assumes the revs climb
	/// evenly across the recording — not strictly true, since an engine pulls hardest in the middle of its range,
	/// but predictable, and wrong by a margin the crossfade absorbs. Neighbouring layers overlap, so a clip
	/// labelled slightly off is played slightly off-pitch rather than jarringly wrong.
	///
	/// They also disambiguate octaves when <see cref="UseDetectedRpm"/> is on, where a rough figure is plenty:
	/// harmonics sit a whole multiple apart, so the prior only has to be within about 40% to pick the right one.
	///
	/// Getting these right matters more than anything else in the options. Read them off the car: idle speed and
	/// redline, or wherever the recording actually starts and stops.
	/// </remarks>
	public float StartRpm = 800.0f;
	public float EndRpm = 7000.0f;

	/// <summary>Seconds to ignore at each end — handy for trimming a key turn or a lift-off.</summary>
	public float SkipStart = 0.0f;
	public float SkipEnd = 0.0f;
}

/// <summary>
/// Cuts one recorded acceleration run — idle to redline in a single pull — into an RPM-indexed bank of loops.
/// </summary>
/// <remarks>
/// Two separate jobs, from two different sources, because they need different things:
///
/// LOOP LENGTH is measured from the audio. Each loop is cut to a whole number of signal periods and crossfaded at
/// the seam, which is what stops it clicking on every wrap. This only needs SOME true period of the waveform, and
/// correlation finds one reliably.
///
/// THE RPM LABEL comes from <see cref="SliceOptions.StartRpm"/> and <see cref="SliceOptions.EndRpm"/>, spread
/// across the recording by position. It does not come from the audio, and the reason is worth recording so nobody
/// re-attempts it: engine sound is strongly periodic at three rates at once — the firing rate, the crank rate
/// below it, and whatever the exhaust rings at, which is fixed and often lands inside the firing range. Nothing in
/// the signal says which is which. Measured against a sweep of known speed, waveform correlation reported the
/// resonance and gave near-identical figures from idle to redline; chaining probes to each other instead made the
/// first reading load-bearing and scrambled every label behind one bad probe; the envelope method that should have
/// stripped the carrier needs carrier and rhythm not to overlap, and they overlap.
///
/// The label matters more than the loop, since the synthesiser derives playback pitch from it — so it is taken
/// from the one thing that is actually known: the engine's idle and redline. That assumes revs climb evenly, which
/// is not quite true, but the error is small and neighbouring layers crossfade over it.
///
/// <see cref="SliceOptions.UseDetectedRpm"/> restores measurement for anyone who wants to try it per-recording.
/// </remarks>
public static class EngineBankSlicer
{
	/// <summary>Combustion events per crank revolution.</summary>
	public static float EventsPerRevolution(SliceOptions _Options)
	{
		return _Options.Cylinders / (_Options.FourStroke ? 2.0f : 1.0f);
	}

	public static WavData LoadWav(byte[] _Bytes)
	{
		if (_Bytes is null || _Bytes.Length < 44)
			return null;

		if (_Bytes[0] != 'R' || _Bytes[1] != 'I' || _Bytes[2] != 'F' || _Bytes[3] != 'F')
			return null;

		int channels = 1;
		int sampleRate = 44100;
		int bitsPerSample = 16;
		int dataOffset = -1;
		int dataLength = 0;

		// Walk the chunks rather than assuming a 44-byte header; plenty of wavs carry extra chunks first.
		int offset = 12;

		while (offset + 8 <= _Bytes.Length)
		{
			string id = System.Text.Encoding.ASCII.GetString(_Bytes, offset, 4);
			int size = BitConverter.ToInt32(_Bytes, offset + 4);

			if (id == "fmt ")
			{
				channels = BitConverter.ToInt16(_Bytes, offset + 10);
				sampleRate = BitConverter.ToInt32(_Bytes, offset + 12);
				bitsPerSample = BitConverter.ToInt16(_Bytes, offset + 22);
			}
			else if (id == "data")
			{
				dataOffset = offset + 8;
				dataLength = size;

				break;
			}

			offset += 8 + size + (size & 1);
		}

		if (dataOffset < 0 || bitsPerSample != 16 || channels < 1)
			return null;

		dataLength = Math.Min(dataLength, _Bytes.Length - dataOffset);

		int frames = dataLength / 2 / channels;

		if (frames <= 0)
			return null;

		float[] samples = new float[frames];

		for (int f = 0; f < frames; f++)
		{
			float sum = 0.0f;

			for (int c = 0; c < channels; c++)
				sum += BitConverter.ToInt16(_Bytes, dataOffset + ((f * channels + c) * 2));

			samples[f] = sum / channels / short.MaxValue;
		}

		return new WavData { Samples = samples, SampleRate = sampleRate };
	}

	/// <summary>
	/// Cuts the bank.
	/// </summary>
	public static List<SlicedLayer> Slice(WavData _Wav, SliceOptions _Options, out string _Report)
	{
		List<SlicedLayer> layers = new();
		System.Text.StringBuilder report = new();

		if (_Wav?.Samples is not { Length: > 0 })
		{
			_Report = "No audio loaded.";

			return layers;
		}

		int rate = _Wav.SampleRate;
		float events = EventsPerRevolution(_Options);

		if (events <= 0.0f)
		{
			_Report = "Cylinder count must be at least 1.";

			return layers;
		}

		// A low-passed copy purely for DETECTION. The firing fundamental is low and the upper harmonics are what
		// confuse autocorrelation into locking an octave high, so they are removed before measuring — but every
		// sample that gets written out comes from the untouched original.
		float[] detect = LowPass(RemoveDc(_Wav.Samples), rate, 600.0f);

		// Lag bounds straight from the RPM bounds, so detection can never report an impossible engine speed.
		int minLag = (int)(rate / (_Options.MaxRpm / 60.0f * events));
		int maxLag = (int)(rate / (_Options.MinRpm / 60.0f * events));

		minLag = Math.Max(minLag, 8);
		maxLag = Math.Min(maxLag, _Wav.Samples.Length / 4);

		if (maxLag <= minLag)
		{
			_Report = "RPM range is too narrow, or the recording is too short to measure.";

			return layers;
		}

		int window = Math.Min(maxLag * 4, _Wav.Samples.Length);

		int start = (int)(_Options.SkipStart * rate);
		int end = _Wav.Samples.Length - (int)(_Options.SkipEnd * rate);

		start = Math.Clamp(start, 0, _Wav.Samples.Length - 1);
		end = Math.Clamp(end, start + window, _Wav.Samples.Length);

		int count = Math.Max(1, _Options.LayerCount);
		int usable = end - start - window;

		if (usable <= 0)
		{
			_Report = "Nothing left to slice after the skip settings.";

			return layers;
		}

		report.AppendLine($"{_Wav.DurationSeconds:0.00}s at {rate}Hz, {events:0.#} firings/rev");

		// PASS ONE: measure every probe before deciding anything, because the correction below needs neighbours.
		int[] positions = new int[count];
		float[] firingHz = new float[count];

		for (int i = 0; i < count; i++)
		{
			// Spread the probes across the usable span. The recording's own shape decides what RPM each lands on,
			// which is why the results are rarely evenly spaced — and why they should not be forced to be.
			positions[i] = start + (count == 1 ? usable / 2 : usable * i / (count - 1));

			float period = DetectPeriod(detect, positions[i], window, minLag, maxLag);

			firingHz[i] = period > 0.0f ? rate / period : 0.0f;
		}

		if (_Options.UseDetectedRpm)
		{
			int corrected = ResolveOctaves(firingHz, _Options, events);

			report.AppendLine($"labelling from DETECTED pitch ({corrected} octave corrections) — verify these figures");
		}
		else
		{
			report.AppendLine($"labelling from the {_Options.StartRpm:0}-{_Options.EndRpm:0}rpm span; detection sets loop length only");
		}

		HashSet<int> seen = new();

		for (int i = 0; i < count; i++)
		{
			if (firingHz[i] <= 0.0f)
			{
				report.AppendLine($"  [{i}] no stable pitch found, skipped");

				continue;
			}

			// Loop length always comes from the MEASURED period — that is what makes the seam seamless, and it only
			// needs some true period of the signal, which detection supplies reliably.
			float period = rate / firingHz[i];

			// The label is a different question, and by default a different source. See UseDetectedRpm.
			float progress = count == 1 ? 0.5f : (float)i / (count - 1);

			int rpm = _Options.UseDetectedRpm
				? (int)MathF.Round(firingHz[i] * 60.0f / events)
				: (int)MathF.Round(_Options.StartRpm + (_Options.EndRpm - _Options.StartRpm) * progress);

			if (rpm < _Options.MinRpm || rpm > _Options.MaxRpm)
			{
				report.AppendLine($"  [{i}] {rpm}rpm out of range, skipped");

				continue;
			}

			// Two probes can land on the same revs when the run pauses or a gearchange flattens it. Duplicate
			// reference speeds break pair selection, which assumes strictly ascending layers.
			if (!seen.Add(rpm))
			{
				report.AppendLine($"  [{i}] {rpm}rpm duplicate, skipped");

				continue;
			}

			SlicedLayer layer = ExtractLoop(_Wav, positions[i], period, _Options.LoopSeconds);

			if (layer is null)
			{
				report.AppendLine($"  [{i}] {rpm}rpm too close to the end to loop, skipped");

				continue;
			}

			layer.Rpm = rpm;

			layers.Add(layer);

			report.AppendLine($"  {rpm,5}rpm  {layer.Cycles} cycles  {layer.Samples.Length / (float)rate:0.000}s");
		}

		layers.Sort((a, b) => a.Rpm.CompareTo(b.Rpm));

		report.AppendLine($"{layers.Count} layers");

		_Report = report.ToString();

		return layers;
	}

	/// <summary>
	/// Cuts a loop that contains a WHOLE number of firing cycles, then crossfades its seam.
	/// </summary>
	/// <remarks>
	/// A loop holding a fractional cycle restarts mid-bang, and a waveform discontinuity is a click — heard once
	/// per loop, which at a third of a second is three clicks a second and utterly damning. Snapping the length to
	/// the detected period is what makes these usable as sustained loops at all.
	///
	/// Whole cycles still leave a small mismatch, because the engine is speeding up throughout and the end of the
	/// loop is fractionally higher than its start. The crossfade takes the material just PAST the loop point and
	/// blends it over the head, so the join is a short overlap rather than a step.
	/// </remarks>
	private static SlicedLayer ExtractLoop(WavData _Wav, int _Position, float _Period, float _LoopSeconds)
	{
		int rate = _Wav.SampleRate;
		int cycles = Math.Max(2, (int)MathF.Round(_LoopSeconds * rate / _Period));
		int length = (int)MathF.Round(cycles * _Period);
		int fade = (int)MathF.Round(_Period);

		if (length <= 0 || _Position + length + fade >= _Wav.Samples.Length)
			return null;

		float[] loop = new float[length];

		Array.Copy(_Wav.Samples, _Position, loop, 0, length);

		// Blend the material just past the end over the head. After this the last sample runs into the first
		// without a step, which is what "seamless" actually means.
		for (int i = 0; i < fade && i < length; i++)
		{
			float t = (float)i / fade;

			loop[i] = loop[i] * t + _Wav.Samples[_Position + length + i] * (1.0f - t);
		}

		short[] output = new short[length];

		for (int i = 0; i < length; i++)
			output[i] = (short)(Math.Clamp(loop[i], -1.0f, 1.0f) * short.MaxValue);

		return new SlicedLayer
		{
			Samples = output,
			SampleRate = rate,
			Period = _Period,
			Cycles = cycles
		};
	}

	/// <summary>
	/// Forces the measured firing rate to rise across the recording, snapping octave errors as it goes.
	/// </summary>
	/// <remarks>
	/// THIS IS WHAT MAKES DETECTION RELIABLE, and it took measuring a known bank to find out. Pitch detection on
	/// engine audio is genuinely hard: a four-cylinder repeats every two firings per crank revolution and a V8
	/// every four, so the signal is strongly periodic at whole multiples of the firing period. Correlation-based
	/// methods lock onto those multiples about as readily as onto the truth, and no threshold separates them —
	/// tested against fifteen clips of known speed, the raw detector was exactly right below 2600rpm and exactly
	/// four times too slow above it, with nothing in the measurement itself to say which was which.
	///
	/// The recording answers it. A single pull only ever speeds UP, so a measured drop is impossible and can only
	/// be an octave error. Multiplying by the smallest whole number that restores the rise recovers the true rate.
	/// On that same bank this took the spread in firings-per-revolution from a factor of four down to 2%.
	///
	/// The catch is the first probe, which has no predecessor to be judged against: an error there shifts every
	/// later value with it. Starting the recording at a steady idle, where detection is easiest, is the defence.
	/// </remarks>
	private static int ResolveOctaves(float[] _FiringHz, SliceOptions _Options, float _EventsPerRev)
	{
		int corrected = 0;
		int count = _FiringHz.Length;

		for (int i = 0; i < count; i++)
		{
			if (_FiringHz[i] <= 0.0f)
				continue;

			// Where the prior says this probe roughly is. Straight-line, which is wrong about real engines — they
			// pull hardest in the middle — but only ever used to choose between candidates a whole multiple apart,
			// so being loose is harmless.
			float t = count == 1 ? 0.5f : (float)i / (count - 1);
			float expectedRpm = _Options.StartRpm + (_Options.EndRpm - _Options.StartRpm) * t;
			float expectedHz = MathF.Max(expectedRpm, 1.0f) / 60.0f * _EventsPerRev;

			float best = _FiringHz[i];
			float bestError = MathF.Abs(MathF.Log(best / expectedHz));

			// Compared in log space so being twice too fast and half too slow count equally — in linear terms the
			// high side would always look worse and the search would drift downward.
			foreach (float multiple in OctaveMultiples)
			{
				float candidate = _FiringHz[i] * multiple;
				float error = MathF.Abs(MathF.Log(candidate / expectedHz));

				if (error < bestError)
				{
					bestError = error;
					best = candidate;
				}
			}

			if (!best.AlmostEqual(_FiringHz[i]))
			{
				_FiringHz[i] = best;
				corrected++;
			}
		}

		return corrected;
	}

	/// <summary>
	/// Whole-number relationships a firing pattern can hide behind, and their reciprocals.
	/// </summary>
	/// <remarks>
	/// Both directions are needed. Detection can land on a subharmonic — the crank period rather than the firing
	/// period — or on an upper harmonic, and a search that could only multiply would leave the second kind wrong.
	/// </remarks>
	private static readonly float[] OctaveMultiples =
	{
		2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 8.0f,
		1.0f / 2.0f, 1.0f / 3.0f, 1.0f / 4.0f, 1.0f / 5.0f, 1.0f / 6.0f, 1.0f / 8.0f
	};

	/// <summary>
	/// Finds the firing period in samples, using the normalised square difference function.
	/// </summary>
	/// <remarks>
	/// NSDF rather than plain autocorrelation, and the FIRST strong peak rather than the tallest. Raw correlation
	/// grows with the number of terms summed and so favours long lags, which is precisely the wrong bias when the
	/// long lags are subharmonics. Normalising by the energy of both windows removes that, and taking the first
	/// peak within 10% of the best prefers the shortest period that explains the signal.
	///
	/// It is still not enough on its own — see SnapOctaves, which is what actually makes this trustworthy.
	/// </remarks>
	private static float DetectPeriod(float[] _Samples, int _Start, int _Window, int _MinLag, int _MaxLag)
	{
		if (_Start + _Window >= _Samples.Length)
			return 0.0f;

		float[] nsdf = new float[_MaxLag + 2];

		for (int lag = _MinLag; lag <= _MaxLag; lag++)
		{
			int overlap = _Window - lag;

			if (overlap <= 0)
				break;

			double correlation = 0.0;
			double energy = 0.0;

			for (int i = 0; i < overlap; i++)
			{
				float a = _Samples[_Start + i];
				float b = _Samples[_Start + i + lag];

				correlation += a * b;
				energy += a * a + b * b;
			}

			nsdf[lag] = energy > 0.000000000001 ? (float)(2.0 * correlation / energy) : 0.0f;
		}

		// Key maxima: the high point of each positive run. Ordinary local maxima are far too noisy to use.
		List<int> peaks = new();

		bool inRun = false;
		int peak = -1;

		for (int lag = _MinLag + 1; lag < _MaxLag; lag++)
		{
			if (!inRun)
			{
				if (nsdf[lag] > 0.0f && nsdf[lag] >= nsdf[lag - 1])
				{
					inRun = true;
					peak = lag;
				}

				continue;
			}

			if (nsdf[lag] > nsdf[peak])
				peak = lag;

			if (nsdf[lag] <= 0.0f)
			{
				peaks.Add(peak);
				inRun = false;
			}
		}

		if (inRun && peak > 0)
			peaks.Add(peak);

		if (peaks.Count == 0)
			return 0.0f;

		float best = 0.0f;

		foreach (int candidate in peaks)
			best = MathF.Max(best, nsdf[candidate]);

		if (best <= 0.0f)
			return 0.0f;

		foreach (int candidate in peaks)
		{
			if (nsdf[candidate] >= best * 0.9f)
				return Refine(nsdf, candidate);
		}

		return 0.0f;
	}

	/// <summary>
	/// Parabolic fit through the peak and its neighbours, for sub-sample precision.
	/// </summary>
	/// <remarks>
	/// Worth the few lines: the lag is an integer, so at high revs where the period is short, being one sample out
	/// is already a percent or two of error — tens of RPM on the label, and a permanently mistuned layer.
	/// </remarks>
	private static float Refine(float[] _Scores, int _Lag)
	{
		if (_Lag <= 0 || _Lag + 1 >= _Scores.Length)
			return _Lag;

		float previous = _Scores[_Lag - 1];
		float current = _Scores[_Lag];
		float next = _Scores[_Lag + 1];

		float denominator = previous - 2.0f * current + next;

		if (MathF.Abs(denominator) < 0.0000001f)
			return _Lag;

		float shift = 0.5f * (previous - next) / denominator;

		return _Lag + Math.Clamp(shift, -1.0f, 1.0f);
	}

	private static float[] RemoveDc(float[] _Samples)
	{
		float mean = 0.0f;

		foreach (float sample in _Samples)
			mean += sample;

		mean /= _Samples.Length;

		float[] result = new float[_Samples.Length];

		for (int i = 0; i < _Samples.Length; i++)
			result[i] = _Samples[i] - mean;

		return result;
	}

	private static float[] LowPass(float[] _Samples, int _SampleRate, float _Cutoff)
	{
		float rc = 1.0f / (MathF.Tau * _Cutoff);
		float dt = 1.0f / _SampleRate;
		float alpha = dt / (rc + dt);

		float[] result = new float[_Samples.Length];
		float value = 0.0f;

		for (int i = 0; i < _Samples.Length; i++)
		{
			value += alpha * (_Samples[i] - value);
			result[i] = value;
		}

		return result;
	}

	/// <summary>Wraps raw samples back up as a 16-bit mono PCM wav.</summary>
	public static byte[] BuildWav(short[] _Samples, int _SampleRate)
	{
		int dataLength = _Samples.Length * 2;

		using System.IO.MemoryStream stream = new();
		using System.IO.BinaryWriter writer = new(stream);

		writer.Write(System.Text.Encoding.ASCII.GetBytes("RIFF"));
		writer.Write(36 + dataLength);
		writer.Write(System.Text.Encoding.ASCII.GetBytes("WAVE"));

		writer.Write(System.Text.Encoding.ASCII.GetBytes("fmt "));
		writer.Write(16);
		writer.Write((short)1);
		writer.Write((short)1);
		writer.Write(_SampleRate);
		writer.Write(_SampleRate * 2);
		writer.Write((short)2);
		writer.Write((short)16);

		writer.Write(System.Text.Encoding.ASCII.GetBytes("data"));
		writer.Write(dataLength);

		foreach (short sample in _Samples)
			writer.Write(sample);

		writer.Flush();

		return stream.ToArray();
	}
}