using MPSGE,  Random, DataFrames, PlotlyJS
import JuMP.Containers:DenseAxisArray,@container 
import Distributions.Uniform
function random_parameter(set; lower = 0, upper = 1)
    return @container([set], rand(Uniform(lower, upper)))
end
function create_data(N::Int; random_seed::Union{Int,Missing} = 1234)
    if !ismissing(random_seed)
        Random.seed!(random_seed)
    end
    regions = Symbol.("r",1:N)

    land_rent = random_parameter(regions) #n0
    capital_rent = random_parameter(regions) #k0
    wage = random_parameter(regions) #l0

    return (regions = regions, land_rent = land_rent, capital_rent = capital_rent, wage = wage)
end

regions, land_rent, capital_rent, wage = create_data(10)
initial_μ = 4
outputY = land_rent .+ capital_rent .+ wage #y0
# Assume all regions have balanced budgets
consumption = land_rent .+ capital_rent .+ wage
consumption_share = consumption./sum(consumption)

M1 = MPSGEModel()

@parameters(M1, begin
    productivity_index[regions], 1, (description = "Productivity index")
    μ, initial_μ, (description = "Logit parameter, elasticity for sorting")
    good_elasticity, 2, (description = "Good elasticity")
    γ, 1, (description = "Congestion spillover elasticity")
end)

@sectors(M1, begin
    Y[regions], (description = "Production in region")
    C[regions], (description = "Consumption in region") # Total consumption activity/Index, U x N
end)

@commodities(M1, begin
    P[regions], (description = "Regional supply")
    PN[regions], (description = "Rental index for land")
    PK, (description = "Captial price")
    PL[regions], (description = "Regional wage")
    PC[regions], (description = "Consumption price")
end)

@consumer(M1, RA[regions], description = "Representative agent for households living in region")

@auxiliaries(M1, begin
    U[regions], (description = "Per-capita utility index (benchmark = 1)")
    N[regions], (description = "Population index (benchmark = 1)")
    congestion_externality[regions], (description = "Congestion spillover")
    price_index, (description = "Price index to fix as numeraire")
end)

set_start_value.(U, 1)
set_start_value.(N, 1)
set_start_value.(congestion_externality, 0)
set_start_value(price_index, 1)

for r∈regions
    @production(M1, Y[r], [s=0.5, kl=>s=1, t=0], begin
        @output(P[r], productivity_index[r] * outputY[r], t)
        @input(PN[r], land_rent[r], s)
        @input(PK, capital_rent[r], kl)
        @input(PL[r], wage[r], kl)
    end)
end

for r∈regions
    @production(M1, C[r], [s=good_elasticity, t=0], begin
        @output(PC[r], (1-congestion_externality[r])*consumption[r], t)
        [@input(P[rr], consumption_share[r]*outputY[rr], s) for rr∈regions]...
    end)
end

for r∈regions
    @demand(M1, RA[r], begin
        @final_demand(PC[r], outputY[r])
        @endowment(PL[r], wage[r]*N[r])
        @endowment(PN[r], land_rent[r]) ## Was an 'error, having N as a multiplier here instead of with PL
        @endowment(PK, capital_rent[r]*N[r])
    end)
end

# Per capital utility (economic well being)
for r∈regions
    @aux_constraint(M1, U[r], 
        U[r] - RA[r]/(PC[r]*N[r]*consumption[r]) )
end

# Logit model determines the number of individuals located in region r:
for r∈regions
    @aux_constraint(M1, N[r], 
        N[r] - exp(μ*U[r])/sum(consumption_share[rr]*exp(μ*U[rr]) for rr in regions))
end
for r∈regions
    @aux_constraint(M1, congestion_externality[r], 
        congestion_externality[r] - (N[r]^γ - 1))
end
@aux_constraint(M1, price_index, price_index - sum(P[r]*outputY[r] for r in regions) / sum(outputY[r] for r in regions))

# set_silent(M1)
solve!(M1, cumulative_iteration_limit=0)
bnchM1 = generate_report(M1)
bnchM1[!,:var] = string.(bnchM1[:,:var])
for r in regions
    push!(bnchM1, ["Prod_Ind$r" value(productivity_index[r])  0])
    push!(bnchM1, ["ConsShare$r" consumption_share[r] * value(M1[:N][r]) 0])
    push!(bnchM1, ["ConsShare0$r" consumption_share[r] 0])
end
price_index_bnch = sum(value(M1[:P][r])*outputY[r] for r in regions) / sum(outputY[r] for r in regions)
solve!(M1, cumulative_iteration_limit=1000000)

# fix(RA[:r1],outputY[:r1])
# fix(P[:r1],1)
fix(price_index,1)
set_value!.(M1[:productivity_index], random_parameter(regions, lower = 0.5, upper = 1.5))
for r in regions
    fix(congestion_externality[r],0)
end
# fix(RA[:r1],consumption[:r1]) Can set the numeraire, or by default it will be the largest consumer (r10)
solve!(M1, cumulative_iteration_limit=100000)

Shock_noextM1=generate_report(M1)
Shock_noextM1[!,:var] = string.(Shock_noextM1[:,:var])
for r in regions
    push!(Shock_noextM1, ["Prod_Ind$r" value(productivity_index[r])  0])
    push!(Shock_noextM1, ["ConsShare$r" consumption_share[r] * value(M1[:N][r]) 0])
    push!(Shock_noextM1, ["ConsShare0$r" consumption_share[r] 0])
end
price_index_noext = sum(value(M1[:P][r])*outputY[r] for r in regions) / sum(outputY[r] for r in regions)

data_frames_noext = []
# set_silent(M)
for μ in 0:.1:5
    set_value!(M1[:μ], μ)
    solve!(M1)

    tmp =  DataFrame( 
        region = regions,
        population = [value(M1[:N][r])*consumption_share[r] for r in regions],
        U = [value(M1[:U][r]) for r in regions],
        N = [value(M1[:N][r]) for r in regions],
        Y = [value(M1[:Y][r]) for r in regions],
        C = [value(M1[:C][r]) for r in regions],
        productivity_index = [value(M1[:productivity_index][r]) for r in regions],
        P = [value(M1[:P][r])/price_index_noext for r in regions],
        PN = [value(M1[:PN][r])/price_index_noext for r in regions],
        PL = [value(M1[:PL][r])/price_index_noext for r in regions],
        RA = [value(M1[:RA][r])/price_index_noext for r in regions],
        μ = value(M1[:μ]),
        good_elasticity = value(M1[:good_elasticity]),
    ) 
    # report(M1, regions, land_rent, capital_rent, wage)
    tmp[!,:μ] .= μ

    push!(data_frames_noext,tmp)
end

df_noext = vcat(data_frames_noext...);

# Finally, we plot the results using PlotlyJS
p = PlotlyJS.plot(df_noext, x=:μ, y=:N, color=:region, kind = "lines", category_orders=attr(region=sort(df_noext[end-10:end,:],:N, rev=true)[:,:region]),Layout(yaxis_title_text = "N: Index of population",xaxis_title_text = "μ: logit parameter"))

p = PlotlyJS.plot(df_noext, x=:μ, y=:Y, color=:region, mode = "lines")

for r in regions
    unfix(congestion_externality[r])
    set_lower_bound(congestion_externality[r], -10)
    print(lower_bound(congestion_externality[r]))
end

solve!(M1, cumulative_iteration_limit=100000)
price_index_shock = sum(value(M1[:P][r])*outputY[r] for r in regions) / sum(outputY[r] for r in regions)
ShockM1 = generate_report(M1)
ShockM1[!,:var] = string.(ShockM1[:,:var])
for r in regions
    push!(ShockM1, ["Prod_Ind$r" value(productivity_index[r])  0])
    push!(ShockM1, ["ConsShare$r" consumption_share[r] * value(M1[:N][r]) 0])
    push!(ShockM1, ["ConsShare0$r" consumption_share[r] 0])
end
    
data_frames = []
# set_silent(M)
for γ in 0:.1:5
    set_value!(M1[:γ], γ)
    solve!(M1)

    tmp =  DataFrame( 
        region = regions,
        population = [value(M1[:N][r])*consumption_share[r] for r in regions],
        U = [value(M1[:U][r]) for r in regions],
        N = [value(M1[:N][r]) for r in regions],
        Y = [value(M1[:Y][r]) for r in regions],
        congestion_externality = [value(M1[:congestion_externality][r]) for r in regions],
        productivity_index = [value(M1[:productivity_index][r]) for r in regions],
        P = [value(M1[:P][r])/price_index_shock for r in regions],
        PN = [value(M1[:PN][r])/price_index_shock for r in regions],
        PL = [value(M1[:PL][r])/price_index_shock for r in regions],
        RA = [value(M1[:RA][r])/price_index_shock for r in regions],
        γ = value(M1[:γ]),
        good_elasticity = value(M1[:good_elasticity]),
    ) 
    # report(M1, regions, land_rent, capital_rent, wage)
    # tmp[!,:γ] .= γ

    push!(data_frames,tmp)
end
df = vcat(data_frames...);

# Finally, we plot the results using PlotlyJS
p = PlotlyJS.plot(df, x=:γ, y=:congestion_externality, color=:region, mode = "lines")
# p = PlotlyJS.plot(df, x=:γ, y=:N, color=:region, mode = "lines")

p = PlotlyJS.plot(df, x=:γ, y=:population, color=:region, mode = "lines")
p = PlotlyJS.plot(df, x=:γ, y=:N, color=:region, mode = "lines", category_orders=attr(region=sort(df_noext[end-10:end,:],:N, rev=true)[:,:region]),
 Layout(yaxis_title_text = "N: Index of population change",xaxis_title_text = "γ: Congestion elasticity parameter",legend=attr(x=0.9, y=.98)))
