########################################################################
# Codes for the paper:
# An Anthropocene Map of Genetic Diversity
#
# Andreia Miraldo, Sen Li, Michael K. Borregaard, Alexander Floréz-Rodriguéz, Shyam Gopalakrishnan, Mirneza Risvanovic, Zhiheng Wang, Carsten Rahbek, Katharine A. Marske & David Nogués-Bravo
#
# Submitted to Science, 2016
# Code in this file by Michael K. Borregaard
#########################################################################

#Load the necessary library
using DataFrames, DataFramesMeta, StatsBase


"""
Calculate the GD (Genetic Diversity) value

**Parameters**
* 'dats' : A master matrix data frame generated by the 'create_master_matrices()' function
"""
function calcGD(dats)
    meanpaspec = by(dats, [:cell, :species], x -> mean(x[:num_per_bp]))
    GD = by(meanpaspec, :cell, x -> mean(x[:x1]))
    names!(GD, [:site, :GDval])
    GD, meanpaspec
end

"""
Calculate the weighted (by total # of sequence pairs) GD (Genetic Diversity) value

**Parameters**
* 'dats' : A master matrix data frame generated by the 'create_master_matrices()' function
"""
function calcGD_weighted(dats)
    meanpaspec = by(dats, [:cell, :species], x -> DataFrame(PIhat = mean(x[:num_per_bp]), numpairs = size(x,1)))
    GD = by(meanpaspec, :cell, x -> mean(x[:PIhat], WeightVec(x[:numpairs])))
    names!(GD, [:site, :GDval])
    GD, meanpaspec
end

"""
Calculate the total number of base pairs in a cell

**Parameters**
* 'dat' :A master matrix data frame generated by the 'create_master_matrices()' function
"""
function totalbasepairs(dat)
    seqs = unique!(dat[[:cell, :species, :seq2, :length_seq2]])
    seq1s = by(dat, [:cell, :species], df->df[[ :seq1, :length_seq1]][1,:])
    names!(seq1s, names(seqs))
    seqs = vcat(seqs, seq1s)
    by(seqs, [:cell], df-> sum(dropna(df[:length_seq2])))
end

"""
Calculate the total number of species in a cell

**Parameters**
* 'dat' :A master matrix data frame generated by the 'create_master_matrices()' function
"""
function totalspecies(dat)
    by(dat, :cell) do df
        sps = unique(df[:species])
        size(sps,1)
    end
end

"""
Calculate the richness and basepairs in a cell

**Parameters**
* 'dat' :A master matrix data frame generated by the 'create_master_matrices()' function
"""
function calc_cellvalues(dats)
    ret, full = calcGD(dats)
    ret[:richness] = by(dats, :cell, totalspecies)[:x1]
    ret[:basepairs] = by(dats, :cell, totalbasepairs)[:x1]
    ret, full
end


function ready_data(filename::AbstractString)
    ret = readtable(filename)
    ret = @where(ret, !isna(:seq1))
    ret = @where(ret, :overlap .> 0.5)
    ret
end
