#@title # <font color = "deeppink">１．パッケージのインストールやライブラリのインポート</font> { display-mode: "form" }

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

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

#@title # <font color = "deeppink">２．ルールの式（QUBO形式）を
#書く</font> { display-mode: "form" }

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール0：メニューごとに必要な量子ビットを用意
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
q1 = symbols_list(2, 'q1_{}')
q2 = symbols_list(4, 'q2_{}')
q3 = symbols_list(3, 'q3_{}')
q4 = symbols_list(1, 'q4_{}')
q5 = symbols_list(2, 'q5_{}')

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール1： メニューは1000 〜 1300kcalの範囲で選ぶ
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
H = 0   # 初期化

# 不等式用の量子ビット（補助ビット）の用意
s = symbols_list(7, 's{}')

# 補助の量子ビットをワンホットエンコーディングにする（強い条件）
k = 1
H += k * (s[0] + s[1] + s[2] + s[3] + s[4] + s[5] + s[6] - 1)**2

# 不等式（強い条件）
# 1000 ≦ 450 * q1[0] + ... + 650 * q5[1] ≦ 1300 （ただし50刻み）
k = 0.01
H += k * (450*q1[0] + 450*q1[1] \
          + 300*q2[0] + 300*q2[1] + 300*q2[2] + 300*q2[3] \
          + 350*q3[0] + 350*q3[1] + 350*q3[2] \
          + 800*q4[0] \
          + 650*q5[0] + 650*q5[1] \
          - (1000*s[0] + 1050*s[1] + 1100*s[2] + 1150*s[3] + 1200*s[4] + 1250*s[5] + 1300*s[6]))**2

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# ルール2： 選んだメニューの幸福度に応じて報酬を与える
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# それぞれのメニュー（量子ビット）を食べる(1)か食べない(0)か、幸福度を係数にして報酬とする（弱い条件）
k = 0.01
H += -k * (4*q1[0] + 4*q1[1] \
          + 2*q2[0] + 2*q2[1] + 2*q2[2] + 2*q2[3] \
          + 3*q3[0] + 3*q3[1] + 3*q3[2]\
          + 6*q4[0] \
          + 5*q5[0] + 5*q5[1])

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# その他の条件
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# 降順
H += (1 - q3[0]) * q3[1]
H += (1 - q3[1]) * q3[2]


#@title # <font color = "deeppink">３．TYTANを使ったコンパイルとサンプリング</font> { display-mode: "form" }

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=10000)         # サンプリング（QUBO行列を上記で選択したサンプラーで解く。shotsの回数分解く。）


#@title # <font color = "deeppink">４．結果の表示</font> { display-mode: "form" }
import numpy as np

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