using Test
using ElectronPhonon
using LinearAlgebra

@testset "transport with screening" begin
    BASE_FOLDER = dirname(dirname(pathof(ElectronPhonon)))
    folder = joinpath(BASE_FOLDER, "test", "data_cubicBN")

    model = load_model(folder, epmat_outer_momentum="el", load_symmetry_operators=true)
    model.el_velocity_mode = :BerryConnection

    # temporary directory to store output data file
    tmp_dir = joinpath(BASE_FOLDER, "test", "tmp")
    mkpath(tmp_dir)

    window_k  = (15.0, 16.0) .* unit_to_aru(:eV)
    window_kq = (15.0, 16.0) .* unit_to_aru(:eV)

    nklist = (15, 15, 15)
    nqlist = (15, 15, 15)

    transport_params = ElectronTransportParams{Float64}(
        Tlist = [300.0, 200.0] .* unit_to_aru(:K),
        nlist = [1.0e19, 1.0e20] .* model.volume ./ unit_to_aru(:cm)^3,
        smearing = (:Gaussian, 50.0 * unit_to_aru(:meV)),
        volume = model.volume,
        nband_valence = 4,
        spin_degeneracy = 2,
    )

    # FIXME: In LindhardScreeningParams, n means carrier per bohr^3.
    #        In RPAScreeningParams, n means carrier per unit cell.

    # Some arbitrary screening parameter
    screening_params_lindhard = LindhardScreeningParams(
        degeneracy = 4,
        m_eff = 0.4,
        nlist = transport_params.nlist ./ model.volume,
        ϵM = model.polar_phonon.ϵ,
        smearing = 0.05 * unit_to_aru(:eV)
    )

    screening_params_rpa = RPAScreeningParams(
        spin_degeneracy = 2,
        Tlist = transport_params.Tlist,
        nlist = transport_params.nlist,
        ϵM = model.polar_phonon.ϵ,
        smearing = 0.05 * unit_to_aru(:eV),
        volume = model.volume,
        nband_valence = 4,
    )

    # Reference data generated by Julia
    σ_serta_reference = Dict(
        :Lindhard => [0.13989211028360177, 3.17537840929968],
        :RPA => [0.13444720985278077, 2.142379681074033],
    )
    σ_reference = Dict(
        :Lindhard => [0.07411494913216994, 1.8045314987187495],
        :RPA => [0.07106564977717034, 1.1656991021860752],
    )

    for method in [:Lindhard, :RPA]
        if method == :Lindhard
            screening_params = screening_params_lindhard
        elseif method == :RPA
            screening_params = screening_params_rpa
        end

        # Calculate matrix elements
        @time ElectronPhonon.run_transport(
            model, nklist, nqlist,
            folder = tmp_dir,
            window_k  = window_k,
            window_kq = window_kq,
            average_degeneracy = false,
            run_for_qme = true,
            screening_params = screening_params,
        )

        filename = joinpath(tmp_dir, "btedata_coherence.rank0.h5")

        qme_model = load_QMEModel(filename, transport_params);
        bte_compute_μ!(qme_model, do_print=true);
        compute_qme_scattering_matrix!(qme_model, compute_Sᵢ=true);

        if method == :RPA
            @test screening_params_rpa.μlist ≈ transport_params.μlist
        end

        # Calculate mobility
        out_linear = solve_electron_linear_conductivity(qme_model, qme_offdiag_cutoff=Inf, rtol=1e-5);
        @test out_linear.σ_serta ≈ cat([x * I(3) for x in σ_serta_reference[method]]..., dims=3)
        @test out_linear.σ ≈ cat([x * I(3) for x in σ_reference[method]]..., dims=3)
    end
end
