--
-- HarvestExtension
-- Class to manage the harvesting (get custom yield, track the yield and apply to yieldMap)
--
-- @author Stefan Maurus
-- @date 21/09/2020
--
-- Copyright (C) GIANTS Software GmbH, Confidential, All Rights Reserved.


HarvestExtension = {}

HarvestExtension.YIELD_DEBUG = false

HarvestExtension.MOD_NAME = g_currentModName

local HarvestExtension_mt = Class(HarvestExtension)

function HarvestExtension.new(pfModule, customMt)
    local self = setmetatable({}, customMt or HarvestExtension_mt)

    self.pfModule = pfModule

    self.debugValues = {
        nActualValue = 0,
        nTargetValue = 0,
        nFactor = 0,
        regularNFactor = 0,
        nRegularValue = 0,
        yieldPotential = 0,
        ignoreOverfertilization = false,
        pHActualValue = 0,
        pHTargetValue = 0,
        pHRegularValue = 0,
        pHFactor = 0,
        regularPHFactor = 0,
        plowFactor = 0,
        weedFactor = 0,
        yieldFactor = 0,
        yieldFactorRegular = 0
    }

    g_messageCenter:subscribe(MessageType.UNLOADING_STATIONS_CHANGED, self.onUnloadingStationsChanged, self)

    addConsoleCommand("paShowYieldDebug", "Displayes precision farming yield debug", "toggleYieldDebug", self)

    return self
end

function HarvestExtension:delete()
    g_messageCenter:unsubscribeAll(self)
    removeConsoleCommand("paShowYieldDebug")
end

function HarvestExtension:loadFromXML(xmlFile, key, baseDirectory, configFileName, mapFilename)
    return true
end

function HarvestExtension:toggleYieldDebug()
    HarvestExtension.YIELD_DEBUG = not HarvestExtension.YIELD_DEBUG
end

function HarvestExtension:update(dt)
    if HarvestExtension.YIELD_DEBUG then
        local debugValues = self.debugValues

        local nitrogenMap = self.pfModule.nitrogenMap
        local pHMap = self.pfModule.pHMap

        renderText(0.01, 0.42, 0.015, string.format("Nitrogen: Actual: %d | Target: %d | Regular: %d => Factor: %d%% | Regular Factor: %d%%",
        nitrogenMap:getNitrogenValueFromInternalValue(debugValues.nActualValue),
        nitrogenMap:getNitrogenValueFromInternalValue(debugValues.nTargetValue),
        nitrogenMap:getNitrogenValueFromInternalValue(debugValues.nRegularValue), debugValues.nFactor * 100, debugValues.regularNFactor * 100))

        renderText(0.01, 0.40, 0.015, string.format("Nitrogen: Yield Potential: %d%% | ignoreOverfertilization: %s",
        debugValues.yieldPotential * 100, debugValues.ignoreOverfertilization))

        renderText(0.01, 0.38, 0.015, string.format("pH: Actual: %.3f | Target: %.3f | Regular: %.3f => Factor: %d%% | Regular Factor: %d%%",
        pHMap:getPhValueFromInternalValue(debugValues.pHActualValue),
        pHMap:getPhValueFromInternalValue(debugValues.pHTargetValue),
        pHMap:getPhValueFromInternalValue(debugValues.pHRegularValue), debugValues.pHFactor * 100, debugValues.regularPHFactor * 100))

        renderText(0.01, 0.36, 0.015, string.format("Plow Factor: %d%% | Weed Factor: %d%%",
        debugValues.plowFactor * 100, debugValues.weedFactor * 100))

        renderText(0.01, 0.34, 0.015, string.format("Yield: %d%% | Regular Yield: %d%%",
        (debugValues.yieldFactor * debugValues.yieldPotential) / 2 * 100, (debugValues.yieldFactorRegular * debugValues.yieldPotential) / 2 * 100))
    end
end

function HarvestExtension:onUnloadingStationsChanged()
    self.unloadingStations = g_currentMission.storageSystem:getUnloadingStations()
end

function HarvestExtension:getYieldFromArea(combine, cutter, inputFruitType, realArea)
    local pixelToSqm = g_currentMission:getFruitPixelsToSqm() -- 4096px are mapped to 2048m
    local literPerSqm = 1
    if inputFruitType ~= FruitType.UNKNOWN then
        local fruitDesc = g_fruitTypeManager:getFruitTypeByIndex(inputFruitType)
        literPerSqm = fruitDesc.literPerSqm
    end
    local sqm = realArea * pixelToSqm
    local yield = sqm * literPerSqm * combine.spec_combine.threshingScale

    local farmId = cutter:getLastTouchedFarmlandFarmId()
    if farmId ~= AccessHandler.EVERYONE then
        local damage = combine:getVehicleDamage()
        if damage > 0 then
            yield = yield * (1 - damage * Combine.DAMAGED_YIELD_REDUCTION)
        end
    end

    return yield
end

function HarvestExtension:getBestPriceForFillType(fillType)
    if self.unloadingStations == nil then
        self.unloadingStations = g_currentMission.storageSystem:getUnloadingStations()
    end

    local maxPrice = 0
    for _, unloadingStation in pairs(self.unloadingStations) do
        if unloadingStation.getEffectiveFillTypePrice ~= nil then
            if unloadingStation:getIsFillTypeAllowed(fillType.index) then
                if unloadingStation.getAppearsOnStats == nil or unloadingStation:getAppearsOnStats() then
                    maxPrice = math.max(maxPrice, unloadingStation:getEffectiveFillTypePrice(fillType.index, ToolType.UNDEFINED))
                end
            end
        end
    end

    if maxPrice == 0 then
        maxPrice = fillType.pricePerLiter or 0
    end

    return maxPrice
end

function HarvestExtension:updateLatestFactors(pfModule, vehicle, requiresPHFactorUpdate, requiresNFactorUpdate)
    local nFactor = vehicle.lastNFactor or 0
    local regularNFactor = vehicle.lastRegularNFactor or 0
    local currentYieldPotential = vehicle.lastYieldPotential or 1
    if requiresPHFactorUpdate then
        local nActualValue, nTargetValue, yieldPotential, nRegularValue, lastIgnoreOverfertilization = pfModule.nitrogenMap:updateLastNitrogenValues()
        if nActualValue > -1 and nTargetValue > -1 then
            nFactor = pfModule.nitrogenMap:getYieldFactorByLevelDifference(nActualValue - nTargetValue, lastIgnoreOverfertilization)
            vehicle.lastNFactor = nFactor
            vehicle.lastYieldPotential = yieldPotential
            currentYieldPotential = yieldPotential

            regularNFactor = pfModule.nitrogenMap:getYieldFactorByLevelDifference(nRegularValue - nTargetValue, lastIgnoreOverfertilization)
            vehicle.lastRegularNFactor = regularNFactor

            if HarvestExtension.YIELD_DEBUG then
                local debugValues = self.debugValues
                debugValues.nActualValue = nActualValue
                debugValues.nTargetValue = nTargetValue
                debugValues.yieldPotential = yieldPotential
                debugValues.nRegularValue = nRegularValue
                debugValues.ignoreOverfertilization = lastIgnoreOverfertilization
                debugValues.nFactor = nFactor
                debugValues.regularNFactor = regularNFactor
            end
        end
    end

    local pHFactor = vehicle.lastPHFactor or 0
    local regularPHFactor = vehicle.lastRegularPHFactor or 0
    if requiresNFactorUpdate then
        local pHActualValue, pHTargetValue, pHRegularValue = pfModule.pHMap:updateLastPhValues()
        if pHActualValue > -1 and pHTargetValue > -1 and pHRegularValue > -1 then
            pHFactor = pfModule.pHMap:getYieldFactorByLevelDifference(pHActualValue - pHTargetValue)
            vehicle.lastPHFactor = pHFactor

            regularPHFactor = pfModule.pHMap:getYieldFactorByLevelDifference(pHRegularValue - pHTargetValue)
            vehicle.lastRegularPHFactor = regularPHFactor

            if HarvestExtension.YIELD_DEBUG then
                local debugValues = self.debugValues
                debugValues.pHFactor = pHFactor
                debugValues.regularPHFactor = regularPHFactor
                debugValues.pHActualValue = pHActualValue
                debugValues.pHTargetValue = pHTargetValue
                debugValues.pHRegularValue = pHRegularValue
            end
        end
    end

    -- regular factors are not allowed to be better than PF factors
    -- if you treat your field not as targeted in PF you wouldn't without as well
    return nFactor, math.min(regularNFactor, nFactor), currentYieldPotential, pHFactor, math.min(regularPHFactor, pHFactor)
end

function HarvestExtension:overwriteGameFunctions(pfModule)
    --
    --  Apply custom yield multiplier and track the yield
    --
    pfModule:overwriteGameFunction(Cutter, "processCutterArea", function(superFunc, vehicle, workArea, dt)
        local specCutter = vehicle.spec_cutter
        local combine = specCutter.workAreaParameters.combineVehicle

        local xs, zs, xw, zw, xh, zh, _
        local farmlandId
        local requiresPHFactorUpdate = false
        local requiresNFactorUpdate = false
        if combine ~= nil then
            xs, _, zs = getWorldTranslation(workArea.start)
            xw, _, zw = getWorldTranslation(workArea.width)
            xh, _, zh = getWorldTranslation(workArea.height)

            for fruitType, state in pairs(specCutter.workAreaParameters.fruitTypesToUse) do
                if state then
                    local phMapUpdated, nMapUpdated = pfModule.soilMap:updateCoverArea(fruitType, xs, zs, xw, zw, xh, zh, specCutter.allowsForageGrowthState)
                    requiresPHFactorUpdate = requiresPHFactorUpdate or phMapUpdated
                    requiresNFactorUpdate = requiresNFactorUpdate or nMapUpdated
                end
            end

            farmlandId = g_farmlandManager:getFarmlandIdAtWorldPosition((xs + xw) * 0.5, (zs + zw) * 0.5)
        end

        local nFactor, regularNFactor, currentYieldPotential, pHFactor, regularPHFactor = self:updateLatestFactors(pfModule, vehicle, requiresPHFactorUpdate, requiresNFactorUpdate)

        self.checkMultiplier = true
        local lastRealArea, lastArea = superFunc(vehicle, workArea, dt)
        self.checkMultiplier = false

        if lastRealArea > 0 and self.lastMultiplier ~= nil and combine ~= nil then
            local inputFruitType = specCutter.workAreaParameters.lastFruitType

            lastRealArea = lastRealArea / self.lastMultiplier -- restore raw value

            lastRealArea = lastRealArea * 1.05 -- apply global 5% more yield to counteract the increased (realistic) fertilizer usage

            local plowWeedMultiplier = 0
            plowWeedMultiplier = plowWeedMultiplier + self.lastPlowFactor * 0.15
            plowWeedMultiplier = plowWeedMultiplier + self.lastWeedFactor * 0.15

            local multiplier = 1
            multiplier = multiplier + nFactor * 0.50
            multiplier = multiplier + pHFactor * 0.20
            multiplier = multiplier + plowWeedMultiplier

            local regularMultiplier = 1
            regularMultiplier = regularMultiplier + regularNFactor * 0.50
            regularMultiplier = regularMultiplier + regularPHFactor * 0.20
            regularMultiplier = regularMultiplier + plowWeedMultiplier

            local regularRealArea = lastRealArea * regularMultiplier * currentYieldPotential

            specCutter.workAreaParameters.lastRealArea = lastRealArea * multiplier * currentYieldPotential

            if pfModule.yieldMap ~= nil then
                pfModule.yieldMap:setAreaYield(xs, zs, xw, zw, xh, zh, multiplier * currentYieldPotential)
            end

            if HarvestExtension.YIELD_DEBUG then
                local debugValues = self.debugValues
                debugValues.yieldFactor = multiplier
                debugValues.yieldFactorRegular = regularMultiplier
                debugValues.plowFactor = self.lastPlowFactor
                debugValues.weedFactor = self.lastWeedFactor
            end

            local fillTypeIndex = g_fruitTypeManager:getFillTypeIndexByFruitTypeIndex(inputFruitType)
            local conversionFactor = 1
            if specCutter.fruitTypeConverters[inputFruitType] ~= nil then
                fillTypeIndex = specCutter.fruitTypeConverters[inputFruitType].fillTypeIndex
                conversionFactor = specCutter.fruitTypeConverters[inputFruitType].conversionFactor
            end

            local fillType = g_fillTypeManager:getFillTypeByIndex(fillTypeIndex)

            local yield = self:getYieldFromArea(combine, vehicle, inputFruitType, specCutter.workAreaParameters.lastRealArea) * conversionFactor
            local yieldWeight = yield * (fillType.massPerLiter / FillTypeManager.MASS_SCALE)
            local regularYield = self:getYieldFromArea(combine, vehicle, inputFruitType, regularRealArea) * conversionFactor

            local farmlandStatistics, _, _ = vehicle:getPAStatisticInfo()
            if farmlandStatistics ~= nil and farmlandId ~= nil then
                farmlandStatistics:updateStatistic(farmlandId, "yield", yield)
                farmlandStatistics:updateStatistic(farmlandId, "yieldWeight", yieldWeight)
                farmlandStatistics:updateStatistic(farmlandId, "yieldRegular", regularYield)

                local maxPrice = self:getBestPriceForFillType(fillType)
                farmlandStatistics:updateStatistic(farmlandId, "yieldBestPrice", maxPrice * yield)
            end
        end

        self.lastMultiplier = nil
        self.lastPlowFactor = nil
        self.lastWeedFactor = nil

        return lastRealArea, lastArea
    end)

    --
    --  Apply custom yield multiplier and track the yield on grass
    --
    pfModule:overwriteGameFunction(Mower, "processMowerArea", function(superFunc, vehicle, workArea, dt)
        local specMower = vehicle.spec_mower

        local xs, zs, xw, zw, xh, zh, _
        local farmlandId
        local requiresPHFactorUpdate = false
        local requiresNFactorUpdate = false

        xs, _, zs = getWorldTranslation(workArea.start)
        xw, _, zw = getWorldTranslation(workArea.width)
        xh, _, zh = getWorldTranslation(workArea.height)

        for inputFruitType, converterData in pairs(specMower.fruitTypeConverters) do
            local phMapUpdated, nMapUpdated = pfModule.soilMap:updateCoverArea(inputFruitType, xs, zs, xw, zw, xh, zh)
            requiresPHFactorUpdate = requiresPHFactorUpdate or phMapUpdated
            requiresNFactorUpdate = requiresNFactorUpdate or nMapUpdated
        end

        farmlandId = g_farmlandManager:getFarmlandIdAtWorldPosition((xs + xw) * 0.5, (zs + zw) * 0.5)

        local nFactor, regularNFactor, currentYieldPotential, pHFactor, regularPHFactor = self:updateLatestFactors(pfModule, vehicle, requiresPHFactorUpdate, requiresNFactorUpdate)

        self.replaceMowerMultiplierCurrentYieldPotential = currentYieldPotential

        self.replaceMowerMultiplierNFactor = nFactor
        self.replaceMowerMultiplierRegularNFactor = regularNFactor
        self.replaceMowerMultiplierPHFactor = pHFactor
        self.replaceMowerMultiplierRegularPHFactor = regularPHFactor

        self.replaceMowerMultiplier = true
        local lastRealArea, lastArea = superFunc(vehicle, workArea, dt)
        self.replaceMowerMultiplier = false

        if lastRealArea > 0 and self.replaceMowerNewMultiplier ~= nil and self.replaceMowerRegularMultiplier ~= nil then
            -- get raw area change
            local rawPickupLiters = workArea.lastPickupLiters / self.replaceMowerNewMultiplier
            local regularPickupLiters = rawPickupLiters * self.replaceMowerRegularMultiplier

            local fillTypeIndex
            local dropArea = vehicle:getDropArea(workArea)
            if dropArea ~= nil then
                fillTypeIndex = dropArea.fillType
            elseif specMower.fillUnitIndex ~= nil then
                fillTypeIndex = vehicle:getFillUnitFillType(specMower.fillUnitIndex)
            end

            if fillTypeIndex ~= nil and fillTypeIndex ~= FillType.UNKNOWN then
                if pfModule.yieldMap ~= nil then
                    pfModule.yieldMap:setAreaYield(xs, zs, xw, zw, xh, zh, self.replaceMowerRegularMultiplier)
                end

                local fillType = g_fillTypeManager:getFillTypeByIndex(fillTypeIndex)

                local pickupWeight = workArea.lastPickupLiters * (fillType.massPerLiter / FillTypeManager.MASS_SCALE)

                local farmlandStatistics, _, _ = vehicle:getPAStatisticInfo()
                if farmlandStatistics ~= nil and farmlandId ~= nil then
                    farmlandStatistics:updateStatistic(farmlandId, "yield", workArea.lastPickupLiters)
                    farmlandStatistics:updateStatistic(farmlandId, "yieldWeight", pickupWeight)
                    farmlandStatistics:updateStatistic(farmlandId, "yieldRegular", regularPickupLiters)

                    local maxPrice = self:getBestPriceForFillType(fillType)
                    farmlandStatistics:updateStatistic(farmlandId, "yieldBestPrice", maxPrice * workArea.lastPickupLiters)
                end
            end
        end

        return lastRealArea, lastArea
    end)

    pfModule:overwriteGameFunction(FSBaseMission, "getHarvestScaleMultiplier", function(superFunc, mission, fruitTypeIndex, sprayFactor, plowFactor, limeFactor, weedFactor)
        local multiplier = superFunc(mission, fruitTypeIndex, sprayFactor, plowFactor, limeFactor, weedFactor)

        if self.checkMultiplier then
            self.lastMultiplier = multiplier
            self.lastPlowFactor = plowFactor
            self.lastWeedFactor = weedFactor
        end

        if self.replaceMowerMultiplier then
            multiplier = 1
            multiplier = multiplier + plowFactor * 0.15
            multiplier = multiplier + weedFactor * 0.15
            multiplier = multiplier + self.replaceMowerMultiplierRegularNFactor * 0.50
            multiplier = multiplier + self.replaceMowerMultiplierRegularPHFactor * 0.20
            multiplier = multiplier * self.replaceMowerMultiplierCurrentYieldPotential

            self.replaceMowerRegularMultiplier = multiplier

            multiplier = 1
            multiplier = multiplier + plowFactor * 0.15
            multiplier = multiplier + weedFactor * 0.15
            multiplier = multiplier + self.replaceMowerMultiplierNFactor * 0.50
            multiplier = multiplier + self.replaceMowerMultiplierPHFactor * 0.20
            multiplier = multiplier * self.replaceMowerMultiplierCurrentYieldPotential

            self.replaceMowerNewMultiplier = multiplier
        end

        return multiplier
    end)

    --
    --  Reset the cover lock bit after sowing
    --
    pfModule:overwriteGameFunction(SowingMachine, "processSowingMachineArea", function(superFunc, vehicle, workArea, dt)
        local spec = vehicle.spec_sowingMachine

        if vehicle:getLastSpeed() > 0.5 then
            if spec.workAreaParameters.isActive and spec.workAreaParameters.canFruitBePlanted then
                local xs,_,zs = getWorldTranslation(workArea.start)
                local xw,_,zw = getWorldTranslation(workArea.width)
                local xh,_,zh = getWorldTranslation(workArea.height)

                pfModule.soilMap:resetCoverLock(xs, zs, xw, zw, xh, zh)
            end
        end

        return superFunc(vehicle, workArea, dt)
    end)

    --
    --  Reset the cover lock bit after cultivating
    --
    pfModule:overwriteGameFunction(Cultivator, "processCultivatorArea", function(superFunc, vehicle, workArea, dt)
        if vehicle:getLastSpeed() > 0.5 then
            local xs,_,zs = getWorldTranslation(workArea.start)
            local xw,_,zw = getWorldTranslation(workArea.width)
            local xh,_,zh = getWorldTranslation(workArea.height)

            pfModule.soilMap:resetCoverLock(xs, zs, xw, zw, xh, zh)
        end

        return superFunc(vehicle, workArea, dt)
    end)

    --
    --  Add nitrogen where we destroy fruits
    --
    pfModule:overwriteGameFunction(FSDensityMapUtil, "updateDestroyCommonArea", function(superFunc, startWorldX, startWorldZ, widthWorldX, widthWorldZ, heightWorldX, heightWorldZ, limitToField, addFertilizerLock)
        pfModule.nitrogenMap:updateDestroyCommonArea(startWorldX, startWorldZ, widthWorldX, widthWorldZ, heightWorldX, heightWorldZ, limitToField, addFertilizerLock)

        superFunc(startWorldX, startWorldZ, widthWorldX, widthWorldZ, heightWorldX, heightWorldZ, limitToField, addFertilizerLock)

        pfModule.nitrogenMap:postUpdateDestroyCommonArea(startWorldX, startWorldZ, widthWorldX, widthWorldZ, heightWorldX, heightWorldZ, limitToField, addFertilizerLock)
    end)
end
