#Used Packages
using Random, Distributions, StatsBase, Plots, CSV, DelimitedFiles

stop(text="Stop.") = throw(StopException(text))

struct StopException{T}
    S::T
end

function Base.showerror(io::IO, ex::StopException, bt; backtrace=true)
    Base.with_output_color(get(io, :color, false) ? :green : :nothing, io) do io
        showerror(io, ex.S)
    end
end


# Parameters and variables used in the model
NumberRun = 100 # the number of runs (replicates)
for run=1:NumberRun

PopSize =1000 #Population size (N)
HalfPopSize = div(PopSize,2)

NumberOrnament = 10 # Number of ornaments (t)

# Initialisation of the traits (O, P, M*)
InitialOrnamentMean = 0; InitialOrnamentSD = 1
InitialPreferenceMean = 0; InitialPreferenceSD = 1
InitialExtraMean = 5; InitialExtraSD = 1

# Initialisation of the mutations
MutationOrnamentProbability = 0.01; MutationOrnamentMean = 0; MutationOrnamentSD = 1 # πO; mean; σO
MutationPreferenceProbability = 0.01; MutationPreferenceMean = 0; MutationPreferenceSD = 1 # πP; mean; σP
MutationExtraProbability = 0.01; MutationExtraMean = 0; MutationExtraSD = 1 # πM; mean; σM

# Costs of the ornaments
CostStructure = 1 # 0 = additive cost strucutre; 1 = multiplicative cost structure
CostSteepness = 0.1
MortalityFunction(ornament) = CostSteepness * ornament^2 # ornament mortality function (cost parameter c), dependent on the ornament's size

StrengthPreference = 0.05 # Preference strength coefficient (s)

MateSearchCost = 0.00001 # Mate-search cost (k)

Pleiotropy = 0 # Degree of pleiotropy affecting ornaments and preferences (p)

NumGens = 100000 # Number of generations


SaveOutputEveryXGens = 10 # Save simulation results for each run, every tenth generation
PrintOutputEveryXGens = 50 # Print simulation results


# Generating the initial trait genetics (O; P; M*)
function InitialGenes()

	OrnamentGenes = rand(Normal(InitialOrnamentMean, InitialOrnamentSD), (HalfPopSize, NumberOrnament))
	PreferenceGenes = rand(Normal(InitialPreferenceMean, InitialPreferenceSD), (HalfPopSize, NumberOrnament))
	ExtraGenes = rand(Normal(InitialExtraMean, InitialExtraSD), (HalfPopSize, 1))
	hcat(OrnamentGenes, PreferenceGenes, ExtraGenes)
end

mGenes = InitialGenes();  fGenes = InitialGenes() #male and female genes

function Mutation() # Generating the mutations for the next generation

		MutationOrnament = rand(Normal(MutationOrnamentMean, MutationOrnamentSD), NumberOrnament) .* rand(Bernoulli(MutationOrnamentProbability), NumberOrnament)
		MutationPreference = rand(Normal(MutationPreferenceMean, MutationPreferenceSD), NumberOrnament) .* rand(Bernoulli(MutationPreferenceProbability), NumberOrnament)
		MutationExtra = rand(Normal(MutationExtraMean, MutationExtraSD), 1) .* rand(Bernoulli(MutationExtraProbability), 1)

		hcat(MutationOrnament', MutationPreference', MutationExtra)

end

GenesOutput = Array{Float64}(undef, 0, 2 * NumberOrnament + 1)
OrnamentOutput = Array{Float64}(undef, 0, NumberOrnament)
ExtraOutput = Array{Float64}(undef, 0, 1)
mViabilityOutput = Array{Float64}(undef, 0, 1)

# Generation loop
for gen=1:NumGens

# Male's viability (v): determined by his ornament expression and the mortality function (cost parameter c)
# Different for the two ornament cost structures: v_additive; v_multiplicative
	mOrnament = (1 - Pleiotropy/2) * mGenes[:,1:NumberOrnament] + Pleiotropy/2 * mGenes[:,(NumberOrnament+1):(2*NumberOrnament)]
	mOrnamentPositive = max.(mOrnament, 0)

	if CostStructure == 0
		OrnamentTotal = sum(mOrnamentPositive, dims=2)
		mViability = vec(max.(1 .- MortalityFunction.(OrnamentTotal), 0))
	else
		OrnamentSurvival = 1 .- min.(MortalityFunction.(mOrnamentPositive), 1)
		mViability = vec(prod(OrnamentSurvival, dims=2))
	end

	if sum(mViability)==0
		stop("All males are dead")
	end


	fPreferences = (Pleiotropy/2) * fGenes[:,1:NumberOrnament] + (1-Pleiotropy/2) * fGenes[:,(NumberOrnament+1):(2*NumberOrnament)] # female preferences (P)
	fExtra = max.(fGenes[:,2*NumberOrnament + 1], 0) # Female’s number of mates (M*)
	fFecundity = max.(1 .- MateSearchCost * fExtra.^2 , 0) # Female fecundity (f)

	AllBabyGenes = Array{Float64}(undef, 0, 2 * NumberOrnament + 1)

# Mating loop
	for i=1:PopSize

		# A female chooses her partner from 1+M* suitors, according to her 'rating' (R)
		fChosenIndex = sample(1:HalfPopSize, Weights(fFecundity))
		fChosenGenes = fGenes[fChosenIndex, :]

		mWanted = 1 + rand(Poisson(max(fChosenGenes[2*NumberOrnament + 1],0)))

		mSuitorIndices = sample(1:HalfPopSize, Weights(mViability), mWanted)
		mSuitorGenes = mGenes[mSuitorIndices, :]

		mSuitorOrnaments = mOrnamentPositive[mSuitorIndices, :]
		fChosenPreferences = fPreferences[fChosenIndex, :]
		fWeights = exp.(StrengthPreference * mSuitorOrnaments * fChosenPreferences)
		mChosenIndex = sample(1:mWanted, Weights(fWeights))
		mChosenGenes = mSuitorGenes[mChosenIndex,:]

		# Generating the offspring
		WhichParentGenes = rand(Bernoulli(), 2*NumberOrnament+1)' # Deciding which genes are passed on from the parents (female, male)
		BabyGenes = WhichParentGenes .* fChosenGenes' + (1 .- WhichParentGenes) .* mChosenGenes' + Mutation() # THe offspring's genes, including mutations

		AllBabyGenes = [AllBabyGenes;BabyGenes]

	end


# The offspring of one generation are assigned as the future females and males of the next generation
		 fGenes = AllBabyGenes[1:HalfPopSize, :]
		 mGenes = AllBabyGenes[HalfPopSize+1:PopSize, :]

# Save the output results for each generation
		if mod(gen, SaveOutputEveryXGens)==0
			AllGenes = vcat(mGenes, fGenes)
					GenesOutput = [GenesOutput ; mean(AllGenes, dims=1)]
					OrnamentOutput = [OrnamentOutput ; mean(mOrnamentPositive, dims=1)]
					ExtraOutput = [ExtraOutput ; mean(fExtra)]
					mViabilityOutput = [mViabilityOutput ; mean(mViability)]

# Print the output results
		if mod(gen, PrintOutputEveryXGens)==0
			println(gen)
			println("mean ornament sizes $(round.(mean(mOrnamentPositive, dims=1), digits=2))")
			println("mean preferences $(round.(mean(fPreferences, dims=1), digits=2))")
			println("mean male viability $(round(mean(mViability),digits=2))")
			println("mean mate-search effort $(round(mean(fExtra), digits=1))")
			println("mean fecundity $(round(mean(fFecundity), digits=2))")

		end

	end

 end

end