import json
from pathlib import Path

import joblib
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.inspection import permutation_importance
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import KFold, cross_validate, train_test_split
from sklearn.neural_network import MLPRegressor
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'
FIG_DIR = OUT_DIR / 'figures'
MODEL_DIR = OUT_DIR / 'models'
for directory in [OUT_DIR, FIG_DIR, MODEL_DIR]:
    directory.mkdir(parents=True, exist_ok=True)

plt.rcParams['font.sans-serif'] = ['Noto Sans CJK SC', 'WenQuanYi Zen Hei', 'SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
sns.set_theme(style='whitegrid')

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


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


def evaluate_model(model, x_train, x_test, y_train, y_test):
    model.fit(x_train, y_train)
    pred = model.predict(x_test)
    return {
        'RMSE': rmse(y_test, pred),
        'MAE': float(mean_absolute_error(y_test, pred)),
        'R2': float(r2_score(y_test, pred)),
    }, pred


def make_preprocessor(scale_numeric=False):
    numeric_transformer = StandardScaler() if scale_numeric else 'passthrough'
    return ColumnTransformer(
        transformers=[
            ('cat', OneHotEncoder(handle_unknown='ignore'), CATEGORICAL_FEATURES),
            ('num', numeric_transformer, NUMERIC_FEATURES),
        ],
        remainder='drop',
        verbose_feature_names_out=False,
    )


def make_models():
    return {
        '线性回归': Pipeline([
            ('preprocess', make_preprocessor(scale_numeric=True)),
            ('model', LinearRegression()),
        ]),
        '随机森林': Pipeline([
            ('preprocess', make_preprocessor()),
            ('model', RandomForestRegressor(n_estimators=500, random_state=42, min_samples_leaf=2, n_jobs=-1)),
        ]),
        '梯度提升树': Pipeline([
            ('preprocess', make_preprocessor()),
            ('model', GradientBoostingRegressor(random_state=42, n_estimators=300, learning_rate=0.05, max_depth=3)),
        ]),
        '支持向量回归': Pipeline([
            ('preprocess', make_preprocessor(scale_numeric=True)),
            ('model', SVR(C=200, gamma='scale', epsilon=50)),
        ]),
        '神经网络': Pipeline([
            ('preprocess', make_preprocessor(scale_numeric=True)),
            ('model', MLPRegressor(hidden_layer_sizes=(80, 40), activation='relu', alpha=0.001, max_iter=3000, random_state=42, early_stopping=True)),
        ]),
    }


def load_data():
    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['日期'])
    return daily, future


def feature_importance(best_model, x_test, y_test):
    model_name = type(best_model.named_steps['model']).__name__
    if hasattr(best_model.named_steps['model'], 'feature_importances_'):
        transformed_names = best_model.named_steps['preprocess'].get_feature_names_out()
        raw_importances = pd.Series(best_model.named_steps['model'].feature_importances_, index=transformed_names)
        collapsed = {}
        for feature in FEATURES:
            collapsed[feature] = raw_importances[[idx == feature or idx.startswith(feature + '_') for idx in raw_importances.index]].sum()
        importance = pd.Series(collapsed).sort_values(ascending=False)
        method = f'{model_name} 内置重要性'
    else:
        result = permutation_importance(best_model, x_test, y_test, n_repeats=30, random_state=42, scoring='neg_root_mean_squared_error', n_jobs=-1)
        importance = pd.Series(result.importances_mean, index=FEATURES).sort_values(ascending=False)
        method = '排列重要性'
    return importance, method


def save_figures(daily, y_test, y_pred, residuals, importance, future_pred_df):
    plt.figure(figsize=(12, 6))
    sns.lineplot(data=daily, x='日期', y=TARGET, color='#2563eb')
    plt.title('2011-2012 年共享单车每日总骑行量趋势')
    plt.xlabel('日期')
    plt.ylabel('总骑行量')
    plt.tight_layout()
    plt.savefig(FIG_DIR / 'historical_trend.png', dpi=180)
    plt.close()

    plt.figure(figsize=(8, 6))
    sns.scatterplot(x=y_test, y=y_pred, color='#16a34a', edgecolor='white')
    low, high = min(y_test.min(), y_pred.min()), max(y_test.max(), y_pred.max())
    plt.plot([low, high], [low, high], '--', color='#ef4444', linewidth=1.5)
    plt.title('测试集真实值-预测值散点图')
    plt.xlabel('真实总骑行量')
    plt.ylabel('预测总骑行量')
    plt.tight_layout()
    plt.savefig(FIG_DIR / 'actual_vs_predicted.png', dpi=180)
    plt.close()

    plt.figure(figsize=(8, 5))
    sns.histplot(residuals, kde=True, color='#f97316')
    plt.axvline(0, color='#111827', linestyle='--', linewidth=1)
    plt.title('测试集残差分布')
    plt.xlabel('残差（真实值 - 预测值）')
    plt.ylabel('频数')
    plt.tight_layout()
    plt.savefig(FIG_DIR / 'residual_distribution.png', dpi=180)
    plt.close()

    plt.figure(figsize=(9, 6))
    plot_imp = importance.sort_values(ascending=True)
    plt.barh(plot_imp.index, plot_imp.values, color='#7c3aed')
    plt.title('特征重要性排序')
    plt.xlabel('重要性')
    plt.tight_layout()
    plt.savefig(FIG_DIR / 'feature_importance.png', dpi=180)
    plt.close()

    plt.figure(figsize=(12, 6))
    sns.lineplot(data=future_pred_df, x='日期', y='预测总骑行量', marker='o', color='#0891b2')
    plt.title('2013 年前 100 天共享单车骑行量预测趋势')
    plt.xlabel('日期')
    plt.ylabel('预测总骑行量')
    plt.tight_layout()
    plt.savefig(FIG_DIR / 'future_prediction_trend.png', dpi=180)
    plt.close()


def main():
    daily, future = load_data()
    missing = daily[FEATURES + [TARGET]].isna().sum().to_dict()
    duplicates = int(daily.duplicated(subset=['日期']).sum())

    x = daily[FEATURES]
    y = daily[TARGET]
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)

    models = make_models()
    initial_rows = []
    predictions = {}
    for name, model in models.items():
        metrics, pred = evaluate_model(model, x_train, x_test, y_train, y_test)
        initial_rows.append({'模型': name, **metrics})
        predictions[name] = pred
        joblib.dump(model, MODEL_DIR / f'{name}.pkl')

    initial_df = pd.DataFrame(initial_rows).sort_values('RMSE')
    best_name = initial_df.iloc[0]['模型']
    best_model = models[best_name]
    best_pred = predictions[best_name]
    residuals = y_test.to_numpy() - best_pred

    cv = KFold(n_splits=5, shuffle=True, random_state=42)
    cv_result = cross_validate(
        best_model,
        x,
        y,
        cv=cv,
        scoring={'rmse': 'neg_root_mean_squared_error', 'mae': 'neg_mean_absolute_error', 'r2': 'r2'},
        n_jobs=-1,
    )
    cv_summary = {
        '模型': best_name,
        'RMSE均值': float(-cv_result['test_rmse'].mean()),
        'RMSE标准差': float(cv_result['test_rmse'].std()),
        'MAE均值': float(-cv_result['test_mae'].mean()),
        'MAE标准差': float(cv_result['test_mae'].std()),
        'R2均值': float(cv_result['test_r2'].mean()),
        'R2标准差': float(cv_result['test_r2'].std()),
    }

    importance, importance_method = feature_importance(best_model, x_test, y_test)

    future_pred = best_model.predict(future[FEATURES])
    future_pred_df = future.copy()
    future_pred_df['预测总骑行量'] = np.maximum(np.round(future_pred), 0).astype(int)

    save_figures(daily, y_test, best_pred, residuals, importance, future_pred_df)

    initial_df.to_csv(OUT_DIR / 'model_comparison.csv', index=False, encoding='utf-8-sig')
    pd.DataFrame([cv_summary]).to_csv(OUT_DIR / 'cross_validation_summary.csv', index=False, encoding='utf-8-sig')
    importance.rename('重要性').to_csv(OUT_DIR / 'feature_importance.csv', encoding='utf-8-sig')
    future_pred_df.to_excel(OUT_DIR / '2013前100天骑行量预测结果.xlsx', index=False)

    test_pred_df = pd.DataFrame({
        '日期': daily.loc[x_test.index, '日期'].values,
        '真实总骑行量': y_test.values,
        '预测总骑行量': np.round(best_pred).astype(int),
        '残差': np.round(residuals, 2),
    }).sort_values('日期')
    test_pred_df.to_csv(OUT_DIR / 'test_set_predictions.csv', index=False, encoding='utf-8-sig')

    joblib.dump(best_model, MODEL_DIR / 'best_bike_demand_model.pkl')

    summary = {
        'data': {
            'rows': int(len(daily)),
            'columns': int(daily.shape[1]),
            'future_rows': int(len(future)),
            'missing': missing,
            'duplicate_dates': duplicates,
            'target_min': int(y.min()),
            'target_max': int(y.max()),
            'target_mean': float(y.mean()),
        },
        'initial_comparison': initial_df.to_dict(orient='records'),
        'best_model': best_name,
        'best_test_metrics': initial_df.iloc[0].to_dict(),
        'cross_validation': cv_summary,
        'importance_method': importance_method,
        'top_features': importance.head(8).to_dict(),
        'future_prediction': {
            'min': int(future_pred_df['预测总骑行量'].min()),
            'max': int(future_pred_df['预测总骑行量'].max()),
            'mean': float(future_pred_df['预测总骑行量'].mean()),
        },
        'outputs': {
            'figures': [str(p) for p in sorted(FIG_DIR.glob('*.png'))],
            'best_model': str(MODEL_DIR / 'best_bike_demand_model.pkl'),
            'prediction_excel': str(OUT_DIR / '2013前100天骑行量预测结果.xlsx'),
        }
    }
    (OUT_DIR / 'analysis_summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()