export genofreq, genofreq_expected, genocount_observed, genocount_expected

"""
    genocount_observed(locus::GenoArray)
Return a `Dict` of genotype counts of a single locus in a
`PopData` object.
"""
@inline function genocount_observed(locus::T) where T<:GenoArray
    # conditional testing if all genos are missing
    all(ismissing.(locus)) && return missing
    countmap(skipmissing(locus))
end

"""
    genocount_expected(locus::GenoArray)
Return a `Dict` of the expected genotype counts of a single locus in a
`PopData` object. Expected counts are calculated as the product of observed
allele frequencies multiplied by the number of non-missing genotypes.
"""
function genocount_expected(locus::T) where T<:GenoArray
    #count number of non-missing genotypes in the locus
    n = nonmissing(locus)

    # Get expected number of genotypes in a locus
    ## get the observed allele frequencies
    allele_dict = allelefreq(locus)

    ## split the appropriate pairs into their own vectors
    alle, freq = collect(keys(allele_dict)), collect(values(allele_dict))
    ## calculate expected genotype frequencies by multiplying all-by-all x n
    expected_genotype_freq = vec(freq * freq' .* n)

    # reform genotype frequencies with same all-by-all approach
    genos = reverse.(Base.Iterators.product(alle, alle) |> collect |> vec)
    expected = Dict{nonmissingtype(eltype(locus)), Float64}()
    for (geno, freq) in zip(genos, expected_genotype_freq)
        expected[geno] = get!(expected, geno, 0.0) + freq
    end

    return expected
end


"""
    genofreq(locus::GenoArray)
Return a `Dict` of genotype frequencies of a single locus in a
`PopData` object.
"""
@inline function genofreq(locus::T) where T<:GenoArray
    # conditional testing if all genos are missing
    all(ismissing.(locus)) && return missing
    proportionmap(locus |> skipmissing |> collect)
end


"""
    genofreq(data::PopData, locus::String; population::Bool = false)
Return a `Dict` of genotype frequencies of a single locus in a `PopData`
object. Use `population = true` to return a table of genotype frequencies
of that locus per population.
### Example
```
cats = @nancycats;
genofreq(cats, "fca8")
genofreq(cats, "fca8", population = true)
```
"""
function genofreq(data::PopData, locus::String; population::Bool=false)
    if !population
        tmp = data.genodata[data.genodata.locus .== locus, :genotype]
        genofreq(tmp)
    else
        tmp = data.genodata[data.genodata.locus .== locus, :]
        DataFrames.combine(groupby(tmp, :population), :genotype => genofreq => :freq)
    end
end

"""
    genofreq_expected(locus::GenoArray)
Return a `Dict` of the expected genotype frequencies of a single locus in a
`PopData` object. Expected frequencies are calculated as the product of
observed allele frequencies.
"""
function genofreq_expected(locus::T) where T<:GenoArray
    # Get expected number of genotypes in a locus
    ## get the observed allele frequencies
    allele_dict = allelefreq(locus)

    ## split the appropriate pairs into their own vectors
    alle, freq = collect(keys(allele_dict)), collect(values(allele_dict))

    ## calculate expected genotype frequencies by multiplying all-by-all
    expected_genotype_freq = vec(freq * freq')

    # reform genotype frequencies with same all-by-all approach
    genos = reverse.(Base.Iterators.product(alle, alle) |> collect |> vec)

    # reform genotype frequencies into a Dict
    expected = Dict{nonmissingtype(eltype(locus)), Float64}()
    for (geno, freq) in zip(genos, expected_genotype_freq)
        expected[geno] = get!(expected, geno, 0.0) + freq
    end

    return expected
end

"""
    genofreq_expected(data::PopData, locus::String; population::Bool = false)
Return a `Dict` of expected genotype frequencies of a single locus in a
`PopData` object. Use `population = true` to return a table of expected genotype
frequencies of that locus per population.
### Example
```
cats = @nancycats;
genofreq_expected(cats, "fca8")
genofreq_expected(cats, "fca8", population = true)
```
"""
function genofreq_expected(data::PopData, locus::String; population::Bool=false)
    if !population
        tmp = data.genodata[data.genodata.locus .== locus, :genotype]
        genofreq_expected(tmp)
    else
        tmp = data.genodata[data.genodata.locus .== locus, :]
        DataFrames.combine(groupby(tmp, :population), :genotype => genofreq_expected => :freq)
    end
end
