import numpy as np
from scipy import stats
import pandas as pd
from numpy.linalg import eig, inv, svd

# 表4.11的数据
data = [
    [191, 155, 179, 145],
    [195, 149, 201, 152],
    [181, 148, 185, 149],
    [183, 153, 188, 149],
    [176, 144, 171, 142],
    [208, 157, 192, 152],
    [189, 150, 190, 149],
    [197, 159, 189, 152],
    [188, 152, 197, 159],
    [192, 150, 187, 151],
    [179, 158, 186, 148],
    [183, 147, 174, 147],
    [174, 150, 185, 152],
    [190, 159, 195, 157],
    [188, 151, 187, 158],
    [163, 137, 161, 130],
    [195, 155, 183, 158],
    [186, 153, 173, 148],
    [181, 145, 182, 146],
    [175, 140, 165, 137],
    [192, 154, 185, 152],
    [174, 143, 178, 147],
    [176, 139, 176, 143],
    [197, 167, 200, 158],
    [190, 163, 187, 150]
]

# 转换为DataFrame以便数据处理
df = pd.DataFrame(data, columns=['x1', 'x2', 'y1', 'y2'])

# 将数据分为X组(父母)和Y组(子女)
X = df[['x1', 'x2']].values
Y = df[['y1', 'y2']].values

# 计算样本数
n = X.shape[0]

# 步骤1: 计算协方差矩阵
# 计算X和Y的协方差矩阵
S_xx = np.cov(X, rowvar=False) * (n-1)/n  # 使用无偏估计
S_yy = np.cov(Y, rowvar=False) * (n-1)/n
S_xy = np.zeros((2, 2))
for i in range(n):
    S_xy += np.outer(X[i] - X.mean(axis=0), Y[i] - Y.mean(axis=0))
S_xy /= n
S_yx = S_xy.T

# 整个协方差矩阵
S = np.block([
    [S_xx, S_xy],
    [S_yx, S_yy]
])

print("协方差矩阵 S_xx:")
print(S_xx)
print("\n协方差矩阵 S_yy:")
print(S_yy)
print("\n协方差矩阵 S_xy:")
print(S_xy)

# 步骤2: 计算相关系数矩阵
# 计算X和Y的相关系数矩阵
D_x = np.sqrt(np.diag(np.diag(S_xx)))
D_y = np.sqrt(np.diag(np.diag(S_yy)))
D_x_inv = inv(D_x)
D_y_inv = inv(D_y)

R_xx = D_x_inv @ S_xx @ D_x_inv
R_yy = D_y_inv @ S_yy @ D_y_inv
R_xy = D_x_inv @ S_xy @ D_y_inv
R_yx = R_xy.T

# 整个相关系数矩阵
R = np.block([
    [R_xx, R_xy],
    [R_yx, R_yy]
])

print("\n相关系数矩阵 R_xx:")
print(R_xx)
print("\n相关系数矩阵 R_yy:")
print(R_yy)
print("\n相关系数矩阵 R_xy:")
print(R_xy)

# 步骤3: 典型相关分析
# 计算 R_xx^(-1/2) * R_xy * R_yy^(-1) * R_yx * R_xx^(-1/2)
R_xx_inv_sqrt = inv(np.linalg.cholesky(R_xx))
R_yy_inv = inv(R_yy)
M = R_xx_inv_sqrt @ R_xy @ R_yy_inv @ R_yx @ R_xx_inv_sqrt

# 计算M的特征值和特征向量
eigvals, eigvecs = eig(M)
# 对特征值排序
idx = np.argsort(eigvals)[::-1]
eigvals = eigvals[idx]
eigvecs = eigvecs[:, idx]

# 典型相关系数是特征值的平方根
canonical_corr = np.sqrt(np.real(eigvals))

print("\n典型相关系数:")
for i, corr in enumerate(canonical_corr):
    print(f"r_{i+1} = {corr:.4f}")

# 计算 a 和 b 系数
a = R_xx_inv_sqrt @ eigvecs
# 计算 b
b = R_yy_inv @ R_yx @ a @ np.diag(1/canonical_corr)

print("\nX组的典型系数 a:")
print(a)
print("\nY组的典型系数 b:")
print(b)

# 典型变量
U = X @ a
V = Y @ b

# 步骤4: 进行显著性检验
# Wilks' Lambda检验
def wilks_lambda_test(canonical_corr, n, p, q):
    """
    进行Wilks' Lambda显著性检验
    canonical_corr: 典型相关系数
    n: 样本数
    p: X变量数
    q: Y变量数
    """
    results = []
    m = min(p, q)
    
    for k in range(m):
        # 计算Wilks' Lambda
        lambda_val = np.prod([1 - cc**2 for cc in canonical_corr[k:]])
        
        # 自由度
        df1 = (p - k) * (q - k)
        df2 = n - p - q + 1 - k
        
        # F统计量
        s = np.sqrt((p**2 * q**2 - 4) / (p**2 + q**2 - 5))
        if s > 1:
            s = 1
        r = (n - 1 - (p + q) / 2) / 2
        u = (p * q - 2) / 4
        
        if p**2 + q**2 - 5 > 0:
            t = np.sqrt((p * q)**2 / (p**2 + q**2 - 5))
            F = ((1 - lambda_val**(1/t)) / (lambda_val**(1/t))) * ((r * t - u) / df1)
        else:
            F = ((1 - lambda_val) / lambda_val) * (df2 / df1)
            
        # P值
        p_value = 1 - stats.f.cdf(F, df1, df2)
        
        results.append({
            'k': k,
            'Wilks Lambda': lambda_val,
            'F': F,
            'df1': df1,
            'df2': df2,
            'p-value': p_value
        })
    
    return results

# 执行Wilks' Lambda检验
p = X.shape[1]  # X变量数
q = Y.shape[1]  # Y变量数
test_results = wilks_lambda_test(canonical_corr, n, p, q)

print("\nWilks' Lambda检验结果:")
for result in test_results:
    print(f"检验第{result['k']+1}个及以后的典型相关：")
    print(f"Wilks' Lambda = {result['Wilks Lambda']:.4f}")
    print(f"F = {result['F']:.4f}, df1 = {result['df1']}, df2 = {result['df2']}")
    print(f"p-value = {result['p-value']:.4f}")
    if result['p-value'] < 0.05:
        print("在α=0.05水平下显著\n")
    else:
        print("在α=0.05水平下不显著\n")

# 计算标准化的典型载荷（结构相关系数）
# X组变量与U的相关系数
r_x_u = np.zeros((p, p))
for i in range(p):
    for j in range(p):
        r_x_u[i, j] = np.corrcoef(X[:, i], U[:, j])[0, 1]

# Y组变量与V的相关系数
r_y_v = np.zeros((q, q))
for i in range(q):
    for j in range(q):
        r_y_v[i, j] = np.corrcoef(Y[:, i], V[:, j])[0, 1]

print("\nX组变量与典型变量U的结构相关系数:")
print(r_x_u)
print("\nY组变量与典型变量V的结构相关系数:")
print(r_y_v)

# 计算典型冗余
# X经U解释Y的比例
redundancy_x_via_u = np.mean(r_x_u**2, axis=0) * canonical_corr**2

# Y经V解释X的比例
redundancy_y_via_v = np.mean(r_y_v**2, axis=0) * canonical_corr**2

print("\nX经典型变量U解释Y的冗余:")
for i, red in enumerate(redundancy_x_via_u):
    print(f"第{i+1}对: {red:.4f}")

print("\nY经典型变量V解释X的冗余:")
for i, red in enumerate(redundancy_y_via_v):
    print(f"第{i+1}对: {red:.4f}")

# 总结结果
print("\n===== 总结 =====")
print(f"在α=0.05水平下，第一个典型相关系数r_1 = {canonical_corr[0]:.4f}，"
      f"p值为{test_results[0]['p-value']:.4f}，{"显著" if test_results[0]['p-value'] < 0.05 else "不显著"}")
if len(canonical_corr) > 1:
    print(f"第二个典型相关系数r_2 = {canonical_corr[1]:.4f}，"
          f"p值为{test_results[1]['p-value']:.4f}，{"显著" if test_results[1]['p-value'] < 0.05 else "不显著"}")