-- Gather all environemental effects in a single table & sort
-- Need to be done only once foe each trait
--
verbose = false

local function collect_EE(envEffFrames)
    local allEE = {}
    -- For each frame
    for _, frame in pairs(envEffFrames) do
        -- For each row & col
        for row = 1, #frame do
            for col = 1, #frame[row] do
                -- We can't use env effect that equal zero (divide by 0)
                if  frame[row][col] ~= 0.0 then
                    allEE[#allEE + 1] = frame[row][col]
                else
		  if verbose then 
                      print("WARNING: ENV VAL IS NULL")
                  end
                  allEE[#allEE + 1] = 0  -- historically this line was not there, forbidding E=0
                end
            end
        end
    end
    table.sort(allEE)
    return allEE
end



-- Here we draw a representative subsample of environmental effects
-- Input "valuesE" is a table containing all env effects sorted
-- Input nal is the number of env. eff. we want to select (= number of alleles at locus)
local function selectE(nal, valuesE)
    -- Return value (selected environmental effects)
    local selE = {}
    -- Tot number of env.eff.
    local nE = #valuesE

    -- if nal == nE we take each E one time
    if nE == nal then
        selE = valuesE
    -- if nal > nE, we take nal/nE times each E (with round precision)
    elseif nal >  nE  then
	    local a = nE/nal
	    for i = 1, nal, 1 do
	        local tIndex = math.ceil(a*i)
	        table.insert(selE, valuesE[tIndex]) -- subsample the table
	    end
	-- if nal  <  nE  we divide  valuesE in nal blocks of identical size
	-- and chose each block mean as E
    else
        -- Init boundaries
	    local borneInf = 0
	    local borneSup = 0
	    -- create nal blocs boundaries dividing nE groups of identical size
	    for i = 1, nal, 1 do
            -- move boundaries
            borneInf = borneSup + 1
            borneSup = math.floor((nE/nal)*i)
            -- Sum values between boundaries
            local sum = 0
            for j = borneInf, borneSup, 1 do
                sum = sum + valuesE[j]
            end
            -- Mean from sum
            table.insert(selE, sum / (borneSup - borneInf + 1))
        end
    end
    table.sort(selE)
    return selE
end
-- End pick of E




------------------------------------------------------
-- Step 1: loci weights pick. Use  var_loci_weights --
------------------------------------------------------
-- may use a gamma law (if gamma param > 0) or identical weights otherwize
local function get_weights(species, nloc, varLociWeights)-- end, gammaParam, config)
    local values = {}
    local sum = 0
    verbose = false --display weights or not
    --Get random weights
    for i = 1, nloc, 1 do
        if varLociWeights <= 0 then
            values[i] = 1
        else
            values[i] = math.abs(1/nloc + gasdev() * math.sqrt(varLociWeights)) --close to gamma distribution
            sum = sum + values[i]
        end
    end


    -- Normalise them [ONLY IF GAMMA]
    if varLociWeights > 0 then
        for i = 1, nloc, 1 do
            values[i] = values[i]/sum*nloc
        end
    end
    if (verbose) then --display weights
        print("GAMMA[1,2]")
        for pos = 1, #values, 1 do
            print("  "..values[pos])
        end
    end
    return values
end


-----------------------------------
-- Step 2 : pick allelic effects --
-----------------------------------
-- Inputs:
-- nAllLt => number of alleles for each locus of the loctype
-- weights => weight of each locus, calculated using "get_weights"
-- envEff => all env. eff. of the trait sorted in a table
--local function get_AE(id, phi, pow, weights, nAllLt, envEff, h2, phiA, ltypes, Eref, nloc, ploidypLoci, Vz)
local function get_AE(trait, ltype, weights, nAllLt, envEff, ltypes, ploidypLoci)
    id = ltype.id
    phi = ltype.phi
    pow = ltype.pow
    nloc = #ltype.lociList
    phiA = trait.refPhiA
    Eref = trait.Eref
    Vz = trait.varZ

    assert(id == "a" or id == "b" or id == "c", "get_all_effects: unsuported type")

    --print("EREF? "..id.." "..Eref)

    --Tableau des effets alléliques (return value)
    local allEff = {}
    verbose = false   --display alleff or not
    allEffText = ""


    isBorC = false --check which alleff formula will be used. If b or c present, use the plasticity model
    for _, lt in pairs(ltypes) do
        if (lt.id == "b" or lt.id == "c") and #(lt.lociList) > 0 then
            isBorC = true
        end
    end

    -- Core of the function
    if(verbose)then print('allEff') end
    for loc = 1, #nAllLt, 1 do
        -- Number of alleles for this locus
        local nal = nAllLt[loc]
        local ploidy = ploidypLoci[loc]
        --Init sub-table
        allEff[loc] = {}
        --Sample E effects
        local selectedEE = {}
        --print(nal)
        --[[print("LENVF"..#envEff)
        for _, lt in pairs(envEff) do
            print(_..lt)
        end
        --]]
        if id ~= "a" then selectedEE = selectE(nal, envEff) end

        -- For each allele of the locus, get allelic effects
        for al = 1, nal, 1 do
            local envEffect = 1 -- for type "a", value is 1
            -- For type b and c, we get E effet to the power (2*pow)
            --print(selectedEE[al])
            --print(pow)
            if id ~= "a" then envEffect = selectedEE[al]^(pow*2) end

            -- Core formula
            if isBorC then -- SEPARATED CASES: ONLY A VS A&B
                --local Veps = Vz*(1-h2)  --TODO:check formula
                if id == "a" then
                    --local Va = phi/(1-phi)*Vz*(1-h2)
                    allEff[loc][al] =   math.sqrt(( weights[loc] * trait.Va) / (nloc*ploidy))* gasdev();

                    local allEffText = allEffText.."    a:"..phi.."->"..allEff[loc][al]
                elseif id == "b" then
                    --Va = phiA/(1-phiA)
                    --local Va = phiA/(1-phiA)*Vz*(1-h2)
                    --Vb = phi*(Va+Veps)/(1-phi*Eref^2)
                    --local Vb = phi*(Va+Veps)/((1-phi)*Eref)
                    --print("Vb", Vb)
                    allEff[loc][al] =   math.sqrt(( weights[loc] * trait.Vb) / (nloc*ploidy))* gasdev();
                    --allEffText = allEffText.."    b:"..phi.."->"..allEff[loc][al]
                end
            else
                --local Vg = h2/(1-h2)
                --local Vg = h2*Vz
                allEff[loc][al] =   math.sqrt(( weights[loc] * trait.Va ) / (nloc*ploidy))* gasdev();
                --allEff[loc][al] =   math.sqrt(Va) * gasdev();
                allEffText = allEffText.."    a:"..trait.Va.."->"..allEff[loc][al]
            end
        end

        -- Shuffeling (Optional for type a)
        local suffleInd = get_k(nal, 1, nal)
        for i = 1, nal, 1 do
            local j = suffleInd[i]
            allEff[loc][i], allEff[loc][j] = allEff[loc][j], allEff[loc][i]
        end
        -- end shuffeling
    end
    if (verbose) then
        print(allEffText)
    end
    return allEff
end


--------------------------------------------------------------------------------
-------------------------------  Module interface ------------------------------
--------------------------------------------------------------------------------
local module = {}
module.get_weights = get_weights
module.get_AE = get_AE
module.collect_EE = collect_EE
return module
