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

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

from tytan import *
import numpy as np
import pandas as pd

from sklearn import datasets                             # サンプルデータセット
from sklearn.model_selection import train_test_split     # データ分割用
from matplotlib import pyplot as plt                     # グラフ描画用
from sklearn.tree import DecisionTreeClassifier as DTC   # 決定木
from sklearn.tree import plot_tree                       # 決定木の可視化
from sklearn.metrics import accuracy_score               # 精度計算用
from sklearn import metrics                              # 精度計算用

# Titanicデータセットを取得

# X: 性別、年齢などのデータ  y: 正解データ
# Data obtained from http://hbiostat.org/data courtesy of the Vanderbilt University Department of Biostatistics.
# X, y = datasets.fetch_openml(data_id=40945, return_X_y=True, as_frame=True)
df = pd.read_csv('https://hbiostat.org/data/repo/titanic3.csv')
X = df.drop(columns='survived')
y = df['survived']

# データの確認

X    # 性別、年齢などのデータ
y    # 正解データ   1: 生還  0: 死亡

# 前処理

# データがないところを中央値などで埋める
X["age"] = X["age"].fillna(X["age"].median())
X["fare"] = X["fare"].fillna(X["fare"].median())
X["embarked"] = X["embarked"].fillna("S")
X["sex"] = pd.get_dummies(X["sex"], drop_first=True)
X = pd.concat([X, pd.get_dummies(X["embarked"], prefix="embarked")], axis=1).drop(columns=["embarked"])

# 今回使わないデータを削除
feature_columns =["name", "sibsp", "parch", "ticket","fare", "cabin", "boat", "body", "home.dest"]
X = X.drop(columns=feature_columns, axis=1)

print(X.columns)
X = X.values
y = y.values
y = y.astype(int)

# 正解値を平均値計算の都合により {0, 1}から{-1, 1}に変換
y = (y - 0.5) * 2

X    # 性別、年齢などのデータ
y    # 正解データ   1: 生還  -1: 死亡

seed=0
np.random.seed(seed)

# データを学習用とテスト用に7:3に分割
train_X, test_X, train_y, test_y  = train_test_split(X, y, test_size=0.3, random_state=seed)

print("X: ", X.shape, X.dtype)
print("y: ", y.shape, y.dtype)
print("train_X: ", train_X.shape, train_X.dtype)
print("train_y: ", train_y.shape, train_y.dtype)
print("test_X: ", test_X.shape, test_X.dtype)
print("test_y: ", test_y.shape, test_y.dtype)

# パラメータ設定
param_lambda = 3                # 正則化調整パラメータ
num_random = 100                # ランダムデータの数
num_learner = 100               # 学習器の数

# 決定木の深さを指定してモデルを構築
models = [DTC(max_depth=3) for i in range(num_learner)]

y_pred_list_train = []
y_pred_list_test = []

for model in models:

    # データをランダムに抽出
    train_idx = np.random.choice(np.arange(train_X.shape[0]), num_random)

    # 決定木作成
    model.fit(X=train_X[train_idx], y=train_y[train_idx])

    # 予測
    y_pred_list_train.append(model.predict(train_X))
    y_pred_list_test.append(model.predict(test_X))

y_pred_list_train = np.asanyarray(y_pred_list_train)
y_pred_list_test = np.array(y_pred_list_test)

# 決定木の可視化（代表で1つ）
plt.figure(figsize=(10, 7))
plot_tree(models[0], filled=True, class_names=['Class 0', 'Class 1'])
plt.show()

# 予測、精度確認 

# 多数決で正解値を決定（重みはすべて1）
y_pred_test = np.sign(np.sum(y_pred_list_test.astype(int), axis=0))

# 予測結果の精度のスコア計算
acc_test_base = metrics.accuracy_score(y_true=test_y, y_pred=y_pred_test)
print("Base Model: {}".format(acc_test_base))

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

# 学習器の数分、量子ビットを用意
q = symbols_list(num_learner, 'q{}')

#+++++++++++++++++++++++++++++++++
# QBoostの定式化
#+++++++++++++++++++++++++++++++++
# 初期化
H_clf = 0
H_norm = 0
H = 0

for i in range(len(train_y)):
    temp_sum = 0
    for j in range(len(q)):
        temp_sum += q[j] * y_pred_list_train[j, i]
    temp_sum = (1/num_learner * temp_sum - train_y[i]) ** 2
    H_clf += temp_sum

# 正則化項
for j in range(len(q)):
    H_norm += q[j]

H = H_clf + param_lambda * H_norm

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

qubo, offset = Compile(H).get_qubo()
# solver = sampler.SASampler(seed=0)
solver = sampler.ArminSampler(seed=1)
result = solver.run(qubo, shots=100)

for r in result:
    print(r)

# <font color = "deeppink">４．結果の表示</font> { display-mode: "form" }

accs_train = []
accs_test = []

decoded_solutions = np.array([record[0] for record in np.array(result)])

for decoded_solution in decoded_solutions:
    idx_clf = []

    for key, val in decoded_solution.items():
        if val == 1:
            idx = int(key.split('q')[1].split(']')[0])
            idx_clf.append(idx)

    np.set_printoptions(threshold=np.inf)

    # 重み付き多数決（アンサンブル）
    y_pred_train = np.sign(np.sum(y_pred_list_train[idx_clf, :], axis=0))
    y_pred_test = np.sign(np.sum(y_pred_list_test[idx_clf, :], axis=0))

    acc_train = metrics.accuracy_score(y_true=train_y, y_pred=y_pred_train)
    acc_test = metrics.accuracy_score(y_true=test_y, y_pred=y_pred_test)

    accs_train.append(acc_train)
    accs_test.append(acc_test)

# ベースとなるモデルとQBoostの精度を確認
print("Base Model: {}".format(acc_test_base))
print("QBoost    : {}".format(max(accs_test)))

accs_test