import sys, subprocess, zipfile
from pathlib import Path as _P
if "google.colab" in sys.modules:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "catboost", "gdown"], check=True)
    if not any(_P("/content").rglob("train.csv")):
        import gdown
        gdown.download(id="1v6Ce6HLFvByXzDdcHQD1NMgA-XnJf5kf", output="/content/data_download", quiet=False)
        z = _P("/content/data_download")
        if zipfile.is_zipfile(z):
            zipfile.ZipFile(z).extractall("/content/data")
            for inner in list(_P("/content/data").rglob("*.zip")):  # zip inside a zip
                zipfile.ZipFile(inner).extractall(inner.parent)
    print("train.csv found at:", [str(p) for p in _P("/content").rglob("train.csv")])

import json
import os
import shutil
import subprocess
import time
import warnings
from pathlib import Path

import numpy as np
import pandas as pd
import lightgbm as lgb
import xgboost as xgb
from catboost import CatBoostClassifier
from scipy import sparse
from scipy.optimize import minimize
from sklearn.cluster import MiniBatchKMeans
from sklearn.metrics import log_loss
from sklearn.model_selection import StratifiedGroupKFold
from sklearn.neighbors import NearestNeighbors

warnings.filterwarnings("ignore")

# ======================= SETTINGS =======================
QUICK_TEST = bool(int(os.environ.get("PT_QUICK", "0")))  # True = 10% of operators + few rounds (pipeline check only)
USE_LIGHTGBM = False                                        # LightGBM runs on CPU only (its GPU mode is slower on Kaggle): ~8 min/fold
RUN_ABLATION = False                                        # test H1/H2 (feature blocks) before the big training
FOLDS, SEED, K = 5, 42, 4
MAX_ROUNDS = 300 if QUICK_TEST else 6000
LOG_EVERY = 100 if QUICK_TEST else 250                     # print training progress every N rounds
# =========================================================

TAG = f"final{'q' if QUICK_TEST else ''}_f{FOLDS}"          # checkpoint names (quick-test checkpoints never mix with full ones)
WORK = Path("/kaggle/working") if Path("/kaggle/working").exists() else Path(".")
CKPT = WORK / "ckpt"
CKPT.mkdir(exist_ok=True)
SEARCH = [CKPT] + (sorted(p for p in Path("/kaggle/input").glob("**/ckpt") if p.is_dir()) if Path("/kaggle/input").exists() else [])


def gpu_available():
    if shutil.which("nvidia-smi") is None:
        return False
    try:
        return subprocess.run(["nvidia-smi"], capture_output=True, timeout=10).returncode == 0
    except Exception:
        return False


GPU = gpu_available()
if GPU:  # all-GPU models; LightGBM only if asked for
    MODELS = ["logreg", "xgb", "xgb_reg", "cat"] + (["lgb_multi", "lgb_reg"] if USE_LIGHTGBM else [])
else:
    MODELS = ["logreg", "lgb_multi", "lgb_reg", "xgb", "xgb_reg"] + (["cat"] if os.environ.get("PT_FORCE_CAT") else [])
    print("No GPU found: running on CPU and skipping CatBoost. Turn on GPU T4 for a much faster run.")
N_THREADS = os.cpu_count()
print(f"GPU: {GPU} | CPU threads: {N_THREADS} | models: {MODELS} | quick test: {QUICK_TEST}")
print("checkpoint folders searched:", [str(p) for p in SEARCH])

def ck_load(name):
    for d in SEARCH:
        p = d / f"{name}.npz"
        if p.exists():
            with np.load(p, allow_pickle=False) as z:
                return {k: z[k] for k in z.files}
    return None


def ck_save(name, **arrays):
    tmp = CKPT / f"{name}.tmp.npz"
    np.savez(tmp, **arrays)
    tmp.replace(CKPT / f"{name}.npz")


class Progress:
    """Overall progress bar over all big steps. ETA = remaining steps x the measured time of the same kind of step
    (e.g. each remaining XGBoost fold is estimated from finished XGBoost folds, not from the slower LightGBM ones)."""

    def __init__(self, plan):
        self.plan, self.total, self.done, self.t0, self.kind_times = plan, len(plan), 0, time.time(), {}

    def bar(self, width=30):
        full = int(width * self.done / self.total)
        return "█" * full + "░" * (width - full) + f" {self.done}/{self.total} ({100 * self.done / self.total:.0f}%)"

    def start(self, label):
        self.label, self.ts = label, time.time()
        print(f"\n▶ {label}   [{self.bar()}]", flush=True)

    def finish(self, from_checkpoint=False):
        self.done += 1
        took = time.time() - self.ts
        if not from_checkpoint:
            self.kind_times.setdefault(self.label.split(":")[0], []).append(took)
        left = [s.split(":")[0] for s in self.plan[self.done:]]
        timed = sum(np.mean(self.kind_times[k]) for k in left if k in self.kind_times)
        untimed = sum(k not in self.kind_times for k in left)
        eta = "-" if not left else f"{timed / 60:.0f} min" + (f" + {untimed} steps not timed yet" if untimed else "")
        how = "loaded from checkpoint" if from_checkpoint else f"took {took / 60:.1f} min"
        print(f"✓ {self.label}: {how} | [{self.bar()}] | elapsed {(time.time() - self.t0) / 60:.1f} min | ETA ~{eta}", flush=True)


def conf(yt, yp, w=None):
    return np.bincount(yt * K + yp, weights=w, minlength=K * K).reshape(K, K)


def f1_from_conf(cm):
    tp = np.diag(cm)
    return float(np.mean(2 * tp / np.maximum(cm.sum(0) + cm.sum(1), 1e-12)))


def fast_f1(yt, yp, w=None):
    """macro F1 (optionally sample-weighted), ~50× faster than sklearn; used inside the tuning loops"""
    return f1_from_conf(conf(yt, yp, w))


PLAN = (["load data + grouped folds", "build features", "ablation of feature blocks"]
        + [f"{m}: fold {f + 1}/{FOLDS}" for m in MODELS for f in range(FOLDS)]
        + ["level-2 operator stacker", "scenario table", "write submissions"])
prog = Progress(PLAN)

prog.start("load data + grouped folds")


DATA_DIR = ""  # optional: paste the folder that contains train.csv, e.g. "/kaggle/input/workstation-rental-price-tier-prediction"


def find_data():
    if DATA_DIR:
        return Path(DATA_DIR)
    if "PT_DATA" in os.environ:
        return Path(os.environ["PT_DATA"])
    for root in ["/kaggle/input", "/content", "."]:
        for dirpath, _, files in os.walk(root, followlinks=True):  # followlinks: Kaggle mounts inputs as symlinks
            if "train.csv" in files and "test.csv" in files:
                return Path(dirpath)
    found = [str(q) for q in Path("/kaggle/input").glob("*")] if Path("/kaggle/input").exists() else []
    raise FileNotFoundError(
        "train.csv / test.csv not found. The competition data is not attached to this notebook. "
        "Fix: right panel -> 'Add Input' -> 'Competitions' -> search 'workstation-rental-price-tier-prediction' -> Add. "
        f"Currently in /kaggle/input: {found}")


DATA = find_data()
train = pd.read_csv(DATA / "train.csv")
test = pd.read_csv(DATA / "test.csv")

# operator_ref alone isn't unique (different operators share names); these columns together identify one operator
HOST = ["operator_ref", "operator_platform_age_years", "operator_platform_age_months", "operator_active_years",
        "operator_active_months", "operator_class", "operator_portfolio_size"]
host = pd.concat([train[HOST], test[HOST]], ignore_index=True).groupby(HOST, dropna=False, sort=False).ngroup().to_numpy()
if QUICK_TEST:  # keep 10% of operators (whole operators, so the split logic stays the same)
    keep = (host * 2654435761 % 1000) < 100
    n0 = len(train)
    train, test = train[keep[:n0]].reset_index(drop=True), test[keep[n0:]].reset_index(drop=True)
    host = host[keep]

n_tr = len(train)
y = train["price_tier"].to_numpy() - 1
folds = list(StratifiedGroupKFold(FOLDS, shuffle=True, random_state=SEED).split(np.zeros(n_tr), y, host[:n_tr]))
fold_id = np.zeros(n_tr, int)
for i, (_, v) in enumerate(folds):
    fold_id[v] = i
print(train.shape, test.shape, "| test rows whose operator is in train:", round(np.isin(host[n_tr:], host[:n_tr]).mean(), 4))
prog.finish()