--------------------------------------------------------------------------------
---------------------------- A -- Pop demographics------------------------------
--------------------------------------------------------------------------------

-- WARNING !!!
-- Demographics are generated as a matrix, but the we use it as vector
-- Conversion needed

-- TODO => add support for usrData type "demo"
-- Check type & dimensions
local function demo_check(demoSpe, mapLength, mapWidth)
    --for key,value in pairs(demoSpe) do print(key,value) end
    assert_type(demoSpe, 'table', "Individuals: demogaphy: not a valid table")
    assert(#demoSpe==mapWidth, "Width pop matrix not consistant")
    for key,value in pairs(demoSpe)
    do
        assert(#value==mapLength, "Length pop matrix not consistant for line "..key)
    end

    return demoSpe
end

--[[ Generate or load
--local function generate_demo(usrDemo, species, mapLength, mapWidth)
local function generate_demo(usrDemo, species, mapLength, mapWidth)
    local demo = {}
    local speId = species.id

    if is_type(species.initDemo, 'table') then
        if species.initDemo.source == "param" then
            --if species.initDemo.vertical then print("vertical") else print("horizontal") end
            assert((species.initDemo.vertical and mapLength<=#species.initDemo.npop)
                or (not species.initDemo.vertical and mapWidth<=#species.initDemo.npop), "Wrong pop param format: npop vector too short")

            demo = make_profile_matrix(species.initDemo.npop, mapLength, mapWidth, species.initDemo.vertical)
        elseif species.initDemo.source == "value" then
            demo = make_matrix(species.initDemo.npop, mapLength, mapWidth)
        elseif species.initDemo.source == "userdata" then
            demo = demo_check(usrDemo[speId], mapWidth, mapLength)
        else
            print("\nError: source of demography not recognized")
            os.exit()
        end
    ---- Failure
    else
        print("\nError: no demographic parameter for species "..speId)
        os.exit()
    end
    return demo
end --]]


--------------------------------------------------------------------------------
------------------- B -- Allelic frequencies -----------------------------------
--------------------------------------------------------------------------------

-- Check type and dimensions of usr allFreq (n locus * n alleles ini)
-- TODO -- Incomplete => structure must include data for each pop
-- TODO -- Add userdata support
local function allFreq_check(usrAllFreq, nAllpLoci)
    assert_type(usrAllFreq, 'table', "Individuals: usr allFreq: invalid table")
    assert(#usrAllFreq == #nAllpLoci, "Individuals: usr allFreq: wrong size")
    for i = 1, #nAllpLoci do
        assert(is_type_vector(usrAllFreq[i], 'numeric',  nAllpLoci[i]),
               "Individuals: usr allFreq: wrong size/type at locus "..i)
    end
    return  usrAllFreq
end


-- Write geneome structure in a file as a matrix loci*alles,
-- If rows of different length, completed with NA,
-- alleles ini observed are represended with a 1, otheres with a 0
-- This file is an input for dirichlet R function
local function mkAllInitFile(species, rDir)
    local id = species.id
    local geneStruct = species.nAllpLoci
    local alleleInit = species.nAllInit
    local max_nallele = max(species.nAllpLoci)
    local file_name = rDir .. id .. "_genome.csv"

    f = assert(io.open(file_name, "w"))
    f:write("# Number of alleles for each loci of species ", id, "\n")
    for j = 1, #geneStruct, 1 do -- For each loci
        for k = 1,  max_nallele, 1 do -- For each Allele -- Avant geneStruct[j] au lieu de max_nallele
            if k <= alleleInit[j] then
                f:write("1 ")
            elseif k <= geneStruct[j] then
                f:write("0 ")
            else
                f:write("NA ")
            end
        end
        f:write("\n")
    end
    f:close()
    f = nil
end


-- Generate or load
local function generate_allFreq(usrAllFreq, nPop, species, rDir, outDir)
    local allFreq = {}
    ---- Userdata -- TODO add support of userdata, define format
    if is_type(usrAllFreq, 'table') and is_type(usrAllFreq[speId], 'table') then
        allFreq = allFreq_check(usrAllFreq[speId])
    ---- Generated -- regular
    else
        assert_type(species.Nm, 'number', "Allelic Frequencies: Nm undefined")
        -- File describing genetic structure of the species
        mkAllInitFile(species, rDir)
        -- Call to the R script that generates dirichlet distribution,
        -- Ok, gsl_ran_dirichlet could be called from c++ instead of this R call.
        -- Doing this would be interesting in term of portability & readability
        local filename = rDir .. species.id .."_genome.csv"
        local fileout = outDir .. species.id .. "_initial_frequencies"
        local Rscript = rDir .. "dirichlet_standard.r "
        local rnd = Random()
        --print("ALLEFF seed:"..rnd)
        --print("fileout:"..fileout)
        local command_line = Rscript ..
                             "filename=\"" .. filename .. "\" , " ..
                             "Npop=" .. nPop .. " , " ..
                             "Nm=" .. species.Nm .. " , " ..
                             "varname=" .. "all_freq" .. " , " ..
                             "Seed=" .. rnd ..
                             " > \"" .. fileout .. "\""
        -- execute command line
        os.execute(command_line)
        allFreq = dofile(fileout) -- We now have frequencies in : all_freq[locus][pop][allele])
    end
    return allFreq
end



-- Transforms allelics frequencies in cumulative frequences
local function mk_cumulFreq(allFreq)
    local cumulFreq = {}
    for i = 1, #allFreq, 1 do -- For each loci
        cumulFreq[i] = {}
        for j = 1, #allFreq[i], 1 do -- For each pop
            cumulFreq[i][j] = {}
            -- Cumulative sum
            local prev = 0
            for k = 1, #allFreq[i][j], 1 do -- For each allele
                cumulFreq[i][j][k] = allFreq[i][j][k] + prev
                prev = cumulFreq[i][j][k]
            end
            -- Correction of cumulative sum if all probas are equal to 0 (to get equiprobable)
            if cumulFreq[i][j][#cumulFreq[i][j]] == 0 -- Cas ou la sommee est égale à zéro : équiprobable
            then
                for k = 1, #allFreq[i][j], 1 do cumulFreq[i][j][k] = k end
            end
            -- Normalisation
            for k = 1, #allFreq[i][j], 1 do
                cumulFreq[i][j][k] = cumulFreq[i][j][k] / cumulFreq[i][j][#cumulFreq[i][j]]
            end
        end
    end
    return cumulFreq
end


--------------------------------------------------------------------------------
-------------------------- C -- Individual generation --------------------------
--------------------------------------------------------------------------------

-- Initialisation des génotypes des individus à partir des fréquences alléliques initiales de Ncoeff
-- Pour tirer les allèles on utilise "Inverse transform sampling":
--[[
          - Calculer les proba cumulées
          - Tirer un nombere aléatoire uniforme entre 0 et 1
          - Faire une recherche binaire du nombre dans le tableau des proba cumulées
--]]


-- generates individuals for all populations in a specie
local function pop_generate(species, demo, cumAllFreq)
    local vectDemo = matrix_to_vect(demo)
    local ploidy = species.ploidy
    local Npop = #cumAllFreq[1] -- same number of pop for each locus
    local Nlocus = #cumAllFreq
    local populations = {}

    for pop = 1, Npop, 1 do  -- For each pop
        populations[pop] = {}
        for ind = 1, vectDemo[pop], 1 do -- For each individual
            local individual = {}
            individual.genome = {}
            individual.specie = species.id
            individual.age = 1 -- TODO => how do we generate ages

            local cnt = 1
            for locus = 1, Nlocus, 1 do
                -- Mat / Pat (1 / 2)
                for p = 1, ploidy[locus], 1 do
                    local rand = Random()
                    -- Third parameter == 2, we are looking for the upper bound
                    individual.genome[cnt] = binarySearch(cumAllFreq[locus][pop], rand, 2)
                    cnt = cnt + 1
                end
            end
            populations[pop][ind] = individual
            -- TODO => locus cytoplasmiques ?
        end
    end
    return populations
end

-- generates individuals for all populations in a specie
local function pop_generate_ageclasses(species, mode, cumAllFreq)
    local ploidy = species.ploidy
    local Npop = #cumAllFreq[1] -- same number of pop for each locus
    local Nlocus = #cumAllFreq
    local populations = {}

    local deltaAge_per_class = species.deltaAge
    local age_init_per_class = {}
    local nInitInd_per_class = species.sizeByAge

    -- pop ratio to keep proportions between age classes
    local vectDemo = {}
    local ageClass = {}
    local sumDelta = 0
    local nAgeClasses = #deltaAge_per_class
    local patchPerClass = {}
    local maxPatchPerClass = {}

    local minForClass = {}

    for ageclass = 1, nAgeClasses, 1 do
        age_init_per_class[ageclass] = sumDelta
        sumDelta = sumDelta + deltaAge_per_class[ageclass]
    end

    local partialSum = 0
    local partialNPatch = 0
    for ageclass = 1, nAgeClasses, 1 do
        minForClass[ageclass] = 1.0*partialSum/sumDelta
        maxPatchPerClass[ageclass] = math.floor(deltaAge_per_class[ageclass]/sumDelta*Npop)
        patchPerClass[ageclass] = 0
        --print(minForClass[ageclass])
        partialSum = partialSum + deltaAge_per_class[ageclass]
        partialNPatch = partialNPatch + maxPatchPerClass[ageclass]
    end
    if (partialNPatch < Npop) then
        maxPatchPerClass[nAgeClasses] = Npop - partialNPatch + maxPatchPerClass[nAgeClasses]
    end
    --[[for ageclass = 1, nAgeClasses, 1 do
        print(maxPatchPerClass[ageclass])
    end
    for ageclass = 1, nAgeClasses, 1 do
        print(patchPerClass[ageclass])
    end--]]


    local all_young =false
    local random_location_for_ageclasses = true

    if (mode=="evenage") then
        print("EVEN-AGE at intra-pop level")
        if (not all_young) then
          print("Asynchronous at inter-pop level")
          if (random_location_for_ageclasses) then
              local pop = 1
              while (pop <= Npop) do
                  local randomPos = Random()
                  local currPos = nAgeClasses
                  while (currPos > 0 and minForClass[currPos] > randomPos) do
                      currPos = currPos - 1
                  end
                  --print(currPos)
                  if (patchPerClass[currPos] + 1 <= maxPatchPerClass[currPos]) then
                      ageClass[pop] = currPos
                      patchPerClass[currPos] = patchPerClass[currPos] + 1
                      --print(randomPos.." "..ageClass[pop])
                      pop = pop + 1
                  --else
                  --print("Fail patch filling")
                  end

              end
          else
              print("Specified location for ageclasses not implemented yet.")
          end
        else
          print("EVEN-AGE at inter-pop level")
          local pop = 1
          while (pop <= Npop) do
            ageClass[pop] = 1
            pop = pop + 1
          end
        end

    else
        print("Uneven age not implemented yet.")
    end

    for pop = 1, Npop, 1 do  -- For each pop
        populations[pop] = {}
        local nInd = nInitInd_per_class[ageClass[pop]]
        local ageVariation = 0;
        if (not all_young) then
          ageVariation = math.floor(Random()*deltaAge_per_class[ageClass[pop]])
        end
        for ind = 1, nInd, 1 do -- For each individual
            local individual = {}
            individual.genome = {}
            individual.specie = species.id
            individual.age = age_init_per_class[ageClass[pop]] + ageVariation -- TODO => how do we generate ages

            local cnt = 1
            for locus = 1, Nlocus, 1 do
                -- Mat / Pat (1 / 2)
                for p = 1, ploidy[locus], 1 do
                    local rand = Random()
                    -- Third parameter == 2, we are looking for the upper bound
                    individual.genome[cnt] = binarySearch(cumAllFreq[locus][pop], rand, 2)
                    cnt = cnt + 1
                end
            end
            populations[pop][ind] = individual
            -- TODO => locus cytoplasmiques ?
        end
    end
    return populations
end

--------------------------------------------------------------------------------
--------------------------- Generation function --------------------------------
--------------------------------------------------------------------------------

-- Checks individual user data, and reformats
-- Format : pop, age, loc1, loc2, loc3
local function ind_check(usrInds, nPop, species)
    local pop = {}
    local tempInd = {}
    local speId = species.id
    local nAllpLoci = species.nAllpLoci
    local nLocus = species.nLocus
    local ploidy = species.ploidy
    local genomeSize = species.genomeSize
    -- Init table of populations
    for i = 1, nPop do pop[i] = {} end

    for _, ind in ipairs(usrInds) do
        local currentPop = ind[1] + 1
        assert(currentPop > 0 and currentPop <= nPop, "usrInd check: pop out of range")

        local nind = #pop[currentPop] + 1
        --print("IND SIZE", #ind)
        -- Checks length and type
        assert(is_type_vector(ind, 'number', genomeSize + 2),
               "User individual check: wrong format")
        -- Checks range of population
        assert(currentPop > 0 and currentPop <= nPop, "usrInd check: pop out of range")
        -- Check range of alleles & copy alleles
        local genome = {}
        local pos = 0
        for l = 1, nLocus do
            for p = 1, ploidy[l] do
                pos = pos + 1
                assert(ind[pos+2] >= 0 and ind[pos+2] < nAllpLoci[l],--, value Allele by locus constant for all loci?
                   "usrInd check: locus "..l.." out of range in spe "..speId)
                genome[pos] = ind[pos+2] + 1
            end
        end
        -- Save individual
        tempInd = {}
        tempInd.genome = genome
        tempInd.age = ind[2]
        --print("Add indiv to pop", currentPop, nind)
        pop[currentPop][nind] = tempInd
    end
    return pop
end


local function make_multi(usrInds, usrDemo, usrAllFreq, config, species, rDir, outDir)
    local allPops = {}
    local mapLength = config.mapLength; local mapWidth = config.mapWidth
    local nPop = mapLength * mapWidth

    ---- For each species
    for _, spe in ipairs(species) do
        local speId = spe.id
        -- Case 1 => userdata contains individuals for this spe
        if spe.source == "injection" then
            if is_type(usrInds, 'table') and is_type(usrInds[speId], 'table') then
                print("Import individuals from ind.txt")
                --for k, v in ipairs(usrInds[speId]) do print(k, v) end
                allPops[spe.id] = ind_check(usrInds[speId], nPop, spe)
            else
                print('initDemo.source = "injection" ; Problem with file load. Check ind.txt')
            end
        -- Case 2 => no userdata: generation needed
        else
            print("Generate individuals following demography")
            --local demo = generate_demo(usrDemo, spe, mapLength, mapWidth)
            local allFreq = generate_allFreq(usrAllFreq, nPop, spe, rDir, outDir)
            local cumAllFreq = mk_cumulFreq(allFreq)
            allPops[spe.id] = pop_generate_ageclasses(spe, "evenage", cumAllFreq)
        end
    end
    return allPops
end



--------------------------------------------------------------------------------
--------------------------- Module interface -----------------------------------
--------------------------------------------------------------------------------
local module = {}
module.make_multi = make_multi
return module
