local flux = require("fluxMatrix")
--local popInit = require("popInit")
local traitConfig = require("traitsConfig") 


-- Helper functions ------------------------------------------------------------
-- Used to complete fields:
-- If field contains a number value or a single value in a table,
-- then the value is repeated n times, othewise, it does nothing
local function fill_values(field, size)
    if type(field) == 'number' then
        return rep(field, size) 
    elseif type(field) == 'table' and #field == 1 then 
        return rep(field[1], size) 
    else return field end
end

--------------------------------------------------------------------------------
--------------------------- Fields Checking functions --------------------------
-- id --------------------------------------------------------------------------
local function id_check(id, newSpe)
    assert_type(id, 'string', "Species id: Wrong type")
    return id
end

-- nLocus ----------------------------------------------------------------------
local function nLocus_check(nLocus)
    assert_type(nLocus, 'integer', "Species nLocus: Wrong type")
    assert(nLocus > 0, "Species nLocus: Minimum value is 1")
    return nLocus
end

-- genomeSize ------------------------------------------------------------------
local function genomeSize_check(genomeSize, nLocus)
    assert_type(genomeSize, 'integer', "Species genomeSize: Wrong type")
    assert(genomeSize >= nLocus, "Species: genomeSize cannot be < nLocus")
    return genomeSize
end

-- ploidy ----------------------------------------------------------------------
local function ploidy_check(ploidy, nLocus, genomeSize)
    ploidy = fill_values(ploidy, nLocus)
    -- Checking type & size of vector
    assert(is_type_vector(ploidy, 'integer', nLocus), 
        "Species: ploidy is not valid a vector of size " .. nLocus)
    -- Checking range of values
    assert(is_range_vector(ploidy, 1, 3), "Species: ploidy must be > 0 and < 3")
    -- Checking consistency with genome size
    local sum = stats.sum_vector(ploidy)
    assert(sum == genomeSize, "Species: Sum of ploidy != genomeSize")
    return ploidy
end

local function chromoSplit_check(chromoSplit)
    assert(is_type_vector(chromoSplit, 'number', #chromoSplit),
        "Species: chromoSplit wrong type")
    -- need check split values are < genomeSize
    -- need check that first value is 0
    return chromoSplit
end

--[[local function gammaParam_check(gammaParam)
    if gammaParam == nil then
        --print("GAMMANIL")
        return -1
    end

    assert_type(gammaParam, 'number', "Species gammaParam: Wrong type")
    --assert(gammaParam > 0, "Species gamma: Minimum value is 0")
    return gammaParam
end--]]

-- nAllpLoci -------------------------------------------------------------------
local function nAllpLoci_check(nAllpLoci, nLocus)
    nAllpLoci = fill_values(nAllpLoci, nLocus)
    assert(is_type_vector(nAllpLoci, 'integer', nLocus), 
        "Species: nAllpLoci is not valid a vector of size " .. nLocus)
    assert(is_range_vector(nAllpLoci, 1, math.huge), "Species: nAllpLoci must be > 0")
    return nAllpLoci
end

-- nAllInit --------------------------------------------------------------------
local function nAllInit_check(nAllInit, nAllpLoci)
    nAllInit = fill_values(nAllInit, #nAllpLoci)
    assert(is_type_vector(nAllInit, 'integer', #nAllpLoci), 
        "Species: nAllInit is not valid a vector of size " .. #nAllpLoci)
    -- Allowed range depends on nAllpLoci
    for i = 1, #nAllInit, 1 do
        assert(nAllInit[i] <= nAllpLoci[i], "Species: nAllInit > nAllpLoci at loc "..i)
        assert(nAllInit[i] >= 1, "Species: nAllInit < 1 at loc "..i)
    end
    return nAllInit
end

-- allEffSource ----------------------------------------------------------------
local function allEffSource_check(allEffSource)
    assert(allEffSource=="userdata" or allEffSource=="generated", 
              "Species: allEffSource is not valid: accepted values are 'userdata' and 'generated'")
    return allEffSource
end


-- crossingRate ----------------------------------------------------------------
local function crossingRate_check(crossingRate, nLocus)
    crossingRate = fill_values(crossingRate, nLocus)
    assert(is_type_vector(crossingRate, 'number', nLocus), 
        "Species: crossingRate wrong type or size in vector")
    assert(is_range_vector(crossingRate, 0.0, 1.0), 
        "Species: crossingRate must be in range [0.0, 1.0[")
    return crossingRate
end

-- mutationRate ----------------------------------------------------------------
local function mutationRate_check(mutationRate, nLocus)
    mutationRate = fill_values(mutationRate, nLocus)
    assert(is_type_vector(mutationRate, 'number', nLocus), 
        "Species: mutationRate wrong type or size in vector")
    assert(is_range_vector(mutationRate, 0.0, 1.0), 
        "Species: mutationRate must be in range [0.0, 1.0[")
    return mutationRate
end

-- inheritance -------------------------------------------------------------------------
local function inheritance_check(inheritance, nLocus, ploidy)
    assert_type(inheritance, 'table', "Species: inheritance is not a table")
    local s = {}
    --local lst = {"pollen", "female"}
    local lst = {"female", "male"}
    for _, i in pairs(lst) do
        s[i] = fill_values(inheritance[i], nLocus)
        assert(is_type_vector(s[i], 'integer', nLocus), 
            "Species: inheritance: "..i..": no a valid integers vector")
        assert(is_range_vector(s[i], 0, math.huge), 
            "Species: inheritance: "..i..": must be >= 0")
    end
    -- Check sum of herited loci
    for i = 1, nLocus do
        local sum = s["male"][i] + s["female"][i]
        assert(sum == ploidy[i], 
            "Specie: inheritance: sum of male + female != ploidy on locus " .. i)
    end
    return inheritance
end

-- ageClasses -- TODO: mostly a placeholder for now ----------------------------
local function ageClasses_check(ageClasses)
    assert(is_type_vector(ageClasses, 'number', #ageClasses),  
        "Species: ageClasses wrong type")
    return ageClasses
end

-- selfingRate -----------------------------------------------------------------
local function selfingRate_check(selfingRate)
    assert_type(selfingRate, 'number', "Species: selfingRate not a number")
    assert(selfingRate >= 0.0 and selfingRate <= 1.0, 
        "Species: selfingRate not in range [0.0, 1.0] ("..selfingRate..")")
    return selfingRate
end

-- maxAge ----------------------------------------------------------------------
local function maxAge_check(maxAge)
    assert_type(maxAge, 'number', "Species: maxAge not a number")
    assert(maxAge > 0.0,  "Species: maxAge not > 0 ("..maxAge..")")
    return maxAge
end

-- fluxSource ------------------------------------------------------------------
local function fluxSource_check(id, fluxSource)
    assert_type(fluxSource, 'table', "fluxSource undefined in "..id)
    assert_type(fluxSource.pollen, 'string', "fluxSource: pollen undefined in "..id)
    assert_type(fluxSource.seed, 'string', "fluxSource: seed undefined in "..id)
    return fluxSource
end

-- traitsId --------------------------------------------------------------------
local function traitsId_check(traitsId)
    assert(traitsId ~= nil, "Species, traitsId not provided (nil)")
     -- If a string convert to a table of one string
    if type(traitsId) == 'string' then traitsId  = { traitsId, } end
    assert_type(traitsId, 'table', "Species, traitsId not a table")
    assert(#traitsId > 0, "At least on trait must be provided")
    assert(is_type_vector(traitsId, 'string'), "Species, a traitsId field not a string")
    return traitsId
end

local function selIntInteractions_check(selIntInteractions, nTraits, minOmega)
    --print("MOM"..minOmega)
    -- TODO SELINT
    --assert_type(selIntInteractions, 'table', "Species, traitsId not a table")
    --assert(#selIntInteractions == math.pow(nTraits, 2), "The selection strength matrix should be defined as a list of nTraits^2 values: for example for two traits: {s1, s1/2, s2/1, s2}")
    for k, v in pairs(selIntInteractions) do
        assert_type(v, 'number', "Species: selInt value not a number")
        assert(v < minOmega, "Interaction values must be less than the lowest selection intensity (here "..minOmega..")")
    end

    return selIntInteractions
end

-- assortMating --------------------------------------------------------------------
local function assortMating_check(assortMating, traitsId)
    if assortMating == nil then return "" end
    assert_type(assortMating, 'string', "assortMating not a string")
    assert(is_in_table(assortMating, traitsId), "assortMating: "..assortMating.." not in traitsId")
    return assortMating
end

-- initDemo
local function initDemo_check(initDemo)
    assert_type(initDemo, 'table', "species: initDemo wrong type")
    return initDemo
end

-- Nm __ OPTIONAL ! 
local function Nm_check(Nm)
    if Nm == nil then return nil end
    assert_type(Nm, 'number', "species: Nm wrong type")
    return Nm
end

-- cap --------------------------------------------------------------------
local function cap_check(cap)
    assert_type(cap, 'integer', "species: cap wrong type")
    assert(cap > 0, "species: cap wrong size")
    return cap
end

-- growth --------------------------------------------------------------------
local function growth_check(growth)
    assert_type(growth, 'number', "species: growth wrong type")
    --assert(growth >= 0, "species: growth wrong size")
    return growth
end


-- rho --------------------------------------------------------------------
local function rho_check(rho, assortMating)
    if assortMating == nil or assortMating == "" then
        --print("assortNIL")
        return nil
    end
    assert_type(rho, 'number', "species: assortThres wrong type")
    assert(rho > 0, "species: assortThres wrong size")
    return rho
end


--------------------------------------------------------------------------------
----------------------- Flux matrix functions ----------------------------------
--------------------------------------------------------------------------------

-- Create the flux matrix for each frame
local function make_multi_FM(fluxPatFrames, nlongi, nlati)
    local fluxMatFrames = {}
    -- For each frame check the actual data size and type

    for time, data in pairs(fluxPatFrames) do
        fluxMatFrames[time] = flux.makeFluxMat(data, nlongi, nlati)
    end
    return fluxMatFrames
end

-- Check flux matrix (type & size) & retunrs a copy
local function flux_matrix_check(fluxMatFrames, dim, id)
    assert_type(fluxMatFrames, 'table', "User flux_matrix type issue in "..id)
    -- For each frame checks the data size and type
    local count = 0
    for _, data in pairs(fluxMatFrames) do
        count = count + 1
        assert(is_type_matrix(data, 'number', dim, dim), 
               "flux matrix: Incorrect format")
    end
    -- Assert than at least on frame exists
    assert(count > 0, "No flux matrix frames in for this propagule")
    return deep_copy(fluxMatFrames)
end


-- Check dimentions and type
local function flux_pattern_check(fluxPatFrames, id)
    -- Type/existence check
    assert_type(fluxPatFrames, 'table', "flux pattern type issue in "..id)
    -- For each frame check the actual data size and type
    local count = 0
    for frame, data in pairs(fluxPatFrames) do
        count = count + 1
        assert(is_type_matrix(data, 'number', nil, nil), 
               "flux pattern: Incorrect format")
        assert(#data%2 == 1 and #data[1]%2 == 1, 
               "Flux pattern must be of odd dimensions")
    end
    -- Assert than at least on frame exists
    assert( count > 0, "No flux pattern frames in for this propagule")
    return fluxPatFrames
end        
     

-- fluxSource is a field defined in species that indicate either:
-- a predifined function (stepping stone, island..)
-- or a type of custom data to load )

-- TODO IMO this is bad design, fluxsouce should only indicate a source and be mandatory
-- The if source == function, a function name should be indicated in  an optional field
local function load_flux(usrSpe, fluxSource, mapWidth, mapLength)
    -- Fields from usrSpe that might be used
    local id = usrSpe.id
    local usrFM = usrSpe.fluxMatrix
    local usrFP = usrSpe.fluxPattern
    local ms = usrSpe.ms -- might be nil (checked on use)
    local mp = usrSpe.mp -- might be nil (checked on use)
    
    local fluxMatrix = {}
    local dim = mapWidth * mapLength

    -- Iteration on both sexes
    -- for _, propagule in pairs({"seed", "female"}) do
    for _, propagule in pairs({"seed", "pollen"}) do
        -- 1 - Loading a flux matrix, we just check existence & format 
        if fluxSource[propagule] == "matrix" then
            assert_type(usrFM, 'table', "flux matrix undefined for "..id)
            fluxMatrix[propagule] = flux_matrix_check(usrFM[propagule], dim, id)
        -- 2 - Loading a flux pattern, check and then convert to flux matrix
        elseif fluxSource[propagule] == "pattern" then
            assert_type(usrFP, 'table', "flux pattern undefined for "..id)
            assert_type(usrFP[propagule], 'table', "flux pattern undefined for "..id.." "..propagule)
            flux_pattern_check(usrFP[propagule], id)
            fluxMatrix[propagule] = make_multi_FM(usrFP[propagule], mapWidth, mapLength)
        -- 3 - Else we assume it is a built in function (Only one frame in this case)
        elseif type(fluxSource[propagule]) == 'string' then
            local pattern = {}
            -- TODO, deal more elegantly with frames number
            pattern["1"] = flux.make_flux_pattern(
                              fluxSource[propagule], ms, mp, mapWidth, mapLength, propagule)
            fluxMatrix[propagule] = make_multi_FM(pattern, mapWidth, mapLength)
        end
    end
    return fluxMatrix
end

--------------------------------------------------------------------------------
-------------------------------  Main functions --------------------------------
--------------------------------------------------------------------------------
-- This Function checks  the basic validity of a species
local function check(usrSpe, config)
    print("Checking Species: "..usrSpe.id)
    --New specie object (empty table)
    local newSpe = {}
    newSpe.generatedDemo = nil
    
    ---- Simple checks ---------------------------------------------------------
    -- Some optional fields are only checked on use (mp, ms, ...)
    -- id
    newSpe.id = id_check(usrSpe.id, newSpe)
    -- nLocus (deduced in C++ for now)
    newSpe.nLocus = nLocus_check(usrSpe.nLocus)
    -- genomeSize (deduced in C++ for now)
    newSpe.genomeSize = genomeSize_check(usrSpe.genomeSize, newSpe.nLocus) 
    -- ploidy
    newSpe.ploidy = ploidy_check(usrSpe.ploidy, newSpe.nLocus, newSpe.genomeSize)
    -- gamma

    --newSpe.gammaParam = gammaParam_check(usrSpe.gammaParam)
    --chromoSplit
    newSpe.chromoSplit = chromoSplit_check(usrSpe.chromoSplit)
    -- nAllpLoci            
    newSpe.nAllpLoci = nAllpLoci_check(usrSpe.nAllpLoci, newSpe.nLocus)
    -- nAllInit (unsused in C++ for now)
    newSpe.nAllInit = nAllInit_check(usrSpe.nAllInit, newSpe.nAllpLoci)
    --allelic effect source
    newSpe.allEffSource = allEffSource_check(usrSpe.allEffSource)
    -- crossingRate 
    newSpe.crossingRate = crossingRate_check(usrSpe.crossingRate, newSpe.nLocus)
    -- mutationRate
    newSpe.mutationRate = mutationRate_check(usrSpe.mutationRate, newSpe.nLocus)
    -- inheritance
    newSpe.inheritance = inheritance_check(usrSpe.inheritance, newSpe.nLocus, newSpe.ploidy)
    -- ageClasses (not implemented in C++ for now)
    newSpe.ageClasses = {0, 10, 11, 20} --ageClasses_check(usrSpe.ageClasses)
    -- selfing rate
    newSpe.selfingRate = selfingRate_check(usrSpe.selfingRate)
    -- Max age ( Temporary way for non overlapping generations)
    --newSpe.maxAge = maxAge_check(usrSpe.maxAge)
    --fluxSource 
    newSpe.fluxSource = fluxSource_check(usrSpe.id, usrSpe.fluxSource)
    -- traitId (only typecheck)
    newSpe.traitsId = traitsId_check(usrSpe.traitsId)

    -- assortMating
    newSpe.assortMating = assortMating_check(usrSpe.assortMating, newSpe.traitsId)
    -- initDemo
    newSpe.initDemo = initDemo_check(usrSpe.initDemo)
    -- Nm -- OPTINAL -- TODO deal with optional
    newSpe.Nm = Nm_check(usrSpe.Nm)
    -- cap
    newSpe.cap = cap_check(usrSpe.cap)
    -- growth
    newSpe.growth = growth_check(usrSpe.growth)
    -- rho
    newSpe.rho = rho_check(usrSpe.rho, newSpe.assortMating)


    ---- Checking & filling ----------------------------------------------------
    -- Check flux matrix source & generate or load data
    -- We put usrSpe as param to pass the following (optional & unchecked) fields: 
    --     fluxMatrix, fluxPattern, ms, mp
    newSpe.fluxMatrix = load_flux(usrSpe, newSpe.fluxSource, 
                                  config.mapWidth, config.mapLength)
                                  
    -- Match & Check traits 
    newSpe.traits = traitConfig.check_all(usrSpe.traits, newSpe, config)
    
    -- check matrix
    minOmega = 1e9
    for pos, id in ipairs(newSpe.traits) do
        if minOmega > newSpe.traits[pos].minSelInt then
            minOmega = newSpe.traits[pos].minSelInt
        end
    end

    --selectMatrix
    if (newSpe.selIntInteractions == nil) then
        newSpe.selIntInteractions = {}
    end
    newSpe.selIntInteractions = selIntInteractions_check(usrSpe.selIntInteractions, #newSpe.traitsId, minOmega)

    -- TODO remove : for compatibility with old CPP
    newSpe.parameters = {}

    return newSpe
end


-- Check a list of species
local function check_all(allUsrSpe, config)
    newSpecies = {}
    for _, spe in ipairs(allUsrSpe) do
        assert(newSpecies[spe.id] == nil, "Same species id used twice") 
        newSpecies[#newSpecies + 1] = check(spe, config)
    end
    return newSpecies
end


--------------------------------------------------------------------------------
-------------------------------  Module interface ------------------------------
--------------------------------------------------------------------------------
local module = {}
module.check = check
module.check_all = check_all
return module

