using JuMP, CairoMakie
import HiGHS

newspaper_cost = 5
customer_price = 7

supply = 42
demand = 40

# Daily Profit is equal to the newspapers sold less the cost we paid for it
function profit(demand, supply, customer_price, newspaper_cost) 
    return (min.(demand, supply) .* customer_price) .- (newspaper_cost .* supply)
end

# one day
single_profit = profit(demand, supply, customer_price, newspaper_cost)
println("Our Profit for one day with supply $supply is: \$ $single_profit")

# two days
three_day_demand = [40, 60, 25]
multi_day_profit = profit(three_day_demand, supply, customer_price, newspaper_cost)

println("Our Profit for three days with supply $supply is: \$ $sum(multi_day_profit)")

"""Plot Demand vs Profits"""
function plot_profit(demand, supply, profit; bar_labels=true)
        N = length(demand)
        fig = Figure(resolution=(800,400))

        ax1 = Axis(fig[1,1], title="$N-day Newspaper Demand")
        ax2 = Axis(fig[1,2], title="$N-day Newspaper Profit with Supply $supply")

        colors = [i >= 0 ? :grey : :crimson for i in profit]

        if bar_labels
                demand_labels = demand
                profit_labels = profit
        else
                demand_labels = ["" for i in demand]
                profit_labels = ["" for i in profit]
        end
        barplot!(ax1, demand, 
                bar_labels=demand_labels, 
                label_size=16,
                label_offset=-30,
                color=colors)
        barplot!(ax2, profit, 
                bar_labels=profit_labels, 
                label_size=16,
                label_offset=-30,
                color=colors)
        return fig
end

plot_profit(three_day_demand, supply, multi_day_profit)

# 0. Params
demand = 40
M = 100

# 1. Establish model
model = Model(HiGHS.Optimizer)
# Option: To Suppress all the printing, use `set_silent(model)`
set_silent(model)

# 2. Variables
@variable(model, s >= 0, Int)
@variable(model, x >= 0)
@variable(model, y, Bin)

# 3. Constraints
# convert the min(d,i) nonlinear constraint to linear
@constraint(model, x <= demand)
@constraint(model, x <= s)
@constraint(model, x >= demand - M*(1-y))
@constraint(model, x >= s - M * y)

# 4. Objective Function
@objective(model, Max, (customer_price * x) - (newspaper_cost * s))

# Optimize the model and print the value of $s$
optimize!(model)
value(s)

# What was the profit at this optimal value?
objective_value(model)


"""Wrap the optimization problem within a function"""
function optimize_problem(demand)
    # 0. Params
    M = 100
    n = length(demand)

    # 1. Establish model
    model = Model(HiGHS.Optimizer)
    # Option: To Suppress all the printing, use `set_silent(model)`
    set_silent(model)

    # 2. Variables
    @variable(model, s >= 0, Int)
    @variable(model, x[i=1:n] >= 0) # Vectorize x 
    @variable(model, y[i=1:n], Bin) # Vectorize y too

    # 3. Constraints
    # convert the min(d,i) nonlinear constraint to linear
    @constraint(model, x .<= demand) 
    @constraint(model, x .<= s)
    @constraint(model, x .>= demand .- M .* (1 .- y))
    @constraint(model, x .>= s .- M .* y)

    # # 4. Objective Function
    @objective(model, Max, sum((customer_price .* x) .- (newspaper_cost .* s)))
    optimize!(model)

    return value(s)
end

demand = [30, 30, 60, 25, 100]

optimal_supply = optimize_problem(demand)

optimal_profits = profit(demand, optimal_supply, customer_price, newspaper_cost)

plot_profit(demand, optimal_supply, optimal_profits)

using Random, Distributions

# set seed for reproducibility
Random.seed!(42)

# Specify the truncated distribution with truncated(dist, lower=0)
dist = truncated(Normal(40, 20), lower=0)
# Randomly sample 1000
sample_demand = rand(dist, 100)
# Plot
hist(sample_demand, 
     figure=(resolution=(500,300),),
     axis=(title=L"\text{100 Random Samples from Truncated Normal}(40,20)",))


random_sample_supply = optimize_problem(sample_demand)

optimal_profits = profit(sample_demand, random_sample_supply, customer_price, newspaper_cost)

plot_profit(sample_demand, random_sample_supply, optimal_profits, bar_labels=false)

# Disable progress bars and such
using Turing; Turing.setprogress!(false);

demand_seen = sample_demand[1:5]
demand_unseen = sample_demand[6:15];

@model function newspaper(demand)
    σ ~ Truncated(TDist(20), 10, 100)
    μ ~ Normal(mean(demand_seen), 20)
    for i in 1:length(demand)
        demand[i] ~ Truncated(Normal(μ, σ), 0, 100)
    end
    return demand
end

idata = sample(newspaper(demand_seen), NUTS(), 1000; progress=false)

typeof(idata)

import Plots
Plots.histogram(idata)

# You can also access the data like so:
samples = idata[:μ].data
hist(samples[1:end], figure=(resolution=(400,200),),)

keys(idata)

describe(idata)

using DataFramesMeta
first(DataFrame(idata), 10)

typeof(demand_seen)

@model function linear_reg(x, y, σ = 0.1)
    β ~ Normal(0, 1)
    for i ∈ eachindex(y)
        y[i] ~ Normal(β * x[i], σ)
    end
end;
σ = 0.1; 
f(x) = 2 * x + 0.1 * randn();
# Create x/y data
Δ = 0.1; 
xs_train = 0:Δ:10; ys_train = f.(xs_train);
xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

# add observations
m_train = linear_reg(xs_train, ys_train, σ);
chain_lin_reg = sample(m_train, NUTS(100, 0.65), 200);

# Set a function parameter to `Missing` to predict it
m_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)), σ);
predictions = predict(m_test, chain_lin_reg)

ys_pred = vec(mean(Array(group(predictions, :y)); dims = 1));
sum(abs2, ys_test - ys_pred) ≤ 0.1

# See predictions for 1 additional day of demand:
demand_preds = predict(newspaper(Vector{Union{Missing, Float64}}(undef, 1)), idata)

# We can plot the histogram directly with Plots, too
Plots.histogram(demand_preds)

# Then the posterior predictive samples would be:
demand_samples = demand_preds[:"demand[1]"].data |> vec

bayesian_optimum_supply = optimize_problem(demand_samples)

bayesian_optimal_profits = profit(demand_samples, bayesian_optimum_supply, customer_price, newspaper_cost)

plot_profit(demand_samples, bayesian_optimum_supply, bayesian_optimal_profits, bar_labels=false)