import numpy as np
import matplotlib.pyplot as plt
from scipy import sparse
import time

# スパース行列の生成（COO形式）
row = np.array([0, 1, 2, 3, 0, 2])
col = np.array([0, 1, 2, 1, 3, 0])
data = np.array([5, 8, 3, 6, 2, 7])

A_coo = sparse.coo_matrix((data, (row, col)), shape=(4, 4))
print("COO形式:")
print(A_coo)
print(f"\n密行列:\n{A_coo.toarray()}")

# CSR形式への変換
A_csr = A_coo.tocsr()
print(f"\nCSR形式:")
print(f"  data:    {A_csr.data}")
print(f"  indices: {A_csr.indices}")
print(f"  indptr:  {A_csr.indptr}")

# ランダムスパース行列の生成と性能比較
n = 5000
density = 0.01  # 1%が非零

# スパース行列の生成
A_sparse = sparse.random(n, n, density=density, format='csr')
A_dense = A_sparse.toarray()
x = np.random.randn(n)

print(f"\n=== 性能比較 ({n}x{n}, density={density}) ===")
print(f"非零要素数: {A_sparse.nnz}")
print(f"スパース率: {1 - A_sparse.nnz / (n*n):.4f}")
print(f"密行列メモリ: {A_dense.nbytes / 1024**2:.1f} MB")
print(f"スパース行列メモリ: {(A_sparse.data.nbytes + A_sparse.indices.nbytes + A_sparse.indptr.nbytes) / 1024**2:.3f} MB")

# 行列-ベクトル積の速度比較
n_trials = 100
t_start = time.time()
for _ in range(n_trials):
    y_dense = A_dense @ x
t_dense = (time.time() - t_start) / n_trials

t_start = time.time()
for _ in range(n_trials):
    y_sparse = A_sparse @ x
t_sparse = (time.time() - t_start) / n_trials

print(f"\n行列-ベクトル積の平均時間:")
print(f"  密行列:      {t_dense*1000:.3f} ms")
print(f"  スパース行列: {t_sparse*1000:.3f} ms")
print(f"  速度比:       {t_dense/t_sparse:.1f}x")

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# 左: スパースパターンの表示
A_small = sparse.random(50, 50, density=0.05, format='csr')
axes[0].spy(A_small, markersize=2, color='navy')
axes[0].set_title(f'Sparsity Pattern (50x50, nnz={A_small.nnz})')
axes[0].set_xlabel('Column')
axes[0].set_ylabel('Row')

# 右: density vs 計算時間
densities = [0.001, 0.005, 0.01, 0.05, 0.1, 0.2, 0.5]
times_sparse = []
times_dense = []
n_test = 2000

for d in densities:
    A_sp = sparse.random(n_test, n_test, density=d, format='csr')
    A_de = A_sp.toarray()
    x_test = np.random.randn(n_test)

    t0 = time.time()
    for _ in range(20):
        A_sp @ x_test
    times_sparse.append((time.time() - t0) / 20 * 1000)

    t0 = time.time()
    for _ in range(20):
        A_de @ x_test
    times_dense.append((time.time() - t0) / 20 * 1000)

axes[1].semilogy(densities, times_sparse, 'bo-', label='Sparse (CSR)')
axes[1].semilogy(densities, times_dense, 'rs-', label='Dense')
axes[1].set_xlabel('Density')
axes[1].set_ylabel('Time (ms)')
axes[1].set_title(f'Matrix-Vector Product Time ({n_test}x{n_test})')
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()