import pandas as pd
fp = "../data/regimed_coffee_prices.csv"
df = pd.read_csv(fp)
df

regimes = df["regime"].unique().tolist()

for r in regimes:
    reg_select = (df["regime"] == r) # select each regime
    df_reg = df[reg_select]
    # the index contains the indices of each regime, discretize the prices with qcut and set
    df.loc[df_reg.index, "price"] = pd.qcut(df_reg["cents_per_lb"], 3, labels=["L", "M", "H"])

    

    

for r in regimes:
    reg_select = (df["regime"] == r) # select each regime
    df_reg = df[reg_select]
    # maintain regime point count to mark the previous price for the first entry in each regime as na
    rpc = 0
    for ri, row in df_reg.iterrows():
        if rpc == 0 :
            rpc += 1
            continue
        else:
            df.loc[ri, "previous_price"] = df.loc[ (ri -1), "price"]
            rpc += 1

sum(df["previous_price"].isna()) 

fp = "../data/regimed_coffee_prices.csv"
df.to_csv(fp, index=False)

matrix_dict = {}
for r in regimes:
    reg_select = (df["regime"] == r) # select each regime
    df_reg = df[reg_select]
    # the index contains the indices of each regime, discretize the prices with qcut and set
    df_sm = pd.crosstab(df_reg.price, df_reg.previous_price)
    # the next step normalizes the entry in each row by the row sum
    df_sm = df_sm.div(df_sm.sum(axis=1), axis=0).round(3)
    matrix_dict[r] = df_sm
    fp = "../data/stochastic_matrix_coffee_price-regime-" + r + ".csv"
    df_sm.to_csv(fp, index=True)
    

matrix_dict["R-5"]

import plotly.express as px

fig = px.imshow(matrix_dict["R-5"], text_auto=True)
fig.update_layout(
    title={
        'text': "Stochastic Matrix for Region 5",
        'y':.95,
        'x':0.5,
        'xanchor': 'center',
        'yanchor': 'top'})

fig.show()