from dotenv import dotenv_values
fp = "../notebooks/.env"
config = dotenv_values(fp)

SOLVER = "cbc"

from amplpy import AMPL, ampl_notebook

ampl = ampl_notebook(
    modules=["cbc"],  # modules to install
    license_uuid=config["AMPL_UUID"],  # license to use
)  # instantiate AMPL object and register magics

%%writefile bakery.mod
param N; # number of samples

param c;
param p;
param h;

set indices := 1..N;
param xi{indices};

# first stage variable: x (amount of baguettes baked)
var x integer >=0;

var first_stage_profit = -c * x;

# second stage variables: y (sold) and z (unsold)
var y{indices} integer >=0;
var z{indices} integer >=0;

# second stage constraints
s.t. cantsellbaguettedonthave {i in indices}: y[i] <= xi[i];
s.t. baguettedonotdisappear {i in indices}: y[i] + z[i] == x;

var second_stage_profit = sum{i in indices}(p * y[i] - h * z[i]) / N;

# objective
maximize total_expected_profit: first_stage_profit + second_stage_profit;

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math
import scipy.stats as stats

from scipy.stats import nbinom
# Two-stage stochastic LP using SAA
# Setting parameters
## TL;DR version on pricing Profit Index. Costs = Selling Price, selling price is known from data, profit index is 1.7 
c = 52
p = 90
h = 1
# Two-stage stochastic LP using SAA

np.random.seed(1)
N = 5000
shape = 9.42
intensity = (1/6.97)
samples = nbinom.rvs(shape,intensity,size=N)
ampl = AMPL()
ampl.read("bakery.mod")

# Load the data
ampl.param["N"] = int(N)
ampl.param["xi"] = samples
ampl.param["c"] = c
ampl.param["p"] = p
ampl.param["h"] = h

ampl.solve(solver=SOLVER)
assert ampl.solve_result == "solved", ampl.solve_result

xval = ampl.var["x"].value()
total_expected_profit = (ampl.obj["total_expected_profit"].value()/100) # The divide by 100 is to convert from cent to euro


print(
    f"Approximate solution with fitted deman distribution using N={N:.0f} samples"
)
print(f"Approximate optimal solution: x = {xval:.2f} baguettes")
print(f"Approximate expected daily profit: {total_expected_profit:.2f}€")

zs = ampl.var["z"]

ys = ampl.var["y"]

ny = ys.numInstances()

yvals = [ ys[i].value() for i in range(1, ny +1)]

zvals = [ zs[i].value() for i in range(1, ny +1)]

df_soln = pd.DataFrame.from_records({"yi": yvals, "zi": zvals})

df_soln