issues.jl

using Test
using UnROOT

SAMPLES_DIR = joinpath(@__DIR__, "samples")


@testset "Issues" begin
    # issue 9 — fixed-size array branches ([N]/D, [M][N]/D, name[N]/D)
    rootfile = UnROOT.samplefile("issue9.root")
    @test_throws "is a TTree" UnROOT.array(rootfile, "arrays")
    @test_throws "is a TTree" UnROOT.array(rootfile, "structs")
    # individual branch reading
    @test UnROOT.array(rootfile, "arrays/nInt") == Int32[1]
    @test UnROOT.array(rootfile, "arrays/6dVec") == [UnROOT.FixLenVector(SVector{6,Float64}(1,2,3,4,5,6))]
    @test UnROOT.array(rootfile, "arrays/2x3Mat") == [UnROOT.FixLenVector(SVector{6,Float64}(1,2,3,4,5,6))]
    @test UnROOT.array(rootfile, "structs/2x3mat") == [UnROOT.FixLenVector(SVector{6,Float64}(1,2,3,4,5,6))]
    # LazyTree and arrays() (plural) work
    t = LazyTree(rootfile, "arrays")
    @test t[1].nInt == 1
    @test t[1].var"6dVec" == UnROOT.FixLenVector(SVector{6,Float64}(1,2,3,4,5,6))
    res = UnROOT.arrays(rootfile, "structs")
    @test res[1] == Int32[1]
    @test res[2] == [UnROOT.FixLenVector(SVector{6,Float64}(1,2,3,4,5,6))]
    close(rootfile)

    rootfile = UnROOT.samplefile("issue7.root")
    @test 2 == length(keys(rootfile))
    @test [1.0, 2.0, 3.0] == UnROOT.array(rootfile, "TreeD/nums")
    @test [1.0, 2.0, 3.0] == UnROOT.array(rootfile, "TreeF/nums")
    close(rootfile)

    # issue #55 and #156
    rootfile = UnROOT.samplefile("cms_ntuple_wjet.root")
    pts1 = UnROOT.array(rootfile, "variable/met_p4/fCoordinates/fCoordinates.fPt"; raw=false)
    pts2 = LazyTree(rootfile, "variable", [r"met_p4/fCoordinates/.*", "mll"])[!, Symbol("met_p4_fPt")]
    pts3 = LazyBranch(rootfile, "variable/good_jets_p4/good_jets_p4.fCoordinates.fPt")
    @test 24 == length(pts1)
    @test Float32[69.96958, 25.149912, 131.66693, 150.56802] == pts1[1:4]
    @test pts1 == pts2
    @test pts3[1:2] == [[454.0, 217.5, 89.5, 30.640625], [184.375, 33.28125, 32.28125, 28.46875]]
    close(rootfile)

    # issue 61
    rootfile = UnROOT.samplefile("issue61.root")
    @test LazyBranch(rootfile, "Events/Jet_pt")[:] == Vector{Float32}[[], [27.324587, 24.889547, 20.853024], [], [20.33066], [], []]
    close(rootfile)

    # issue 78
    rootfile = UnROOT.samplefile("issue61.root")
    arr = LazyTree(rootfile,"Events").Jet_pt;
    _ = length.(arr);
    @test length.(arr.buffer) == length.(arr.buffer_range)
    close(rootfile)

    # issue 108
    # unsigned short -> Int16, ulong64 -> UInt64
    # file minified with `rooteventselector --recreate -l 2 "trackntuple.root:trackingNtuple/tree" issue108_small.root`
    rootfile = UnROOT.samplefile("issue108_small.root")
    @test LazyBranch(rootfile, "tree/trk_algoMask")[2] == [0x0000000000004000, 0x0000000000004000, 0x0000000000004000, 0x0000000000004000]
    @test LazyBranch(rootfile, "tree/pix_ladder")[3][1:5] == UInt16[0x0001, 0x0001, 0x0001, 0x0001, 0x0003]
    close(rootfile)

    # issue 116
    rootfile = UnROOT.samplefile("issue116.root")
    @test length(rootfile["fTree"].fBranches.elements) == 112
    close(rootfile)

    # issue 246
    arr = LazyTree(joinpath(SAMPLES_DIR, "issue246.root"), "tree_NOMINAL").v_mcGenWgt
    @test all(reduce(vcat, arr) .== 1.0)

    # issue 323
    f = UnROOT.samplefile("issue323.root")
    t = LazyTree(f, "sim", [r"ghost/ghost\.(.*)" => s"\1"])
    @test 1200 == length(t)
    @test t[1].time[2] ≈ 36.396744f0
    @test t[end].xpos[end] ≈ 788.35144f0

    # issue 265 — RecoveredTBasket offsets were stored as Vector{UInt32} instead of
    # Vector{Int32}, causing a MethodError when indexing jagged LazyBranch entries
    # in files that contain embedded (recovered) TBaskets.
    f = UnROOT.samplefile("nanoAOD_2015_CMS_Open_Data_ttbar.root")
    t = LazyTree(f, "Events")
    @test length(t) == 200
    @test t.Jet_pt[1] == Float32[17.921875, 15.734375]
    @test t.Electron_pt[2] == Float32[32.679607]
    @test t.MET_pt[1:3] == Float32[33.261875, 33.803047, 70.08252]
    @test count(x -> length(x) > 0, t.Electron_pt[:]) == 63
    close(f)

    # issue 241 — TTree v5 / TBranch v8 old format: all data stored as embedded baskets
    # (fWriteBasket=0, fEntryNumber as Int32, fBasketEntry as Int32, fEntries/fTotBytes/fZipBytes as Float64)
    f = UnROOT.samplefile("issue241.root")
    t = LazyTree(f, "proton")
    @test 462 == length(t)
    @test 76.55430958116864 ≈ t.ekin[1]
    @test 0.012379961254578768 ≈ t.edep[1]
    @test 6.220370248201434 ≈ t.trackx[1]
    @test propertynames(t) ⊇ (:ekin, :edep, :trackx, :tracky, :trackz, :id, :process)
    close(f)

    # issue 377
    f = UnROOT.samplefile("issue377.root")
    arr = UnROOT.array(f, "podio_metadata/events___CollectionTypeInfo/events___CollectionTypeInfo.dataType")
    t = LazyTree(f, "podio_metadata", ["events___CollectionTypeInfo"])
    @test 1 == length(t.events___CollectionTypeInfo_dataType)
    @test 26 == length(t.events___CollectionTypeInfo_dataType[1])
    @test "edm4hep::CaloHitContributionCollection" == t.events___CollectionTypeInfo_dataType[1][1]
    @test "podio::LinkCollection<edm4hep::Vertex,edm4hep::ReconstructedParticle>" == t.events___CollectionTypeInfo_dataType[1][end]
    @test arr == t.events___CollectionTypeInfo_dataType
end

function _test_clean_GC(fname)
    for i in 1:5
        f = UnROOT.samplefile(fname)
        t = LazyTree(f, "t1")
        f = t = nothing
    end
end

@testset "Clean GC issue #260" begin
    fname = "tree_with_large_array_lzma.root"

    _test_clean_GC(fname)
    GC.gc()
    GC.gc()
    sleep(2)
    GC.gc()
    sleep(2)

    pid = getpid()
    @static if Sys.islinux()
        @test isempty(filter(contains(fname), readlines("/proc/$pid/smaps")))
    elseif Sys.isapple()
        @test isempty(filter(contains(fname), readlines(`vmmap $(getpid())`)))
    elseif Sys.iswindows()
        # TODO: add test for windows
    end
end

@testset "PR 266" begin
    f = UnROOT.samplefile("edm4hep_266.root")
    tree = LazyTree(f, "events", r"PandoraPFOs/PandoraPFOs.[(a-z)(A-Z))]")

    @test length(tree.PandoraPFOs_energy[1]) == 79
    @test length(tree.var"PandoraPFOs_covMatrix[10]"[1]) == 790
end

@testset "PR 342 TLeafC" begin
    df = LazyTree(UnROOT.samplefile("TLeafC_pr342.root"), "G4Sim")
    @test all(df.Process[1:10] .== ["Radioactivation", "msc", "eIoni", "Transportation", "ionIoni", "Radioactivation", "msc", "eIoni", "ionIoni", "Radioactivation"])
end