df.dropna(inplace=True)

import pandas as pd
import numpy as np
import xgboost as xgb
from IPython.core.display import display
from typing import List, Union

df = pd.read_csv('sample_nan.csv')
display(df)

df_mean = df.copy()
df_mean.fillna(df_mean.mean(), inplace=True)  # 平均値
display(df_mean)

df_median = df.copy()
df_median.fillna(df_median.median(), inplace=True)  # 中央値
display(df_median)

df_mode = df.copy()
df_mode.fillna(df_mode.mode().iloc[0], inplace=True)  # 最頻値
display(df_mode)

class BayesianAverage:
    """
    Bayesian Averageを計算する
    """
    def __init__(self, m: Union[float, int], C: int):
        """
        イニシャライザ

        @param m: あらかじめの値
        @param C: 値mの観測回数
        """
        self.m = m
        self.C = C

    def predict(self, xs: List[Union[float, int]]):
        """
        Bayesian Averageを計算する関数
        あらかじめ値mのデータをC回観測した経験の上で、実測値の平均を計算する

        @param xs: 実測値のリスト
        @return: Bayesian Average
        """
        return (np.sum(xs) + (self.C * self.m)) / (len(xs) + self.C)

df_state = df.copy()
ba_age, ba_point = BayesianAverage(40, 3), BayesianAverage(80, 3)  # <- ハイパーパラメータ
# ba_age: 40歳の人を３人追加する設定 (本来であれば、クラスごとに傾向が異なるので、クラスごとに異なるオブジェクトを生成する必要がある)
# ba_point: 80点の人を3人追加する設定 (本来であれば、クラスごとに傾向が異なるので、クラスごとに異なるオブジェクトを生成する必要がある)
df_state = df_state.groupby('state').agg({'age': ba_age.predict, 'point': ba_point.predict})
display(df_state)  # 一つ一つの値はBayesian Average
#         age  point
# state
# CA     40.0   82.0
# NY     36.0   80.0
# TX     47.0   77.5
# これらを埋める
df_bayes_ave = df.copy()
for column in ['age', 'point']:
    for state, data in df_bayes_ave.groupby('state'):
        bayes_ave = df_state.loc[state, column]
        df_bayes_ave.loc[df_bayes_ave['state'] == state, column] = data[column].fillna(value=bayes_ave)
display(df_bayes_ave)

# ====Step1: データ準備==============
# [前提] df_bayes_aveを用いて、stateの最後の行NaNを予測する
# state -> indexへの変換辞書
state_s2i = {state: i for i, state in enumerate(df_bayes_ave['state'].dropna().unique())}
print(f'stateの種類: {state_s2i}')  # stateの種類: {'NY': 0, 'CA': 1, 'TX': 2}

X_train = df_bayes_ave[['age', 'point']].dropna()  # 訓練データ(説明変数)には欠損を含めない
y_train = df_bayes_ave['state'].dropna().map(state_s2i)  # 訓練データ(目的変数)には欠損を含めない
X_test = df_bayes_ave[df_bayes_ave.isnull().any(axis=1)][['age', 'point']]  # テストデータは欠損が含まれているレコードのみ抽出
# ================================

# ======Step2: 学習=================
# xgboostのオブジェクトに変換する
dtrain = xgb.DMatrix(X_train, label=y_train)
# ハイパーパラメータ設定
param = {'max_depth': 2, 'eta': 0.1, 'objective': 'multi:softmax', 'num_class': len(state_s2i)}
# 学習
bst = xgb.train(param, dtrain, 20)
# ===============================

# ======Step3: 予測=================
# 予測
dtest = xgb.DMatrix(X_test)
pred_test = bst.predict(dtest)
print(pred_test)
# [1.]
# ================================

# index -> stateへの変換辞書
state_i2s = {i: state for state, i in state_s2i.items()}
# 予測値を欠損値に代入する
X_test['state'] = [state_i2s[int(pred)] for pred in pred_test]
# 欠損がない訓練データと欠損を埋めたテストデータを結合
train_test = pd.concat([X_train.join(y_train.map(state_i2s)), X_test])
df_bayes_ave2 = pd.concat([df_bayes_ave['name'], train_test], axis=1)[['name', 'age', 'state', 'point']]
display(df_bayes_ave2)

df_isnull = df.copy()
df_isnull['point_isnull'] = df_isnull.isnull()['point'].map({True: 1, False: 0})
df_isnull['point'].fillna(np.mean(df_isnull['point']), inplace=True)  # 平均値で埋める
display(df_isnull)

df_null_num = df.copy()
df_null_num['null_num'] = df_null_num.isnull().sum(axis=1)
display(df_null_num)

# null_combの説明
0: 欠損を含まない
1: pointのみ欠損
2: age, pointが欠損
3: ageのみ欠損
4: state, pointが欠損