from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import ExtraTreesRegressor, GradientBoostingRegressor, HistGradientBoostingRegressor, RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.svm import SVR

BASE_DIR = Path(__file__).resolve().parent
DATA_PATH = Path('/root/.hermes-web-ui/upload/default/28752a235eb6ef9d.xlsx')
OUT_DIR = BASE_DIR / 'outputs' / 'time_split_experiment'
OUT_DIR.mkdir(parents=True, exist_ok=True)

TARGET = '总骑行量'
BASE_FEATURES = ['年', '季节', '月', '节假日', '星期', '工作日', '天气', '气温', '体感温度', '湿度', '风速']
BASE_CAT = ['年', '季节', '月', '节假日', '星期', '工作日', '天气']


def rmse(y_true, y_pred):
    return float(np.sqrt(mean_squared_error(y_true, y_pred)))


def add_features(df, min_date):
    df = df.copy()
    df['日期数值'] = (df['日期'] - min_date).dt.days
    df['一年中的第几天'] = df['日期'].dt.dayofyear
    df['月内日'] = df['日期'].dt.day
    df['周序号'] = df['日期'].dt.isocalendar().week.astype(int)
    df['日期_sin'] = np.sin(2 * np.pi * df['一年中的第几天'] / 365.25)
    df['日期_cos'] = np.cos(2 * np.pi * df['一年中的第几天'] / 365.25)
    df['周序号_sin'] = np.sin(2 * np.pi * df['周序号'] / 52.18)
    df['月内日_sin'] = np.sin(2 * np.pi * df['月内日'] / 31)
    df['年内进度'] = df['一年中的第几天'] / 365.25
    df['季度内月份'] = ((df['月'] - 1) % 3) + 1
    df['是否月初'] = (df['月内日'] <= 7).astype(int)
    df['是否月末'] = (df['月内日'] >= 24).astype(int)
    df['日期分箱4'] = pd.qcut(df['日期数值'], q=4, labels=False, duplicates='drop').astype(int)
    df['日期分箱8'] = pd.qcut(df['日期数值'], q=8, labels=False, duplicates='drop').astype(int)
    return df


def make_pipeline(model_name, estimator, features, categorical):
    numeric = [col for col in features if col not in categorical]
    need_scale = model_name == 'SVR'
    preprocessor = ColumnTransformer(
        transformers=[
            ('cat', OneHotEncoder(handle_unknown='ignore', sparse_output=False), categorical),
            ('num', StandardScaler() if need_scale else 'passthrough', numeric),
        ],
        verbose_feature_names_out=False,
    )
    return Pipeline([('preprocess', preprocessor), ('model', estimator)])


def main():
    daily = pd.read_excel(DATA_PATH, sheet_name='每天数据')
    future = pd.read_excel(DATA_PATH, sheet_name='预测集')
    daily['日期'] = pd.to_datetime(daily['日期'])
    future['日期'] = pd.to_datetime(future['日期'])
    daily = daily.sort_values('日期').reset_index(drop=True)
    future = future.sort_values('日期').reset_index(drop=True)
    min_date = daily['日期'].min()
    daily = add_features(daily, min_date)
    future = add_features(future, min_date)

    split_idx = int(len(daily) * 0.8)
    train = daily.iloc[:split_idx].copy()
    test = daily.iloc[split_idx:].copy()

    feature_schemes = {
        '12_日期数值': (['日期数值'] + BASE_FEATURES, BASE_CAT),
        '12_日期sin': (['日期_sin'] + BASE_FEATURES, BASE_CAT),
        '12_日期cos': (['日期_cos'] + BASE_FEATURES, BASE_CAT),
        '12_年内第几天': (['一年中的第几天'] + BASE_FEATURES, BASE_CAT),
        '12_年内进度': (['年内进度'] + BASE_FEATURES, BASE_CAT),
        '12_周序号': (['周序号'] + BASE_FEATURES, BASE_CAT + ['周序号']),
        '12_周序号sin': (['周序号_sin'] + BASE_FEATURES, BASE_CAT),
        '12_月内日': (['月内日'] + BASE_FEATURES, BASE_CAT + ['月内日']),
        '12_月内日sin': (['月内日_sin'] + BASE_FEATURES, BASE_CAT),
        '12_季度内月份': (['季度内月份'] + BASE_FEATURES, BASE_CAT + ['季度内月份']),
        '12_是否月初': (['是否月初'] + BASE_FEATURES, BASE_CAT + ['是否月初']),
        '12_是否月末': (['是否月末'] + BASE_FEATURES, BASE_CAT + ['是否月末']),
        '12_日期分箱4': (['日期分箱4'] + BASE_FEATURES, BASE_CAT + ['日期分箱4']),
        '12_日期分箱8': (['日期分箱8'] + BASE_FEATURES, BASE_CAT + ['日期分箱8']),
    }

    models = {
        '梯度提升树': lambda: GradientBoostingRegressor(random_state=42, n_estimators=300, learning_rate=0.05, max_depth=3),
        '随机森林': lambda: RandomForestRegressor(n_estimators=500, random_state=42, min_samples_leaf=2, n_jobs=-1),
        '极端随机树': lambda: ExtraTreesRegressor(n_estimators=500, random_state=42, min_samples_leaf=2, n_jobs=-1),
        'HistGBDT': lambda: HistGradientBoostingRegressor(random_state=42, max_iter=300, learning_rate=0.05, max_leaf_nodes=31),
        'SVR': lambda: SVR(C=200, gamma='scale', epsilon=50),
    }

    rows = []
    for scheme_name, (features, categorical) in feature_schemes.items():
        for model_name, factory in models.items():
            pipe = make_pipeline(model_name, factory(), features, categorical)
            pipe.fit(train[features], train[TARGET])
            pred = pipe.predict(test[features])
            future_pred = np.maximum(np.round(pipe.predict(future[features])), 0).astype(int)
            rows.append({
                '方案': scheme_name,
                '模型': model_name,
                'RMSE': rmse(test[TARGET], pred),
                'MAE': float(mean_absolute_error(test[TARGET], pred)),
                'R2': float(r2_score(test[TARGET], pred)),
                '未来预测最小值': int(future_pred.min()),
                '未来预测最大值': int(future_pred.max()),
                '未来预测均值': float(future_pred.mean()),
            })

    result = pd.DataFrame(rows)
    result.to_csv(OUT_DIR / 'extended_12_feature_nonlinear_comparison.csv', index=False, encoding='utf-8-sig')
    print('各模型最佳前5：')
    for model_name, group in result.groupby('模型'):
        print('\n[' + model_name + ']')
        print(group.sort_values('RMSE').head(5).to_string(index=False, float_format=lambda x: f'{x:.3f}'))
    print('\n总体前15：')
    print(result.sort_values('RMSE').head(15).to_string(index=False, float_format=lambda x: f'{x:.3f}'))
    print('\n保存：', OUT_DIR / 'extended_12_feature_nonlinear_comparison.csv')


if __name__ == '__main__':
    main()