import numpy as np
import math
import matplotlib.pyplot as plt
import seaborn as sns
plt.style.use('seaborn-white')
import pandas as pd
from matplotlib import animation, rc
import torch.nn.functional as F
import torch
import torch.nn as nn
import torch.optim as optim
plt.rcParams.update({'pdf.fonttype': 'truetype'})
import pickle
pc2 = pickle.load(open('price.pkl','rb'))
from PIL import Image 
import streamlit as st

st.title("Price Optimization")

def to_tensor(x):
    return torch.from_numpy(np.array(x).astype(np.float32))
def prediction(price_max,price_step,policy_net):
    price_grid = np.arange(price_step, price_max, price_step)
    sample_state = [0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0., \
                1.,     0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.]
    Q_s = policy_net(to_tensor(sample_state))
    a_opt = Q_s.max(0)[1].detach()
    plt.figure(figsize=(16, 5))
    plt.xlabel("Price action ($)")
    plt.ylabel("Q ($)")
    plt.bar(price_grid, Q_s.detach().numpy(), color='crimson',  width=6, alpha=0.8)
    plt.savefig('price.png')    
    return price_grid[a_opt]
    
def fun():
    st.header("Optimal Price Action")
    st.subheader(str(a))
    
    return 
st.header("Enter the Specification")
max_value = st.number_input('Enter the Maximum Value of Price',min_value=50,value = 500,step=1)
step = st.number_input('Enter the Price step',min_value = 10,value = 10,step=1)
a = prediction(max_value,step,pc2)
if st.button('Predict'):
    fun()
    image = Image.open('price.png')
    st.image(image,caption = 'Price Optimization',width =1000)