using FileTrees, OnlineStats, Serialization


function runtest(ctx, imgsize...)
  mktempdir() do tmpdir
      tree = FileTree(tmpdir)
      params = [(x, y) for x=-1:0.037:1, y=-1:0.037:1]
      for i = 1:size(params,1)
          for j = 1:size(params,2)
              tree = touch(tree, "$i/$j.png"; value=params[i, j])
          end
      end

      t1 = FileTrees.mapvalues(tree, lazy=true) do params
          randn(imgsize...)
      end

      res = FileTrees.save(t1; lazy=true, exec=false) do file
          serialize(path(file), file[])
      end
      @time exec(ctx, res)

      t2 = FileTrees.load(tree; lazy=true) do f
          deserialize(path(f))
      end

      t3 = FileTrees.mapvalues(t2) do img
          map(sin, sum(img; dims=ndims(img)))
      end

      t4 = FileTrees.mapvalues(t3) do simg
          fit!(Hist(0:0.05:1), simg)
      end
         
      @time h = exec(reducevalues(merge, t4))
  end
end

## Current main:

julia> runtest(FileTrees.Dagger.Context(), 1, 1, 1); # deprecation spam removed
4.447815 seconds (9.74 M allocations: 728.701 MiB, 8.65% gc time, 110.86% compilation time: 3% of which was recompilation)
10.924932 seconds (11.41 M allocations: 733.204 MiB, 1.82% gc time, 35.11% compilation time)

julia> runtest(FileTrees.Dagger.Context(), 300, 300, 3);
10.110492 seconds (7.98 M allocations: 6.523 GiB, 11.12% gc time, 0.12% compilation time)
20.979808 seconds (10.45 M allocations: 10.732 GiB, 10.96% gc time)

## This PR:

julia> runtest(FileTrees.Dagger.Context(), 1, 1, 1);
6.149975 seconds (13.90 M allocations: 871.720 MiB, 8.53% gc time, 69.51% compilation time: 2% of which was recompilation)
4.329147 seconds (15.66 M allocations: 957.089 MiB, 8.32% gc time, 74.08% compilation time)

julia> runtest(FileTrees.Dagger.Context(), 300, 300, 3);
11.772658 seconds (11.74 M allocations: 6.846 GiB, 13.51% gc time, 0.07% compilation time)
17.383311 seconds (14.18 M allocations: 10.991 GiB, 17.50% gc time)

julia> runtest(FileTrees.SingleTreadedContext(), 1, 1, 1);
3.805192 seconds (1.81 M allocations: 116.478 MiB, 0.62% gc time, 13.30% compilation time)
13.108113 seconds (15.73 M allocations: 986.425 MiB, 4.47% gc time, 39.17% compilation time: 2% of which was recompilation)

julia> runtest(FileTrees.SingleTreadedContext(), 300, 300, 3);
17.225328 seconds (1.02 M allocations: 6.146 GiB, 4.74% gc time)
18.537098 seconds (14.18 M allocations: 11.064 GiB, 15.04% gc time, 0.03% compilation time)

julia> runtest(FileTrees.ThreadContext(), 1, 1, 1);
1.733510 seconds (1.97 M allocations: 136.489 MiB, 0.72% gc time, 140.78% compilation time)
13.661712 seconds (15.70 M allocations: 993.805 MiB, 4.57% gc time, 35.86% compilation time: 2% of which was recompilation)

julia> runtest(FileTrees.ThreadContext(), 300, 300, 3);
6.914679 seconds (1.70 M allocations: 6.268 GiB, 11.00% gc time, 0.12% compilation time)
13.141357 seconds (13.87 M allocations: 10.997 GiB, 19.63% gc time)

julia> versioninfo()
Julia Version 1.10.3
Commit 0b4590a550 (2024-04-30 10:59 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Windows (x86_64-w64-mingw32)
CPU: 8 × 11th Gen Intel(R) Core(TM) i5-1145G7 @ 2.60GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, tigerlake)
Threads: 4 default, 0 interactive, 2 GC (on 8 virtual cores)
Environment:
JULIA_EDITOR = code
JULIA_NUM_THREADS = 4