using CSV, NamedArrays, DataFrames

datayear=2021 ## only 2021 and 2017 are avaiable
include("HHdataparse.jl")

using MPSGE

TRN = String.(HHdata["sets"]["trn"])
G =  String.(HHdata["sets"]["g"])# Commodities
GM = String.(HHdata["sets"]["gm"])# Subset of goods
M =  String.(HHdata["sets"]["m"]) #transport and trade margins
R =  String.(HHdata["sets"]["r"]) # Regions
Q =  String.(HHdata["sets"]["q"]) # Alias of Regions 
H =  String.(HHdata["sets"]["h"]) # HHs, hh1-5
S =  String.(HHdata["sets"]["s"]) # Sectors (same as commodities)

## Data
a0 = HHdata["parameters"]["a0"]
c0_h = HHdata["parameters"]["c0_h"] # Total consumption demand per hh, region
cd0_h = HHdata["parameters"]["cd0_h"] # Consumption demand of each good per hh, region
dd0 = HHdata["parameters"]["dd0"] # local commodity demand
dm0 = HHdata["parameters"]["dm0"]# Margin demand per region, per regionally traded good
esubL = HHdata["parameters"]["esubL"] # Demand elasticity between leisure and consuption per hh, per region
etaK = HHdata["parameters"]["etaK"] # Transformation elasticity between regions and sectors.
fint0 = HHdata["parameters"]["fint0"]# 0 so not used. Is for ...?
fsav0 = HHdata["parameters"]["fsav0"] # 0 foreign savings?
g0 = HHdata["parameters"]["g0"] # Government demand for goods from each sector/region 
govdef0 = HHdata["parameters"]["govdef0"] # Government foriegn deficit
hhtrn0 = HHdata["parameters"]["hhtrn0"] # Household transfers:: hucval:unemployment compensation, hwcval:workers compensation, hssval:social security, 
#hssival:supplemental security, hpawval:public assistance or welfare, hvetval:veterans benefits, hsurval:survivors income, hdisval:disability, 
#hedval:educational assistance, hcspval:child support, hfinval:financial assistance, medicare, medicaid	
i0 = HHdata["parameters"]["i0"] # Investment demand
id0 = HHdata["parameters"]["id0"] # Intermediate demand
kd0 = HHdata["parameters"]["kd0"] # Capital demand
ke0 = HHdata["parameters"]["ke0"] # Household interest payments
ld0 = HHdata["parameters"]["ld0"] # Labour demand
le0 = HHdata["parameters"]["le0"] # Household labor endowment [region, region, hh] ? includes cross-border labor? 
ls0 = HHdata["parameters"]["ls0"] # Labor supply (net) [regions, hh]
lsr0 = HHdata["parameters"]["lsr0"] # Leisure demand [regions, hh]
m0 = HHdata["parameters"]["m0"] # Imports [regions, sectors]
md0 = HHdata["parameters"]["md0"] # Total margin demand [regions, margins, sectors]
nd0 = HHdata["parameters"]["nd0"] # Regional demand from national market [regions, sectors]
nm0 = HHdata["parameters"]["nm0"] # Margin demand from national market [regions, sectors, margins]
rx0 = HHdata["parameters"]["rx0"] # Re-exports of goods and services [regions, sectors]
s0 = HHdata["parameters"]["s0"] # Aggregate supply [regions, sectors]
sav0 = HHdata["parameters"]["sav0"] # Household saving [regions, hh]
ta0 = HHdata["parameters"]["ta0"] # Tax net subsidy rate on intermediate demand [regions, sectors]
tfica0 = HHdata["parameters"]["tfica0"] # Household FICA labor tax rate [regions, hh]
tk0 = HHdata["parameters"]["tk0"] # Capital tax rate [regions]
tl_avg0 = HHdata["parameters"]["tl_avg0"] # Household average labor tax rate [regions, hh]
tl0 = HHdata["parameters"]["tl0"] # Household marginal labor tax rate [regions, hh]
tm0 = HHdata["parameters"]["tm0"] # Import tariff [regions, sectors]
totsav0 = HHdata["parameters"]["totsav0"] # Aggregate savings
trn0 = HHdata["parameters"]["trn0"] # Household transfer payments [regions, hh]
ty0 = HHdata["parameters"]["ty0"] # Output tax on production [regions, sectors]
x0 = HHdata["parameters"]["x0"] # 	Exports of goods and services [regions, sectors]
xd0 = HHdata["parameters"]["xd0"] # Regional supply to local market
xn0 = HHdata["parameters"]["xn0"] # Regional supply to national market
yh0 = HHdata["parameters"]["yh0"] # 0: Household production
ys0 = HHdata["parameters"]["ys0"] # Sectoral/intermediate supply

HH = MPSGEModel()

@parameters(HH, begin
    ta[R,G], 0, (description = "Consumption Tax")
    ty[R,S], 0, (description = "Production Tax")
    tm[R,G], 0, (description = "Import Tax")
    tk[R,S], 0, (description = "Capital Tax")
    tfica[R,H], 0, (description = "FICA Labor Shares")
    tl[R,H], 0, (description = "Marginal Labor Tax")
end)

for r∈R, g∈G
    set_value!(ta[r,g], ta0[r,g])
    set_value!(ty[r,g], ty0[r,g])
    set_value!(tm[r,g], tm0[r,g])
    set_value!(tk[r,g], tk0[r])
end

for r∈R,h∈H
    set_value!(tfica[r,h], tfica0[r,h])
    set_value!(tl[r,h], tl0[r,h])
end

@sectors(HH, begin
    Y[R,S], (description = "Production")
    X[R,G], (description = "Disposition")
    A[R,G], (description = "Absorption")
    LS[R,H], (description = "Labor Supply")
    KS, (description = "Aggregate Capital Stock")
    C[R,H], (description = "Household Consumption")
    MS[R,M], (description = "Margin Supply")
    INVEST_DEMAND
    GOVT_DEMAND
end)

@commodities(HH, begin
    PA[R,G], (description = "Regional Market (input)")
    PY[R,G], (description = "Regional Market (output)")
    PD[R,G], (description = "Local Market Price")
    RK[R,S], (description = "Sectoral Rental Rate")
    RKS, (description = "Capital Stock")
    PM[R,M], (description = "Margin Price")
    PC[R,H], (description = "Consumption Index")
    PN[G], (description = "National Market Price for goods")
    PLS[R,H], (description = "Leisure Price")
    PL[R], (description = "Regional Wage Rate")
    PK, (description = "Aggregate return to capital")
    PFX, (description = "Foreign Exchange Rate")
    INVEST_COMMODITY
    GOVT_COMMODITY
end)

@consumers(HH, begin
    RA[R,H], (description = "Representative Agent")
    NYSE, (description = "Aggregate Capital Owner")
    INVEST, (description = "Aggregate Investor")
    GOVT, (description = "Aggregate Government")
    #ROW # if fint0 !=0
end)

@auxiliaries(HH, begin
    SAVRATE, (description = "Domestic Savings Rate")
    TRANS, (description = "Budget balance rationing variable")
    SSK, (description = "Steady-state capital stock")
    CPI, (description = "Consumer Price Index")
end)

set_start_value(SAVRATE, 1)
set_start_value(TRANS, 1)
set_start_value(SSK, 1)
set_start_value(CPI, 1)

for r∈R,ŝ∈S
    @production(HH, Y[r,ŝ], [t=0, s=0, va=>s=1], begin
        [@output(PY[r,g], ys0[r,ŝ,g], t, taxes = [Tax(GOVT, ty[r,ŝ])], reference_price = 1-ty0[r,ŝ]) for g∈G]...
        [@input(PA[r,g], id0[r,g,ŝ], s) for g∈G]...
        @input(PL[r], ld0[r,ŝ], va)
        @input(RK[r,ŝ], kd0[r,ŝ], va, taxes=[Tax(GOVT, tk[r,ŝ])], reference_price = 1+tk0[r])
    end)
end

for r∈R,g∈G
    @production(HH, X[r,g], [t=4, s=1], begin
        @output(PFX, x0[r,g] - rx0[r,g], t)
        @output(PN[g], xn0[r,g], t)
        @output(PD[r,g], xd0[r,g], t)
        @input(PY[r,g], s0[r,g], s)
    end)
end

for r∈R, g∈G
    @production(HH, A[r,g], [t=0, s=0, dm=>s=4, d=>dm=2], begin
        @output(PA[r,g], a0[r,g], t, taxes = [Tax(GOVT, ta[r,g])], reference_price = 1-ta0[r,g])
        @output(PFX, rx0[r,g], t)
        @input(PN[g], nd0[r,g], d)
        @input(PD[r,g], dd0[r,g], d)
        @input(PFX, m0[r,g], dm, taxes = [Tax(GOVT, tm[r,g])], reference_price = 1+tm0[r,g])
        [@input(PM[r,m], md0[r,m,g], s) for m∈M]...
    end)
end

for r∈R, m∈M
    @production(HH, MS[r,m], [t=0, s=0], begin
        @output(PM[r,m], sum(md0[r,m,gm] for gm∈GM), t)
        [@input(PN[gm], nm0[r,gm,m], s) for gm∈GM]...
        [@input(PD[r,gm], dm0[r,gm,m], s) for gm∈GM]...
    end)
end

for r∈R,h∈H
    @production(HH, C[r,h], [t=0, s=1], begin
        @output(PC[r,h], c0_h[r,h], t)
        [@input(PA[r,g], cd0_h[r,g,h], s) for g∈G]...
    end)
end

for r∈R, h∈H
    @production(HH, LS[r,h], [t=0, s=1], begin
        [@output(PL[q], le0[r,q,h], t, taxes=[Tax(GOVT, tl[r,h] + tfica[r,h])], reference_price = 1-tl0[r,h]-tfica0[r,h]) for q∈Q]...
        @input(PLS[r,h], ls0[r,h], s)
    end)
end

@production(HH, KS, [t=etaK, s=1], begin
    [@output(RK[r,s], kd0[r,s], t) for r∈R, s∈S]...
    @input(RKS, sum(kd0[r,s] for r∈R, s∈S), s)
end)


for r∈R,h∈H
    @demand(HH, RA[r,h], begin
        @final_demand(PC[r,h], c0_h[r,h])
        @final_demand(PLS[r,h], lsr0[r,h])
        @endowment(PLS[r,h], ls0[r,h]+lsr0[r,h])
        @endowment(PFX, TRANS*sum(hhtrn0[r,h,trn] for trn∈TRN))
        @endowment(PLS[r,h], (tl[r,h] - tl_avg0[r,h])*sum(le0[r,q,h] for q∈Q))
        @endowment(PK, ke0[r,h])
        @endowment(PFX, -sav0[r,h]*SAVRATE)
    end, elasticity = esubL[r,h])
end



@demand(HH, NYSE, begin
    @final_demand(PK, sum(yh0[r,g] for r∈R, g∈G)+ sum(kd0[r,s] for r∈R, s∈S)*SSK)#)
    [@endowment(PY[r,g], yh0[r,g]) for r∈R, g∈G]...
    @endowment(RKS, SSK*sum(kd0[r,s] for r∈R, s∈S))
end)


@production(HH, INVEST_DEMAND, [t=0, s=0], begin
    @output(INVEST_COMMODITY, sum(i0[r,g] for r∈R, g∈G), t)
    [@input(PA[r,g], i0[r,g], s) for r∈R, g∈G]...
end)

@demand(HH, INVEST, begin
    #[@final_demand(PA[r,g], i0[r,g]) for r∈R, g∈G]...
    @final_demand(INVEST_COMMODITY, sum(i0[r,g] for r∈R, g∈G))
    @endowment(PFX, totsav0*SAVRATE)
    @endowment(PFX, fsav0)
end)#, elasticity = 0)


@production(HH, GOVT_DEMAND, [t=0, s= 1], begin
    @output(GOVT_COMMODITY, sum(g0[r,g] for r∈R,g∈G), t)
    [@input(PA[r,g], g0[r,g], s) for r∈R, g∈G]...

end)    


@demand(HH, GOVT, begin
    #[@final_demand(PA[r,g], g0[r,g]) for r∈R, g∈G]...
    @final_demand(GOVT_COMMODITY, sum(g0[r,g] for r∈R,g∈G))
    @endowment(PFX, -TRANS*sum(trn0[r,h] for r∈R, h∈H))
    @endowment(PFX, govdef0)
    [@endowment(PLS[r,h], -(tl[r,h] - tl_avg0[r,h])*sum(le0[r,q,h] for q∈Q)) for r∈R, h∈H]...
end)


#@demand(HH, ROW, begin
#    @final_demand(PFX, fint0)
#    @endowment(PK, fint0)
#end)

@aux_constraint(HH, SSK,
    sum(i0[r,g]*PA[r,g] for r∈R, g∈G) - sum(i0[r,g] for r∈R, g∈G) *RKS
)

@aux_constraint(HH, SAVRATE,
    INVEST - sum(PA[r,g]*i0[r,g] for r∈R, g∈G)*SSK
)

@aux_constraint(HH, TRANS,
    GOVT - sum(PA[r,g]*g0[r,g] for r∈R, g∈G)
)

@aux_constraint(HH, CPI, 
    CPI - sum(PC[r,h]*c0_h[r,h] for r∈R, h∈H)/sum(c0_h[r,h] for r∈R, h∈H)
)


fix(HH[:PFX], 1)

sttime=time()
solve!(HH, cumulative_iteration_limit=0)
endtime=time()
println(round((endtime-sttime),digits=4)," seconds= ",round((endtime-sttime)/60,digits=2)," minutes")
bnchHH = generate_report(HH)
set_silent(HH)
# unset_silent(HH)
# for tf in 0.26:0.01:0.35
# println(tf)
# end
tf = 0.1
    for r∈R,g∈G
        if tm0[r,g] < tf
        set_value!(HH[:tm][r,g],  tf)
        end
    end
sttime=time()
solve!(HH)
endtime=time()
println(tf," ",round((endtime-sttime),digits=4)," seconds= ",round((endtime-sttime)/60,digits=2)," minutes")
# push!(CPIarray2, [tf value(CPI)])
# end
# CPIarray2

tarrif_1 = generate_report(HH)

function variable_to_dataframe(var::MPSGE.MPSGEIndexedVariable, columns::Symbol...)
    X = Iterators.product(axes(var)...)

    L = []
    for a in X
        push!(L, a)
    end

    df = DataFrame(
        domain = L,
        value = [value(var[a...]) for a in L]
    ) |>
    x -> transform(x,
        :domain => ByRow(identity) => [columns...]
    ) |>
    x -> select(x, [columns..., :value]) 

    return df
end


import PlotlyJS: plot as plotly
import PlotlyJS.savefig
# import PlotlyJS: plot


var = HH[:PC]
df = variable_to_dataframe(var, :state, :household)
PLwage = variable_to_dataframe(HH[:PL], :state)

### Set up for the HH by state income chart
X = df |> 
    x -> innerjoin(
        x,
        PLwage,
        on = :state,
        renamecols = "" => "_wage"
    ) |>
    x -> transform(x,
        [:value, :value_wage] => ByRow((x,y) -> (y/x-1)*100) => :real_value
    ) #

# =
    bounds = X |>
    x -> combine(x,
        :real_value => maximum => :max,
        :real_value => minimum => :min
    ) |>
    x -> transform(x,
    :max => ByRow(y -> 0) => :max,     # :max => ByRow(y -> round(y,digits = 2)+5*y/1_000) => :max, alternative with different y axes range
    :min => ByRow(y -> y-.1) => :min   # :min => ByRow(y -> round(y,digits = 2)-5*y/1_000) => :min, alternative with different y axes range
    )
using PlotlyJS

### Bar chart of real income effect per HH per state
plt = X |>
    x -> sort(x, :real_value) |>
    x ->
     transform(x,
        :household => ByRow(y -> replace(y,"hh1"=>"HH1: lowest income quintile","hh2"=>"HH2","hh3"=>"HH3","hh4"=>"HH4","hh5"=>"HH5: highest income quintile"))=>:household) |>
    X -> plotly(
        # X -> PlotlyJS.plot(
    X, 
    
    x = :state, 
    y = :real_value, 
    color = :household, 
    type = :bar,     
    labels = attr(household="Household types"),
    Layout(
        font_family="Garamond",
        font_size=14,
        title = var.description,
        legend=attr(x=0.05, y=.95),
        legend_title_font_color="black",
        xaxis=attr(tickangle=90),
        xaxis_title = "States",
        yaxis_title = "Real Value % (local price / local wage)",
        yaxis_range = [bounds[1,:min], bounds[1,:max]]
    ))
  
### Set up for trade exposure plots  
    getorder1 =  X |>
    x -> sort(x, :real_value)
    getorder =  X |>
    x -> filter(x-> x.household=="hh1",getorder1 )
    getorder[:,:state]

### BEA data, with State post office abbreviations added to CSV    
GSP = CSV.read(joinpath(@__DIR__,"GSPtotals.csv"), DataFrame)
# ### Convert to NamedArray, select just datayear. String on States key to avoid string3 mismatch
# GSP = NamedArray(Matrix(GSPdf[:,["GeoName","TableName","Description",string(datayear)]]),(String.(GSPdf[:,"States"]),names(GSPdf[!,["GeoName","TableName","Description",string(datayear)]])),("regions","data"))
GDPyear = DataFrame([GSP[:,"States"] GSP[:,string(datayear)].*10^-3],[:States, Symbol("GSP_"*string(datayear))])
    mExposure1 = DataFrame(States=names(m0)[1], imports=convert(Array,sum(m0[R,:], dims=2)[:,"sum(s)"]))
    xExposure1 = DataFrame(States=names(x0)[1], exports=convert(Array,sum(x0[R,:], dims=2)[:,"sum(s)"]))
    NetTradeExposure = innerjoin(mExposure1,xExposure1,GDPyear, on = :States)
    NetTradeExposure.net_imports = NetTradeExposure.imports .- NetTradeExposure.exports
    NetTradeExposure.net_imports_GSP = NetTradeExposure.net_imports ./ NetTradeExposure.GSP_2021

    mExposure =              mExposure1[indexin(getorder[:,:state], mExposure1.States),:]
    NetTradeExposure = NetTradeExposure[indexin(getorder[:,:state],NetTradeExposure.States),:]


### Plot of import value by state, order from 
plt = mExposure |>
    # x -> sort(x, :imports) |>
    X -> 
    plotly(
    X, 
    x = :States, 
    y = :imports, 
    type = :bar, 
    marker_color="lightslategrey",
    Layout(
        # color = :green,     # title = var.description,
        yaxis_title_text="Imports \$BillUS",
        barmode = "relative"
    #    yaxis_range = [1, 1.1]
       ) )  

### Plot of Net trade/imports (imports - exports) value by state, order from 
plt = NetTradeExposure |>
       # x -> sort(x, :imports) |>
       X -> 
   plotly(
       X, 
       x = :States, 
       y = :net_imports, 
       type = :bar, 
       marker_color="darkblue",
       Layout(
           color = :green,     # title = var.description,
           yaxis_title_text="Imports \$BillUS",
           barmode = "relative"
       #    yaxis_range = [1, 1.1]
          ) )  

using Plots.PlotMeasures
### Plot of Net trade/imports (imports - exports) value by state, order from 
xticks = 0.5:length(NetTradeExposure.States)-0.5; xticknames = NetTradeExposure.States; 
Plots.bar(NetTradeExposure.States,NetTradeExposure.net_imports_GSP, color=:darkblue,ylim=(-.1,.1), 
size=(1700,800),
 bar_width=0.2,
   legend=false, 
   margin = 11mm,
   ylabel="Trade as % of Gross State Product",
xlabel= "States",guidefont=font(16,"Palatino Roman"),
   yticks=([-0.1, -0.05, 0, 0.05, 0.1], ["-10%", "-5%", "0%", "5%", "10%"]),ytickfont=font(12,"Palatino Roman"), xlim=(0,length(NetTradeExposure.States)))
# Set xticks and xticknames to ensure all labels display
xticks!(xticks, xticknames, xtickfont=font(12,"Palatino Roman"),xrotation=-90)

# savefig(plt,joinpath("C:/Users/Eli/OneDrive - Stanford/Box_2024-11-15/CGE/MPSGE-JL/First MPSGE paper","./LocalCommodPrices.png"))

Gamsresult = CSV.read(joinpath(@__DIR__,"HHgams_output.csv"), DataFrame) |>
    x -> rename(x, :Column1 => :state) |>
    x -> stack(x, Not(:state), variable_name = :household, value_name = :gams)

outerjoin(
    X |>
        x -> select(x, :state, :household, :real_value=> :julia),
    Gamsresult,
    on = [:state, :household]
) |>
x -> transform(x,
    [:julia, :gams] => ByRow((x,y) -> (y-x)) => :diff
) |>
x -> sort(x, :diff) |>
x -> CSV.write("CompareHHjlandGAMSoutput.csv", x)

# plotly(
#     df, 
#     x = :state, 
#     y = :value, 
#     color = :household, 
#     type = :bar, 
#     Layout(
#         title = var.description,
#         yaxis_range = [bounds[1,:min], bounds[1,:max]])
#     )
using CategoricalArrays

var = HH[:PD]
df = variable_to_dataframe(var, :state, :good)
df.good = Symbol.(df.good)
goods = [:pet, :agr,:fbp]

df.state = CategoricalArray(df.state, levels = getorder[:,:state])
sort!(df, :state)
plt = df |>
    x -> subset(x,
        :good => ByRow(in(goods))) |>
    x -> transform(x, :good => ByRow(y -> string(y))=>:good) |>
    x -> transform(x, :good => ByRow(y -> replace(y,"agr"=>"Agricultural goods","fbp"=>"Food, beverage and tobacco","pet"=>"Petroleum and coal products"))=>:good) |>
    x -> rename(x, :good => :Commodities) |>
    x -> transform(x, :value => ByRow(y -> y-1)=>:value) |>

    X -> 
plotly(
    X, 
    x = :state, 
    y = :value, 
    color = :Commodities, 
    type = :bar,
        Layout(
            font_family="Garamond",
        font_size=14,
        title = "Local Market Prices",
        sort = getorder[:,:state],
        xaxis_title_text="States",
        yaxis_title_text="Change in Price",
        marker_color=["green", "brown", "darkgrey"],
        legend=attr(x=0.05, y=.95),
       yaxis_range = (0, maximum(X.value)+.019),
    yaxis_tickvals = 0:0.01:0.10,
    yaxis_ticktext = [string(Int64(round(v * 100, digits=0))) * "%" for v in 0:0.01:0.1],        

       )
    )  
    restyle!(plt, marker_color=["green", "lightbrown", "darkgrey"])
	plt
## Testing for commodities with all or most prices non-zero for use as an example
# prods_nonzerolocalprice = DataFrame(state=String[], commods=String[],bnchprice=Float64[])
# # testdf = DataFrame(state=String[], commods=String[],bnchprice=Float64[])
# for r in R
#        for g in G
#         # g = "gas"
#        if dd0[r,g]+sum(dm0[r,g,M])+xd0[r,g] >0
#        push!(prods_nonzerolocalprice,[r g dd0[r,g]+sum(dm0[r,g,M])+xd0[r,g]])
#     #    push!(testdf,[r g dd0[r,g]+sum(dm0[r,g,M])+xd0[r,g]])
#        println(r,",",g,": ",dd0[r,g])
#        end
#        end
#        end
#        show(stdout,"text/plain",sort(prods_nonzerolocalprice,:commods))
#        show(stdout,"text/plain",sort(leftjoin(combine(groupby(prods_nonzerolocalprice,  :commods), nrow => :count),Sectors, on = :commods=>:index),:count))

oilimports=DataFrame(States=names(m0)[1],imports=convert(Array,m0[:,"oil"]))
oilexports=DataFrame(States=names(x0)[1],exports=convert(Array,x0[:,"oil"]))
plt = innerjoin(oilimports,oilexports, on = :States)|>
    x -> transform(x,[:imports, :exports] => ByRow((a,b) -> a - b) => :net_imports) |>
    X ->
plotly(
    X, 
    x = :States, 
    y = :net_imports, 
    
    type = :bar, 
    marker_color="darkgrey",
    Layout(
        yaxis_title_text="Oil Imports \$10BUS",
        
       )
    ) 