import pandas as pd
import numpy as np
from sklearn.linear_model import ElasticNet
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import joblib

# ---------- 基本インターフェース ----------

def preprocess(df: pd.DataFrame):
    """
    Excelから読み込んだDataFrameを前処理
    - o: が目的変数
    - -: が除外列
    - 文字列はカテゴリ変数として one-hot 化
    """
    # 目的変数列
    obj_cols = [c for c in df.columns if c.startswith("o:")]
    # 使用しない列を除外
    drop_cols = [c for c in df.columns if c.startswith("-:")]
    use_cols = [c for c in df.columns if c not in obj_cols + drop_cols]

    X = df[use_cols]
    y = df[obj_cols]

    # カテゴリ変数を one-hot 化
    X = pd.get_dummies(X, drop_first=True)

    # 欠損値除去
    X = X.dropna()
    y = y.loc[X.index]

    return X, y


def model(args=None):
    """ElasticNet モデルを返す"""
    if args is None:
        args = {"alpha":1.0, "l1_ratio":0.5, "random_state":0}
    return ElasticNet(**args)


def fit(model, X, y):
    """モデルを学習"""
    return model.fit(X, y)


# ElasticNet版
def predict(model, X):
    mean = model.predict(X)
    return mean, None   # 分散なし → None


def metrics(y_true, y_pred):
    """評価指標を返す"""
    return {
        "MAE": mean_absolute_error(y_true, y_pred),
        "MSE": mean_squared_error(y_true, y_pred),
        "R2": r2_score(y_true, y_pred)
    }


def save_model(model, file_path):
    """モデルを保存"""
    joblib.dump(model, file_path)


def load_model(file_path):
    """モデルを読み込み"""
    return joblib.load(file_path)


# ---------- 拡張インターフェース ----------

def set_params(model, params: dict):
    """ハイパーパラメータを設定"""
    model.set_params(**params)


def get_params(model):
    """ハイパーパラメータを取得"""
    return model.get_params()


def cross_validate(model, X, y, cv=5, scoring="r2"):
    """交差検証スコアを返す"""
    scores = cross_val_score(model, X, y, cv=cv, scoring=scoring)
    return {"mean": np.mean(scores), "std": np.std(scores), "scores": scores}


def feature_importance(model, X):
    """特徴量重要度を返す (ElasticNetでは係数)"""
    return dict(zip(X.columns, model.coef_))


# ---------- 使用例 ----------

if __name__ == "__main__":
    # Excelファイル読み込み
    df = pd.read_excel("sample.xlsx")

    # 前処理
    X, y = preprocess(df)

    # モデル生成
    m = model()

    # 学習
    m = fit(m, X, y)

    # 予測
    y_pred = predict(m, X)

    # 評価
    print("Metrics:", metrics(y, y_pred))

    # パラメータ確認
    print("Params:", get_params(m))

    # 交差検証
    print("CV:", cross_validate(m, X, y))

    # 特徴量重要度
    print("Feature importance:", feature_importance(m, X))

    # 保存・読み込み
    save_model(m, "elasticnet_model.pkl")
    loaded = load_model("elasticnet_model.pkl")
    print("Loaded model params:", get_params(loaded))