# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
from datetime import timedelta

# ==================== 配置区域 ====================
FILES = [
    "四点共震数据集1.csv",
    "四点共震数据集2.csv",
    "四点共震数据集3.csv",
    "四点共震数据集4.csv",
]
START_DATE = "2000-01-01"   # 使用最高质量的数据段

# 粗扫描窗口（小时）
COARSE_WINDOWS = [24, 48, 72, 96, 120, 144, 168]
# 细化步长（小时）
FINE_STEP = 6
# =================================================


def load_data():
    """加载数据，过滤 START_DATE 之后，统一去掉时区"""
    print("正在加载数据...")
    start_dt = pd.Timestamp(START_DATE)
    dfs = []
    for i, f in enumerate(FILES, 1):
        df = pd.read_csv(f, parse_dates=["time"])
        if df["time"].dt.tz is not None:
            df["time"] = df["time"].dt.tz_localize(None)
        df = df[df["time"] >= start_dt]
        df.sort_values("time", inplace=True)
        df.reset_index(drop=True, inplace=True)
        dfs.append(df)
        print(f"  地点 {i} 加载完成，共 {len(df)} 条记录（{START_DATE} 后）")
    return dfs


def compute_stats(trigger_hours, dfs):
    """
    计算给定时间窗口（小时）下的统计结果。
    返回：{ "den": 总三链数, "num": 总四链数, "ratio": 条件概率, "rel_mags": 相对震级列表 }
    """

    def traceback_in_allowed(current_time, start_loc, allowed_locs):
        """在 allowed_locs 范围内接力溯源"""
        used_locs = {start_loc}
        cur = current_time
        while len(used_locs) < 4:
            candidates = []
            for loc in allowed_locs:
                if loc in used_locs:
                    continue
                df = dfs[loc]
                mask = (df["time"] < cur) & (df["time"] >= cur - timedelta(hours=trigger_hours))
                sub = df.loc[mask]
                if sub.empty:
                    continue
                idx = sub["time"].idxmin()
                row = sub.loc[idx]
                candidates.append((loc, row["time"], row["mag"]))
            if not candidates:
                break
            best_loc, best_time, _ = min(candidates, key=lambda x: x[1])
            used_locs.add(best_loc)
            cur = best_time
        return len(used_locs)

    def has_three_chain_among_others(others, max_time):
        """检查 others 三地内部是否能形成三链"""
        for start_loc in others:
            df = dfs[start_loc]
            sub = df[df["time"] < max_time]
            if sub.empty:
                continue
            row = sub.loc[sub["time"].idxmax()]
            reached = traceback_in_allowed(row["time"], start_loc, allowed_locs=others)
            if reached >= 3:
                return True
        return False

    def traceback_full_four(response_idx, row):
        """完整四点溯源"""
        cur = row["time"]
        used_locs = {response_idx}
        chain_mags = [row["mag"]]
        while len(used_locs) < 4:
            candidates = []
            for loc in range(4):
                if loc in used_locs:
                    continue
                df = dfs[loc]
                mask = (df["time"] < cur) & (df["time"] >= cur - timedelta(hours=trigger_hours))
                sub = df.loc[mask]
                if sub.empty:
                    continue
                idx = sub["time"].idxmin()
                row_cand = sub.loc[idx]
                candidates.append((loc, row_cand["time"], row_cand["mag"]))
            if not candidates:
                break
            best_loc, best_time, best_mag = min(candidates, key=lambda x: x[1])
            used_locs.add(best_loc)
            chain_mags.append(best_mag)
            cur = best_time
        if len(used_locs) == 4:
            rel_mag = chain_mags[0] - np.mean(chain_mags[1:])
            return True, rel_mag
        return False, None

    total_three = 0
    total_four = 0
    all_rel_mags = []

    for resp_idx in range(4):
        others = [i for i in range(4) if i != resp_idx]
        resp_df = dfs[resp_idx]
        for _, row in resp_df.iterrows():
            if not has_three_chain_among_others(others, row["time"]):
                continue
            total_three += 1
            hit, rel_mag = traceback_full_four(resp_idx, row)
            if hit:
                total_four += 1
                all_rel_mags.append(rel_mag)

    return {
        "den": total_three,
        "num": total_four,
        "ratio": total_four / total_three if total_three > 0 else 0,
        "rel_mags": all_rel_mags,
    }


def main():
    print("=" * 80)
    print("【滑动窗口分析】寻找最优时间间隔")
    print(f"数据范围：{START_DATE} 之后")
    print("=" * 80)

    dfs = load_data()

    # ---- 第一阶段：粗扫描 ----
    print("\n【第一阶段：粗扫描（24~168小时，步长24小时）】")
    print("-" * 80)
    print(f"{'窗口(小时)':<12} | {'三链分母':<10} | {'四链分子':<10} | {'条件概率':<10} | {'相对震级均值':<12}")
    print("-" * 80)

    coarse_results = {}
    for h in COARSE_WINDOWS:
        stats = compute_stats(h, dfs)
        coarse_results[h] = stats
        rel_mean = np.mean(stats["rel_mags"]) if stats["rel_mags"] else 0
        print(f"{h:<12} | {stats['den']:<10} | {stats['num']:<10} | {stats['ratio']:.2%}     | {rel_mean:.3f}")

    # 找出粗扫描中的最优窗口
    best_h = max(coarse_results, key=lambda x: coarse_results[x]["ratio"])
    best_ratio = coarse_results[best_h]["ratio"]
    print("-" * 80)
    print(f"✅ 粗扫描最优窗口：{best_h} 小时，条件概率 {best_ratio:.2%}")

    # ---- 第二阶段：细化扫描（在最优窗口附近 ±18 小时，步长 6 小时） ----
    print("\n【第二阶段：细化扫描（围绕最优窗口，步长6小时）】")
    print("-" * 80)
    print(f"{'窗口(小时)':<12} | {'三链分母':<10} | {'四链分子':<10} | {'条件概率':<10} | {'相对震级均值':<12}")
    print("-" * 80)

    fine_windows = range(max(6, best_h - 18), best_h + 19, FINE_STEP)
    fine_results = {}
    for h in fine_windows:
        # 如果该窗口已经在粗扫描中算过，直接复用，避免重复计算
        if h in coarse_results:
            stats = coarse_results[h]
        else:
            stats = compute_stats(h, dfs)
        fine_results[h] = stats
        rel_mean = np.mean(stats["rel_mags"]) if stats["rel_mags"] else 0
        mark = " ★" if h == best_h else ""
        print(f"{h:<12} | {stats['den']:<10} | {stats['num']:<10} | {stats['ratio']:.2%}     | {rel_mean:.3f}{mark}")

    # ---- 最终最优窗口 ----
    final_best_h = max(fine_results, key=lambda x: fine_results[x]["ratio"])
    final_stats = fine_results[final_best_h]

    print("-" * 80)
    print("\n" + "=" * 80)
    print("【最终结论】")
    print("=" * 80)
    print(f"最优时间窗口：{final_best_h} 小时")
    print(f"对应条件概率：{final_stats['ratio']:.2%}")
    print(f"三链分母：{final_stats['den']}")
    print(f"四链分子：{final_stats['num']}")

    if final_stats["rel_mags"]:
        arr = np.array(final_stats["rel_mags"])
        print(f"相对震级均值：{arr.mean():.3f}")
        print(f"相对震级中位数：{np.median(arr):.3f}")
        print(f"相对震级标准差：{arr.std():.3f}")


if __name__ == "__main__":
    main()