---

# ML: uplift по продуктам (T-learner), политика по ожидаемой инкрементальной прибыли
def load_data(data_dir: Path) -> dict[str, pd.DataFrame]:
    return {
        "history_clients": pd.read_csv(data_dir / "history_clients.csv"),
        "campaign_history": pd.read_csv(data_dir / "campaign_history.csv"),
        "products": pd.read_csv(data_dir / "products.csv", encoding="utf-8-sig"),
        "region_profiles": pd.read_csv(data_dir / "region_profiles.csv"),
    }


def feature_columns(history_clients: pd.DataFrame, public_clients: pd.DataFrame) -> list[str]:
    common = [c for c in history_clients.columns if c in public_clients.columns]
    drop = {"client_id", "evaluation_month"} | {f"eligible_p{j}" for j in PRODUCT_IDS}
    return [c for c in common if c not in drop]


def make_model() -> HistGradientBoostingClassifier:
    return HistGradientBoostingClassifier(
        max_iter=200,
        learning_rate=0.03,
        max_leaf_nodes=15,
        min_samples_leaf=1000,
        l2_regularization=1.0,
        random_state=SEED,
    )


def fit_uplift_models(
    history: pd.DataFrame,
    campaigns: pd.DataFrame,
    features: list[str],
    product_ids: list[int],
) -> dict[int, tuple[HistGradientBoostingClassifier, HistGradientBoostingClassifier]]:
    models = {}
    for j in product_ids:
        table = build_product_table(history, campaigns, j)
        table = table[table[f"eligible_p{j}"] == 1]
        treated = table[table["offered"] == 1]
        control = table[table["offered"] == 0]
        m1 = make_model().fit(treated[features], treated["purchased"])
        m0 = make_model().fit(control[features], control["purchased"])
        models[j] = (m1, m0)
    return models


def score_candidates(
    candidates: pd.DataFrame,
    clients: pd.DataFrame,
    features: list[str],
    models: dict,
    products: pd.DataFrame,
) -> pd.DataFrame:
    margin = dict(zip(products["product_id"], products["margin_rub"]))
    cost = dict(zip(products["product_id"], products["contact_cost_rub"]))
    feats = clients.set_index("client_id")[features]
    scored = []
    for j, (m1, m0) in models.items():
        cand = candidates[candidates["product_id"] == j]
        if cand.empty:
            continue
        x = feats.loc[cand["client_id"].to_numpy()]
        tau = m1.predict_proba(x)[:, 1] - m0.predict_proba(x)[:, 1]
        part = cand.copy()
        part["score"] = tau * margin[j] - cost[j]
        scored.append(part)
    return pd.concat(scored, ignore_index=True)


def select_offers(scored: pd.DataFrame, n_clients: int, contact_fraction: float) -> pd.DataFrame:
    best = scored.sort_values(["score", "client_id"], ascending=[False, True])
    best = best.drop_duplicates("client_id", keep="first")
    best = best[best["score"] > 0]
    max_contacts = int(math.floor(contact_fraction * n_clients))
    best = best.head(max_contacts)
    return pd.DataFrame(
        {"client_id": best["client_id"].to_numpy(), "offer_id": best["product_id"].astype(int).to_numpy()}
    )


def build_policy(
    public_clients: pd.DataFrame,
    data_dir: Path,
) -> pd.DataFrame:
    np.random.seed(SEED)
    data = load_data(data_dir)
    history = add_region_features(data["history_clients"], data["region_profiles"])
    public = add_region_features(public_clients, data["region_profiles"])
    features = feature_columns(history, public)
    product_ids = [j for j in PRODUCT_IDS if j not in EXCLUDED_PRODUCTS]
    models = fit_uplift_models(history, data["campaign_history"], features, product_ids)
    candidates = make_product_candidates(public, data["products"])
    candidates = candidates[candidates["product_id"].isin(product_ids)]
    scored = score_candidates(candidates, public, features, models, data["products"])
    return select_offers(scored, len(public_clients), CONTACT_FRACTION)


def run(data_dir: Path, output_path: Path) -> None:
    np.random.seed(SEED)

    public_clients_path = data_dir / "public_clients.csv"
    if not public_clients_path.exists():
        raise FileNotFoundError(f"File not found: {public_clients_path}")

    public_clients = pd.read_csv(public_clients_path)

    offers = build_policy(
        public_clients=public_clients,
        data_dir=data_dir,
    )

    submission = make_submission(
        clients=public_clients,
        offers=offers,
    )

    diagnostics = validate_submission_frame(
        submission=submission,
        clients=public_clients,
        contact_fraction=CONTACT_FRACTION,
    )

    if not diagnostics["is_valid"]:
        raise ValueError(f"Generated submission is invalid: {diagnostics}")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    submission[["client_id", "offer_id"]].to_csv(output_path, index=False)

    print(f"Submission saved to: {output_path}")
    print(f"Validation: {diagnostics}")
    print(f"Contacts: {contact_statistics(submission)}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data-dir",
        type=Path,
        required=True,
        help="Directory containing competition CSV files",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("submission.csv"),
        help="Output submission path",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    run(
        data_dir=args.data_dir,
        output_path=args.output,
    )