--------------------------------------------------------------------------------
------------------------- Locus Types ------------------------------------------
--------------------------------------------------------------------------------
-- Note: For now only types a, b and c are supported.
-- This limitation could be restricted to the generation of allelic effects
-- for the lua part in cpp, it is a, b c only anyway


-- Dependencies
local confAE = require("init_allelic_effect")


--------------------------------------------------------------------------------
------------------------- Fields checking functions ----------------------------
--------------------------------------------------------------------------------

-- id --------------------------------------------------------------------------
local function id_check(id)
    assert_type(id, 'string', "Trait: ltypes: id: wrong type")
    return id
end

-- lociList --------------------------------------------------------------------
-- All locus listed must exist in the species (nLocus)
local function lociList_check(lociList, nLocus)
    -- Check existence & type of all values
    assert_type(lociList, 'table', "Trait: ltypes: lociList: wrong type")
    if (#lociList > 0) then
        assert(is_type_vector(lociList, 'number'), "Trait: ltypes: lociList: wrong type")
        --assert(#lociList > 0, "Trait: ltypes: lociList: no values")
        -- Values need to be sorted
        quicksort(lociList)
        -- Assert range (begin & end, beacause sorted)
        assert(lociList[1] > 0, "Trait: ltype: locilist: loci number cannot be < 0")
        assert(lociList[#lociList] <= nLocus,
                "Trait: ltype: locilist: loci number cannot be > "..nLocus)
        -- Assert unique values
        assert(is_unique(lociList), "Trait: ltype: locilist: repetition detected")
        return lociList
    else
        --print("Empty locus list.")
        return {}
    end
end

-- Phi -------------------------------------------------------------------------
local function phi_check(phi)
    assert_type( phi, 'number', "Trait: ltypes: phi: wrong type")
    return phi
end

-- Pow -------------------------------------------------------------------------
local function pow_check(pow)
    assert_type( pow, 'number', "Trait: ltypes: pow: wrong type")
    return pow
end



--------------------------------------------------------------------------------
----------------------------- generated fields ---------------------------------
--------------------------------------------------------------------------------
-- Checks dimensions & size for allelic effects
local function allEff_check(allEff, nAllLt)
    assert_type(allEff, 'table', "ltype: allEff_check: Wrong type")
    assert(is_type_vector(allEff, 'table', #nAllLt),
           "ltype: allEff_check: Wrong size of type")
    for i = 1, #allEff do
        assert(is_type_vector(allEff[i], 'number', nAllLt[i]),
               "ltype: allEff_check: Wrong type/size at locus "..i)
    end
    return allEff
end


--  make_allelic_effects => either usrdata or generate
local function allEff_make(allEff, allEnvEff, newLt, trait, species, ltypes)
    -- Shortcuts for var from newLt, trait and species
    local lociList = newLt.lociList; local id = newLt.id
    --local phi = newLt.phi; local pow = newLt.pow
    local varLociWeights = trait.varLociWeights
    --local iTotVa = trait.initTotVa
    --local h2 = trait.h2
    --local Vz = trait.varZ
    local nAllpLoci = species.nAllpLoci
    local ploidy = species.ploidy
    --print("ploidy"..ploidy[1])
    --local gammaParam = species.gammaParam

    -- First, get number of alleles for each locus of "lociList"
    local nAllLt = {}
    for i = 1, #lociList do nAllLt[i] = nAllpLoci[lociList[i]] end
    local ploidypLoci = {}
    for i = 1, #lociList do ploidypLoci[i] = ploidy[lociList[i]] end

    -- In case of no userdata allelic effect, we generate data
    if species.allEffSource == "generated" then --or allEff == nil or (type(allEff) == 'table' and #allEff == 0) then 
        -- Get weights
        local weights = confAE.get_weights(species, #lociList, varLociWeights)
	print("Generate new allelic effects...")

        --allEff = confAE.get_AE(id, phi, pow, weights,  nAllLt, allEnvEff, h2, phiA, ltypes, trait.Eref, #lociList, ploidypLoci, Vz)
        allEff = confAE.get_AE(trait, newLt, weights,  nAllLt, allEnvEff, ltypes, ploidypLoci)

    -- In case userdata is provided, we check format
    else
	if allEff == nil or (type(allEff) == 'table' and #allEff == 0) then
		print("All Eff load failed. Abort simulation.")
	else
                print("Import allelic effects from allEff.txt")
	end
        allEff = allEff_check(allEff, nAllLt)
    end
    return allEff
end


--------------------------------------------------------------------------------
----------------------------- Core functions -----------------------------------
--------------------------------------------------------------------------------
local function check(allEnvEff, usrLtype, trait, species, ltypes)
    local newLt = {}
    -- Fields check
    newLt.id = id_check(usrLtype.id)
    newLt.lociList = lociList_check(usrLtype.lociList, species.nLocus)
    if trait.isGE then
        newLt.phi = phi_check(usrLtype.phi)
    else
        newLt.phi = phi_check(trait.h2) -- isGE = false <=> only ltype a, phiA = h2
    end

    if newLt.id =="a" then
        newLt.pow = 0
    elseif newLt.id =="b" then
        newLt.pow = 1
    end


    if (newLt.lociList ~= {}) then
        assert(newLt.phi > 0, "Trait: ltype: phi must be > 0")
        if newLt.id=="a" then
            trait.phiA = newLt.phi
        end
        -- Generation / loading of allelic effects
        local allelEff = usrLtype.allEff
        newLt.allelicEff = allEff_make(allelEff, allEnvEff, newLt, trait, species, ltypes)
        return newLt
    else
        return {}
    end

end


-- Locus types check
local function check_all(ltypes, trait, species)
    assert_type(ltypes, 'table', "Trait: ltypes: wrong type")

    -- Gets a sorted table of all allelic effects for this trait
    -- Used for generations of allelic effects only:
    --[[print("envEff"..type(trait.envEff))
    for _, lt in pairs(trait.envEff) do
      print("1) ")--..lt)
      for _, ltt in pairs(lt) do
        print("2) ")--..ltt)
        for _, lttt in pairs(ltt) do
          print("3) "..lttt)
        end
      end
    end--]]
    local allEE = confAE.collect_EE(trait.envEff)
    --print("CHECK LT"..type(allEE)..#allEE)
    --for _, lt in pairs(allEE) do
    --  print(_..lt)
    --end

    --local sumPhi = 0
    local newLtypes = {}
    for _, lt in pairs(ltypes) do
        newLt = check(allEE, lt, trait, species, ltypes)
        if (newLt ~= {}) then
            newLtypes[#newLtypes+1] =newLt
            --sumPhi = sumPhi + newLtypes[#newLtypes].phi
        else
            print("newLt empty "..lt.id)
        end
    end
    --assert(sumPhi == 1, "Trait: ltypes: phi sum != 1")
    return newLtypes
end


--------------------------------------------------------------------------------
-------------------------------  Module interface ------------------------------
--------------------------------------------------------------------------------

local module = {}
module.check = check
module.check_all = check_all
return module
