module TestIAEA
using Test
using PhaseSpaceIO
using PhaseSpaceIO: position
using PhaseSpaceIO.Testing
using Accessors
using PhaseSpaceIO: ptype, read_particle, write_particle, CompressedIAEAParticle
using StaticArrays
using LinearAlgebra
using CoordinateTransformations

@testset "CompressedIAEAParticle" begin
    path = testdatapath("some_file.IAEAphsp")
    p = first(PhspVector(path))

    # does not throw
    show(IOBuffer(), p)
    show(IOBuffer(), [p])
    p2 = IAEAParticle(p)
    @test p isa CompressedIAEAParticle
    @test p2 isa IAEAParticle
    @test propertynames(p) == propertynames(p2)
    for prop in propertynames(p)
        @test getproperty(p, prop) == getproperty(p2, prop)
    end
    @test p == p2
    @test [p] == [p2]
    pos = @SVector randn(3)
    dir = normalize(@SVector randn(3))
    @test (@set position(p) = pos) == (@set position(p2) = pos)
    @test (@set direction(p) = dir) == (@set direction(p2) = dir)

    trafo = AffineMap(@SMatrix(randn(3,3)), @SVector(randn(3)))
    @test trafo(p) == trafo(p2)
end

@testset "IAEA save load with many constants" begin
    mktempdir() do dir
        h = IAEAHeader{0,0}((x=1, y=2, z=3, u=0, v=1, weight=5))
        p_truth = IAEAParticle(x=1,y=2,z=3,u=0,v=1,w=0,weight=5, typ=photon, E=6)
        path = IAEAPath(joinpath(dir, "phsp"))
        iaea_writer(path, h) do w
            write(w, p_truth)
        end

        r_reloaded, ps_reloaded = phsp_iterator(path) do iter
            iter.header.record_contents, collect(iter)
        end
        @test r_reloaded == h.record_contents
        @test length(ps_reloaded) == 1
        @test first(ps_reloaded) == p_truth
        @test eltype(ps_reloaded) <: IAEAParticle

        v = PhspVector(path)
        @test length(v) == 1
        @test eltype(v) <: CompressedIAEAParticle
        @test first(v) == p_truth
    end
end

@testset "similar_header" begin
    path = testdatapath("some_file.IAEAphsp")
    h = phsp_iterator(PhaseSpaceIO.similar_header, path)
    @test typeof(h) == typeof(IAEAHeader{0,1}())
    ps = phsp_iterator(collect, path)
    h  = PhaseSpaceIO.similar_header(ps)
    @test typeof(h) == typeof(IAEAHeader{0,1}())
end

@testset "finalizer" begin
    function write_sloppy(path, ps, r)
        w = iaea_writer(path, r)
        for p in ps
            write(w, p)
        end
    end
    path = IAEAPath(tempname())
    ps = [arbitrary(IAEAParticle{0,1})]
    r = IAEAHeader{0,1}()
    write_sloppy(path, ps, r)
    GC.gc()
    ps2 = phsp_iterator(collect, path)
    @test length(ps) == length(ps2) == 1
    @test first(ps) ≈ first(ps2)
    rm(path)
end

@testset "read write single particle" begin
    h = IAEAHeader{0,1}()
    P = ptype(h)
    @test P == IAEAParticle{0,1}
    p_ref = P(photon, 
        1.0f0, 2.0f0, 
        3.0f0, 4.0f0, 5.0f0, 
        0.53259337f0, 0.3302265f0, -0.7792912f0, 
        true, (), (13,))
    
    path = testdatapath("some_file.IAEAphsp")
    ps = iaea_iterator(collect, path)
    @test length(ps) == 1
    @test first(ps) == p_ref
    iaea_iterator(path) do iter
        @test iter.header.attributes[:INSTITUTION] == "hello world"
        @test iter.header.attributes[:ORIG_HISTORIES] == "18446744073709551615"
    end
    
    io = IOBuffer()
    write_particle(io, p_ref, h)
    seekstart(io)
    p = @inferred read_particle(io, h)
    @test p === p_ref
    @test eof(io)
end

@testset "test IAEAPhspIterator" begin
    path = testdatapath("some_file.IAEAphsp")
    phsp = iaea_iterator(path)
    @test length(phsp) == 1
    @test eltype(phsp) === IAEAParticle{0,1}
    @test collect(phsp) == collect(phsp)
    @test length(collect(phsp)) == 1
    close(phsp)
end

function test_header_contents(path)
    s = read(path, String)
    for key in [
        :IAEA_INDEX,
        :TITLE,
        :FILE_TYPE,
        :CHECKSUM,
        :RECORD_LENGTH,
        :BYTE_ORDER,
        :ORIG_HISTORIES,
        :PARTICLES,
        :TRANSPORT_PARAMETERS,
        :MACHINE_TYPE,
        :MONTE_CARLO_CODE_VERSION,
        :GLOBAL_PHOTON_ENERGY_CUTOFF,
        :GLOBAL_PARTICLE_ENERGY_CUTOFF,
        :COORDINATE_SYSTEM_DESCRIPTION,
        :BEAM_NAME,
        :FIELD_SIZE,
        :NOMINAL_SSD,
        :MC_INPUT_FILENAME,
        :VARIANCE_REDUCTION_TECHNIQUES,
        :INITIAL_SOURCE_DESCRIPTION,
        :PUBLISHED_REFERENCE,
        :AUTHORS,
        :INSTITUTION,
        :LINK_VALIDATION,
        :ADDITIONAL_NOTES,]
        @test occursin(string(key), s)
    end
end

@testset "test iaea_iterator iaea_writer" begin
    for _ in 1:100
        f = rand(1:3)
        i = rand(1:3)
        n = rand(1:1000)
        ps = [arbitrary(IAEAParticle{f,i}) for _ in 1:n]
        r = IAEAHeader{f,i}()
        dir = tempname()
        mkpath(dir)
        path = IAEAPath(joinpath(dir, "hello"))
        iaea_writer(path,r) do w
            for p in ps
                write(w, p)
            end
        end
        @test ispath(path.header)
        @test ispath(path.phsp)
        test_header_contents(path.header)

        ps_reload = iaea_iterator(collect, path)
        @test all(ps_reload .≈ ps)
        @test eltype(ps_reload) === eltype(ps)
        rm(path)

        @test !ispath(path.header)
        @test !ispath(path.phsp)
    end
end

end #module
