import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import roc_auc_score, roc_curve
from scipy import stats

np.random.seed(42)

# 合成データの生成: 正常データ + 点異常 + 文脈異常
T = 1000
t = np.arange(T)

# 正常データ: 正弦波 + ノイズ
normal_data = np.sin(2 * np.pi * t / 50) + 0.3 * np.random.randn(T)

# 異常ラベル（0: 正常, 1: 異常）
labels = np.zeros(T, dtype=int)

# 点異常を挿入（ランダムな位置にスパイクを追加）
anomaly_data = normal_data.copy()
point_anomaly_idx = [150, 300, 450, 600, 750, 850]
for idx in point_anomaly_idx:
    anomaly_data[idx] += np.random.choice([-1, 1]) * (4 + np.random.rand())
    labels[idx] = 1

# 文脈異常を挿入（急激なレベルシフト）
anomaly_data[500:520] += 2.5
labels[500:520] = 1

# データの可視化
fig, axes = plt.subplots(2, 1, figsize=(14, 6))
axes[0].plot(t, anomaly_data, linewidth=0.7)
axes[0].scatter(t[labels == 1], anomaly_data[labels == 1],
                c='red', s=20, zorder=5, label='Anomaly')
axes[0].set_title('Synthetic Time Series with Anomalies', fontsize=14)
axes[0].set_xlabel('Time', fontsize=12)
axes[0].set_ylabel('Value', fontsize=12)
axes[0].legend(fontsize=10)
axes[0].grid(True, alpha=0.3)

axes[1].plot(t, labels, 'r-', linewidth=1.5)
axes[1].set_title('Ground Truth Anomaly Labels', fontsize=14)
axes[1].set_xlabel('Time', fontsize=12)
axes[1].set_ylabel('Label', fontsize=12)
axes[1].set_ylim(-0.1, 1.1)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

def moving_average_detector(data, window=50, k=3.0):
    """移動平均 + σルールによる異常検知"""
    n = len(data)
    scores = np.zeros(n)
    for i in range(window, n):
        window_data = data[i - window:i]
        mu = np.mean(window_data)
        sigma = np.std(window_data, ddof=1)
        if sigma > 0:
            scores[i] = np.abs(data[i] - mu) / sigma
        else:
            scores[i] = 0.0
    return scores

# 移動平均法の実行
ma_scores = moving_average_detector(anomaly_data, window=50, k=3.0)

def cusum_detector(data, mu_0=None, k=0.5, h=4.0):
    """CUSUM（累積和制御図）による異常検知"""
    n = len(data)
    if mu_0 is None:
        mu_0 = np.mean(data[:100])  # 最初の100点から目標値を推定

    C_plus = np.zeros(n)
    C_minus = np.zeros(n)
    scores = np.zeros(n)

    for i in range(1, n):
        # 上方CUSUM: 平均の上昇を検出
        C_plus[i] = max(0, C_plus[i-1] + data[i] - mu_0 - k)
        # 下方CUSUM: 平均の低下を検出
        C_minus[i] = max(0, C_minus[i-1] - data[i] + mu_0 - k)
        # 異常スコア: 上方と下方の最大値
        scores[i] = max(C_plus[i], C_minus[i])

    return scores

# CUSUM法の実行
cusum_scores = cusum_detector(anomaly_data, k=0.5, h=4.0)

from statsmodels.tsa.seasonal import STL

def stl_detector(data, period=50):
    """STL分解ベースの異常検知"""
    stl = STL(data, period=period, robust=True)
    result = stl.fit()
    residuals = result.resid

    # 残差のIQRに基づくスコア
    q1, q3 = np.percentile(residuals, [25, 75])
    iqr = q3 - q1
    if iqr > 0:
        scores = np.abs(residuals) / iqr
    else:
        scores = np.abs(residuals)
    return scores

# STL分解法の実行
stl_scores = stl_detector(anomaly_data, period=50)

from sklearn.ensemble import IsolationForest

def isolation_forest_detector(data, window=20, contamination=0.05):
    """Isolation Forestの時系列適用（特徴量エンジニアリング）"""
    n = len(data)
    features = []
    valid_idx = []

    for i in range(window, n):
        w = data[i - window:i]
        feat = [
            np.mean(w),                    # 移動平均
            np.std(w, ddof=1),             # 移動標準偏差
            data[i] - np.mean(w),          # 偏差
            data[i] - data[i-1],           # 差分
            np.abs(data[i] - data[i-1]),   # 差分の絶対値
            np.max(w),                     # ウィンドウ最大値
            np.min(w),                     # ウィンドウ最小値
        ]
        features.append(feat)
        valid_idx.append(i)

    features = np.array(features)

    # Isolation Forestの学習と予測
    clf = IsolationForest(contamination=contamination, random_state=42, n_estimators=200)
    clf.fit(features)

    # スコアの計算（-1 * decision_function で異常ほど大きい値）
    raw_scores = -clf.decision_function(features)

    # 全時刻のスコアに変換
    scores = np.zeros(n)
    for idx, s in zip(valid_idx, raw_scores):
        scores[idx] = s

    return scores

# Isolation Forestの実行
if_scores = isolation_forest_detector(anomaly_data, window=20)

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

class LSTMAutoencoder(nn.Module):
    """LSTM Autoencoder"""
    def __init__(self, input_dim=1, hidden_dim=32, latent_dim=16):
        super().__init__()
        # エンコーダ
        self.encoder = nn.LSTM(input_dim, hidden_dim, batch_first=True)
        self.enc_to_latent = nn.Linear(hidden_dim, latent_dim)

        # デコーダ
        self.latent_to_dec = nn.Linear(latent_dim, hidden_dim)
        self.decoder = nn.LSTM(hidden_dim, hidden_dim, batch_first=True)
        self.output_layer = nn.Linear(hidden_dim, input_dim)

    def forward(self, x):
        # x: (batch, seq_len, input_dim)
        batch_size, seq_len, _ = x.shape

        # エンコード
        _, (h_enc, _) = self.encoder(x)
        latent = self.enc_to_latent(h_enc.squeeze(0))  # (batch, latent_dim)

        # デコード
        dec_input = self.latent_to_dec(latent).unsqueeze(1).repeat(1, seq_len, 1)
        dec_output, _ = self.decoder(dec_input)
        recon = self.output_layer(dec_output)  # (batch, seq_len, input_dim)
        return recon

def lstm_ae_detector(data, window=30, epochs=50, lr=1e-3):
    """LSTM-AEによる異常検知"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # ウィンドウで分割
    n = len(data)
    windows = []
    for i in range(n - window + 1):
        windows.append(data[i:i + window])
    windows = np.array(windows)

    # 正規化
    mean_val = np.mean(data[:int(n * 0.7)])
    std_val = np.std(data[:int(n * 0.7)])
    windows_norm = (windows - mean_val) / (std_val + 1e-8)

    # テンソル変換
    X = torch.tensor(windows_norm, dtype=torch.float32).unsqueeze(-1)

    # 訓練データ（最初の70%を正常とみなす）
    train_size = int(len(X) * 0.7)
    X_train = X[:train_size]

    train_loader = DataLoader(TensorDataset(X_train), batch_size=32, shuffle=True)

    # モデルの学習
    model = LSTMAutoencoder(input_dim=1, hidden_dim=32, latent_dim=16).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()

    model.train()
    for epoch in range(epochs):
        for (batch,) in train_loader:
            batch = batch.to(device)
            optimizer.zero_grad()
            recon = model(batch)
            loss = criterion(recon, batch)
            loss.backward()
            optimizer.step()

    # 全データの再構成誤差を計算
    model.eval()
    with torch.no_grad():
        X_all = X.to(device)
        recon_all = model(X_all)
        errors = torch.mean((X_all - recon_all) ** 2, dim=(1, 2)).cpu().numpy()

    # 各時刻のスコアに変換（ウィンドウの最後の時刻に割り当て）
    scores = np.zeros(n)
    for i in range(len(errors)):
        t_idx = i + window - 1
        scores[t_idx] = max(scores[t_idx], errors[i])

    return scores

# LSTM-AEの実行
lstm_scores = lstm_ae_detector(anomaly_data, window=30, epochs=50)

def bocpd_detector(data, hazard_lambda=200, mu_prior=0.0, sigma_prior=1.0):
    """ベイズオンライン変化点検出（BOCPD）"""
    n = len(data)
    # ラン長確率行列 (最大ラン長 x 時刻)
    max_run = n
    R = np.zeros((max_run + 1, n + 1))
    R[0, 0] = 1.0  # 初期条件: r_0 = 0

    # 十分統計量
    mu_params = np.zeros(max_run + 1)    # 各ラン長ごとの平均
    var_params = np.ones(max_run + 1)    # 各ラン長ごとの分散
    count = np.zeros(max_run + 1)        # データ数
    sum_x = np.zeros(max_run + 1)        # 合計
    sum_x2 = np.zeros(max_run + 1)       # 二乗合計

    # ハザード関数（定数: 1/lambda）
    H = 1.0 / hazard_lambda

    changepoint_prob = np.zeros(n)

    for t in range(n):
        x = data[t]

        # 予測確率の計算
        predprobs = np.zeros(t + 1)
        for r in range(t + 1):
            if count[r] < 2:
                # データが少ない場合は事前分布から予測
                predprobs[r] = stats.norm.pdf(x, mu_prior, sigma_prior)
            else:
                mu_r = sum_x[r] / count[r]
                var_r = (sum_x2[r] - sum_x[r]**2 / count[r]) / (count[r] - 1)
                var_r = max(var_r, 1e-6)
                # Student-t近似（正規近似で代替）
                pred_var = var_r * (1 + 1.0 / count[r])
                predprobs[r] = stats.norm.pdf(x, mu_r, np.sqrt(pred_var))

        # 成長確率
        growth = R[:t+1, t] * predprobs * (1 - H)

        # 変化点確率
        cp = np.sum(R[:t+1, t] * predprobs * H)

        # 更新
        R[1:t+2, t+1] = growth
        R[0, t+1] = cp

        # 正規化
        evidence = np.sum(R[:t+2, t+1])
        if evidence > 0:
            R[:t+2, t+1] /= evidence

        # 変化点確率（r=0の確率）
        changepoint_prob[t] = R[0, t+1]

        # 十分統計量の更新
        new_count = count[:t+1] + 1
        new_sum = sum_x[:t+1] + x
        new_sum2 = sum_x2[:t+1] + x**2

        count[1:t+2] = new_count
        sum_x[1:t+2] = new_sum
        sum_x2[1:t+2] = new_sum2

        # r=0（新しいラン）をリセット
        count[0] = 0
        sum_x[0] = 0
        sum_x2[0] = 0

    return changepoint_prob

# BOCPDの実行
bocpd_scores = bocpd_detector(anomaly_data, hazard_lambda=200)

# 有効範囲（全手法が値を持つ範囲）
valid_start = 50  # 移動平均のウィンドウ幅の影響

# 各手法のスコアを有効範囲で切り出し
scores_dict = {
    'Moving Avg + σ': ma_scores[valid_start:],
    'CUSUM': cusum_scores[valid_start:],
    'STL Decomposition': stl_scores[valid_start:],
    'Isolation Forest': if_scores[valid_start:],
    'LSTM-AE': lstm_scores[valid_start:],
    'BOCPD': bocpd_scores[valid_start:],
}
labels_valid = labels[valid_start:]

# ROC-AUCの計算と表示
print("=" * 50)
print(f"{'Method':<25} {'ROC-AUC':>10}")
print("=" * 50)

auc_results = {}
for name, scores in scores_dict.items():
    try:
        auc = roc_auc_score(labels_valid, scores)
        auc_results[name] = auc
        print(f"{name:<25} {auc:>10.4f}")
    except ValueError:
        print(f"{name:<25} {'N/A':>10}")
print("=" * 50)

# ROC曲線の可視化
plt.figure(figsize=(10, 8))
colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
for (name, scores), color in zip(scores_dict.items(), colors):
    try:
        fpr, tpr, _ = roc_curve(labels_valid, scores)
        auc = auc_results.get(name, 0)
        plt.plot(fpr, tpr, color=color, linewidth=2,
                 label=f'{name} (AUC={auc:.3f})')
    except ValueError:
        pass

plt.plot([0, 1], [0, 1], 'k--', linewidth=1, alpha=0.5)
plt.xlabel('False Positive Rate', fontsize=12)
plt.ylabel('True Positive Rate', fontsize=12)
plt.title('ROC Curves: Time Series Anomaly Detection Methods', fontsize=14)
plt.legend(fontsize=10, loc='lower right')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

fig, axes = plt.subplots(7, 1, figsize=(14, 20))

# 元データ
axes[0].plot(t, anomaly_data, linewidth=0.7)
axes[0].scatter(t[labels == 1], anomaly_data[labels == 1],
                c='red', s=15, zorder=5)
axes[0].set_title('Original Data with Anomalies', fontsize=12)
axes[0].set_ylabel('Value', fontsize=10)
axes[0].grid(True, alpha=0.3)

# 各手法のスコア
method_scores = [
    ('Moving Avg + σ Rule', ma_scores),
    ('CUSUM', cusum_scores),
    ('STL Decomposition', stl_scores),
    ('Isolation Forest', if_scores),
    ('LSTM-AE', lstm_scores),
    ('BOCPD', bocpd_scores),
]

for i, (name, scores) in enumerate(method_scores):
    ax = axes[i + 1]
    ax.plot(t, scores, linewidth=0.7, color=colors[i])
    # 異常位置をハイライト
    for idx in range(T):
        if labels[idx] == 1:
            ax.axvline(x=idx, color='red', alpha=0.1, linewidth=1)
    ax.set_title(f'{name} Anomaly Score', fontsize=12)
    ax.set_ylabel('Score', fontsize=10)
    ax.grid(True, alpha=0.3)

axes[-1].set_xlabel('Time', fontsize=12)
plt.tight_layout()
plt.show()