-- Dependencies
local initGradient = require("gradient")
local initE = require("initE")
local initZopt = require("initZopt")
local confLT = require("ltypeConfig")

--------------------------------------------------------------------------------
--------------------------- Fields Checking functions --------------------------
--------------------------------------------------------------------------------
-- Each fuction in designed to check validity of each field
-- thus preventing as many errors as possible in configuration files

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

-- species : checl that the trait is for this species --------------------------
local function species_check(traitSpe, speId, traitId)
    assert(traitSpe == speId, "Wrong species association for trait"..traitId)
    return traitSpe
end

-- Var loci weight -------------------------------------------------------------
local function varLociWeights_check(varLociWeights)
    assert_type(varLociWeights, 'number', "Trait: varLociWeights: wrong type")
    return varLociWeights
end

-- Var Z -------------------------------------------------------------
local function varZ_check(varZ)
    assert_type(varZ, 'number', "Trait: varZ: wrong type")
    return varZ
end

-- Mean Z -------------------------------------------------------------
local function varZ_check(meanZ)
    assert_type(meanZ, 'number', "Trait: meanZ: wrong type")
    return meanZ
end


-- Environmental reference for the trait ---------------------------------------
local function check_Eref(Eref)
    assert_type(Eref, 'number', "Cfg: Eref Wrong type")
    assert(Eref > 0, "Cfg: Eref Wrong size")
    return Eref
end

-- non heritable pphenotypic plasticity ----------------------------------------
local function nhpp_check(nhpp)
    assert_type(nhpp, 'boolean', "Trait: nhpp: wrong type for nhpp")
    return nhpp
end

-- Indexing --------------------------------------------------------------------
local function indexing_check(indexing)
    assert_type(indexing, 'boolean', "Trait: wrong type for indexing")
    return indexing
end

--[[-- Var Dominance ---------------------------------------------------------------
local function varD_check(varD)
    assert_type(varD, 'number', "Trait: wrong type for varD")
    return varD
end

-- Var Epistasy ----------------------------------------------------------------
local function varE_check(varE)
    assert_type(varE, 'number', "Trait: wrong type for varE")
    return varE
end--]]

-- Heritability ----------------------------------------------------------------
local function h2_check(h2)
    assert_type(h2, 'number', "Trait: wrong type for h2")
    return h2
end



--------------------------------------------------------------------------------
------------------------------ Matching & Generation  --------------------------
--------------------------------------------------------------------------------


-- InitTotVaParam ---------------------------------------------------------------
-- used only if ltype a is present, otherwise an inblock formula is used
local function initTotVa_make(h2)
    return h2 / (1 - h2)
end



-- selection intensity ---------------------------------------------------------
-- Check user provided frames (number of frames > 0, size of each frame, type)
local function selInt_check(selInt, mapLength, mapWidth)
    assert_type(selInt, 'table', "trait: selInt: wrong type")
    local count = 0
    local hasStepOne = false
    for frame, data in pairs(selInt) do
        if (frame=="1") then hasStepOne = true end
        count = count + 1
        assert(is_type_matrix(data, 'number', mapLength, mapWidth),
                "Trait: selInt: Incorrect matrix size")
    end
    assert( hasStepOne, "Step 1 state for E for this trait is mandatory")
    assert( count > 0, "No selInt frames in for this propagule")
    return selInt
end

-- Function making selection intensity frames acdcording to options
local function selInt_make(selInt, selIntParam, mapLength, mapWidth)
    assert_type(selIntParam, 'table', "Trait: wrong type for selIntParam")
    assert_type(selIntParam.source, 'string', "Trait: selIntParam: wrong type for source")
    -- if userdata, check user frames
    if selIntParam.source == "userdata" then
        selInt_check(selInt, mapLength, mapWidth)
    -- if value, make one frame with a matrix filled with this value
    elseif selIntParam.source == "value" then
        assert_type(selIntParam.omega, 'number', "Trait: selIntParam: value: wrong type")
        selInt = {}
        selInt["1"] = make_matrix(selIntParam.omega, mapLength, mapWidth)
    -- Param could be used to generate a specific matrix (ex gradient) with parameters
    elseif selIntParam.source == "param" then
        -- No implementation provided here
        print("Not built-in generation function for selection intensity")
        os.exit()
    else
        print("Trait: selInt: source: Not a valid source ("..selInt.source..")")
        os.exit()
    end
    return selInt
end




-- Environmental effect  -------------------------------------------------------
local function envEff_check(envEff, mapLength, mapWidth)
    assert_type(envEff, 'table', "trait: envEff: wrong type")
    local count = 0
    local hasStepOne = false;
    for frame, data in pairs(envEff) do
        if (frame=="1") then hasStepOne = true end
        count = count + 1
        --print("DATA:"..#data.." "..mapLength.." "..mapWidth.." "..frame)
        assert(is_type_matrix(data, 'number', mapLength, mapWidth),
                "Trait: envEff: Incorrect matrix size")
    end
    assert( hasStepOne, "Step 1 state for E for this trait is mandatory")
    assert( count > 0, "No envEff frames in for this trait")
    return envEff
end


local function generate_envEff(envEffParam, lati, initTotV)
    -- Check envEff param type
    -- Check params
    local Ke = envEffParam.Ke
    assert_type(Ke, 'number', "Trait: envEffParam: Ke: wrong type")
    assert(Ke>0, 'Ke has to be positive, automatic slope generation for environmental effect is always following the larger map dimension')
    -- Generate env eff gradient
    return initE.make(lati, Ke)
end


local function envEff_make(envEff, envEffParam, mapWidth, mapLength, initTotV)
    assert_type(envEffParam, 'table', "Trait: wrong type for envEffParam")
    assert_type(envEffParam.source, 'string', "Trait: envEffParam: source: wrong type")
    if envEffParam.source == "userdata" then
        envEff_check(envEff, mapLength, mapWidth)

    elseif envEffParam.source == "value" then
        assert_type(envEffParam.E, 'number', "Trait: envEffParam: value: wrong type")
        envEff = {}
        envEff["1"] = make_matrix(envEffParam.E, mapLength, mapWidth)
    elseif envEffParam.source == "param" then
        -- Generate latitudes gradient
        local gradLati = 1 --envEffParam.gradLati
        assert_type(gradLati, 'number', "Trait: gradLati wrong type")
        local lati = initGradient.make_latitudes(mapLength, mapWidth, gradLati)
        -- Generate env Eff
        envEff = {}
        envEff["1"] = generate_envEff(envEffParam, lati, initTotV)
    elseif envEffParam.source == "disabled" then
        envEff = {}
        envEff["1"] = make_matrix(0, mapLength, mapWidth)
    else
        print("Trait: envEff: source: Not a valid source ("..envEff.source..")")
        os.exit()
    end
    return envEff
end



-- Phenotypic optimum ----------------------------------------------------------
-- Check user frames
local function zopt_check(zopt, mapLength, mapWidth)
    assert_type(zopt, 'table', "trait: zopt: wrong type")
    local count = 0
    local hasStepOne = false
    for frame, data in pairs(zopt) do
        if (frame=="1") then hasStepOne = true end
        count = count + 1
        assert(is_type_matrix(data, 'number', mapLength, mapWidth),
                "Trait: zopt: Incorrect matrix size")
    end
    assert( hasStepOne, "Step 1 state for E for this trait is mandatory")
    assert( count > 0, "No zopt frames in for this trait")
    return zopt
end


local function generate_zopt(zoptParam, lati, envEff, iniTotVa, envEffSource)
    -- Check zopt param type
    local Kzopt = zoptParam.Kzopt
    assert_type(Kzopt, 'number', "Trait: zopt param: Kzopt: wrong type")
    assert(Kzopt > 0, 'Kzopt must be positive. To have a negative gradient, set opposed=true')
    local opposed = zoptParam.opposed
    if (envEffSource == "param") then
        assert_type(opposed, 'boolean', "Trait: zopt param: opposed: wrong type")
    else
        assert((opposed ~= true), "Trait: zopt param: opposed can't be defined since E is not a gradient of source 'param'")
    end

    --local longActiv = zoptParam.longActiv
    --assert_type(longActiv, 'boolean', "Trait: zopt param: longActiv: wrong type")
    --local longStd = zoptParam.longStd
    --assert_type(longStd, 'number', "Trait: zopt param: longStd: wrong type")
    -- calculate zopts gradient
    return initZopt.make(lati, envEff, iniTotVa, Kzopt, opposed)--, longActiv, longStd)
end


-- Function making selection zopt according to options
local function zopt_make(zopt, zParam, envEff, mapWidth, mapLength, iniTotVa, envEffSource)
    assert_type(zParam, 'table', "Trait: wrong type for zoptParam")
    assert_type(zParam.source, 'string', "Trait: zoptParam: wrong type for source")
    if zParam.source == "userdata" then
        zopt_check(zopt, mapLength, mapWidth, mapLength, mapWidth)
    elseif zParam.source == "value" then
        assert_type(zParam.zopt, 'number', "Trait: zoptParam: value: wrong type")
        zopt = {}
        zopt["1"] = make_matrix(zParam.zopt, mapLength, mapWidth)
    elseif zParam.source == "param" then
        -- Generate latitudes
        local gradLati = 1 --zParam.gradLati
        assert_type(gradLati, 'number', "Trait: gradLati wrong type")
        local lati = initGradient.make_latitudes(mapLength, mapWidth, gradLati)
        -- Generate Zopt
        zopt = {}
        zopt["1"] = generate_zopt(zParam, lati, envEff["1"], iniTotVa, envEffSource)
    else
        print("Trait: zopt: source: Not a valid source ("..zopt.source..")")
        os.exit()
    end
    return zopt
end



--------------------------------------------------------------------------------
-------------------------------- Matching --------------------------------------
--------------------------------------------------------------------------------
local function match_ltypes(ltypes, allelic_effects)
    -- In case no userData is present, aboard
    if allelic_effects == nil then
      print("Match ltypes: no alleff")
      return
    end
    print("Match ltypes: alleff found")
    -- Assert ltypes existence & type
    assert_type(ltypes, 'table', "Trait: ltypes: Wrong type")
    -- Iterate on each allelic effect frame
    for id, allEff in pairs(allelic_effects) do
        local flag = false
        for _, lt in pairs(ltypes) do
            if id == lt.id then
                lt.allEff = allEff
                flag = true
            end
        end
        -- If no match is found execution is aborded
        local s1 = "Trait: user's allelic effect \""
        local s2 = "\" could not be matched with existing ltypes."
        local s3 = "Check trait id and loctype id"
        assert(flag == true, s1..id..s2..s3)
    end
    return -- no return value
end

--------------------------------------------------------------------------------
--------------------------- Core function --------------------------------------
--------------------------------------------------------------------------------


local function check(usrTrait, species, config)

    local newTrait = {}
    local mapLength = config.mapLength
    local mapWidth = config.mapWidth

    ---- Simple type checking

    -- HERITABILITY MODE: A (no G*E), or B (G*E)
    newTrait.isGE = false
    newTrait.refPhiA = nil
    newTrait.refPhiB = nil
    newTrait.Va = nil
    newTrait.Vb = nil
    newTrait.Veps = nil
    for _, lt in pairs(usrTrait.ltypes) do
        if lt.id == "a" then
            newTrait.refPhiA = lt.phi
        elseif lt.id == "b" and #lt.lociList>0 then
            newTrait.refPhiB = lt.phi
            newTrait.isGE = true
        end
    end
    --print("ISGE", isGE)

    -- id
    newTrait.id = id_check(usrTrait.id)
    
    print("Checking Trait: "..newTrait.id)
    -- species : check that the trait is for this species
    newTrait.species = species_check(usrTrait.species, species.id, newTrait.id)
    -- Var loci weight
    newTrait.varLociWeights = varLociWeights_check(usrTrait.varLociWeights)
    -- var Z
    newTrait.varZ = varZ_check(usrTrait.phenoVar)
    newTrait.meanZ = varZ_check(usrTrait.phenoMean)
    if (not newTrait.isGE) then
        -- non heritable phenotypic plasticity
        newTrait.nhpp = nhpp_check(usrTrait.nhpp)
        -- Heritability
        newTrait.h2 = h2_check(usrTrait.h2)
        newTrait.refPhiA = newTrait.h2

        newTrait.Va = newTrait.h2*newTrait.varZ
        newTrait.varEps = newTrait.varZ*(1-newTrait.h2)
    else
        -- environmental reference value
        for _, lt in pairs(usrTrait.ltypes) do
            assert_type(lt.phi, 'number', "Trait: wrong type for phi (phi defined for each ltype is necessary when plasticity")
        end

        newTrait.nhpp = false
        if (newTrait.nhpp ~= usrTrait.nhpp) then
            print("Trait: warning: trait nhpp manually set to true, but b ltype (G*E) defined, so nhpp reset to false.")
        end

        newTrait.Eref = check_Eref(usrTrait.Eref)

        -- DEFINE VARIANCES FOR A AND B LOCI AND E
        h2 = newTrait.refPhiA + newTrait.refPhiB
        newTrait.varEps = newTrait.varZ*(1-h2)  --check formula : CHECKED
        print("VEPS F1: "..newTrait.varEps)
        for _, lt in pairs(usrTrait.ltypes) do
          if lt.id == "a" then
              newTrait.Va = lt.phi/(1-lt.phi)*newTrait.varEps
              --newTrait.varEps = newTrait.Va*(1-h2)/h2--newTrait.varZ*(1-newTrait.h2)
              --print("VEPS F2: "..newTrait.varEps)
          elseif lt.id == "b" then
              newTrait.Vb = lt.phi*(newTrait.Va+newTrait.varEps)/((1-lt.phi)*newTrait.Eref)
          end
        end
    end


    -- Indexing
    newTrait.indexing = true --indexing_check(usrTrait.indexing)



    ---- Generated or userData values
    -- initTotVa
    newTrait.initTotVa = initTotVa_make(newTrait.refPhiA)
    -- Selection intensity (checks both selInt AND selIntParam)
    newTrait.selInt = selInt_make(usrTrait.selInt,  usrTrait.selIntParam,
                                  mapLength, mapWidth)
    -- Environmental effects (checks both envEff AND envEffParam)
    --newTrait.envEff = envEff_make(usrTrait.envEff, usrTrait.envEffParam,
    --                              mapWidth, mapLength, newTrait.initTotVa)
    --for key,value in pairs(usrTrait.e) do print(key,value) end
    local initTotV = newTrait.initTotVa
    --[[if (newTrait.isGE) then
        --TODO:define va and vb as trait arg instead of two separated computations
        local Va = newTrait.refPhiA/(1-newTrait.refPhiA)
        local Vb = newTrait.refPhiB*(Va+1)/((1-newTrait.refPhiB)*usrTrait.Eref)
        local h2 = (Va+Vb*usrTrait.Eref)/(Va+Vb*usrTrait.Eref+1)
        initTotV = h2/(1-h2) -- ???
        --print("ITV?", initTotV)
    end--]]

    newTrait.envEff = envEff_make(usrTrait.e, usrTrait.envEffParam,
                                  mapWidth, mapLength, initTotV)
    -- Phenotypic optimump (checks both zopt AND zoptParam)
    newTrait.zopt = zopt_make(usrTrait.zopt, usrTrait.zoptParam, newTrait.envEff,
                              mapWidth, mapLength, initTotV, usrTrait.envEffParam.source)

    ---- Loading ltypes
    -- match allelic_effects with ltypes
    match_ltypes(usrTrait.ltypes, usrTrait.allelicEffects)
    newTrait.ltypes = confLT.check_all(usrTrait.ltypes, newTrait, species)

    -- TODO : remove and load matrix frome cpp
    newTrait.zopt_l = matrix_to_vect(newTrait.zopt["1"])
    newTrait.selInt_l = matrix_to_vect(newTrait.selInt["1"])
    newTrait.e_l = matrix_to_vect(newTrait.envEff["1"])

    newTrait.selInt_l_k = {}
    newTrait.selInt_l_t = {}

    newTrait.e_l_k = {}
    newTrait.e_l_t = {}

    newTrait.zopt_l_k = {}
    newTrait.zopt_l_t = {}

    newTrait.minSelInt = 1e9

    for key,value in pairs(newTrait.selInt) do
        --print("V"..#value.."->"..#matrix_to_vect(value))
        minV = math.min(matrix_to_vect(value))[1]
        table.insert(newTrait.selInt_l_k, tonumber(key))
        if minV < newTrait.minSelInt then
            newTrait.minSelInt = minV
            --print("new V "..minV)
        end

        append_table(newTrait.selInt_l_t, matrix_to_vect(value))
    end

    for key,value in pairs(newTrait.envEff) do
        --print("V"..#value.."->"..#matrix_to_vect(value))
        table.insert(newTrait.e_l_k, tonumber(key))
        append_table(newTrait.e_l_t, matrix_to_vect(value))
    end

    for key,value in pairs(newTrait.zopt) do
        --print("V"..#value.."->"..#matrix_to_vect(value))
        table.insert(newTrait.zopt_l_k, tonumber(key))
        append_table(newTrait.zopt_l_t, matrix_to_vect(value))
    end

    newTrait.config = {}
    newTrait.config.a = 1
    newTrait.config.b = 2
    -- end TODO

    return newTrait
end

-- Check a list of traits
local function check_all(allUsrTraits, species, config)
    local traitsId = species.traitsId
    local newTraits = {}
    for pos, id in ipairs(traitsId) do
        local trait = allUsrTraits[id]
        assert_type(trait, 'table', "User trait "..id.." not found")
        newTraits[pos] = check(trait, species, config)
    end
    return newTraits
end
--------------------------------------------------------------------------------
-------------------------------  Module interface ------------------------------
--------------------------------------------------------------------------------

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