from flask import Flask, request, jsonify
from sklearn.preprocessing import MinMaxScaler
import pandas as pd
import os
from app2 import predict_stock_codes

app = Flask(__name__)

# Load the prediction model
model = CustomModel()

# Define a function to prepare the data for prediction
def prepare_data(date):
    # Get the historical data for the given date
    data = bs.query_history_k_data_plus(
        "sz.000001",  # Shanghai Composite Index
        "date,open,high,low,close,volume",
        start_date="2005-05-30",
        end_date=date,
        frequency="d"
    )
    data_list = []
    while (data.error_code == '0') & data.next():
        data_list.append(data.get_row_data())
    data_df = pd.DataFrame(data_list, columns=data.fields)

    # Convert 'open' and 'close' columns to numeric type
    data_df['open'] = pd.to_numeric(data_df['open'])
    data_df['close'] = pd.to_numeric(data_df['close'])

    # Filter out stocks that meet the conditions
    data_df = data_df[(data_df["open"] >= 0.98 * data_df["close"].shift(1).fillna(0)) & (data_df["open"] <= 1.02 * data_df["close"].shift(1).fillna(0))] 
    data_df = data_df[(data_df["high"] == data_df["close"]) & (data_df["low"] == data_df["close"])]  # limit-up condition
    data_df = data_df[(data_df["open"]!= 0) & (data_df["close"]!= 0)]  # exclude zero prices

    # Scale the data using MinMaxScaler
    scaler = MinMaxScaler()
    data_df[['open', 'high', 'low', 'close', 'volume']] = scaler.fit_transform(data_df[['open', 'high', 'low', 'close', 'volume']])

    return data_df

# Define a route to predict the top 5 stock codes
@app.route('/predict', methods=['POST'])
def predict():
    date = request.json['date']
    data_df = prepare_data(date)
    if data_df.empty:
        return jsonify({'error': 'No data available for the given date'}), 400
    y_pred = model.predict(data_df)
    top_5_stocks = predict_stock_codes(y_pred, data_df)
    return jsonify({'top_5_stocks': top_5_stocks})

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