local FUN = require("__pycoalprocessing__/prototypes/functions/functions")

--------------------------------------------------------------------------------

local function ends_with(str, ending)
    return ending == "" or str:sub(-#ending) == ending;
end

local function withoutLast(str, count)
    return string.sub(str, 1, string.len (str) - count);
end


local function getBaseAssemblerTryDif (baseAssemblerName, withDif)

     local baseAssembler = data.raw["assembling-machine"][baseAssemblerName];

     if baseAssembler == nil then

         local nameWithMK;

         if withDif then
             nameWithMK = baseAssemblerName .. "-mk01"
         else
             nameWithMK = baseAssemblerName .. "mk01"
         end

         return data.raw["assembling-machine"][nameWithMK];
     else
         return baseAssembler;
     end

end

local function getBaseAssembler(baseAssemblerName)

    log ("--------------------------: " .. baseAssemblerName);

    local withDif = ends_with(baseAssemblerName, "-");

    local realBaseAssemblerName;

    if withDif then
        realBaseAssemblerName = withoutLast(baseAssemblerName, 1)
    else
        realBaseAssemblerName = baseAssemblerName
    end

    return getBaseAssemblerTryDif (realBaseAssemblerName, withDif);

end

for _, asm in pairs(data.raw["assembling-machine"]) do

    -- Check Only Py Assemblers?

    if ends_with(asm.name, "mk02") or ends_with(asm.name, "mk03") or ends_with(asm.name, "mk04") then

        local baseAssemblerName = withoutLast (asm.name, 4);

        local baseAssembler = getBaseAssembler(baseAssemblerName);

        if baseAssembler ~= nil then

            asm.module_specification.module_slots = math.max (asm.module_specification.module_slots, baseAssembler.module_specification.module_slots);

        else

            log ("base assembler for " .. asm.name .. " not found");

        end
    end
end

--------------------------------------------------------------------------------

data.raw["recipe"]["nichrome-2"].category = "casting";

--------------------------------------------------------------------------------

data.raw["assembling-machine"]["advanced-foundry-mk01"].module_specification.module_slots = 2;
data.raw["assembling-machine"]["advanced-foundry-mk02"].module_specification.module_slots = 2;
data.raw["assembling-machine"]["advanced-foundry-mk03"].module_specification.module_slots = 2;
data.raw["assembling-machine"]["advanced-foundry-mk04"].module_specification.module_slots = 2;

data.raw["assembling-machine"]["automated-factory-mk01"].module_specification.module_slots = 1;
data.raw["assembling-machine"]["automated-factory-mk02"].module_specification.module_slots = 2;
data.raw["assembling-machine"]["automated-factory-mk03"].module_specification.module_slots = 4;
data.raw["assembling-machine"]["automated-factory-mk04"].module_specification.module_slots = 4;

--------------------------------------------------------------------------------

local restoreProducts =
{
"tar-processing-unit",
"tar-processing-unit-mk02",
"tar-processing-unit-mk03",
"tar-processing-unit-mk04",
"washer",
"washer-mk02",
"washer-mk03",
"washer-mk04"
};


for index, asmName in pairs(restoreProducts) do

    local assembler = data.raw["assembling-machine"][asmName];

    assembler.allowed_effects = {"consumption", "speed", "productivity", "pollution"};

end

--------------------------------------------------------------------------------

FUN.productivity({

"proton-donor",
"proton-receiver",
"crystallographic-substrate",

"pa-sulfur",
"pa-aromatics",
"pa-benzene",
"pa-propene",
"pa-silicon",
"pa-coal",
"pa-niobium",
"pa-chromium",
"pa-molybdenum",
"pa-vanadium",
"pa-uranium",
"pa-uranium-235",
"pa-diamond",
"pa-proton",
"pa-oxygen",
"pa-helium",
"pa-hydrogen",
"pa-boron",
"pa-coal2",
"pa-aluminium",
"pa-nickel",
"pa-zinc",
"pa-silver",
"pa-gold",
"pa-vanadium2",



"clean-nexelit",
"nexelit-ore-1",
"aramid-2",
"biofilm-pyht",
"ceramic",
"ceramic-2",
"sodium-sulfate",
"sodium-bisulfate",
"sand-classification2",
"kerosene",
"filtration-media",
"carbon-filter",
"carbon-filter-mk02",
"carbon-filter-mk03",
"carbon-filter-mk04",
"bolts",
"small-parts-01",
"small-parts-02",
"small-parts-03",
"rubber-01",
"rubber-02",
"rubber-03",
"rubber-04",
"styrene",
"btx-to-styrene",
"carbon-black",
"carbon-black-2",
"belt",
"guar-01",
"guar-02",
"guar-03",
"guar-04",
"guar-05",
"guar-separation",
"guar-gum",
"drilling-fluid-1",
"drilling-fluid-2",
"drilling-fluid-3"})

--------------------------------------------------------------------------------