from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.linear_model import LinearRegression
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 = ['年', '季节', '月', '节假日', '星期', '工作日', '天气', '气温', '体感温度', '湿度', '风速']
CATEGORICAL_BASE = ['年', '季节', '月', '节假日', '星期', '工作日', '天气']


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


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


def add_date_features(frame, min_date):
    frame = frame.copy()
    frame['日期数值'] = (frame['日期'] - min_date).dt.days
    day_of_year = frame['日期'].dt.dayofyear
    frame['日期_sin'] = np.sin(2 * np.pi * day_of_year / 365.25)
    frame['日期_cos'] = np.cos(2 * np.pi * day_of_year / 365.25)
    return frame


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_date_features(daily, min_date)
    future = add_date_features(future, min_date)

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

    feature_sets = {
        '11特征_无日期': {
            'features': BASE_FEATURES,
            'categorical': CATEGORICAL_BASE,
        },
        '12特征_日期数值': {
            'features': ['日期数值'] + BASE_FEATURES,
            'categorical': CATEGORICAL_BASE,
        },
        '13特征_日期sin_cos': {
            'features': ['日期_sin', '日期_cos'] + BASE_FEATURES,
            'categorical': CATEGORICAL_BASE,
        },
        '12特征_仅日期sin': {
            'features': ['日期_sin'] + BASE_FEATURES,
            'categorical': CATEGORICAL_BASE,
        },
    }

    model_factories = {
        '线性回归': lambda: LinearRegression(),
        '随机森林': lambda: RandomForestRegressor(n_estimators=500, random_state=42, min_samples_leaf=2, n_jobs=-1),
        '梯度提升树': lambda: GradientBoostingRegressor(random_state=42, n_estimators=300, learning_rate=0.05, max_depth=3),
        'SVR': lambda: SVR(C=200, gamma='scale', epsilon=50),
    }

    rows = []
    future_rows = []
    for fs_name, spec in feature_sets.items():
        features = spec['features']
        categorical = spec['categorical']
        for model_name, factory in model_factories.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({
                '特征方案': fs_name,
                '模型': model_name,
                '特征数': len(features),
                '训练集样本数': len(train),
                '测试集样本数': len(test),
                '训练日期起': train['日期'].min().date().isoformat(),
                '训练日期止': train['日期'].max().date().isoformat(),
                '测试日期起': test['日期'].min().date().isoformat(),
                '测试日期止': test['日期'].max().date().isoformat(),
                'RMSE': rmse(test[TARGET], pred),
                'MAE': float(mean_absolute_error(test[TARGET], pred)),
                'R2': float(r2_score(test[TARGET], pred)),
                '测试预测最小值': float(pred.min()),
                '测试预测最大值': float(pred.max()),
                '未来预测最小值': int(future_pred.min()),
                '未来预测最大值': int(future_pred.max()),
                '未来预测均值': float(future_pred.mean()),
            })
            if model_name == '梯度提升树':
                temp = future[['日期'] + BASE_FEATURES].copy()
                temp['特征方案'] = fs_name
                temp['模型'] = model_name
                temp['预测总骑行量'] = future_pred
                future_rows.append(temp)

    result = pd.DataFrame(rows).sort_values(['RMSE', 'R2'], ascending=[True, False])
    result.to_csv(OUT_DIR / 'time_split_feature_model_comparison.csv', index=False, encoding='utf-8-sig')
    pd.concat(future_rows, ignore_index=True).to_excel(OUT_DIR / 'gbdt_future_predictions_by_feature_set.xlsx', index=False)

    print('时间顺序划分：')
    print(f"训练集 {len(train)} 条：{train['日期'].min().date()} 至 {train['日期'].max().date()}")
    print(f"测试集 {len(test)} 条：{test['日期'].min().date()} 至 {test['日期'].max().date()}")
    print('\n按 RMSE 排序：')
    print(result[['特征方案', '模型', 'RMSE', 'MAE', 'R2', '未来预测最小值', '未来预测最大值', '未来预测均值']].to_string(index=False, float_format=lambda x: f'{x:.3f}'))
    print(f"\n结果已保存：{OUT_DIR}")


if __name__ == '__main__':
    main()