from typing import List, Callable, Set
from copy import deepcopy
from collections import defaultdict
import pandas as pd
from pandas import DataFrame, Series
from sklearn.preprocessing import LabelEncoder, PowerTransformer, MultiLabelBinarizer
from sklearn.metrics import accuracy_score
from sklearn.model_selection import StratifiedKFold,GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.compose import make_column_transformer
from sklearn.utils.validation import check_is_fitted
from lightgbm import LGBMClassifier
import matplotlib.pyplot as plt

def load_data():
    train_df: DataFrame = pd.read_csv("train_data.csv", index_col='id')
    train_X: DataFrame = train_df.drop("y",axis=1)
    train_y: Series = train_df.loc[:,"y"]
    test_X: DataFrame = pd.read_csv("test_data.csv", index_col='id')
    return train_X,train_y,test_X
train_X,train_y,test_X = load_data()
train_X.head()

#utils
def _unique_values_from_df(X: DataFrame)->list:
    cols = X.columns
    unique_values: Set = set()
    for col in cols:
        unique_values.update(pd.unique(X[col]))
    return list(unique_values)

def _label_encode(X:DataFrame, cols:list, le: LabelEncoder)->DataFrame:
    for col in cols:
        X[col] = le.transform(X[col])
    return X

def team_feat(team:str = "A", item:str = "weapon")->List[str]:
    """
    param: team:  A | B | both
    param: item:  weapon | rank | level
    return: list of columns name.
    """
    assert item in ["weapon","rank","level"]
    assert team in ["A","B","both"]
    def _list_embed(team:str,item:str)->List[str]:
        dst = []
        for i in range(1,5):
            col_name = f"{team}{i}-{item}"
            dst.append(col_name)
        return dst
    
    A_feats = _list_embed("A",item)
    B_feats = _list_embed("B",item)
    both_feats = A_feats + B_feats
    team_feats_dict = {"A":A_feats,"B":B_feats,"both":both_feats}
    return team_feats_dict[team]
    

from sklearn.base import TransformerMixin,BaseEstimator

## BaseEstimator,TransformerMixinを継承した前処理のためのクラスを作る．
## fitやtransformなどのメソッドが共通するので，sklearnの便利なモジュールを組み合わせることで生産性が高まります．
class BaseTransformer:
    def fit(self,X, y=None, **fit_params):
        return self
    def transform(self,X,y=None):
        return X
    def _valid_cols_exist(self,X: DataFrame, col_names: List[str]):
        """
        カラムの存在確認
        """
        inclusion: set = set(col_names).issubset(set(X.columns))
        if inclusion:
            pass
        else:
            not_found_cols: set = set(col_names) - set(col_names).intersection(set(X.columns))
            raise Exception(f"{list(not_found_cols)} cannot be found...")
        
# pandasのDatetime型に変換
# このように，あまり変わらない(=どの変換にも必須)の前処理はクラス化しておくと後々便利
class PeriodTransformer(BaseTransformer,TransformerMixin, BaseEstimator):
    def __init__(self,period_format:str = '%Y-%m-%dT%H:%M:%S+00:00'):
        self.period_format = period_format
    def transform(self,X,y=None):
        self._valid_cols_exist(X,['period'])
        X['period'] = pd.to_datetime(X['period'], format=self.period_format)
        print(X.shape)
        return X

class DeleteColumnTransformer(BaseTransformer,TransformerMixin, BaseEstimator):
    def __init__(self,deleted_columns:List[str]=["game-ver","lobby","period"]):
        self.deleted_columns = deleted_columns
    def transform(self,X,y=None):
        self._valid_cols_exist(X,self.deleted_columns)
        X = X.drop(self.deleted_columns,axis=1)
        print(X.shape)
        return X

class FloatTransformer(BaseTransformer,TransformerMixin, BaseEstimator):
    def transform(self,X,y=None):
        X.astype("float")
        return X
        
        

# どのモデルにも共通するけど，試行錯誤でコロコロ変えたい前処理は以下の関数に記述
def _preprocess(X: DataFrame,y=None)->DataFrame:
    #時間・日時に関する特徴量
    X["day_of_the_week"] = X["period"].dt.dayofweek
    X["hour"] = X["period"].dt.hour
    return X

class NormalLabelEncoder(BaseTransformer,TransformerMixin, BaseEstimator):
    def __init__(self,cols:list,fillna_word:str="none"):
        self.cols = cols
        self.fillna_word = fillna_word
        self.encoders = defaultdict(LabelEncoder)
    def fit(self,X,y=None):
        for col in self.cols:
            self.encoders[col].fit(X[col].fillna(self.fillna_word))
        return self
    def transform(self,X,y=None):
        for col in self.cols:
            X[col] = self.encoders[col].transform(X[col].fillna(self.fillna_word))
        print(X.shape)
        return X
        
class LevelPreprocessor(BaseTransformer,TransformerMixin, BaseEstimator):
    def __init__(self,fillna:str="mean"):
        """
        param fillna: how to fill nan cell.
                "mean" | "median" | (本当は"predict"も入れて，nanの値を予測することもしたかったのですが，
                ちょっと実装に時間がかかりそうなので，このチュートリアルでは取り上げません)
        """
        self.fillna = fillna
        self.level_cols = team_feat("both","level")
        self.fillna_dict = dict()
    def fit(self,X,y=None):
        for level_col in self.level_cols:
            fillna_val = eval(f"X[level_col].{self.fillna}()")
            self.fillna_dict[level_col] = fillna_val
        return self
    def transform(self,X,y=None):
        for level_col in self.level_cols:
            fillna_val = self.fillna_dict[level_col]
            X[level_col].fillna(fillna_val,inplace=True)
        print(X.shape)
        return X
            
            
    
    
# モデルごとに前処理を分ける．今回はGBT系のアルゴリズム用の前処理クラスを書く．
class GBTPreprocessor(BaseTransformer,TransformerMixin, BaseEstimator):
    def __init__(self,preprocess_fn:Callable,le:LabelEncoder = LabelEncoder()):
        self.le = le
        self.preprocess_fn = preprocess_fn
    def transform(self,X,y=None):
        X = self.preprocess_fn(X)
        print(X.shape)
        return X
        
# 以下の関数は公式チュートリアルを参考にした．
def trans_weapon(df, cols=team_feat("A","weapon"))->DataFrame:
    le = LabelEncoder()
    weapon = df.fillna('none')[cols]
    le.fit(weapon[cols[3]])#weaponの１~3には"none"が含まれないので
    for col in cols:
        weapon[col] = le.transform(weapon[col])    
    return weapon,le

def weapon_feature_df(df)->DataFrame:
    a_weapon, a_weapon_transformer = trans_weapon(df, cols=team_feat("A","weapon"))
    b_weapon, b_weapon_transformer = trans_weapon(df, cols=team_feat("B","weapon"))
    X = pd.concat([a_weapon, b_weapon], axis=1)
    return X

class WeaponLabelEncoder(BaseTransformer,TransformerMixin, BaseEstimator):
    def __init__(self,le=LabelEncoder()):
        self.le = le
        self.weapon_cols = team_feat("both","weapon")
        
    def fit(self,X,y=None):
        weapons = _unique_values_from_df(X[self.weapon_cols].fillna('none'))
        self.le.fit(weapons)      
        return self
    
    def transform(self,X,y=None):
        X[self.weapon_cols] = X[self.weapon_cols].fillna('none')
        X = _label_encode(X, self.weapon_cols, self.le)
        print(X.shape)
        return X
        

#カラムごとに処理を加えたい場合，
#↑のWeaponLabelEncoderの実装のように説明変数を返せばいいのですが，
#カラムごとに行いたい処理がある場合,make_column_transformerを用いることで実装できます．
#参考:https://qiita.com/R0w0/items/3b3d8e660b8abc1f804d

pt = make_column_transformer(
    (PowerTransformer(), team_feat("both","level")),remainder="passthrough" # (procedure, cols)f
    (PowerTransformer(), team_feat("both","level")),remainder="passthrough"
#     (OneHotEncoder(sparse=False, drop="first"), cat_cols)
)

categorical_cols = ['game-ver', 'lobby-mode', 'lobby', 'mode', 'stage', 'A1-rank', 'A2-rank', 'A3-rank', 'A4-rank', 'B1-rank', 'B2-rank', 'B3-rank','B4-rank']
clf = Pipeline(
    [
        ("PeriodTransformer",PeriodTransformer()),
        ("LevelPreprocessor",LevelPreprocessor()),
        ("NormalLabelEncoder",NormalLabelEncoder(categorical_cols)),
        ("WeaponLabelEncoder",WeaponLabelEncoder()),
        ("GBTPreprocessor",GBTPreprocessor(_preprocess)),
        ("DeleteColumnTransformer",DeleteColumnTransformer()),
        ("FloatTransformer",FloatTransformer()),
        ("pt",pt),
        ("lgbm",LGBMClassifier())
    ]
)

clf.fit(train_X,train_y)

train_period_ser = PeriodTransformer().transform(train_X)["period"]
test_period_ser = PeriodTransformer().transform(test_X)["period"]

plt.figure(figsize=(7,7))
train_period_ser.hist(bins=20)
plt.ylabel("train_data_nume")
plt.xlabel("datetime")
plt.show()

plt.figure(figsize=(7,7))
test_period_ser.hist(bins=20)
plt.ylabel("test_data_nume")
plt.xlabel("datetime")
plt.show()

def cv_score(clf,X,y,n_splits=5,kf=StratifiedKFold(n_splits=5))->float:
    score = 0.
    counter = 1
    for train_index, valid_index in kf.split(X, y):
        train_X,valid_X = X.iloc[train_index,:].copy(),X.iloc[valid_index,:].copy()
        train_y,valid_y = y.iloc[train_index],y.iloc[valid_index]
        c_clf = deepcopy(clf)
        c_clf.fit(train_X,train_y)
        preds = c_clf.predict(valid_X)
        acc_score = accuracy_score(valid_y,preds)
        print(f"fold{counter} score is :{acc_score}")
        score += acc_score
        counter += 1
    return score / n_splits

cv_score(clf,train_X,train_y)