import numpy as np
import random
import math
import matplotlib.pyplot as plt

def get_related_cells(row, col):
    """指定されたマスに関連する（同じ行、列、3x3ブロックにある）すべてのマスのインデックスを返す"""
    related = set()
    # 同じ行と列
    for i in range(9):
        related.add((row, i))
        related.add((i, col))
    # 同じ3x3ブロック
    start_row, start_col = 3 * (row // 3), 3 * (col // 3)
    for i in range(3):
        for j in range(3):
            related.add((start_row + i, start_col + j))
    related.remove((row, col))
    return list(related)

def calculate_energy(board):
    """
    エネルギー関数 H(X) を計算する。
    エネルギーは、制約を破っているペア（同じ行・列・ブロックに同じ数字がある）の数。
    """
    energy = 0
    # 行
    for i in range(9):
        counts = np.bincount(board[i, :][board[i, :] > 0])
        energy += sum(c * (c - 1) // 2 for c in counts)
    # 列
    for i in range(9):
        counts = np.bincount(board[:, i][board[:, i] > 0])
        energy += sum(c * (c - 1) // 2 for c in counts)
    # 3x3ブロック
    for br in range(3):
        for bc in range(3):
            block = board[br*3:(br+1)*3, bc*3:(bc+1)*3].flatten()
            counts = np.bincount(block[block > 0])
            energy += sum(c * (c - 1) // 2 for c in counts)
            
    return energy

def gibbs_sampler_for_sudoku(initial_board, M, B):
    """
    ギブスサンプラーを用いて数独を解く (アルゴリズム 11)

    :param initial_board: 初期盤面 (0は空マス)
    :param M: サンプルサイズ (イテレーション回数)
    :param B: 逆温度 (beta)
    :return: (解かれた盤面, イテレーション履歴, エネルギー履歴)
    """
    # 状態の初期化 (X <- Xo)
    board = np.copy(initial_board)
    fixed_cells = np.where(board > 0)
    empty_cells = list(zip(*np.where(board == 0)))

    # 空のマスをランダムな数字で埋める
    for r, c in empty_cells:
        board[r, c] = random.randint(1, 9)

    print("初期盤面（ランダムに充填）:")
    print(board)
    print("-" * 20)

    # プロット用の履歴
    iteration_history = []
    energy_history = []

    # for t <- 1 to M do
    for t in range(M):
        # j ~ Unif({1, ..., |V|}) (ランダムに頂点を選択)
        # ここでは変更可能な空きマスからランダムに一つ選ぶ
        if not empty_cells:
            print("変更可能なマスがありません。")
            break
        
        r, c = random.choice(empty_cells)
        
        weights = []
        # W <- 0
        W = 0.0

        # for k = 1 to q do (q=9)
        for k in range(1, 10):
            original_val = board[r, c]
            board[r, c] = k # Xj <- k

            # H(X(k)) の計算
            # 全体を再計算する代わりに、変更したマス周辺のエネルギー変化を計算する
            conflicts = 0
            for nr, nc in get_related_cells(r, c):
                if board[nr, nc] == k:
                    conflicts += 1
            
            # pk <- exp(-BH(X(k))) (重みの計算)
            pk = math.exp(-B * conflicts)
            weights.append(pk)
            # W <- W + pk
            W += pk
        
        # (pk)k <- (pk)k / W (重みの正規化)
        probabilities = [w / W for w in weights]

        # Xj ~ Cat(p1, ..., pq) (カテゴリカル分布からのサンプリング)
        new_val = np.random.choice(np.arange(1, 10), p=probabilities)
        board[r, c] = new_val

        # イテレーションの状況を出力し、履歴を保存
        if (t + 1) % 100 == 0: # 100回ごとに記録
            current_energy = calculate_energy(board)
            iteration_history.append(t + 1)
            energy_history.append(current_energy)
            if (t + 1) % 1000 == 0: # 1000回ごとに表示
                print(f"イテレーション {t+1}/{M}, 現在のエネルギー (違反数): {current_energy}")
            if current_energy == 0:
                print(f"イテレーション {t+1} で解が見つかりました！")
                # 最後の状態を記録
                iteration_history.append(t + 1)
                energy_history.append(current_energy)
                break
    
    return board, iteration_history, energy_history

if __name__ == '__main__':
    # 数独の問題 (0は空マス)
    # 'Easy'レベルの問題
    problem = np.array([
        [5, 3, 0, 0, 7, 0, 0, 0, 0],
        [6, 0, 0, 1, 9, 5, 0, 0, 0],
        [0, 9, 8, 0, 0, 0, 0, 6, 0],
        [8, 0, 0, 0, 6, 0, 0, 0, 3],
        [4, 0, 0, 8, 0, 3, 0, 0, 1],
        [7, 0, 0, 0, 2, 0, 0, 0, 6],
        [0, 6, 0, 0, 0, 0, 2, 8, 0],
        [0, 0, 0, 4, 1, 9, 0, 0, 5],
        [0, 0, 0, 0, 8, 0, 0, 7, 9]
    ])

    # パラメータ設定
    M = 50000  # サンプルサイズ (イテレーション回数)
    B = 5.0    # 逆温度

    print("元の問題:")
    print(problem)
    print("-" * 20)

    # ギブスサンプラーを実行
    solution, iterations, energies = gibbs_sampler_for_sudoku(problem, M, B)

    print("-" * 20)
    print("最終結果:")
    print(solution)
    print(f"最終エネルギー (違反数): {calculate_energy(solution)}")

    # 結果をプロット
    if iterations:
        plt.figure(figsize=(10, 6))
        plt.plot(iterations, energies, marker='.', linestyle='-')
        plt.title('Energy (Violations) over Iterations')
        plt.xlabel('Iteration')
        plt.ylabel('Number of Violations (Energy)')
        plt.grid(True)
        plt.show()