# app.py
from flask import Flask, render_template, request
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression

app = Flask(__name__)

# Load data from the CSV file
data = pd.read_csv("student-mat.csv")

# Map the 'health' variable to three classes
data['health_class'] = pd.cut(data['health'], bins=[-1, 2, 3, 5], labels=['плохое', 'среднее', 'хорошее'])

# Define features and target variable
features = ['Pstatus', 'guardian', 'internet', 'romantic', 'famrel', 'freetime', 'goout', 'Dalc', 'Walc', 'absences']
X = pd.get_dummies(data[features])
y = data['health_class']

# Split the data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.01, random_state=42)

# Create and train the logistic regression model
model = LogisticRegression(max_iter=1000)
model.fit(X_train, y_train)

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        # Get data from the form
        form_data = request.form
        input_data = pd.DataFrame({
            'Pstatus': [form_data['Pstatus']],
            'guardian': [form_data['guardian']],
            'internet': [form_data['internet']],
            'romantic': [form_data['romantic']],
            'famrel': [int(form_data['famrel'])],
            'freetime': [int(form_data['freetime'])],
            'goout': [int(form_data['goout'])],
            'Dalc': [int(form_data['Dalc'])],
            'Walc': [int(form_data['Walc'])],
            'absences': [int(form_data['absences'])]
        })

        # One-hot encode the categorical variables consistently
        input_data = pd.get_dummies(input_data)
        # Ensure the input features match the features seen during training
        input_data = input_data.reindex(columns=X.columns, fill_value=0)

        # Make prediction
        prediction = model.predict(input_data)[0]

        return render_template("index.html", prediction_result=prediction)

    return render_template("index.html")

if __name__ == "__main__":
    app.run(host="localhost", port=5000)