beta2.lua

-- Bugs
-- --------------------------------------------------------------------------
-- Constants
	traitBinValues = 
	{ -- lLimit, uLimit, l_cumulativeProbability, u_cumulativeProbability, groupProbability
		{0	,9		,0			,0.003875	,0.003875},
		{10	,24		,0.003875	,0.023125	,0.01925},
		{25	,39		,0.023125	,0.10775	,0.084625},
		{40	,60		,0.10775	,0.89225	,0.7845},
		{61	,75		,0.89225	,0.976875	,0.084625},
		{76	,90		,0.976875	,0.996125	,0.01925},
		{91	,100	,0.996125	,1			,0.003875}
	}
-- --------------------------------------------------------------------------
-- Classes
--TraitArrayBins = defclass(TraitArrayBins)
TraitArrayBins = {}

TraitArrayBins.__index = TraitArrayBins

function TraitArrayBins.init(selectedUnitCaste, int)
	local trait = {}
	setmetatable(trait,TraitArrayBins)
	trait.name = df.personality_facet_type[int-1]
	trait.traitMin = selectedUnitCaste.personality.a[int-1]
	trait.traitMed = selectedUnitCaste.personality.b[int-1]
	trait.traitMax = selectedUnitCaste.personality.c[int-1]
	return trait
end

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

-- --------------------------------------------------------------------------
-- Functions
function boostAttributes(selectedUnitCaste, unit)

--[[ these were not easy to get, see code...
df.physical_attribute_type
0       STRENGTH
1       AGILITY
2       TOUGHNESS
3       ENDURANCE
4       RECUPERATION
5       DISEASE_RESISTANCE

df.mental_attribute_type
0       ANALYTICAL_ABILITY
1       FOCUS
2       WILLPOWER
3       CREATIVITY
4       INTUITION
5       PATIENCE
6       MEMORY
7       LINGUISTIC_ABILITY
8       SPATIAL_SENSE
9       MUSICALITY
10      KINESTHETIC_SENSE
11      EMPATHY
12      SOCIAL_AWARENESS
]]

handleBoostType("Physical", selectedUnitCaste, unit)
handleBoostType("Mental", selectedUnitCaste, unit)

end

function handleBoostType(string_type, selectedUnitCaste, unit)

	attributeBin = {}
	
	if string_type == "Physical" then
		print (string_type)
		numberOfAttr = #selectedUnitCaste.attributes.phys_att_range
	else
		print (string_type)
		numberOfAttr = #selectedUnitCaste.attributes.ment_att_range
	end
	
	print ("Number of Attributes: ", numberOfAttr)

-- cycle through attributes
-- for i,j in ipairs(df.physical_attribute_type) do print(i,j) end
	for attrNumber = 1, numberOfAttr do

	if string_type == "Physical" then
			numberOfBins = #selectedUnitCaste.attributes.phys_att_range[attrNumber-1]
			attribute_name = df.physical_attribute_type[attrNumber-1]
		else
			numberOfBins = #selectedUnitCaste.attributes.ment_att_range[attrNumber-1]
			attribute_name = df.mental_attribute_type[attrNumber-1]
	end
	
	print("Attribute Name: ", attribute_name)
	print("Number of bins: ", numberOfBins)
	
	-- cycle through bins, print each bin
	for line = 1, numberOfBins do
	
		if string_type == "Physical" then
			attributeBin[line] = selectedUnitCaste.attributes.phys_att_range[attrNumber-1][line-1]
		else
			attributeBin[line] = selectedUnitCaste.attributes.ment_att_range[attrNumber-1][line-1]
		end
		
		print("Bin: " , line, attributeBin[line])
	end
	
	-- Figure out Percent
	if string_type == "Physical" then
			rawValue = unit.body.physical_attrs[attrNumber-1].value
		else
			rawValue = unit.status.current_soul.mental_attrs[attrNumber-1].value
		end
										
	if string_type == "Physical" then	
		print("Raw Value: ", unit.body.physical_attrs[attrNumber-1].value)
	else
		print("Raw Value: ", unit.status.current_soul.mental_attrs[attrNumber-1].value)
	end
	
	-- If value is above final bin, skip
	if rawValue >= attributeBin[#attributeBin] then
		print("Value meets or exceeds max embark value, not attempting")
	else
		percent = calcPercent(rawValue, attributeBin)
		print ("Percent: ", percent)
		
		-- Boost this attributes value
		if string_type == "Physical" then
			boostValue(attrNumber, rawValue, percent, attributeBin, unit, type_string, "Physical")
		else
			boostValue(attrNumber, rawValue, percent, attributeBin, unit, type_string, "Mental")
		end
	end		
	
end

end

function figureValueUL_Limit (rawValue, attributeBin)

	local counter = #attributeBin
	
	while rawValue < attributeBin[counter]  do
      l_limit = attributeBin[counter-1]
      u_limit = attributeBin[counter]
	  counter = counter-1
    end

	return l_limit, u_limit, counter
end

function calcPercent(rawValue, attributeBin)

	local l_limit, u_limit, counter = figureValueUL_Limit(rawValue, attributeBin)
	
	print ("l_limit", l_limit)
	print ("u_limit", u_limit)
	print ("counter: ", counter)
	
	local value = (rawValue - l_limit + 1)
	
	print("# of bins: ", #attributeBin)
	
	local numOfBins = #attributeBin
	
	
	value = value / (u_limit - l_limit + 1)
	value = value * ( 1/ (numOfBins-1) )
	value = value + ( ( counter - 1 ) * 1/(numOfBins-1))
	
	return value
end

--would like to clean this up, don't like type_string, should be able to assign something to a var	
function boostValue(attrNumber, rawValue, percent, attributeBin, unit, line, type_string)

	-- Check if we're going to boost the value
	
	local randomRoll = math.random()
	print("RNG: ", randomRoll)
	
	if randomRoll > percent then
		print ("Boosting!")
			
		local l_limit, u_limit, counter = figureValueUL_Limit(rawValue, attributeBin)
		
		local binsLeft = #attributeBin - counter 
		
		print ("Bins Left: ", binsLeft)
		
		percentWithinBin = (rawValue - l_limit) / (u_limit - l_limit)
		print("Percent within bin: ", percentWithinBin)
		
		binPercents = {}
		
		for position = 1, binsLeft do
			if position == 1 then
				binPercents[position] = percentWithinBin
			else
				binPercents[position] = 1
			end
		end
		
		-- get sum
		sumOfPercents = sum(binPercents)
		print("Sum of Percents: ", sumOfPercents)
			
		for position = 1, #binPercents do
			binPercents[position] = binPercents[position]/sumOfPercents
		end

		-- last = 0 ; for ... ; last = last + value ; binPc[i] = last ; end
		-- cumulative sum array
		for position = 1, #binPercents do
			if position == 1 then
			  -- do nothing
			else
				binPercents[position] = binPercents[position-1]+binPercents[position]
				newCumPosition = position
			end
		end
		
		--printall(binPercents)
		
		-- new idea, the new RNG will be based on the original roll

		-- do a second RNG to determine what bin we're going to land in
		secondRoll = (randomRoll - percent)/(1 - percent)
		print("secondRoll: ", secondRoll)
		
		-- found flag
		local found = 0
		
		local binPosition = 1
		
		-- search for matching range
		while found == 0  do
			if secondRoll > binPercents[binPosition] then
				binPosition = binPosition + 1
			else
				found = 1
				print("BinPosition: ", binPosition)
				print("BinPositionFound: ", binPercents[binPosition])
				print("LowerBinPosition: ", binPercents[binPosition-1])
				if binPosition == 1 then
					thirdRoll = (secondRoll - 0) / (binPercents[binPosition] - 0)
					else
					thirdRoll = (secondRoll - binPercents[binPosition-1]) / (binPercents[binPosition]-binPercents[binPosition-1])
				end
				print ("Third roll: ", thirdRoll)
			end
		end
		


		binP = findBinPosition (l_limit, attributeBin)
		
		l_limit = attributeBin[binP + binPosition - 1]
		u_limit = attributeBin[binP + binPosition]
		
		print("new l_limit: ", l_limit)
		print("new u_limit: ", u_limit)
		
	-- gen new value!
	
		--if value is greater than l_limit, increase l_limit to raw_value + 1
		if rawValue > l_limit then
			l_limit = rawValue + 1
		end
		
		range = u_limit - l_limit + 1
		
		
		--thirdRoll = math.random()
		--print ("Position: ", newCumPosition)
		--print ("BinPercent(Position): ", binPercents[newCumPosition])
		--thirdRoll = (secondRoll - binPercents[binPosition-1]) / (binPercents[binPosition]-binPercents[binPosition-1])
		
		newValue = (range * thirdRoll ) + l_limit -1
		
		print("New Value: ", newValue)
		
		-- need to specify whether physical or mental here.
		if type_string == "Physical" then
			unit.body.physical_attrs[attrNumber-1].value = round(newValue,0)
			if (unit.body.physical_attrs[attrNumber-1].value * 2) > 5000 then
				unit.body.physical_attrs[attrNumber-1].max_value = 5000
			else
				unit.body.physical_attrs[attrNumber-1].max_value = (unit.body.physical_attrs[attrNumber-1].value * 2)
			end
		else
			unit.status.current_soul.mental_attrs[attrNumber-1].value = round(newValue,0)
			if (unit.status.current_soul.mental_attrs[attrNumber-1].value * 2) > 5000 then
				unit.status.current_soul.mental_attrs[attrNumber-1].max_value = 5000
			else
				unit.status.current_soul.mental_attrs[attrNumber-1].max_value = (unit.status.current_soul.mental_attrs[attrNumber-1].value * 2)
			end
		end
		
	else
		print ("Value was determined too high on this pass")
	end

end

function round(num, idp)
  local mult = 10^(idp or 0)
  return math.floor(num * mult + 0.5) / mult
end

function findBinPosition (binValue, attributeBin)
	for tempValue = 1, #attributeBin do
		if binValue == attributeBin[tempValue] then
			return tempValue
		end
	end
end

function sum(t)
    local sum = 0
    for k,v in pairs(t) do
        sum = sum + v
    end

    return sum
end

function getCitizen()
	local ret={}
	local units={}
	local vs = dfhack.gui.getCurViewscreen()
	if df.viewscreen_setupdwarfgamest:is_instance(vs) then
		print("Embark!")
		units=df.global.world.units.all
	else
		print("Not Embark!")
		units=df.global.world.units.active
	end
	for num,unit in ipairs(units) do
		if dfhack.units.isCitizen(unit) then
			table.insert(ret,unit)
		end
	end
	return ret
end

function InitTraits(selectedUnitCaste, unit)
	
	--[traitNames = {"ANXIETY","ANGER","DEPRESSION","SELF_CONSCIOUSNESS","IMMODERATION","VULNERABILITY","FRIENDLINESS","GREGARIOUSNESS","ASSERTIVENESS","ACTIVITY_LEVEL","EXCITEMENT_SEEKING","CHEERFULNESS","IMAGINATION","ARTISTIC_INTEREST","EMOTIONALITY","ADVENTUROUSNESS","INTELLECTUAL_CURIOSITY","LIBERALISM","TRUST","STRAIGHTFORWARDNESS","ALTRUISM","COOPERATION","MODESTY","SYMPATHY","SELF_EFFICACY","ORDERLINESS","DUTIFULNESS","ACHIEVEMENT_STRIVING","SELF_DISCIPLINE","CAUTIOUSNESS"}   --]
	
		traits = {}
	
		for int = 1, 30 do
			--table.insert(traits,unit)
			traits[int] = TraitArrayBins.init(selectedUnitCaste, int )
		end
	return traits
end

function CalcPercentTraits(selectedUnitCaste, unit, traits)

	-- percent for each trait
	local percent = {}
	
	local trait=unit.status.current_soul.traits
	
	-- for each trait
	for int = 1, 30 do
	
		local originalValue = trait[int-1]

		local tempValue = trait[int-1]

		-- handles conversion from non normal (ie 50) mean traits to 50 mean
		-- =IF(Value<=Med;((Value-Min)/(Med-Min))*50;(((Value-Med)/(Max-Med)*50)+50))
		--printall(unit.)
		--personality = status.current_soul.traits
		if (selectedUnitCaste.personality.a[int-1] ~= 0) or (selectedUnitCaste.personality.b[int-1] ~= 50) or (selectedUnitCaste.personality.c[int-1] ~= 100) then 
		
			if trait[int] <= selectedUnitCaste.personality.b[int] then
				tempValue = ((tempValue - selectedUnitCaste.personality.a[int]) / (selectedUnitCaste.personality.b[int] - selectedUnitCaste.personality.a[int])) * 50
			else
				tempValue = ((((tempValue - selectedUnitCaste.personality.b[int]) / (selectedUnitCaste.personality.c[int] - selectedUnitCaste.personality.b[int])) * 50 ) + 50)
			end
		end
		-- based on #traitBinValues, 
		
		local counter = 7
		local range = 0
		local l_limit = 0
		local u_limit = 0
		local l_range = 0
		local l_cum = 0

		-- find matching l_value, make into own function, so can be used for BoostTrait
		while traitBinValues[counter][1] > tempValue do
			counter = counter - 1
		end
	
		l_limit = traitBinValues[counter][1]
		u_limit = traitBinValues[counter][2]
		l_cum = traitBinValues[counter][3]
		u_cum = traitBinValues[counter][4]
		grp_prob = traitBinValues[counter][5]
		
		range = u_limit - l_limit + 1
		l_range = tempValue - l_limit + 1
		
		percent[int] = ((l_range/range)*(u_cum - l_cum)) + l_cum
		
		print(percent[int])
		--print(l_limit)
		--print(originalValue)
		--print(u_limit)
	end
	
	return percent
end

function ConvertTraitPercentToValue()

end

function InitCaste (unit)
	race = unit.race
	caste = unit.caste
	selectedUnitCaste = df.global.world.raws.creatures.all[race].caste[caste]
	return selectedUnitCaste
end

function BoostTrait(unit, traitPercent, int, randomRoll)
-- [[
	local trait=unit.status.current_soul.traits
	
	local originalValue = trait[int-1]
	
	local newValue
	
	local normedValue =  trait[int-1]
	
	print ("original Value: ", originalValue)
	
	-- find matching l_value, make into own function, so can be used for BoostTrait
	local counter = #traitBinValues
	while traitBinValues[counter][1] > normedValue do
		counter = counter - 1
	end
	
	local binsLeft = #traitBinValues-counter
	print("Bins Left: ", binsLeft)

	l_limit = traitBinValues[counter][1]
	u_limit = traitBinValues[counter][2]
	l_cum = traitBinValues[counter][3]
	u_cum = traitBinValues[counter][4]
	grp_prob = traitBinValues[counter][5]
	
	range = u_limit - l_limit + 1
	l_range = normedValue - l_limit + 1

	local t_min = selectedUnitCaste.personality.a[int-1]
	local t_med = selectedUnitCaste.personality.b[int-1]
	local t_max = selectedUnitCaste.personality.c[int-1]

	--check if not normal, if so, skew to normal
	if (t_min ~= 0) or (t_med ~= 50) or (t_max ~= 100) then
		if originalValue >= t_med then
			normedValue = (((normedValue - 50)/(100-50))*(t_max - t_med)) + t_med
		else
			normedValue = ((normedValue - 0) / (50 - 0) * (t_med - t_min)) + t_min
		end
	else
		--value is normal, no need to norm it
	end
	
	local overage = (randomRoll - traitPercent[int]) / (1 - traitPercent[int])
	print ("Overage: ", overage)

	--create new bins based on new floor
	new_bins = {}
	
	-- lLimit, uLimit, l_cumulativeProbability, u_cumulativeProbability, groupProbability
	l_limit = normedValue
	l_cum = traitPercent[int]
	
	for b_counter = 1, binsLeft+1 do
	
		local t_l_limit
		local t_l_cum
		local t_u_limit
		local t_u_cum
		local t_grp_prob
		
		print("b_counter: ", b_counter)
		
		if b_counter == 1 then
			t_l_limit = l_limit
			t_l_cum = 0
		else
			t_l_limit = traitBinValues[b_counter+#traitBinValues-binsLeft-1][1]
			t_l_cum = ((traitBinValues[b_counter+#traitBinValues-binsLeft-1][3] - traitPercent[int]) / (1-traitPercent[int]))
		end
		t_u_limit = traitBinValues[b_counter+#traitBinValues-binsLeft-1][2]
		t_u_cum = ((traitBinValues[b_counter+#traitBinValues-binsLeft-1][4] - traitPercent[int]) / (1-traitPercent[int]))
		t_grp_prob = (t_u_cum - t_l_cum)/(t_u_limit - t_l_limit + 1)
		--print("l_cum: ", t_l_cum)
		--print("u_cum: ", t_u_cum)
		new_bins[b_counter] = {t_l_limit, t_u_limit, t_l_cum, t_u_cum, t_grp_prob}
	end
	
	-- user overage to come to a new trait value!
	
	local b_counter = 1
	
	local found = 0
	
	while found == 0 do
		
		local t_l_limit = new_bins[b_counter][1]
		local t_u_limit = new_bins[b_counter][2]
		local t_l_cum = new_bins[b_counter][3]
		local t_u_cum = new_bins[b_counter][4]
		local t_grp_prob = new_bins[b_counter][5]
		local spread = t_u_cum - t_l_cum
		local range = t_u_limit - t_l_limit + 1
		local distFromLow = overage - t_l_cum
		
		if overage > t_u_cum then
			
			b_counter = b_counter + 1
		else
			found = 1
			print ("spread: ", spread, "range: ", range)
			print ("distFromL: ", distFromLow, "l_cum:", t_l_cum)
			newValue = (distFromLow / (spread/range) ) + t_l_limit -1
		end
	end
	
	print("New Value: ", newValue)
	
	--now if skewed, skew back, then round, then set
	--check if not normal, if so, skew to normal
	if (t_min ~= 0) or (t_med ~= 50) or (t_max ~= 100) then
		if newValue >= t_med then
				newValue = (newValue - 50)/(100-50)
				newValue = (newValue * (t_max - t_med)) + t_med
			else
				newValue = (newValue - 0)/(50 - 0)
				newValue = (newValue * (t_med - t_min) ) + t_min
		end
	else
	  -- do nothing
	end
	
	print ("After skew check value: ", newValue)
	
	newValue = round(newValue,0)
	
	-- set new value
	trait[int-1] = newValue

end

--Boosts attributes, traits, skills for a specific unit
function ProcessFunction (unit)
				selectedUnitCaste = InitCaste(unit)
				
				--printall(selectedUnitCaste)
				
				boostAttributes(selectedUnitCaste, unit)
				
				-- Traits
				traits = {}
				traits = InitTraits(selectedUnitCaste, unit)
				local traitPercent = {}
				traitPercent = CalcPercentTraits(selectedUnitCaste, unit, traits)
				
				for int = 1, #traitPercent do
					local randomRoll = math.random()
					print("Trait RNG: ", randomRoll)
					if randomRoll >= traitPercent[int] then
						print("boosting trait!")
						BoostTrait(unit, traitPercent, int, randomRoll)
					else
						print("value too high, not boosting")
					
					end
				end
				
				-- Skills

end
-- --------------------------------------------------------------------------

-- --------------------------------------------------------------------------
-- Main
-- Initilization
local opt = ...

if opt == nil then
		
	print("Options: ")
	print("boost_dwarf (SU)Selected Unit, (AU)All Units")
	
	-- AU All Units
	else if opt == "AU" then
		if df.global.world.units.all == nil then 
			print("No Units!  Aborting!")
			else
			-- get all citizens!
			citizens = getCitizen()
		
			for num, unit in ipairs(citizens) do
				ProcessFunction(unit)
			end
		end
		else if opt == "SU" then
			if dfhack.gui.getSelectedUnit() == nil then
				print("No Unit selected!  Aborting!")
			else
					unit = dfhack.gui.getSelectedUnit()
					ProcessFunction(unit)
			end
			
				
		end
			
	end	
end

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