using System;
using System.Linq;
using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Running;
using LanguageExt; // Your Namespace
using LanguageExt;
using LanguageExt;
using PersistentMap; // NuGet: LanguageExt.Core

[MemoryDiagnoser]
public class ImmutableCollectionBenchmarks
{
    [Params(100, 1000, 100_000)]
    public int N;

    private int[] _keys;
    private int[] _values;

    // --- 1. Your Collections ---
    private PersistentMap<int, int, IntStrategy> _niceMap;
    private IntStrategy _strategy;

    // --- 2. Microsoft Collections ---
    private System.Collections.Immutable.ImmutableSortedDictionary<int, int> _msSortedMap;
    private System.Collections.Immutable.ImmutableDictionary<int, int> _msHashMap;

    // --- 3. LanguageExt Collections ---
    private LanguageExt.Map<int, int> _leMap;       // AVL Tree (Sorted)
    private LanguageExt.HashMap<int, int> _leHashMap; // Hash Array Mapped Trie (Unsorted)

    [GlobalSetup]
    public void Setup()
    {
        // Generate Data
        var rand = new Random(42);
        _keys = Enumerable.Range(0, N).Select(x => x * 2).ToArray(); 
        rand.Shuffle(_keys); 
        _values = _keys.Select(k => k * 100).ToArray();
        _strategy = new IntStrategy();

        // 1. Setup NiceBTree
        var transient = BaseOrderedMap<int, int, IntStrategy>.CreateTransient(_strategy);
        for (int i = 0; i < N; i++) transient.Set(_keys[i], _values[i]);
        _niceMap = transient.ToPersistent();

        // 2. Setup MS Immutable
        var msBuilder = System.Collections.Immutable.ImmutableSortedDictionary.CreateBuilder<int, int>();
        for (int i = 0; i < N; i++) msBuilder.Add(_keys[i], _values[i]);
        _msSortedMap = msBuilder.ToImmutable();

        _msHashMap = System.Collections.Immutable.ImmutableDictionary.CreateRange(
            _keys.Zip(_values, (k, v) => new System.Collections.Generic.KeyValuePair<int, int>(k, v)));

        // 3. Setup LanguageExt
        // Note: LanguageExt performs best when bulk-loaded from tuples
        var tuples = _keys.Zip(_values, (k, v) => (k, v));
        _leMap = new LanguageExt.Map<int,int>(tuples); 
        _leHashMap = new HashMap<int, int>(tuples);
    }

    // =========================================================
    // 1. BUILD (Item by Item)
    // Note: LanguageExt has no "Mutable Builder", so this tests
    // the cost of pure immutable inserts vs your Transient/Builder.
    // =========================================================

    [Benchmark(Description = "Build: NiceBTree (Transient)")]
    public int Build_NiceBTree()
    {
        var t = BaseOrderedMap<int, int, IntStrategy>.CreateTransient(_strategy);
        for (int i = 0; i < N; i++) t.Set(_keys[i], _values[i]);
        return t.Count;
    }

    [Benchmark(Description = "Build: MS Sorted (Builder)")]
    public int Build_MsSorted()
    {
        var b = System.Collections.Immutable.ImmutableSortedDictionary.CreateBuilder<int, int>();
        for (int i = 0; i < N; i++) b.Add(_keys[i], _values[i]);
        return b.Count;
    }

    [Benchmark(Description = "Build: LanguageExt Map (AVL)")]
    public int Build_LanguageExt_Map()
    {
        var map = LanguageExt.Map<int, int>.Empty;
        for (int i = 0; i < N; i++) 
        {
            // Pure immutable add
            map = map.Add(_keys[i], _values[i]);
        }
        return map.Count;
    }

    [Benchmark(Description = "Build: LanguageExt HashMap")]
    public int Build_LanguageExt_HashMap()
    {
        var map = LanguageExt.HashMap<int, int>.Empty;
        for (int i = 0; i < N; i++) 
        {
            map = map.Add(_keys[i], _values[i]);
        }
        return map.Count;
    }

    // =========================================================
    // 2. READ (Lookup)
    // =========================================================

    [Benchmark(Description = "Read: NiceBTree")]
    public int Read_NiceBTree()
    {
        int found = 0;
        for (int i = 0; i < N; i++)
        {
            if (_niceMap.TryGetValue(_keys[i], out _)) found++;
        }
        return found;
    }

    [Benchmark(Description = "Read: MS Sorted")]
    public int Read_MsSorted()
    {
        int found = 0;
        for (int i = 0; i < N; i++)
        {
            if (_msSortedMap.ContainsKey(_keys[i])) found++;
        }
        return found;
    }

    [Benchmark(Description = "Read: LanguageExt Map")]
    public int Read_LanguageExt_Map()
    {
        int found = 0;
        for (int i = 0; i < N; i++)
        {
            // Find returns Option<V>, IsSome checks if it exists
            if (_leMap.Find(_keys[i]).IsSome) found++;
        }
        return found;
    }

    [Benchmark(Description = "Read: LanguageExt HashMap")]
    public int Read_LanguageExt_HashMap()
    {
        int found = 0;
        for (int i = 0; i < N; i++)
        {
            if (_leHashMap.Find(_keys[i]).IsSome) found++;
        }
        return found;
    }

    // =========================================================
    // 3. ITERATE (Foreach)
    // =========================================================

    [Benchmark(Description = "Iterate: NiceBTree")]
    public int Iterate_NiceBTree()
    {
        int sum = 0;
        foreach (var kvp in _niceMap) sum += kvp.Key;
        return sum;
    }

    [Benchmark(Description = "Iterate: MS Sorted")]
    public int Iterate_MsSorted()
    {
        int sum = 0;
        foreach (var kvp in _msSortedMap) sum += kvp.Key;
        return sum;
    }

    [Benchmark(Description = "Iterate: LanguageExt Map")]
    public int Iterate_LanguageExt_Map()
    {
        int sum = 0;
        // LanguageExt Map is IEnumerable<(Key, Value)>
        foreach (var item in _leMap) sum += item.Key;
        return sum;
    }

    [Benchmark(Description = "Iterate: LanguageExt HashMap")]
    public int Iterate_LanguageExt_HashMap()
    {
        int sum = 0;
        foreach (var item in _leHashMap) sum += item.Key;
        return sum;
    }

    // =========================================================
    // 4. SET (Persistent / Immutable Update)
    // =========================================================

    [Benchmark(Description = "Set: NiceBTree")]
    public PersistentMap<int, int, IntStrategy> Set_NiceBTree()
    {
        return _niceMap.Set(_keys[N / 2], -1); 
    }

    [Benchmark(Description = "Set: MS Sorted")]
    public System.Collections.Immutable.ImmutableSortedDictionary<int, int> Set_MsSorted()
    {
        return _msSortedMap.SetItem(_keys[N / 2], -1);
    }

    [Benchmark(Description = "Set: LanguageExt Map")]
    public LanguageExt.Map<int, int> Set_LanguageExt_Map()
    {
        return _leMap.SetItem(_keys[N / 2], -1);
    }

    [Benchmark(Description = "Set: LanguageExt HashMap")]
    public LanguageExt.HashMap<int, int> Set_LanguageExt_HashMap()
    {
        return _leHashMap.SetItem(_keys[N / 2], -1);
    }
}