# 返回图片测试
from sklearn.linear_model import LinearRegression
from sklearn.neural_network import MLPRegressor
import lightgbm as lgb
from xgboost import XGBRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import pandas as pd
from fastapi.middleware.cors import CORSMiddleware  # 跨域
from fastapi import FastAPI, Response, BackgroundTasks
import json
import matplotlib.pyplot as plt
import io
import matplotlib
matplotlib.use('AGG')

app = FastAPI()

# 配置跨域白名单
origins = [
    "http://127.0.0.1:5500"
]

# 图片测试
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["POST"],
    allow_headers=["*"],
)

def create_img():
    plt.rcParams['figure.figsize'] = [7.50, 3.50]
    plt.rcParams['figure.autolayout'] = True
    plt.plot([1, 2])
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    plt.close()
    return img_buf

@app.get('/png')
async def get_img(background_tasks: BackgroundTasks):
    img_buf = create_img()
    # get the entire buffer content
    # because of the async, this will await the loading of all content
    bufContents: bytes = img_buf.getvalue()
    background_tasks.add_task(img_buf.close)
    headers = {'Content-Disposition': 'inline; filename="out.png"'}
    return Response(bufContents, headers=headers, media_type='image/png')


# 机器学习测试
# 多元线性回归
@app.post("/mlr")
async def mlr():
    # 引入excel数据
    # 可能存在空行问题，dropna(axis=0)删除空行
    df = pd.read_csv("./1.csv").dropna(axis=0)
    # 取xy
    x = df.iloc[:, :12]
    # y = df.loc[:, "BSR"]
    y = df.loc[:, "SBR"]
    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.3, random_state=0)

    # 标准化
    standardscaler = StandardScaler()
    standardscaler.fit(x_train)
    x_train = standardscaler.transform(x_train)
    x_test = standardscaler.transform(x_test)

    # 多元线性回归
    model = LinearRegression()
    model.fit(x_train, y_train)

    # 测试样本预测
    y_pred = model.predict(x_test)
    return json.dumps(y_pred.tolist())

# 随机森林
@app.post("/rf")
async def rf():
    # 引入excel数据
    # 可能存在空行问题，dropna(axis=0)删除空行
    df = pd.read_csv("./1.csv").dropna(axis=0)
    # 取xy
    x = df.iloc[:, :12]
    # y = df.loc[:, "BSR"]
    y = df.loc[:, "SBR"]
    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.3, random_state=0)

    # 标准化
    standardscaler = StandardScaler()
    standardscaler.fit(x_train)
    x_train = standardscaler.transform(x_train)
    x_test = standardscaler.transform(x_test)

    # 随机森林
    model = RandomForestRegressor()
    model.fit(x_train, y_train)

    # 测试样本预测
    y_pred = model.predict(x_test)
    return json.dumps(y_pred.tolist())

# BP神经网络
@app.post("/bpn")
async def rf():
    # 引入excel数据
    # 可能存在空行问题，dropna(axis=0)删除空行
    df = pd.read_csv("./1.csv").dropna(axis=0)
    # 取xy
    x = df.iloc[:, :12]
    # y = df.loc[:, "BSR"]
    y = df.loc[:, "SBR"]
    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.3, random_state=0)

    # 标准化
    standardscaler = StandardScaler()
    standardscaler.fit(x_train)
    x_train = standardscaler.transform(x_train)
    x_test = standardscaler.transform(x_test)

    # BP神经网络
    model = MLPRegressor(hidden_layer_sizes=(10,), random_state=10,learning_rate_init=0.1)  # BP神经网络回归模型
    model.fit(x_train,y_train)  # 训练模型

    # 测试样本预测
    y_pred = model.predict(x_test)
    return json.dumps(y_pred.tolist())

# XGBoost
@app.post("/xgboost")
async def rf():
    # 引入excel数据
    # 可能存在空行问题，dropna(axis=0)删除空行
    df = pd.read_csv("./1.csv").dropna(axis=0)
    # 取xy
    x = df.iloc[:, :12]
    # y = df.loc[:, "BSR"]
    y = df.loc[:, "SBR"]
    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.3, random_state=0)

    # 标准化
    standardscaler = StandardScaler()
    standardscaler.fit(x_train)
    x_train = standardscaler.transform(x_train)
    x_test = standardscaler.transform(x_test)

    # XGBoost
    model = XGBRegressor(max_depth=5, learning_rate=0.1, n_estimators=160, objective='reg:gamma')
    model.fit(x_train, y_train)

    # 测试样本预测
    y_pred = model.predict(x_test)
    return json.dumps(y_pred.tolist())

# LightGBM
@app.post("/lightgbm")
async def rf():
    # 引入excel数据
    # 可能存在空行问题，dropna(axis=0)删除空行
    df = pd.read_csv("./1.csv").dropna(axis=0)
    # 取xy
    x = df.iloc[:, :12]
    # y = df.loc[:, "BSR"]
    y = df.loc[:, "SBR"]
    # 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.3, random_state=0)

    # 标准化
    standardscaler = StandardScaler()
    standardscaler.fit(x_train)
    x_train = standardscaler.transform(x_train)
    x_test = standardscaler.transform(x_test)

    # LightGBM
    model = lgb.LGBMRegressor(objective='regression',num_leaves=31,learning_rate=0.05,n_estimators=20)
    model.fit(x_train, y_train)

    # 测试样本预测
    y_pred = model.predict(x_test)
    return json.dumps(y_pred.tolist())