import datetime
import os
import pickle
import pandas as pd
from flask_sqlalchemy.session import Session
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split, cross_val_score
from sqlalchemy import text
from xgboost import XGBRegressor

from .database_models import Models, Datasets
from .config import Config


# 加载模型
def load_model(session, model_id):
    model = session.query(Models).filter(Models.ModelID == model_id).first()
    
    if not model:
        raise ValueError(f"Model with ID {model_id} not found.")
    with open(model.ModelFilePath, 'rb') as f:
        return pickle.load(f)


# 模型预测
def predict(session, input_data: pd.DataFrame, model_id):
    # 初始化模型
    ML_model = load_model(session, model_id)  # 根据指定的模型名加载模型
    # model = load_model(model_id)  # 根据指定的模型名加载模型
    predictions = ML_model.predict(input_data)
    return predictions.tolist()

# 计算模型评分
def calculate_model_score(model_info):
    # 加载模型
    with open(model_info.ModelFilePath, 'rb') as f:
        ML_model = pickle.load(f)
    # print("Model requires the following features:", model.feature_names_in_)
    # 数据准备
    if model_info.Data_type == 'reflux':  # 反酸数据集
        # 加载保存的 X_test 和 Y_test
        X_test = pd.read_csv('uploads/data/X_test_reflux.csv')
        Y_test = pd.read_csv('uploads/data/Y_test_reflux.csv')
        print(X_test.columns)  # 在测试时使用的数据的列名
        y_pred = ML_model.predict(X_test)
    elif model_info.Data_type == 'reduce':  # 降酸数据集
        # 加载保存的 X_test 和 Y_test
        X_test = pd.read_csv('uploads/data/X_test_reduce.csv')
        Y_test = pd.read_csv('uploads/data/Y_test_reduce.csv')
        print(X_test.columns)  # 在测试时使用的数据的列名
        y_pred = ML_model.predict(X_test)


    # 计算 R² 分数
    r2 = r2_score(Y_test, y_pred)
    return r2


def train_and_save_model(session, model_type, model_name, model_description, data_type, dataset_id=None):
    try:
        if not dataset_id:
            # 创建新的数据集并复制数据，此过程将不立即提交
            dataset_id = save_current_dataset(session, data_type, commit=False)

            if data_type == 'reflux':
                current_table = 'current_reflux'
            elif data_type == 'reduce':
                current_table = 'current_reduce'

            # 从current数据集表中加载数据
            dataset = pd.read_sql_table(current_table, session.bind)

        elif dataset_id:
            # 从新复制的数据集表中加载数据
            dataset_table_name = f"dataset_{dataset_id}"
            dataset = pd.read_sql_table(dataset_table_name, session.bind)

            if dataset.empty:
                raise ValueError(f"Dataset {dataset_id} is empty or not found.")

        if data_type == 'reflux':
            X = dataset.iloc[:, 1:-1]
            y = dataset.iloc[:, -1]
        elif data_type == 'reduce':
            X = dataset.iloc[:, 2:]
            y = dataset.iloc[:, 1]

        # 训练模型
        model = train_model_by_type(X, y, model_type)

        # 保存模型到数据库
        model_id = save_model(session, model, model_name, model_type, model_description, dataset_id, data_type)

        # 所有操作成功后，手动提交事务
        session.commit()
        return model_name, model_id, dataset_id
    except Exception as e:
        # 如果在任何阶段出现异常，回滚事务
        session.rollback()
        raise e  # 可选择重新抛出异常或处理异常



def save_current_dataset(session, data_type, commit=True):
    """
    创建一个新的数据集条目，并复制对应的数据类型表的数据，但不立即提交事务。

    Args:
    session (Session): SQLAlchemy session对象。
    data_type (str): 数据集的类型。
    commit (bool): 是否在函数结束时提交事务。

    Returns:
    int: 新保存的数据集的ID。
    """
    new_dataset = Datasets(
        Dataset_name=f"{data_type}_dataset_{datetime.datetime.now():%Y%m%d_%H%M%S}",
        Dataset_description=f"Automatically generated dataset for type {data_type}",
        Row_count=0,
        Status='pending',
        Dataset_type=data_type
    )

    session.add(new_dataset)
    session.flush()

    dataset_id = new_dataset.Dataset_ID
    source_table = data_type_table_mapping(data_type)
    new_table_name = f"dataset_{dataset_id}"
    session.execute(text(f"CREATE TABLE {new_table_name} AS SELECT * FROM {source_table};"))

    session.execute(text(f"UPDATE datasets SET status='Datasets upgraded success', row_count=(SELECT count(*) FROM {new_table_name}) WHERE dataset_id={dataset_id};"))

    if commit:
        session.commit()

    return dataset_id

def data_type_table_mapping(data_type):
    """映射数据类型到对应的数据库表名"""
    if data_type == 'reduce':
        return 'current_reduce'
    elif data_type == 'reflux':
        return 'current_reflux'
    else:
        raise ValueError("Invalid data type provided.")


def train_model_by_type(X, y, model_type):
    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    if model_type == 'RandomForest':
        # 随机森林的参数优化
        return train_random_forest(X_train, y_train)
    elif model_type == 'XGBR':
        # XGBoost的参数优化
        return train_xgboost(X_train, y_train)
    elif model_type == 'GBSTR':
        # 梯度提升树的参数优化
        return train_gradient_boosting(X_train, y_train)
    else:
        raise ValueError(f"Unsupported model type: {model_type}")


def train_random_forest(X_train, y_train):
    best_score = -float('inf')
    best_n_estimators = None
    best_max_depth = None
    random_state = 43

    # 筛选最佳的树的数量
    for n_estimators in range(1, 20, 1):
        model = RandomForestRegressor(n_estimators=n_estimators, random_state=random_state)
        score = cross_val_score(model, X_train, y_train, cv=5).mean()
        if score > best_score:
            best_score = score
            best_n_estimators = n_estimators

    print(f"Best number of trees: {best_n_estimators}, Score: {best_score}")

    # 在找到的最佳树的数量基础上，筛选最佳的最大深度
    best_score = 0  # 重置最佳得分，为最大深度优化做准备
    for max_depth in range(1, 5, 1):
        model = RandomForestRegressor(n_estimators=best_n_estimators, max_depth=max_depth, random_state=random_state)
        score = cross_val_score(model, X_train, y_train, cv=5).mean()
        if score > best_score:
            best_score = score
            best_max_depth = max_depth

    print(f"Best max depth: {best_max_depth}, Score: {best_score}")

    # 使用最佳的树的数量和最大深度训练最终模型
    best_model = RandomForestRegressor(n_estimators=best_n_estimators, max_depth=best_max_depth,
                                       random_state=random_state)

    # 传入列名进行训练
    best_model.fit(X_train, y_train)
    # 指定传入的特征名
    best_model.feature_names_in_ = X_train.columns
    return best_model


def train_xgboost(X_train, y_train):
    best_score = -float('inf')
    best_params = {'learning_rate': None, 'max_depth': None}
    random_state = 43

    for learning_rate in [0.01, 0.05, 0.1, 0.2]:
        for max_depth in range(3, 10):
            model = XGBRegressor(learning_rate=learning_rate, max_depth=max_depth, random_state=random_state)
            score = cross_val_score(model, X_train, y_train, cv=5).mean()
            if score > best_score:
                best_score = score
                best_params['learning_rate'] = learning_rate
                best_params['max_depth'] = max_depth

    print(f"Best parameters: {best_params}, Score: {best_score}")

    # 使用找到的最佳参数训练最终模型
    best_model = XGBRegressor(learning_rate=best_params['learning_rate'], max_depth=best_params['max_depth'],
                              random_state=random_state)
    best_model.fit(X_train, y_train)

    return best_model


def train_gradient_boosting(X_train, y_train):
    best_score = -float('inf')
    best_params = {'learning_rate': None, 'max_depth': None}
    random_state = 43

    for learning_rate in [0.01, 0.05, 0.1, 0.2]:
        for max_depth in range(3, 10):
            model = GradientBoostingRegressor(learning_rate=learning_rate, max_depth=max_depth, random_state=random_state)
            score = cross_val_score(model, X_train, y_train, cv=5).mean()
            if score > best_score:
                best_score = score
                best_params['learning_rate'] = learning_rate
                best_params['max_depth'] = max_depth

    print(f"Best parameters: {best_params}, Score: {best_score}")

    # 使用找到的最佳参数训练最终模型
    best_model = GradientBoostingRegressor(learning_rate=best_params['learning_rate'], max_depth=best_params['max_depth'],
                                           random_state=random_state)
    best_model.fit(X_train, y_train)

    return best_model

def save_model(session, model, model_name, model_type, model_description, dataset_id, data_type, commit=False):
    """
    保存模型到数据库，并将模型文件保存到磁盘。
    
    Args:
        session: 数据库会话
        model: 要保存的模型对象
        model_name: 模型的名称
        model_type: 模型的类型
        model_description: 模型的描述信息
        dataset_id: 数据集ID
        data_type: 数据类型
        commit: 是否提交事务
    
    Returns:
        int: 返回保存的模型ID
    """
    prefix_dict = {
        'RandomForest': 'rf_model_',
        'XGBR': 'xgbr_model_',
        'GBSTR': 'gbstr_model_'
    }
    prefix = prefix_dict.get(model_type, 'default_model_')

    try:
        # 从配置中获取保存路径
        model_save_path = Config.MODEL_SAVE_PATH
        
        # 确保路径存在
        os.makedirs(model_save_path, exist_ok=True)

        # 获取当前时间戳
        timestamp = datetime.datetime.now().strftime('%m%d_%H%M')

        # 拼接完整的文件名
        file_name = os.path.join(model_save_path, f'{prefix}{timestamp}.pkl')

        # 保存模型到文件
        with open(file_name, 'wb') as f:
            pickle.dump(model, f)
        print(f"模型已保存至: {file_name}")

        # 创建模型数据库记录
        new_model = Models(
            Model_name=model_name,
            Model_type=model_type,
            Description=model_description,
            DatasetID=dataset_id,
            Created_at=datetime.datetime.now(),
            ModelFilePath=file_name,
            Data_type=data_type
        )

        # 添加记录到数据库
        session.add(new_model)
        session.flush()

        return new_model.ModelID

    except Exception as e:
        print(f"保存模型时发生错误: {str(e)}")
        raise




if __name__ == '__main__':
    # 反酸模型预测
    # 测试 predict 函数
    input_data = pd.DataFrame([{
        "organic_matter": 5.2,
        "chloride": 3.1,
        "cec": 25.6,
        "h_concentration": 0.5,
        "hn": 12.4,
        "al_concentration": 0.8,
        "free_alumina": 1.2,
        "free_iron": 0.9,
        "delta_ph": -0.2
    }])

    model_name = 'RF_filt'

    Acid_reflux_result = predict(input_data, model_name)
    print("Acid_reflux_result:", Acid_reflux_result)  # 预测结果

    # 降酸模型预测
    # 测试 predict 函数
    input_data = pd.DataFrame([{
        "pH": 5.2,
        "OM": 3.1,
        "CL": 25.6,
        "H": 0.5,
        "Al": 12.4
    }])

    model_name = 'rf_model_1214_1008'
    Acid_reduce_result = predict(input_data, model_name)
    print("Acid_reduce_result:", Acid_reduce_result)  # 预测结果