"""Fold-safe feature engineering for the Titanic classroom.

No target is read during feature extraction or imputation.
"""
import numpy as np
import pandas as pd
from pathlib import Path
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression, Ridge


def extract_title(names):
    title = names.str.extract(r"([A-Za-z]+)\.", expand=False)
    title = title.replace({"Mlle": "Miss", "Ms": "Miss", "Mme": "Mrs"})
    return title.where(title.isin(["Mr", "Miss", "Mrs", "Master"]), "Rare")


def course_splits(X, kind="stratified", seed=42):
    """Replay the published validation allocation, aligned by PassengerId.

    Persisting the group allocation avoids platform-dependent tie ordering
    when many Ticket groups have the same size. All 891 course IDs are required.
    """
    manifest = pd.read_csv(Path(__file__).with_name("validation-folds.csv"))
    ids = X["PassengerId"]
    if not ids.is_unique or set(ids) != set(manifest["PassengerId"]):
        raise ValueError("固定划分需要完整的891名乘客及唯一PassengerId；子样本请另设验证设计。")
    column = "ticket_5" if kind == "ticket" else f"stratified_{seed}"
    if kind not in {"ticket", "stratified"} or column not in manifest:
        raise ValueError("Unsupported course validation allocation")
    fold_id = manifest.set_index("PassengerId").loc[ids, column].to_numpy()
    return [(np.flatnonzero(fold_id != k), np.flatnonzero(fold_id == k))
            for k in range(5)]


class FeatureBuilder(BaseEstimator, TransformerMixin):
    def __init__(self, age_strategy="median"):
        self.age_strategy = age_strategy

    def fit(self, X, y=None):
        self.age_median_ = X["Age"].median()
        self.age_groups_ = X.groupby(["Sex", "Pclass"])["Age"].median()
        self.fare_median_ = X["Fare"].median()
        # A learned frequency map: it is fitted only on this training fold.
        self.ticket_counts_ = X["Ticket"].value_counts()
        return self

    def transform(self, X):
        out = X.copy()
        out["AgeMissing"] = out["Age"].isna().astype(int)
        if self.age_strategy == "group":
            keys = pd.MultiIndex.from_frame(out[["Sex", "Pclass"]])
            fill = pd.Series(self.age_groups_.reindex(keys).to_numpy(), index=out.index)
            out["Age"] = out["Age"].fillna(fill)
        out["Age"] = out["Age"].fillna(self.age_median_)
        out["Fare"] = out["Fare"].fillna(self.fare_median_)
        out["FamilySize"] = out["SibSp"] + out["Parch"] + 1
        out["IsAlone"] = (out["FamilySize"] == 1).astype(int)
        out["FamilyBand"] = pd.cut(out["FamilySize"], [0, 1, 4, np.inf], labels=["1", "2-4", "5+"]).astype(str)
        out["HasCabin"] = out["Cabin"].notna().astype(int)
        out["Deck"] = out["Cabin"].str[0].fillna("Unknown")
        out["Title"] = extract_title(out["Name"])
        out["FareLog"] = np.log1p(out["Fare"])
        out["SexClass"] = out["Sex"] + "_" + out["Pclass"].astype(str)
        out["FemaleClass2"] = ((out["Sex"] == "female") & (out["Pclass"] == 2)).astype(int)
        out["FemaleClass3"] = ((out["Sex"] == "female") & (out["Pclass"] == 3)).astype(int)
        out["TicketCountTrain"] = out["Ticket"].map(self.ticket_counts_).fillna(0)
        out["FarePerRecordedTicket"] = out["Fare"] / out["TicketCountTrain"].clip(lower=1)
        return out


def make_model(mode="engineered", age_strategy="median", C=1.0):
    if mode == "basic":
        numeric = ["Age", "Fare", "SibSp", "Parch"]
        categorical = ["Pclass", "Sex", "Embarked"]
    elif mode in {"family", "cabin", "engineered"}:
        numeric = ["Age", "Fare", "FamilySize", "HasCabin"]
        categorical = ["Pclass", "Sex", "Embarked", "Title"]
        if mode == "family":
            numeric.remove("HasCabin")
        if mode != "engineered":
            categorical.remove("Title")
    elif mode == "advanced":
        numeric = ["Age", "FareLog", "AgeMissing"]
        categorical = ["Pclass", "Sex", "Embarked", "Title", "Deck", "FamilyBand"]
    else:
        raise ValueError("mode must be basic, engineered or advanced")
    preprocess = ColumnTransformer([
        ("num", Pipeline([("fill", SimpleImputer(strategy="median")),
                          ("scale", StandardScaler())]), numeric),
        ("cat", Pipeline([("fill", SimpleImputer(strategy="most_frequent")),
                          ("encode", OneHotEncoder(drop="first", handle_unknown="ignore"))]), categorical)
    ])
    if mode == "advanced":
        # Main effects above plus two explicit interactions, with class 1 as
        # the interaction reference. No redundant six-level SexClass encoding.
        preprocess.transformers.append(("interactions", "passthrough", ["FemaleClass2", "FemaleClass3"]))
    return Pipeline([("features", FeatureBuilder(age_strategy)),
                     ("preprocess", preprocess),
                     ("logit", LogisticRegression(max_iter=2000, C=C))])


def ridge_age_impute(train, other):
    """Optional age model; train must be a training fold, never all CV data."""
    numeric = ["Fare", "SibSp", "Parch"]
    categorical = ["Sex", "Pclass"]
    pre = ColumnTransformer([
        ("num", Pipeline([("fill", SimpleImputer(strategy="median")), ("scale", StandardScaler())]), numeric),
        ("cat", OneHotEncoder(handle_unknown="ignore"), categorical)
    ])
    known = train["Age"].notna()
    model = Pipeline([("pre", pre), ("ridge", Ridge(alpha=10))])
    model.fit(train.loc[known, numeric + categorical], train.loc[known, "Age"])
    out = other["Age"].copy()
    missing = out.isna()
    if missing.any():
        out.loc[missing] = np.clip(model.predict(other.loc[missing, numeric + categorical]), 0, 100)
    return out