# １．パッケージのインストールやライブラリのインポート

# TYTANパッケージをインストール
!pip install -U git+https://github.com/tytansdk/tytan

# TYTANパッケージの全てのライブラリをインポート
from tytan import *


# ２．ルールの式（QUBO形式）を書く

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール1：都市と順番の6x5の計30個の量子ビットを用意
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
q = symbols_list([6, 5], 'q{}_{}')
print("都市と順番の6x5の計30個の量子ビット：\n", q, "\n")

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール2：N番目に訪れる都市は1つだけ（絶対に守ってほしい：重み係数は大きめ）
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
H = 0
k = 1    # 重み係数
H += k * (q[0][0] + q[0][1] + q[0][2] + q[0][3] + q[0][4] - 1)**2   # 1番目
H += k * (q[1][0] + q[1][1] + q[1][2] + q[1][3] + q[1][4] - 1)**2   # 2番目
H += k * (q[2][0] + q[2][1] + q[2][2] + q[2][3] + q[2][4] - 1)**2   # 3番目
H += k * (q[3][0] + q[3][1] + q[3][2] + q[3][3] + q[3][4] - 1)**2   # 4番目
H += k * (q[4][0] + q[4][1] + q[4][2] + q[4][3] + q[4][4] - 1)**2   # 5番目
H += k * (q[5][0] + q[5][1] + q[5][2] + q[5][3] + q[5][4] - 1)**2   # 6番目

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール3：ある都市に訪れるタイミングは1回だけ（絶対に守ってほしい）
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
k = 1    # 重み係数
H += k * (q[0][0] + q[1][0] + q[2][0] + q[3][0] + q[4][0]  - 1)**2   # A都市
H += k * (q[0][1] + q[1][1] + q[2][1] + q[3][1] + q[4][1]  - 1)**2   # B都市
H += k * (q[0][2] + q[1][2] + q[2][2] + q[3][2] + q[4][2]  - 1)**2   # C都市
H += k * (q[0][3] + q[1][3] + q[2][3] + q[3][3] + q[4][3]  - 1)**2   # D都市
H += k * (q[0][4] + q[1][4] + q[2][4] + q[3][4] + q[4][4]  - 1)**2   # E都市

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール4：6番目と1番目の都市が同じ場合に報酬（絶対に守ってほしい）
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
k = -1    # 重み係数
H += k * (q[5][0] * q[0][0])   # A都市
H += k * (q[5][1] * q[0][1])   # B都市
H += k * (q[5][2] * q[0][2])   # C都市
H += k * (q[5][3] * q[0][3])   # D都市
H += k * (q[5][4] * q[0][4])   # E都市

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール5：都市間の距離に比例したペナルティ
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# 重み係数について：
# 合計係数値が絶対に守ってほしい係数を超えない程度にうまく調整が必要！
# 具体的には、距離の値が大きめなので係数はかなり小さくする必要がある。
k = 0.001
# k = 0.1
# k = 1

# ペナルティーの式は計125ケース！
# 1番目から2番目の都市への移動（5 * 5 = 25ケース）
H +=  0 * k * (q[0][0] * q[1][0])    # A都市 → A都市
H += 30 * k * (q[0][1] * q[1][0])    # B都市 → A都市
H += 20 * k * (q[0][2] * q[1][0])    # C都市 → A都市
H += 50 * k * (q[0][3] * q[1][0])    # D都市 → A都市
H += 40 * k * (q[0][4] * q[1][0])    # E都市 → A都市

H += 20 * k * (q[0][0] * q[1][1])    # A都市 → B都市
H +=  0 * k * (q[0][1] * q[1][1])    # B都市 → B都市
H += 10 * k * (q[0][2] * q[1][1])    # C都市 → B都市
H += 30 * k * (q[0][3] * q[1][1])    # D都市 → B都市
H += 60 * k * (q[0][4] * q[1][1])    # E都市 → B都市

H += 20 * k * (q[0][0] * q[1][2])    # A都市 → C都市
H += 10 * k * (q[0][1] * q[1][2])    # B都市 → C都市
H +=  0 * k * (q[0][2] * q[1][2])    # C都市 → C都市
H += 20 * k * (q[0][3] * q[1][2])    # D都市 → C都市
H += 20 * k * (q[0][4] * q[1][2])    # E都市 → C都市

H += 50 * k * (q[0][0] * q[1][3])    # A都市 → D都市
H += 30 * k * (q[0][1] * q[1][3])    # B都市 → D都市
H += 30 * k * (q[0][2] * q[1][3])    # C都市 → D都市
H +=  0 * k * (q[0][3] * q[1][3])    # D都市 → D都市
H += 10 * k * (q[0][4] * q[1][3])    # E都市 → D都市

H += 40 * k * (q[0][0] * q[1][4])    # A都市 → E都市
H += 20 * k * (q[0][1] * q[1][4])    # B都市 → E都市
H += 20 * k * (q[0][2] * q[1][4])    # C都市 → E都市
H += 10 * k * (q[0][3] * q[1][4])    # D都市 → E都市
H +=  0 * k * (q[0][4] * q[1][4])    # E都市 → E都市

# 2番目から3番目の都市への移動（5 * 5 = 25ケース）
H +=  0 * k * (q[1][0] * q[2][0])
H += 30 * k * (q[1][1] * q[2][0])
H += 20 * k * (q[1][2] * q[2][0])
H += 50 * k * (q[1][3] * q[2][0])
H += 40 * k * (q[1][4] * q[2][0])

H += 20 * k * (q[1][0] * q[2][1])
H +=  0 * k * (q[1][1] * q[2][1])
H += 10 * k * (q[1][2] * q[2][1])
H += 30 * k * (q[1][3] * q[2][1])
H += 60 * k * (q[1][4] * q[2][1])

H += 20 * k * (q[1][0] * q[2][2])
H += 10 * k * (q[1][1] * q[2][2])
H +=  0 * k * (q[1][2] * q[2][2])
H += 20 * k * (q[1][3] * q[2][2])
H += 20 * k * (q[1][4] * q[2][2])

H += 50 * k * (q[1][0] * q[2][3])
H += 30 * k * (q[1][1] * q[2][3])
H += 30 * k * (q[1][2] * q[2][3])
H +=  0 * k * (q[1][3] * q[2][3])
H += 10 * k * (q[1][4] * q[2][3])

H += 40 * k * (q[1][0] * q[2][4])
H += 20 * k * (q[1][1] * q[2][4])
H += 20 * k * (q[1][2] * q[2][4])
H += 10 * k * (q[1][3] * q[2][4])
H +=  0 * k * (q[1][4] * q[2][4])

# 3番目から4番目の都市への移動（5 * 5 = 25ケース）
H +=  0 * k * (q[2][0] * q[3][0])
H += 30 * k * (q[2][1] * q[3][0])
H += 20 * k * (q[2][2] * q[3][0])
H += 50 * k * (q[2][3] * q[3][0])
H += 40 * k * (q[2][4] * q[3][0])

H += 20 * k * (q[2][0] * q[3][1])
H +=  0 * k * (q[2][1] * q[3][1])
H += 10 * k * (q[2][2] * q[3][1])
H += 30 * k * (q[2][3] * q[3][1])
H += 60 * k * (q[2][4] * q[3][1])

H += 20 * k * (q[2][0] * q[3][2])
H += 10 * k * (q[2][1] * q[3][2])
H +=  0 * k * (q[2][2] * q[3][2])
H += 20 * k * (q[2][3] * q[3][2])
H += 20 * k * (q[2][4] * q[3][2])

H += 50 * k * (q[2][0] * q[3][3])
H += 30 * k * (q[2][1] * q[3][3])
H += 30 * k * (q[2][2] * q[3][3])
H +=  0 * k * (q[2][3] * q[3][3])
H += 10 * k * (q[2][4] * q[3][3])

H += 40 * k * (q[2][0] * q[3][4])
H += 20 * k * (q[2][1] * q[3][4])
H += 20 * k * (q[2][2] * q[3][4])
H += 10 * k * (q[2][3] * q[3][4])
H +=  0 * k * (q[2][4] * q[3][4])

# 4番目から5番目の都市への移動（5 * 5 = 25ケース）
H +=  0 * k * (q[3][0] * q[4][0])
H += 30 * k * (q[3][1] * q[4][0])
H += 20 * k * (q[3][2] * q[4][0])
H += 50 * k * (q[3][3] * q[4][0])
H += 40 * k * (q[3][4] * q[4][0])

H += 20 * k * (q[3][0] * q[4][1])
H +=  0 * k * (q[3][1] * q[4][1])
H += 10 * k * (q[3][2] * q[4][1])
H += 30 * k * (q[3][3] * q[4][1])
H += 60 * k * (q[3][4] * q[4][1])

H += 20 * k * (q[3][0] * q[4][2])
H += 10 * k * (q[3][1] * q[4][2])
H +=  0 * k * (q[3][2] * q[4][2])
H += 20 * k * (q[3][3] * q[4][2])
H += 20 * k * (q[3][4] * q[4][2])

H += 50 * k * (q[3][0] * q[4][3])
H += 30 * k * (q[3][1] * q[4][3])
H += 30 * k * (q[3][2] * q[4][3])
H +=  0 * k * (q[3][3] * q[4][3])
H += 10 * k * (q[3][4] * q[4][3])

H += 40 * k * (q[3][0] * q[4][4])
H += 20 * k * (q[3][1] * q[4][4])
H += 20 * k * (q[3][2] * q[4][4])
H += 10 * k * (q[3][3] * q[4][4])
H +=  0 * k * (q[3][4] * q[4][4])

# 5番目から6番目の都市への移動（5 * 5 = 25ケース）
H +=  0 * k * (q[4][0] * q[5][0])
H += 30 * k * (q[4][1] * q[5][0])
H += 20 * k * (q[4][2] * q[5][0])
H += 50 * k * (q[4][3] * q[5][0])
H += 40 * k * (q[4][4] * q[5][0])

H += 20 * k * (q[4][0] * q[5][1])
H +=  0 * k * (q[4][1] * q[5][1])
H += 10 * k * (q[4][2] * q[5][1])
H += 30 * k * (q[4][3] * q[5][1])
H += 60 * k * (q[4][4] * q[5][1])

H += 20 * k * (q[4][0] * q[5][2])
H += 10 * k * (q[4][1] * q[5][2])
H +=  0 * k * (q[4][2] * q[5][2])
H += 20 * k * (q[4][3] * q[5][2])
H += 20 * k * (q[4][4] * q[5][2])

H += 50 * k * (q[4][0] * q[5][3])
H += 30 * k * (q[4][1] * q[5][3])
H += 30 * k * (q[4][2] * q[5][3])
H +=  0 * k * (q[4][3] * q[5][3])
H += 10 * k * (q[4][4] * q[5][3])

H += 40 * k * (q[4][0] * q[5][4])
H += 20 * k * (q[4][1] * q[5][4])
H += 20 * k * (q[4][2] * q[5][4])
H += 10 * k * (q[4][3] * q[5][4])
H +=  0 * k * (q[4][4] * q[5][4])

# ３．TYTANを使ったコンパイルとサンプリング

qubo, offset = Compile(H).get_qubo()          # コンパイル（H式からQUBO行列とオフセットを取得）
# solver = sampler.SASampler(seed=1)         # サンプラー選択 (SASampler: CPU シミュレーテッドアニーリング)
solver = sampler.ArminSampler(seed=1)      # サンプラー選択 (ArminSampler: GPU シミュレーテッドアニーリング)
result = solver.run(qubo, shots=1000)         # サンプリング（QUBO行列を上記で選択したサンプラーで解く。shotsの回数分解く。）

print(f"offset: {offset:.1f}\n")


# ４．結果の表示

print(f"offset: {offset:.1f}\n")
print("結果, 計算式Hの値（小さい値のほうがよい）, 試行回数")
print("--------------------------------------------------------")
for r in result[:10]:
    arr, subs = Auto_array(r[0]).get_ndarray('q{}_{}')
    print(f'{arr},')
    print(f'H: {r[1]:.2f}, 試行回数: {r[2]}\n')