using System;
using System.Collections.Generic;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;

namespace Mosaic.Colors;

internal class VboxBuilder
{
    private int rMin = (int)MmcqQuantizer.MaxValue;
    private int rMax = 0;
    private int gMin = (int)MmcqQuantizer.MaxValue;
    private int gMax = 0;
    private int bMin = (int)MmcqQuantizer.MaxValue;
    private int bMax = 0;

    public void Populate(PixelAccessor<Rgba32> accessor)
    {
        for (var y = 0; y < accessor.Height; y++)
        {
            var row = accessor.GetRowSpan(y);
            for (var x = 0; x < row.Length; x++)
            {
                var pixel = row[x];

                var r = pixel.R >> MmcqQuantizer.Shift;
                rMin = int.Min(rMin, r);
                rMax = int.Max(rMax, r);

                var g = pixel.G >> MmcqQuantizer.Shift;
                gMin = int.Min(gMin, g);
                gMax = int.Max(gMax, g);

                var b = pixel.B >> MmcqQuantizer.Shift;
                bMin = int.Min(bMin, b);
                bMax = int.Max(bMax, b);
            }
        }
    }

    public VBox Build(ReadOnlyMemory<int> histogram) => new(rMin, rMax, gMin, gMax, bMin, bMax, histogram);
}

internal class VBox
{
    public static IComparer<VBox> CountComparer { get; } = new CountComparerImplementation();

    private sealed class CountComparerImplementation : IComparer<VBox>
    {
        public int Compare(VBox? x, VBox? y) => x is null || y is null ? -1 : y.Count.CompareTo(x.Count);
    }

    public static IComparer<VBox> CountVolumeComparer { get; } = new CountVolumeComparerImplementation();

    private sealed class CountVolumeComparerImplementation : IComparer<VBox>
    {
        public int Compare(VBox? x, VBox? y) =>
            x is null || y is null ? -1 : (y.Count * y.Volume).CompareTo(x.Count * x.Volume);
    }

    public readonly int RMin;
    public readonly int RMax;
    public readonly int GMin;
    public readonly int GMax;
    public readonly int BMin;
    public readonly int BMax;

    public readonly ReadOnlyMemory<int> Histogram;

    public readonly int Volume;

    private readonly Lazy<int> count;
    public int Count => count.Value;

    private readonly Lazy<Rgba32> average;
    public Rgba32 Average => average.Value;

    public VBox(int rMin, int rMax, int gMin, int gMax, int bMin, int bMax, ReadOnlyMemory<int> histogram)
    {
        RMin = rMin;
        RMax = rMax;
        GMin = gMin;
        GMax = gMax;
        BMin = bMin;
        BMax = bMax;
        Histogram = histogram;

        Volume = (RMax - RMin + 1) * (GMax - GMin + 1) * (BMax - BMin + 1);
        count = new(CalculateCount);
        average = new(CalculateAverage);
    }

    private int CalculateCount()
    {
        var histogram = Histogram.Span;
        var count = 0;
        for (var r = RMin; r <= RMax; r++)
        for (var g = GMin; g <= GMax; g++)
        for (var b = BMin; b <= BMax; b++)
        {
            count += histogram[MmcqQuantizer.GetColorIndex(r, g, b)];
        }
        return count;
    }

    private Rgba32 CalculateAverage()
    {
        if (RMin == RMax && GMin == GMax && BMin == BMax)
        {
            return new(RMin << MmcqQuantizer.Shift, GMin << MmcqQuantizer.Shift, BMin << MmcqQuantizer.Shift);
        }

        var histogram = Histogram.Span;
        var total = 0;
        var rSum = 0.0d;
        var gSum = 0.0d;
        var bSum = 0.0d;

        for (var r = RMin; r <= RMax; r++)
        for (var g = GMin; g <= GMax; g++)
        for (var b = BMin; b <= BMax; b++)
        {
            var value = histogram[MmcqQuantizer.GetColorIndex(r, g, b)];
            total += value;
            rSum += value * (r + 0.5d) * MmcqQuantizer.Multiplier;
            gSum += value * (g + 0.5d) * MmcqQuantizer.Multiplier;
            bSum += value * (b + 0.5d) * MmcqQuantizer.Multiplier;
        }

        if (total == 0.0d)
        {
            return new(
                (byte)double.Floor(MmcqQuantizer.Multiplier * (RMin + RMax + 1) / 2.0d),
                (byte)double.Floor(MmcqQuantizer.Multiplier * (RMin + RMax + 1) / 2.0d),
                (byte)double.Floor(MmcqQuantizer.Multiplier * (RMin + RMax + 1) / 2.0d)
            );
        }

        return new(
            (byte)double.Floor(rSum / total),
            (byte)double.Floor(gSum / total),
            (byte)double.Floor(bSum / total)
        );
    }
}