!pip install openjij

import numpy as np
import matplotlib.pyplot as plt
import openjij as oj
import itertools
from scipy.special import logsumexp

def generate_bas_dataset():
    #画像のサイズ
    size = 8

    # 全ての 8ビットのパターンを生成(256通り)
    all_patterns = list(itertools.product([0, 1], repeat=size))
    all_patterns = np.array(all_patterns)

    dataset = []

    # Bars (横縞) の生成(各行が同じ値を持つパターン)
    for pat in all_patterns:

        # 行ごとにON/OFFが決まる
        #patは一次元配列、pat.reshape(-1, 1)で縦ベクトルの2次元配列(8, 1)を作成
        #np.tileで同じタイルを敷き詰める。(1, size=8)より縦方向に8回同じものをコピーして並べる。
        img = np.tile(pat.reshape(-1, 1), (1, size))

        # フラット化して1次元配列 (64ビット) にする
        flat_img = img.flatten()

        # 64ビット画像から先頭62ビットを取得（実質的に最後の2ピクセルを削除）
        input_bits = flat_img[:62]

        # ラベルBars = 01
        label = np.array([0, 1])

        # 62ビットのデータ + 2ビットのラベル
        record = np.concatenate([input_bits, label])
        dataset.append(record)

    # Stripes (縦縞) の生成(各列が同じ値を持つパターン)
    for pat in all_patterns:

        img = np.tile(pat.reshape(1, -1), (size, 1))

        flat_img = img.flatten()

        input_bits = flat_img[:62]

        # ラベルStripes = 10
        label = np.array([1, 0])

        record = np.concatenate([input_bits, label])
        dataset.append(record)

    return np.array(dataset)

def visualize_samples(dataset, num_samples=5):

    # ランダムにサンプル抽出
    indices = np.random.choice(len(dataset), num_samples, replace=False)

    fig, axes = plt.subplots(1, num_samples, figsize=(15, 3))

    for i, idx in enumerate(indices):
        record = dataset[idx]

        # 最後の2ビットはラベル
        label_bits = record[-2:]
        if np.array_equal(label_bits, [0, 1]):
            label_name = "Bar (01)"
        elif np.array_equal(label_bits, [1, 0]):
            label_name = "Stripe (10)"
        else:
            label_name = "Unknown"

        # ラベル付きのBASデータの画像
        img_reconstructed = record.reshape(8, 8)

        axes[i].imshow(img_reconstructed, cmap='cividis') # 論文の黄色・青に近い色味
        axes[i].set_title(f"Idx: {idx}\n{label_name}")
        axes[i].axis('off')

    plt.tight_layout()
    plt.show()

bas_dataset = generate_bas_dataset()

print(f"データセット形状: {bas_dataset.shape}") # 期待値: (512, 64)
print(f"内訳: Bars 256個 + Stripes 256個")

# データの可視化
visualize_samples(bas_dataset)

def __init__(self, n_visible, n_hidden, learning_rate=0.05):
    """
    n_visible: 可視層のユニット数
    n_hidden: 隠れ層のユニット数
    learning_rate: 学習率
    W: 重み行列(0付近の値で初期化)
    b: 可視層のバイアス(0で初期化)
    c: 隠れ層のバイアス(0で初期化)
    """
    self.n_visible = n_visible
    self.n_hidden = n_hidden
    self.lr = learning_rate
    self.W = np.random.normal(0, 0.01, (n_visible, n_hidden))
    self.b = np.zeros(n_visible)
    self.c = np.zeros(n_hidden)

    # 実験1、実験3でエポック毎の正解率と対数尤度をもとめるために保存用のリストを作る
    self.history = {'bars_acc': [], 'stripes_acc': [], 'll': []}

# シグモイド関数
def sigmoid(self, x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

# 可視層の状態から隠れ層の値をサンプリングする関数
def sample_h_given_v(self, v):
    # 隠れ層のユニットが1になる確率を可視層の値を基に計算
    prob_h = self.sigmoid(np.dot(v, self.W) + self.c)
    # 0以上1以下の一様分布に従う乱数を生成し、prob_hのユニット毎の値と比較して0か1を決める
    h_sample = (np.random.rand(*prob_h.shape) < prob_h).astype(float)
    return prob_h, h_sample

# 隠れ層の状態から可視層の値をサンプリングする関数
def sample_v_given_h(self, h):
    prob_v = self.sigmoid(np.dot(h, self.W.T) + self.b)
    v_sample = (np.random.rand(*prob_v.shape) < prob_v).astype(float)
    return prob_v, v_sample

def train_step_cd(self, batch):
    n_batch = batch.shape[0]
    # 可視層には訓練データを入力
    pos_v = batch
    prob_h0, pos_h = self.sample_h_given_v(pos_v) # 式(4)
    _, neg_v = self.sample_v_given_h(pos_h) # 式(5)
    prob_h1, _ = self.sample_h_given_v(neg_v) # 式(6)

    self.W += self.lr * (np.dot(pos_v.T, prob_h0) - np.dot(neg_v.T, prob_h1)) / n_batch # 式(1)
    self.b += self.lr * np.mean(pos_v - neg_v, axis=0) # 式(2)
    self.c += self.lr * np.mean(prob_h0 - prob_h1, axis=0) # 式(3)

def train_step_sa(self, batch, num_reads, num_sweeps=100):
    n_batch = batch.shape[0]
    pos_v = batch
    prob_h0, _ = self.sample_h_given_v(pos_v)

    sampler = oj.SASampler()
    qubo = {}
    for i in range(self.n_visible): qubo[(i, i)] = -self.b[i] # 対角成分に可視バイアスを代入
    for j in range(self.n_hidden): qubo[(self.n_visible+j, self.n_visible+j)] = -self.c[j] # 対角成分に隠れバイアスを代入
    for i in range(self.n_visible):
        for j in range(self.n_hidden): qubo[(i, self.n_visible+j)] = -self.W[i, j] # 対角成分より上側に重みを代入

    response = sampler.sample_qubo(qubo, num_reads=num_reads, num_sweeps=num_sweeps, beta_min=0.2, beta_max=1.0)
    neg_v = response.record.sample[:, :self.n_visible].astype(float)
    prob_h_model, _ = self.sample_h_given_v(neg_v)

    # パラメータをサンプル数で割って平均化する
    self.W += self.lr * (np.dot(pos_v.T, prob_h0) - np.dot(neg_v.T, prob_h_model)) / num_reads # 重みを平均化
    self.b += self.lr * (np.mean(pos_v, axis=0) - np.mean(neg_v, axis=0)) # 可視バイアスを平均化
    self.c += self.lr * (np.mean(prob_h0, axis=0) - np.mean(prob_h_model, axis=0)) # 隠れバイアスを平均化

def get_accuracy(self, test_subset, gibbs_steps=50):
    if len(test_subset) == 0: return 0.0
    correct = 0
    for rec in test_subset:
        true_label = rec[62:]
        v_curr = np.concatenate([rec[:62], np.random.randint(0, 2, size=2)])
        for _ in range(gibbs_steps):
            _, h_s = self.sample_h_given_v(v_curr)
            _, v_s = self.sample_v_given_h(h_s)
            v_curr[62:] = v_s[62:] # Label部分のみ更新
        if np.array_equal(v_curr[62:], true_label): correct += 1
    return correct / len(test_subset)

def reconstruct_image(self, original_vector, mask_indices, gibbs_steps=50):
    """
    指定したインデックスを破損させ、RBM学習器で復元する
    original_vector: 元の画像のビット配列
    mask_indices: 破損させるビットを示すインデックスの配列
    gibbs_steps: ギブスサンプリングのステップ数
    """
    corrupted_vector = original_vector.copy()
    # 破損部分を0と1でランダムに埋める
    noise = np.random.randint(0, 2, size=len(mask_indices))
    corrupted_vector[mask_indices] = noise

    v_curr = corrupted_vector.copy()
    all_indices = np.arange(len(original_vector))

    # np.setdiff1d(配列A, 配列B): 配列Aから配列Bに含まれる要素を除いたもの
    fixed_indices = np.setdiff1d(all_indices, mask_indices)

    for _ in range(gibbs_steps):
        _, h_s = self.sample_h_given_v(v_curr)
        _, v_s = self.sample_v_given_h(h_s)

        # 固定するビット
        v_curr[fixed_indices] = original_vector[fixed_indices]
        v_curr[mask_indices] = v_s[mask_indices]

    return corrupted_vector, v_curr

def get_incorrect_bits(self, dataset, mask_indices, gibbs_steps=50):
    """
    データセットに対し、破損された領域の再構成エラー（誤ビット数）を計算する
    """
    error_list = []
    for original_vector in dataset:
        # 再構成を実行
        _, restored_vector = self.reconstruct_image(original_vector, mask_indices, gibbs_steps)

        # 破損された範囲のビット列において、正解値と再構成後の値を比較
        actual = original_vector[mask_indices] # 正解値
        predicted = restored_vector[mask_indices] # 再構成後の値

        # 異なるビットの数をカウント
        incorrect_count = np.sum(actual != predicted)
        error_list.append(incorrect_count)

    return np.array(error_list)

def get_log_likelihood(self, data, all_hidden_states):
    # 分配関数 Z の計算
    term_h = np.dot(all_hidden_states, self.c)
    vis_act = self.b + np.dot(all_hidden_states, self.W.T)
    # logaddexp(x, y) = log(e^x + e^y)
    log_Z = logsumexp(term_h + np.sum(np.logaddexp(0, vis_act), axis=1))

    # 非正規化尤度の計算
    data_term = np.dot(data, self.b)
    hid_act = self.c + np.dot(data, self.W)
    log_unnorm = data_term + np.sum(np.logaddexp(0, hid_act), axis=1)

    return np.mean(log_unnorm - log_Z)

class IntegratedRBM:
    def __init__(self, n_visible, n_hidden, learning_rate=0.05):
        self.n_visible = n_visible
        self.n_hidden = n_hidden
        self.lr = learning_rate
        self.W = np.random.normal(0, 0.01, (n_visible, n_hidden))
        self.b = np.zeros(n_visible)
        self.c = np.zeros(n_hidden)

        # 履歴保存用：bars正解率、stripes正解率、対数尤度
        self.history = {'bars_acc': [], 'stripes_acc': [], 'll': []}

    def sigmoid(self, x):
        return 1.0 / (1.0 + np.exp(-np.clip(x, -500, 500)))

    def sample_h_given_v(self, v):
        prob_h = self.sigmoid(np.dot(v, self.W) + self.c)
        h_sample = (np.random.rand(*prob_h.shape) < prob_h).astype(float)
        return prob_h, h_sample

    def sample_v_given_h(self, h):
        prob_v = self.sigmoid(np.dot(h, self.W.T) + self.b)
        v_sample = (np.random.rand(*prob_v.shape) < prob_v).astype(float)
        return prob_v, v_sample

    # CD-1のサンプリングによる学習
    def train_step_cd(self, batch):
        n_batch = batch.shape[0]
        pos_v = batch
        prob_h0, pos_h = self.sample_h_given_v(pos_v)
        _, neg_v = self.sample_v_given_h(pos_h)
        prob_h1, _ = self.sample_h_given_v(neg_v)

        self.W += self.lr * (np.dot(pos_v.T, prob_h0) - np.dot(neg_v.T, prob_h1)) / n_batch
        self.b += self.lr * np.mean(pos_v - neg_v, axis=0)
        self.c += self.lr * np.mean(prob_h0 - prob_h1, axis=0)

    # SAのサンプリングによる学習
    def train_step_sa(self, batch, num_reads, num_sweeps=100):
        n_batch = batch.shape[0]
        pos_v = batch
        prob_h0, _ = self.sample_h_given_v(pos_v)

        sampler = oj.SASampler()
        qubo = {}
        for i in range(self.n_visible): qubo[(i, i)] = -self.b[i]
        for j in range(self.n_hidden): qubo[(self.n_visible+j, self.n_visible+j)] = -self.c[j]
        for i in range(self.n_visible):
            for j in range(self.n_hidden): qubo[(i, self.n_visible+j)] = -self.W[i, j]

        response = sampler.sample_qubo(qubo, num_reads=num_reads, num_sweeps=num_sweeps, beta_min=0.2, beta_max=1.0)
        neg_v = response.record.sample[:, :self.n_visible].astype(float)
        prob_h_model, _ = self.sample_h_given_v(neg_v)

        self.W += self.lr * (np.dot(pos_v.T, prob_h0) - np.dot(neg_v.T, prob_h_model)) / num_reads
        self.b += self.lr * (np.mean(pos_v, axis=0) - np.mean(neg_v, axis=0))
        self.c += self.lr * (np.mean(prob_h0, axis=0) - np.mean(prob_h_model, axis=0))

    # 画像分類の正解率
    def get_accuracy(self, test_subset, gibbs_steps=50):
        if len(test_subset) == 0: return 0.0
        correct = 0
        for rec in test_subset:
            true_label = rec[62:]
            v_curr = np.concatenate([rec[:62], np.random.randint(0, 2, size=2)])
            for _ in range(gibbs_steps):
                _, h_s = self.sample_h_given_v(v_curr)
                _, v_s = self.sample_v_given_h(h_s)
                v_curr[62:] = v_s[62:]
            if np.array_equal(v_curr[62:], true_label): correct += 1
        return correct / len(test_subset)

    # 画像再構成
    def reconstruct_image(self, original_vector, mask_indices, gibbs_steps=50):

        corrupted_vector = original_vector.copy()

        noise = np.random.randint(0, 2, size=len(mask_indices))
        corrupted_vector[mask_indices] = noise

        v_curr = corrupted_vector.copy()
        all_indices = np.arange(len(original_vector))
        fixed_indices = np.setdiff1d(all_indices, mask_indices)

        for _ in range(gibbs_steps):
            _, h_s = self.sample_h_given_v(v_curr)
            _, v_s = self.sample_v_given_h(h_s)

            v_curr[fixed_indices] = original_vector[fixed_indices]
            v_curr[mask_indices] = v_s[mask_indices]

        return corrupted_vector, v_curr

    # 画像再構成のエラー数
    def get_incorrect_bits(self, dataset, mask_indices, gibbs_steps=50):

        error_list = []
        for original_vector in dataset:
            _, restored_vector = self.reconstruct_image(original_vector, mask_indices, gibbs_steps)

            actual = original_vector[mask_indices]
            predicted = restored_vector[mask_indices]

            incorrect_count = np.sum(actual != predicted)
            error_list.append(incorrect_count)

        return np.array(error_list)

    # 対数尤度
    def get_log_likelihood(self, data, all_hidden_states):

        term_h = np.dot(all_hidden_states, self.c)
        vis_act = self.b + np.dot(all_hidden_states, self.W.T)
        log_Z = logsumexp(term_h + np.sum(np.logaddexp(0, vis_act), axis=1))

        data_term = np.dot(data, self.b)
        hid_act = self.c + np.dot(data, self.W)
        log_unnorm = data_term + np.sum(np.logaddexp(0, hid_act), axis=1)

        return np.mean(log_unnorm - log_Z)

# データセットの生成
dataset = generate_bas_dataset()

# ラベルごとにデータを分ける
# Label: Bars=[0,1], Stripes=[1,0]
all_bars = dataset[(dataset[:, 62] == 0) & (dataset[:, 63] == 1)]
all_stripes = dataset[(dataset[:, 62] == 1) & (dataset[:, 63] == 0)]

# それぞれのデータセットをシャッフル
np.random.shuffle(all_bars)
np.random.shuffle(all_stripes)

# 訓練データに割り当てる数を決める
n_train_per_class = 200

# 各データセットから200個ずつ訓練データへ、残りをテストデータへ
train_bars = all_bars[:n_train_per_class]
test_bars = all_bars[n_train_per_class:]

train_stripes = all_stripes[:n_train_per_class]
test_stripes = all_stripes[n_train_per_class:]

# 訓練データとテストデータを結合して、さらに全体をシャッフル
train_data = np.concatenate([train_bars, train_stripes])
test_data = np.concatenate([test_bars, test_stripes])

np.random.shuffle(train_data)
np.random.shuffle(test_data)

# 結果の確認
print(f"Train Set: Bars={len(train_bars)}, Stripes={len(train_stripes)} (Total: {len(train_data)})")
print(f"Test Set: Bars={len(test_bars)}, Stripes={len(test_stripes)} (Total: {len(test_data)})")

"""
epochs: 学習回数
batch_size: バッチサイズ
rbm_cd: CD-1のサンプリングによる学習器
rbm_sa: SAのサンプリングによる学習器
"""

# モデルの初期化
epochs = 400
batch_size = 64
rbm_cd = IntegratedRBM(64, 64, learning_rate=0.10)
rbm_sa = IntegratedRBM(64, 64, learning_rate=0.10)

print("Training starting...")
for epoch in range(1, epochs + 1):
    np.random.shuffle(train_data)
    for i in range(0, len(train_data), batch_size):
        batch = train_data[i:i+batch_size]
        rbm_cd.train_step_cd(batch)
        rbm_sa.train_step_sa(batch, num_reads=batch_size)

    # エポック毎に正解率を記録
    rbm_cd.history['bars_acc'].append(rbm_cd.get_accuracy(test_bars))
    rbm_cd.history['stripes_acc'].append(rbm_cd.get_accuracy(test_stripes))
    rbm_sa.history['bars_acc'].append(rbm_sa.get_accuracy(test_bars))
    rbm_sa.history['stripes_acc'].append(rbm_sa.get_accuracy(test_stripes))

    if epoch % 50 == 0:
        print(f"Epoch {epoch} completed.")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Barsの正解率
axes[0].plot(range(epochs), rbm_cd.history['bars_acc'], label='CD-1', color='black')
axes[0].plot(range(epochs), rbm_sa.history['bars_acc'], label='SA (OpenJij)', color='orange')
axes[0].set_title('Classification Accuracy (Bars)')
axes[0].set_xlabel('Epoch'); axes[0].set_ylabel('Accuracy')
axes[0].set_ylim(0, 1.05); axes[0].legend(); axes[0].grid(True, linestyle='--', alpha=0.6)

# Stripesの正解率
axes[1].plot(range(epochs), rbm_cd.history['stripes_acc'], label='CD-1', color='black')
axes[1].plot(range(epochs), rbm_sa.history['stripes_acc'], label='SA (OpenJij)', color='orange')
axes[1].set_title('Classification Accuracy (Stripes)')
axes[1].set_xlabel('Epoch'); axes[1].set_ylabel('Accuracy')
axes[1].set_ylim(0, 1.05); axes[1].legend(); axes[1].grid(True, linestyle='--', alpha=0.6)

plt.tight_layout()
plt.show()

# サンプル画像の用意
sample_bar = dataset[10]      # Barパターン
sample_stripe = dataset[-46]  # Stripeパターン

# ケースA: ラベルの破損
mask_A = [62, 63]
target_A = sample_bar.copy()
# target_A = sample_stripe.copy()
# クラスメソッドとして呼び出し
corp_A, cd_res_A = rbm_cd.reconstruct_image(target_A, mask_A)
_, sa_res_A      = rbm_sa.reconstruct_image(target_A, mask_A)

# ケースB: 右下4×4(16ビット)の破損
mask_B = []
for r in range(4, 8):
    for c in range(4, 8):
        mask_B.append(r * 8 + c)
mask_B = np.array(mask_B)

target_B = sample_bar.copy()
# target_B = sample_stripe.copy()
corp_B, cd_res_B = rbm_cd.reconstruct_image(target_B, mask_B)
_, sa_res_B      = rbm_sa.reconstruct_image(target_B, mask_B)

# ケースC: すべてのビットの破損
mask_C = np.arange(64)
target_C = sample_bar.copy()
# target_C = sample_stripe.copy()

corp_C, cd_res_C = rbm_cd.reconstruct_image(target_C, mask_C)
_, sa_res_C      = rbm_sa.reconstruct_image(target_C, mask_C)

# 描写
fig, axes = plt.subplots(3, 3, figsize=(10, 10))
cmap = 'viridis'
def to_img(vec): return vec.reshape(8, 8)

titles = [
    ("A: Label Masked(2 bits)", corp_A, cd_res_A, sa_res_A),
    ("B: Bottom-Right Mask(16 bits)", corp_B, cd_res_B, sa_res_B),
    ("C: Full Noise(64 bits)", corp_C, cd_res_C, sa_res_C)
]

for i, (title, corp, cd_out, sa_out) in enumerate(titles):
    # 破損した画像の表示
    axes[i, 0].imshow(to_img(corp), cmap=cmap, vmin=0, vmax=1)
    axes[i, 0].set_title(f"{title}\nInput")
    axes[i, 0].axis('off')

    # rbm_cdで再構成した画像の表示
    axes[i, 1].imshow(to_img(cd_out), cmap=cmap, vmin=0, vmax=1)
    axes[i, 1].set_title("CD-1 Output")
    axes[i, 1].axis('off')

    # rbm_saで再構成した画像の表示
    axes[i, 2].imshow(to_img(sa_out), cmap=cmap, vmin=0, vmax=1)
    axes[i, 2].set_title("SA (OpenJij) Output")
    axes[i, 2].axis('off')

plt.tight_layout()
plt.show()

# 画像の右下4×4(16ビット)に対応するインデックスを配列に加える
mask_indices = []
for r in range(4, 8):
    for c in range(4, 8):
        mask_indices.append(r * 8 + c)

mask_indices = np.array(mask_indices)
mask_indices = mask_indices[mask_indices < 62]

# 誤りビット数の計算
print("Calculating reconstruction errors for CD and SA models...")
# クラスメソッドとして呼び出し
cd_train_errors = rbm_cd.get_incorrect_bits(train_data, mask_indices)
cd_test_errors  = rbm_cd.get_incorrect_bits(test_data, mask_indices)
sa_train_errors = rbm_sa.get_incorrect_bits(train_data, mask_indices)
sa_test_errors  = rbm_sa.get_incorrect_bits(test_data, mask_indices)

# 描写
fig, axes = plt.subplots(2, 2, figsize=(12, 10))

max_bits = len(mask_indices)
bins = np.arange(0, max_bits + 2) - 0.5

def plot_hist(ax, data, title, color):
    ax.hist(data, bins=bins, rwidth=0.8, color=color, edgecolor='black', alpha=0.7)
    ax.set_title(title, fontsize=14, fontweight='bold')
    ax.set_xlabel("No. of incorrectly predicted bits")
    ax.set_ylabel("No. of instances")
    ax.set_xticks(range(0, max_bits + 1, 2))
    ax.set_ylim(0, 150)
    ax.grid(axis='y', linestyle='--', alpha=0.5)

# CD学習RBMの誤りビット数(訓練データ)
plot_hist(axes[0, 0], cd_train_errors, "D: CD; Training Data", 'tab:blue')

# CD学習RBMの誤りビット数(テストデータ)
plot_hist(axes[0, 1], cd_test_errors, "E: CD; Test Data", 'tab:blue')

# SA学習RBMの誤りビット数(訓練データ)
plot_hist(axes[1, 0], sa_train_errors, "F: SA; Training Data", 'tab:orange')

# SA学習RBMの誤りビット数(テストデータ)
plot_hist(axes[1, 1], sa_test_errors, "G: SA; Test Data", 'tab:orange')

plt.tight_layout()
plt.show()

# パラメータ設定
n_visible = 64
n_hidden = 20 # 隠れ層のノードを減らす
epochs = 400
batch_size = 64

# 尤度計算用の全隠れ状態の生成 (2^20通り)
print(f"Generating {1 << n_hidden} hidden states...")
indices = np.arange(1 << n_hidden, dtype=np.uint32)[:, np.newaxis]
bit_masks = 1 << np.arange(n_hidden, dtype=np.uint32)[::-1]
all_hidden_states = (indices & bit_masks > 0).astype(float)
print("Done.")

# 比較用の2つのモデルを初期化
rbm_cd2 = IntegratedRBM(n_visible, n_hidden, learning_rate=0.10)
rbm_sa2 = IntegratedRBM(n_visible, n_hidden, learning_rate=0.10)

# 履歴保存用のリストを追加
ll_history_cd = []
ll_history_sa = []
epoch_list = []

print("Starting Training Comparison...")

for epoch in range(1, epochs + 1):
    np.random.shuffle(train_data)

    # バッチ学習
    for i in range(0, len(train_data), batch_size):
        batch = train_data[i:i+batch_size]

        # それぞれ独立したモデルとして学習
        rbm_cd2.train_step_cd(batch)
        rbm_sa2.train_step_sa(batch, num_reads=batch_size)

    # 10エポックごとに尤度を計算して比較
    if epoch % 10 == 0 or epoch == 1:
        # 両モデルの対数尤度を算出
        ll_cd = rbm_cd2.get_log_likelihood(train_data, all_hidden_states)
        ll_sa = rbm_sa2.get_log_likelihood(train_data, all_hidden_states)

        # データの保存
        epoch_list.append(epoch)
        ll_history_cd.append(ll_cd)
        ll_history_sa.append(ll_sa)

        print(f"Epoch {epoch:3d}: CD LL={ll_cd:.2f}, SA LL={ll_sa:.2f}")

# 5. 対数尤度の推移を描写

plt.figure(figsize=(8, 6))
plt.plot(epoch_list, ll_history_cd, 'o-', color='black', label='CD-1', markersize=4)
plt.plot(epoch_list, ll_history_sa, 's-', color='red', label='SA (OpenJij)', markersize=4)

plt.xlabel("Epoch")
plt.ylabel("Log-likelihood")
plt.title(f"Log-likelihood Comparison")
plt.legend()
plt.grid(True, linestyle='--', alpha=0.5)
plt.tight_layout()
plt.show()