import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

np.random.seed(42)

class MTMethod:
    """MT法(マハラノビス・タグチ法)による異常検知"""

    def __init__(self):
        self.mu = None
        self.sigma_inv = None
        self.p = None

    def fit(self, X_normal):
        """正常データから単位空間を構築"""
        self.p = X_normal.shape[1]
        self.mu = X_normal.mean(axis=0)
        sigma = np.cov(X_normal.T)
        self.sigma_inv = np.linalg.inv(sigma)
        return self

    def mahalanobis_distance(self, X):
        """正規化マハラノビス距離(MD)を計算"""
        diff = X - self.mu
        md = np.array([d @ self.sigma_inv @ d for d in diff]) / self.p
        return md

    def predict(self, X, threshold=None):
        """異常判定（threshold=Noneの場合はMD>4を閾値とする）"""
        md = self.mahalanobis_distance(X)
        if threshold is None:
            threshold = 4.0  # 一般的な経験的閾値
        return (md > threshold).astype(int), md

# --- データ生成 ---
n_normal = 200
n_anomaly = 15
p = 3

# 正常データ
mean_normal = np.array([5.0, 3.0, 7.0])
cov_normal = np.array([
    [1.0, 0.5, 0.3],
    [0.5, 1.5, 0.2],
    [0.3, 0.2, 2.0]
])
X_normal = np.random.multivariate_normal(mean_normal, cov_normal, n_normal)

# 異常データ（平均からずれたデータ）
X_anomaly = np.random.multivariate_normal(
    mean_normal + np.array([3.0, 2.0, 4.0]),
    cov_normal * 0.5,
    n_anomaly
)

# --- MT法の適用 ---
mt = MTMethod()
mt.fit(X_normal)

# 正常データのMD
md_normal = mt.mahalanobis_distance(X_normal)

# 異常データのMD
md_anomaly = mt.mahalanobis_distance(X_anomaly)

# 全データの異常判定
X_all = np.vstack([X_normal, X_anomaly])
labels_true = np.array([0] * n_normal + [1] * n_anomaly)
predictions, md_all = mt.predict(X_all)

print(f"=== MT法による異常検知結果 ===")
print(f"正常データのMD: 平均={md_normal.mean():.3f}, 標準偏差={md_normal.std():.3f}")
print(f"異常データのMD: 平均={md_anomaly.mean():.3f}, 標準偏差={md_anomaly.std():.3f}")
print(f"検出された異常数: {predictions.sum()} / 実際の異常数: {labels_true.sum()}")

# --- 可視化 ---
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# MDのヒストグラム
ax1 = axes[0]
ax1.hist(md_normal, bins=30, alpha=0.6, label='Normal', color='blue', density=True)
ax1.hist(md_anomaly, bins=15, alpha=0.6, label='Anomaly', color='red', density=True)
ax1.axvline(x=4.0, color='green', linestyle='--', linewidth=2, label='Threshold=4')
ax1.set_xlabel('Mahalanobis Distance (MD)')
ax1.set_ylabel('Density')
ax1.set_title('Distribution of MD')
ax1.legend()
ax1.grid(True, alpha=0.3)

# MDの時系列的なプロット
ax2 = axes[1]
ax2.scatter(range(n_normal), md_normal, c='blue', s=10, alpha=0.5, label='Normal')
ax2.scatter(range(n_normal, n_normal + n_anomaly), md_anomaly,
            c='red', s=30, alpha=0.8, label='Anomaly')
ax2.axhline(y=4.0, color='green', linestyle='--', linewidth=2, label='Threshold')
ax2.set_xlabel('Sample Index')
ax2.set_ylabel('MD')
ax2.set_title('MD for Each Sample')
ax2.legend()
ax2.grid(True, alpha=0.3)

# 正常データのMDとカイ二乗分布の比較
ax3 = axes[2]
md_sorted = np.sort(md_normal * p)  # p倍してカイ二乗分布と比較
x_chi2 = np.linspace(0, md_sorted.max(), 200)
ax3.hist(md_sorted, bins=30, density=True, alpha=0.6, color='blue', label='Data')
ax3.plot(x_chi2, stats.chi2.pdf(x_chi2, df=p), 'r-', linewidth=2,
         label=f'$\\chi^2({p})$')
ax3.set_xlabel('$D^2$')
ax3.set_ylabel('Density')
ax3.set_title('$D^2$ vs $\\chi^2$ Distribution')
ax3.legend()
ax3.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()