import pickle
import numpy as np
import gym
from flask import Flask, render_template, request

app = Flask(__name__)

# Load the Q-learning model
with open("q-learning.pkl", "rb") as f:
    model = pickle.load(f)

# Initialize the Taxi-v3 environment
env = gym.make("Taxi-v3")

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    try:
        # Get the state from the user input
        state = int(request.form['state'])

        # Get the action from the model
        action = model.predict(np.array([state]))  # This assumes model.predict can work with the state as input

        # Return the result
        return render_template('index.html', state=state, action=action)
    except Exception as e:
        return f"Error: {str(e)}"

if __name__ == '__main__':
    app.run(debug=True)