%reload_ext autoreload
%autoreload 2
import itertools
import numpy as np
import pandas  as pd
from pandas.io import gbq
from pyirt import irt
#BigQueryからデータ抽出するクエリを管理しているモジュール
import queries
from tqdm import tqdm

_train = pd.read_gbq(queries.train_agg(), PROJECT_ID, dialect='standard', location="asia-northeast1")
#pyirtの場合、下記のカラム名に統一し、 [(user_id, item_id, ans_boolean)]という構造(タプルのリスト)にする必要がある
train = _train.rename(columns={
    "anon_student_id": "user_id",
    "question_unique_key": "item_id",
    "is_correct": "ans_boolean"
})

item_param, user_param = {}, {}
problem_hierarchies = train["problem_hierarchy"].unique()
#problem_hierarchy毎にデータを分け、周辺最尤法により推定
for _problem_hierarchy in tqdm(problem_hierarchies, position=0):
    train_by_problem_hierarchy = train.query("problem_hierarchy == @_problem_hierarchy").drop("problem_hierarchy", axis=1)
    train_by_problem_hierarchy = train_by_problem_hierarchy[["user_id", "item_id", "ans_boolean"]].values
    #irtが推定用の関数
    _item_param, _user_param = irt(train_by_problem_hierarchy)
    item_param.update(_item_param)
    user_param.update(_user_param)

user_param_df = pd.DataFrame(user_param.items(), columns=["anon_student_id", "theta"])
item_param_dict = []
for tmp_question_unique_key, param_dict in item_param.items():
    item_param_dict.append({
        "question_unique_key": tmp_question_unique_key,
        "alpha": param_dict["alpha"],
        "beta": param_dict["beta"],
        "c": param_dict["c"]
    })
item_param_df = pd.io.json.json_normalize(item_param_dict)

#推定結果をBigQueryに送る
gbq.to_gbq(item_param_df,'{}.item_param'.format(DATASET), project_id=PROJECT_ID, if_exists='append', location="asia-northeast1")
gbq.to_gbq(user_param_df,'{}.user_param'.format(DATASET), project_id=PROJECT_ID, if_exists='append', location="asia-northeast1")