"""
dir024 cross-fit ElasticNet stacker (public LB 0.3832).

This script uses precomputed parent OOF/test predictions from a companion dataset,
then fits GroupKFold ElasticNet and writes submission.csv.
"""

from itertools import combinations
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_squared_log_error
from sklearn.model_selection import GroupKFold
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

SEED = 42
N_SPLITS = 5


def rmsle(y_true, y_pred):
    y_true = np.clip(np.asarray(y_true, dtype=float), 0.0, None)
    y_pred = np.clip(np.asarray(y_pred, dtype=float), 0.0, None)
    return float(np.sqrt(mean_squared_log_error(y_true, y_pred)))


def build_features(log_mat):
    centered = log_mat - np.mean(log_mat, axis=0, keepdims=True)
    parts = [
        log_mat,
        np.std(log_mat, axis=1, keepdims=True),
        np.max(log_mat, axis=1, keepdims=True) - np.min(log_mat, axis=1, keepdims=True),
    ]
    inter = []
    for i, j in combinations(range(log_mat.shape[1]), 2):
        inter.append((centered[:, i] * centered[:, j]).reshape(-1, 1))
    if inter:
        parts.extend(inter)
    return np.hstack(parts)


def find_input_dir(required_file, hint=""):
    input_root = Path("/kaggle/input")

    candidates = [
        input_root / "bidding-dir024-parent-meta-features-public",
        input_root / "datasets" / "oscarletech" / "bidding-dir024-parent-meta-features-public",
    ]
    for c in candidates:
        if (c / required_file).is_file():
            return c

    matches = [p.parent for p in input_root.rglob(required_file) if p.is_file()]
    if hint:
        hinted = [d for d in matches if hint in str(d)]
        if hinted:
            return hinted[0]
    if matches:
        return matches[0]

    raise FileNotFoundError(f"Could not find {required_file} under /kaggle/input")


def main():
    stack_input = find_input_dir("train_meta_inputs.csv", hint="bidding-dir024-parent-meta-features-public")
    print(f"Using stack input root: {stack_input}")

    stack_train = pd.read_csv(stack_input / "train_meta_inputs.csv")
    stack_test = pd.read_csv(stack_input / "test_meta_inputs.csv")

    assert stack_train["row_id"].is_unique
    assert stack_test["row_id"].is_unique

    pred_cols = [c for c in stack_train.columns if c.startswith("pred_")]
    print(f"train_rows={len(stack_train)} test_rows={len(stack_test)} parents={len(pred_cols)}")

    y_true = np.clip(stack_train["total_bid"].to_numpy(dtype=float), 0.0, None)
    y_log = np.log1p(y_true)

    log_mat = np.log1p(np.clip(stack_train[pred_cols].to_numpy(dtype=float), 0.0, None))
    tlog_mat = np.log1p(np.clip(stack_test[pred_cols].to_numpy(dtype=float), 0.0, None))

    X = build_features(log_mat)
    X_test = build_features(tlog_mat)
    groups = stack_train["job_id"].astype(str).to_numpy()

    cv = GroupKFold(n_splits=N_SPLITS)
    oof_log = np.zeros(len(stack_train), dtype=float)
    test_log = np.zeros(len(stack_test), dtype=float)

    fold_scores = []
    for fold, (tr_idx, va_idx) in enumerate(cv.split(X, y_log, groups=groups), start=1):
        model = make_pipeline(
            StandardScaler(),
            ElasticNet(alpha=0.0018, l1_ratio=0.22, max_iter=9000, random_state=SEED + fold),
        )
        model.fit(X[tr_idx], y_log[tr_idx])

        pred_va = model.predict(X[va_idx])
        pred_te = model.predict(X_test)

        oof_log[va_idx] = pred_va
        test_log += pred_te / N_SPLITS

        f = rmsle(y_true[va_idx], np.expm1(pred_va))
        fold_scores.append(f)
        print(f"fold={fold} rmsle={f:.6f}")

    oof_score = rmsle(y_true, np.expm1(oof_log))
    print(f"oof_rmsle={oof_score:.6f}")

    submission = pd.DataFrame(
        {
            "row_id": stack_test["row_id"].values,
            "total_bid": np.clip(np.expm1(test_log), 0.0, None),
        }
    )

    if len(submission) != 1447:
        raise RuntimeError(f"Unexpected row count: {len(submission)}")
    if submission["total_bid"].isna().any():
        raise RuntimeError("NaN found in predictions")

    out_path = Path("submission.csv")
    submission.to_csv(out_path, index=False)
    print(f"Saved {out_path} rows={len(submission)}")


if __name__ == "__main__":
    main()