#!/usr/bin/env python3
"""
Apply user prior correction rules (切换风险 + 雪天严寒) to test set predictions.
Reads the v002-adjusted predictions + original CSV, applies additional rules,
outputs a new adjusted_predictions.csv with all rules combined.

Usage:
    python apply_prior_rules_test.py <adjusted_predictions_csv> <full_prior_csv> <output_csv>

Rules follow user-prior-mirror skill:
    Layer 1 (highest): 切换风险 — 节假日/环保/政策
    Layer 2: 雪天严寒 — 雪天+极端天气 / 前日雪+严寒
    Time constraint: only adjust 07:00+ (keep 00:00-07:00 unchanged)
    Magnitude constraints: bounds per rule type
"""

import sys
import os
import json
import pandas as pd
import numpy as np
from pathlib import Path


def load_data(adjusted_predictions_path, prior_csv_path):
    """Load and merge adjusted predictions with original CSV features."""
    adj = pd.read_csv(adjusted_predictions_path)
    adj["timestamp"] = pd.to_datetime(adj["timestamp"])

    prior = pd.read_csv(prior_csv_path, low_memory=False)
    prior["date"] = pd.to_datetime(prior["date"])

    # Merge: attach exogenous features from prior CSV
    merged = adj.merge(prior, left_on="timestamp", right_on="date", how="left")

    # Add derived features
    merged["hour"] = merged["timestamp"].dt.hour
    merged["year"] = merged["timestamp"].dt.year
    merged["month"] = merged["timestamp"].dt.month
    merged["day"] = merged["timestamp"].dt.day

    # Normalize categorical text columns
    for col in ["节假日", "环保", "政策", "天气", "前一日天气"]:
        merged[col] = merged[col].astype(str)

    # Normalize 是否极端天气
    merged["是否极端天气_bool"] = (
        merged["是否极端天气"].astype(str).isin(["True", "True", "1", True])
    )

    return adj, merged


def classify_rules(df):
    """Classify each row by which rule(s) apply, respecting priority.

    Returns a DataFrame with columns:
        - rule_group: 'layer1_holiday', 'layer1_env', 'layer1_policy', 'layer2_snow_extreme',
                       'layer2_snow_prev_cold', or None
        - direction: 'up' or 'down'
        - rule_detail: specific matching rule
    """

    # Initialize result columns
    df = df.copy()
    df["rule_group"] = None
    df["direction"] = None
    df["suggested_add"] = 0.0

    # ---- Layer 1: 节假日 ----
    is_holiday = df["节假日"] != "0"

    # Sub-classify 节假日
    holiday_resume = df["节假日"].str.contains("复工|开始", na=False)
    holiday_other = is_holiday & ~holiday_resume

    # 周日保留 (may also be holiday), 周六 is not in holidays

    # ---- Layer 1: 环保 ----
    is_env = df["环保"] != "0"
    env_relax = df["环保"].str.contains("结束|解除|停止", na=False)
    env_control = is_env & ~env_relax

    # ---- Layer 1: 政策 ----
    is_policy = df["政策"] != "0"
    policy_end = df["政策"].str.contains("结束|停止", na=False)
    policy_other = is_policy & ~policy_end

    # Layer 1 priority: holiday > env > policy
    # Holiday
    mask = holiday_resume
    df.loc[mask, "rule_group"] = "layer1_holiday"
    df.loc[mask, "direction"] = "up"
    df.loc[mask, "suggested_add"] = 1200  # +1200 (mid-range of 800-2000)

    mask = holiday_other
    df.loc[mask, "rule_group"] = "layer1_holiday"
    df.loc[mask, "direction"] = "down"
    df.loc[mask, "suggested_add"] = -800  # -800 (mid-range of 500-1000)

    # Env (only if no holiday already set)
    mask = is_env & df["rule_group"].isna()
    df.loc[mask & env_relax, "rule_group"] = "layer1_env"
    df.loc[mask & env_relax, "direction"] = "up"
    df.loc[mask & env_relax, "suggested_add"] = 400  # +400

    df.loc[mask & env_control, "rule_group"] = "layer1_env"
    df.loc[mask & env_control, "direction"] = "down"
    df.loc[mask & env_control, "suggested_add"] = -350  # -350

    # Policy (only if no holiday/env already set)
    mask = is_policy & df["rule_group"].isna()
    df.loc[mask & policy_end, "rule_group"] = "layer1_policy"
    df.loc[mask & policy_end, "direction"] = "up"
    df.loc[mask & policy_end, "suggested_add"] = 500  # +500

    df.loc[mask & policy_other, "rule_group"] = "layer1_policy"
    df.loc[mask & policy_other, "direction"] = "down"
    df.loc[mask & policy_other, "suggested_add"] = -500  # -500

    # Special: 北京、天津供暖 → this is a heating start, direction up
    heating_start = df["政策"].str.contains("供暖", na=False)
    mask = heating_start & df["rule_group"].isna()
    df.loc[mask, "rule_group"] = "layer1_policy"
    df.loc[mask, "direction"] = "up"
    df.loc[mask, "suggested_add"] = 500  # +500

    # ---- Layer 2: 雪天严寒 (only if no Layer 1) ----
    layer2_base = df["rule_group"].isna()

    # Condition 4: 雪天 + 极端天气
    is_snow = df["天气"].str.contains("雪", na=False)
    is_extreme = df["是否极端天气_bool"]
    mask = layer2_base & is_snow & is_extreme
    df.loc[mask, "rule_group"] = "layer2_snow_extreme"
    df.loc[mask, "direction"] = "up"
    df.loc[mask, "suggested_add"] = 1500  # +1500 (mid-range of 800-2000)

    # Condition 5: 前一日雪 + 当前日严寒
    # Recalc layer2_base to include rows still unassigned
    layer2_base = df["rule_group"].isna()
    is_prev_snow = df["前一日天气"].str.contains("雪", na=False)
    is_max_low = pd.to_numeric(df["最高温_数值"], errors="coerce") < -5
    is_min_low = pd.to_numeric(df["最低温_数值"], errors="coerce") < -10
    mask = layer2_base & is_prev_snow & is_max_low & is_min_low
    df.loc[mask, "rule_group"] = "layer2_snow_prev_cold"
    df.loc[mask, "direction"] = "up"
    df.loc[mask, "suggested_add"] = 1000  # +1000 (mid-range of 500-1500)

    return df


def apply_time_constraint(
    rules_df,
    predicted_col="adjusted_predicted_original",
    raw_col="raw_predicted_original",
    exclude_hours_before=7,
):
    """Keep 00:00-07:00 (hour < 7) unchanged - reset to raw_predicted_original."""
    result = rules_df.copy()
    night_mask = result["hour"] < exclude_hours_before
    result.loc[night_mask, predicted_col] = result.loc[night_mask, raw_col]
    result.loc[night_mask, "rule_group"] = result.loc[night_mask, "rule_group"].fillna(
        "none"
    )
    result.loc[night_mask, "suggested_add"] = 0.0
    return result


def apply_adjustments(df):
    """Apply rule adjustments with magnitude bounds."""
    result = df.copy()
    result["adjusted_original_final"] = result["adjusted_predicted_original"].copy()

    # Only apply where a rule is active AND hour >= 7
    has_rule = result["rule_group"].notna() & (result["hour"] >= 7)

    for idx in result[has_rule].index:
        add_val = result.loc[idx, "suggested_add"]
        raw_val = result.loc[idx, "raw_predicted_original"]
        current_val = result.loc[idx, "adjusted_original_final"]

        new_val = current_val + add_val

        # Apply magnitude bounds based on direction
        if add_val > 0:
            # Up: delta >= 0 and delta <= 3000
            delta = new_val - current_val
            delta = max(0, min(delta, 3000))
            new_val = current_val + delta
        elif add_val < 0:
            # Down: delta <= 0 and delta >= -1000
            delta = new_val - current_val
            delta = min(0, max(delta, -1000))
            new_val = current_val + delta

        result.loc[idx, "adjusted_original_final"] = new_val

    return result


def compute_metrics(actuals, predictions):
    """Compute MSE, MAE, WAPE, MASE.
    MASE: naive seasonal=pred_len horizon, scaled by mean absolute seasonal difference.
    """
    actuals = np.array(actuals, dtype=float)
    predictions = np.array(predictions, dtype=float)

    mse = np.mean((actuals - predictions) ** 2)
    mae = np.mean(np.abs(actuals - predictions))
    wape = (
        np.sum(np.abs(actuals - predictions)) / np.sum(np.abs(actuals))
        if np.sum(np.abs(actuals)) > 0
        else 0
    )

    # MASE: naive forecast = previous value at same time (pred_len steps back)
    # For 15-min data with pred_len=96, that's 24h back
    if len(actuals) > 96:
        naive_errors = np.abs(actuals[96:] - actuals[:-96])
        mase = mae / np.mean(naive_errors) if np.mean(naive_errors) > 0 else 0
    else:
        mase = 0

    return {"mse": mse, "mae": mae, "wape": wape, "mase": mase}


def main():
    if len(sys.argv) < 4:
        print(
            f"Usage: {sys.argv[0]} <adjusted_predictions_csv> <full_prior_csv> <output_csv>"
        )
        sys.exit(1)

    adj_path = sys.argv[1]
    prior_path = sys.argv[2]
    output_path = sys.argv[3]

    print(f"[1/5] Loading adjusted predictions from {adj_path}...")
    adj_orig, merged = load_data(adj_path, prior_path)

    print(f"[2/5] Classifying rows by rule priority...")
    classified = classify_rules(merged)
    print(
        f"  Layer 1 节假日 (复工 ↑): {(classified['rule_group'] == 'layer1_holiday').sum()}"
    )
    print(
        f"  Layer 1 节假日 (其他 ↓): {(classified['rule_group'] == 'layer1_holiday').sum()}"
    )
    print(
        f"  Layer 1 环保 (解除 ↑): {(classified['rule_group'] == 'layer1_env').sum()}"
    )
    print(
        f"  Layer 1 环保 (管控 ↓): {(classified['rule_group'] == 'layer1_env').sum()}"
    )
    print(
        f"  Layer 1 政策 (结束 ↑): {(classified['rule_group'] == 'layer1_policy').sum()}"
    )
    print(
        f"  Layer 1 政策 (其他 ↓): {(classified['rule_group'] == 'layer1_policy').sum()}"
    )
    print(
        f"  Layer 2 雪天+极端 (↑): {(classified['rule_group'] == 'layer2_snow_extreme').sum()}"
    )
    print(
        f"  Layer 2 前日雪+严寒 (↑): {(classified['rule_group'] == 'layer2_snow_prev_cold').sum()}"
    )

    print(f"[3/5] Applying time constraint (keep 00:00-07:00 unchanged)...")
    constrained = apply_time_constraint(classified)

    print(f"[4/5] Applying magnitude-bounded adjustments...")
    adjusted = apply_adjustments(constrained)

    # Compute metrics
    raw_metrics = compute_metrics(
        adjusted["actual_original"], adjusted["raw_predicted_original"]
    )
    v002_metrics = compute_metrics(
        adjusted["actual_original"], adjusted["adjusted_predicted_original"]
    )
    final_metrics = compute_metrics(
        adjusted["actual_original"], adjusted["adjusted_original_final"]
    )

    print(f"\n[5/5] Results:")
    print(
        f"{'Metric':>10} | {'Raw UserPrior':>14} | {'V002 Adjusted':>14} | {'V002+PriorRules':>16} | {'V002 Delta':>12} | {'Final Delta':>12}"
    )
    print(
        f"{'-' * 10}-+-{'-' * 14}-+-{'-' * 14}-+-{'-' * 16}-+-{'-' * 12}-+-{'-' * 12}"
    )
    for metric in ["mse", "mae", "wape", "mase"]:
        r = raw_metrics[metric]
        v = v002_metrics[metric]
        f = final_metrics[metric]
        vd = (v - r) / r * 100
        fd = (f - r) / r * 100
        print(
            f"{metric:>10} | {r:>14.6f} | {v:>14.6f} | {f:>16.6f} | {vd:>+11.4f}% | {fd:>+11.4f}%"
        )

    # Save output: add columns for rule_group, direction, suggested_add, adjusted_original_final
    output = adjusted[
        [
            "timestamp",
            "actual_original",
            "raw_predicted_original",
            "adjusted_predicted_original",
            "adjusted_original_final",
            "rule_group",
            "direction",
            "suggested_add",
            "节假日",
            "环保",
            "政策",
            "天气",
            "是否极端天气",
            "最高温_数值",
            "最低温_数值",
            "前一日天气",
        ]
    ].copy()

    output.to_csv(output_path, index=False)
    print(f"\nOutput saved to {output_path}")
    print(f"Total rows: {len(output)}")

    # Summary by date
    print(f"\n--- Date-level summary ---")
    output["date"] = pd.to_datetime(output["timestamp"]).dt.date
    date_groups = output[output["rule_group"].notna()].groupby("date")
    summary = date_groups.agg(
        rows=("rule_group", "count"),
        groups=("rule_group", lambda x: x.value_counts().to_dict()),
        avg_add=("suggested_add", "mean"),
    )
    for date, row in summary.iterrows():
        print(
            f"  {date}: {row['rows']} rows, groups={row['groups']}, avg_add={row['avg_add']:.0f}"
        )

    # Save metrics summary
    metrics_out = {
        "raw": raw_metrics,
        "v002": v002_metrics,
        "v002_plus_prior_rules": final_metrics,
        "delta_v002_vs_raw": {k: v - raw_metrics[k] for k, v in v002_metrics.items()},
        "delta_final_vs_raw": {
            k: final_metrics[k] - raw_metrics[k] for k, v in final_metrics.items()
        },
        "delta_final_vs_v002": {
            k: final_metrics[k] - v002_metrics[k] for k, v in final_metrics.items()
        },
    }
    metrics_path = output_path.replace(".csv", "_metrics.json")
    with open(metrics_path, "w") as f:
        json.dump(metrics_out, f, indent=2)
    print(f"Metrics summary saved to {metrics_path}")


if __name__ == "__main__":
    main()