# KAGGLE: {"gpu": false, "internet": true, "datasets": ["daniilantonov5/3d-mesh-quality-control"], "kernels": ["alekstop/00-eda", "alekstop/01-geom", "alekstop/02-imgstats", "alekstop/03-embed", "alekstop/04-tabular"]}

# %% [markdown]
# # 06 — Стекинг, отдельная работа над `quality` и отправляемый submission
#
# `05-frozen` показал лучшую конфигурацию (concat 4 бэкбонов по `mean` +
# табличные признаки + LightGBM, OOF 13.931), но считал только OOF и не трогал
# тест. Здесь мы:
#
# 1. **воспроизводим первый уровень с предсказанием теста** — фолд-модели
#    обучаются на train, предсказания теста усредняются по фолдам;
# 2. добавляем к нему две другие точки зрения: логрег по эмбеддингам (другой
#    индуктивный сдвиг) и уже готовую табличную модель из `04-tabular`;
# 3. **стекинг второго уровня** — мета-модель на 33 OOF-колонках. Именно она
#    должна вытащить `quality`: «объект хороший» — это функция от вероятностей
#    всех десяти дефектов, и мета-модель учит её явно, вместо нашего ручного
#    бленда с одним весом;
# 4. подбираем пороги пер-лейбл (усреднение по фолдам) и пишем `submission.csv`;
# 5. **adversarial validation** train vs test — проверяем, что пороги, настроенные
#    на train, вообще применимы к тесту.
#
# Всё на CPU: GPU-квота не тратится.

# %%
import os, sys, json, time, warnings
from pathlib import Path

import numpy as np
import pandas as pd

warnings.filterwarnings('ignore')
SEED = 42
np.random.seed(SEED)
T0 = time.time()
WORK = Path(os.environ.get('MQC_WORK', '/kaggle/working'))   # в Colab задаётся окружением
IN = Path(os.environ.get('MQC_IN', '/kaggle/input'))
NFOLD = 5

find = lambda pat: next(p for p in IN.rglob(pat))
csvs = {p.name: p for p in IN.rglob('*.csv')}
pick = lambda pref: sorted([v for k, v in csvs.items() if k.lower().startswith(pref)],
                           key=lambda p: len(p.name))[0]

train = pd.read_csv(pick('train'))
test = pd.read_csv(pick('test'))
folds = pd.read_csv(find('folds.csv'))
LABELS = [c for c in train.columns if c not in ('item_id', 'quality')]
TARGETS = LABELS + ['quality']
NL = len(LABELS)

EMB_DIR = find('emb_meta.json').parent
EIDS = json.load(open(EMB_DIR / 'emb_meta.json'))['ids']
EMB = {p.stem[4:]: np.load(p) for p in sorted(EMB_DIR.glob('emb_*.npy'))}
GEOM = pd.read_parquet(find('geom_features.parquet'))
IMG = pd.read_parquet(find('img_features.parquet'))
print('embeddings:', {k: v.shape for k, v in EMB.items()})

# %%
# --- выравнивание порядка объектов (тот же, что в 04 и 05) ---
full = pd.concat([train.assign(split='train'), test.assign(split='test')], ignore_index=True)
full = full.merge(folds, on='item_id', how='left')
pos = {i: k for k, i in enumerate(EIDS)}
row = full.item_id.map(pos).values.astype(int)

istr = (full.split == 'train').values
ite = ~istr
FOLD = full.fold.fillna(-1).astype(int).values
Y = full[TARGETS].fillna(0).astype(int).values
Ytr, Ftr = Y[istr], FOLD[istr]

TAB = full[['item_id']].merge(GEOM, on='item_id', how='left').merge(IMG, on='item_id', how='left')
TABF = [c for c in TAB.columns if c not in ('item_id', 'err')]
TABX = TAB[TABF].replace([np.inf, -np.inf], np.nan).astype(np.float32).fillna(0).values

# mean по 6 ракурсам — победитель E2 у всех четырёх бэкбонов
XEMB = np.concatenate([E[row].astype(np.float32).mean(1) for E in EMB.values()], 1)
XFULL = np.concatenate([XEMB, TABX], 1)
print('X emb %s | tab %s | full %s' % (XEMB.shape, TABX.shape, XFULL.shape))

# %%
# --- метрика и пороги ---
from sklearn.metrics import f1_score, roc_auc_score

GRID = np.arange(0.02, 0.96, 0.005)


def best_thr(y, p):
    s = np.array([f1_score(y, p > t, zero_division=0) for t in GRID])
    i = int(s.argmax())
    return float(GRID[i]), float(s[i])


def fold_thr(y, p, f):
    """Порог усредняем по фолдам: одиночный порог по всему OOF переобучается."""
    return float(np.mean([best_thr(y[f == k], p[f == k])[0] for k in range(NFOLD)]))


def score_oof(P, tag='', verbose=False):
    """P: (n_train, 11) -> (метрика, F1 quality, weighted F1 дефектов, пороги)"""
    thr, PA, per = {}, np.zeros((len(P), NL), int), {}
    for i, l in enumerate(LABELS):
        thr[l] = fold_thr(Ytr[:, i], P[:, i], Ftr)
        PA[:, i] = P[:, i] > thr[l]
        per[l] = round(f1_score(Ytr[:, i], PA[:, i], zero_division=0), 4)
    fa = f1_score(Ytr[:, :NL], PA, average='weighted', zero_division=0)
    thr['quality'] = fold_thr(Ytr[:, -1], P[:, -1], Ftr)
    fq = f1_score(Ytr[:, -1], P[:, -1] > thr['quality'], zero_division=0)
    tot = 10 * fq + 10 * fa
    if tag:
        print('%-40s f1_final=%6.3f  (q %.4f | a %.4f)' % (tag, tot, fq, fa), flush=True)
    if verbose:
        print('   per-label F1:', per)
    return tot, fq, fa, thr


# %%
# --- уровень 1: три разные точки зрения на объект ---
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from joblib import Parallel, delayed
import lightgbm as lgb

NJOB = os.cpu_count() or 4


def fit_predict(head, A, y, B, C):
    """Обучить одну бинарную голову, вернуть (valid_proba, test_proba)."""
    if y.sum() < 5:
        return np.zeros(len(B), np.float32), np.zeros(len(C), np.float32)
    if head == 'logreg':
        m = LogisticRegression(C=0.05, max_iter=800, class_weight='balanced')
    else:
        m = lgb.LGBMClassifier(objective='binary', learning_rate=0.05, num_leaves=31,
                               n_estimators=400, feature_fraction=0.25,
                               bagging_fraction=0.85, bagging_freq=1, lambda_l2=5.0,
                               verbosity=-1, n_jobs=NJOB, seed=SEED)
    m.fit(A, y)
    return (m.predict_proba(B)[:, 1].astype(np.float32),
            m.predict_proba(C)[:, 1].astype(np.float32))


def run_level1(X, head, tag):
    """5-фолдовый OOF + предсказание теста усреднением по фолд-моделям."""
    Xtr, Xte = X[istr], X[ite]
    oof = np.zeros((len(Xtr), len(TARGETS)), np.float32)
    pte = np.zeros((len(Xte), len(TARGETS)), np.float32)
    t = time.time()
    for k in range(NFOLD):
        tr, va = Ftr != k, Ftr == k
        if head == 'logreg':
            sc = StandardScaler().fit(Xtr[tr])
            A, B, C = sc.transform(Xtr[tr]), sc.transform(Xtr[va]), sc.transform(Xte)
            out = Parallel(n_jobs=NJOB, backend='loky')(
                delayed(fit_predict)(head, A, Ytr[tr, j], B, C) for j in range(len(TARGETS)))
        else:
            A, B, C = Xtr[tr], Xtr[va], Xte
            out = [fit_predict(head, A, Ytr[tr, j], B, C) for j in range(len(TARGETS))]
        for j, (pv, pt) in enumerate(out):
            oof[va, j] = pv
            pte[:, j] += pt / NFOLD
    print('  %s: %.1f min' % (tag, (time.time() - t) / 60), flush=True)
    return oof, pte


L1 = {}
L1['emb_tab_lgb'] = run_level1(XFULL, 'lgbm', 'lgbm(emb+tab)')
np.save(WORK / 'oof_embtab.npy', L1['emb_tab_lgb'][0])
np.save(WORK / 'pte_embtab.npy', L1['emb_tab_lgb'][1])
score_oof(L1['emb_tab_lgb'][0], 'L1 lgbm(emb+tab)')

L1['emb_lr'] = run_level1(XEMB, 'logreg', 'logreg(emb)')
score_oof(L1['emb_lr'][0], 'L1 logreg(emb)')

# табличная модель уже посчитана ядром 04 — переиспользуем, а не считаем заново
TAB_DIR = find('oof_tabular.npy').parent
L1['tab_lgb'] = (np.load(TAB_DIR / 'oof_tabular.npy'), np.load(TAB_DIR / 'pte_tabular.npy'))
score_oof(L1['tab_lgb'][0], 'L1 lgbm(tab) [из 04]')

# %%
# --- уровень 2: стекинг ---
# Мета-признаки — все 33 OOF-колонки (3 модели x 11 голов) плюс производная
# «ни одного дефекта». Мета-модель учится тому, что мы раньше задавали руками:
# согласованности между дефектами и quality.
#
# Оговорка про утечку: OOF-признаки объекта получены моделью, не видевшей этот
# объект, но видевшей другие фолды, на которых учится мета-модель. Это обычная
# «мягкая» утечка стекинга; она общепринята, а мы дополнительно держим тот же
# разбиение фолдов на обоих уровнях, чтобы не усиливать её.
def meta_matrix(idx):
    cols = [L1[k][idx] for k in ('emb_tab_lgb', 'emb_lr', 'tab_lgb')]
    der = [1.0 - c[:, :NL].max(1, keepdims=True) for c in cols]
    return np.concatenate(cols + der, 1).astype(np.float32)


MTR, MTE = meta_matrix(0), meta_matrix(1)
print('meta features:', MTR.shape)


def run_meta(head):
    oof = np.zeros((len(MTR), len(TARGETS)), np.float32)
    pte = np.zeros((len(MTE), len(TARGETS)), np.float32)
    for k in range(NFOLD):
        tr, va = Ftr != k, Ftr == k
        if head == 'logreg':
            sc = StandardScaler().fit(MTR[tr])
            A, B, C = sc.transform(MTR[tr]), sc.transform(MTR[va]), sc.transform(MTE)
        else:
            A, B, C = MTR[tr], MTR[va], MTE
        for j in range(len(TARGETS)):
            y = Ytr[tr, j]
            if y.sum() < 5:
                continue
            if head == 'logreg':
                m = LogisticRegression(C=1.0, max_iter=800, class_weight='balanced')
            else:
                m = lgb.LGBMClassifier(objective='binary', learning_rate=0.03, num_leaves=7,
                                       n_estimators=300, min_child_samples=40,
                                       feature_fraction=0.7, bagging_fraction=0.8,
                                       bagging_freq=1, lambda_l2=10.0, verbosity=-1,
                                       n_jobs=NJOB, seed=SEED)
            m.fit(A, y)
            oof[va, j] = m.predict_proba(B)[:, 1]
            pte[:, j] += m.predict_proba(C)[:, 1] / NFOLD
    return oof, pte


CAND = {}
for head in ('logreg', 'lgbm'):
    o, p = run_meta(head)
    CAND['stack_' + head] = (o, p)
    score_oof(o, 'L2 stack %s' % head)

# %%
# --- бленд уровня 1 с единым весом (устойчивая базовая линия) ---
A_oof, A_pte = L1['emb_tab_lgb']
T_oof, T_pte = L1['tab_lgb']
bw, bs = 0.0, -1
for w in np.arange(0, 1.01, 0.1):
    s = score_oof(w * A_oof + (1 - w) * T_oof)[0]
    if s > bs:
        bw, bs = w, s
print('blend w=%.2f -> %.3f' % (bw, bs))
CAND['blend'] = (bw * A_oof + (1 - bw) * T_oof, bw * A_pte + (1 - bw) * T_pte)
CAND['l1_embtab'] = L1['emb_tab_lgb']

# %%
# --- отдельная работа над quality (половина метрики) ---
# Три источника: прямая голова стека, производная «1 - max по дефектам» и
# среднее ранга. Вес и порог подбираем на OOF.
def quality_variants(oof):
    d = oof[:, -1]
    v = 1.0 - oof[:, :NL].max(1)
    return d, v


print('\nquality по кандидатам:')
QBEST = None
for nm, (o, p) in CAND.items():
    d, v = quality_variants(o)
    for w in np.arange(0, 1.01, 0.05):
        t, f1 = best_thr(Ytr[:, -1], w * v + (1 - w) * d)
        if QBEST is None or f1 > QBEST[0]:
            QBEST = (f1, nm, float(w), t)
    print('  %-16s direct F1=%.4f | derived F1=%.4f' %
          (nm, best_thr(Ytr[:, -1], d)[1], best_thr(Ytr[:, -1], v)[1]))
print('лучший quality: F1=%.4f  (%s, w_derived=%.2f, thr=%.3f)' % QBEST)

# %%
# --- выбор финальной конфигурации ---
print('\nсводка кандидатов:')
ROWS = []
for nm, (o, p) in CAND.items():
    tot, fq, fa, thr = score_oof(o, 'FINAL? ' + nm)
    ROWS.append(dict(cand=nm, f1_final=tot, f1_quality=fq, f1_artefacts=fa))
R = pd.DataFrame(ROWS).sort_values('f1_final', ascending=False)
R.to_csv(WORK / 'experiments_stack.csv', index=False)

BEST = R.iloc[0]['cand']
oof_b, pte_b = CAND[BEST]
tot, fq, fa, THR = score_oof(oof_b, verbose=True)

# quality берём от лучшего варианта, если он заметно лучше прямой головы стека
qf1, qnm, qw, qthr = QBEST
q_oof = qw * (1.0 - CAND[qnm][0][:, :NL].max(1)) + (1 - qw) * CAND[qnm][0][:, -1]
q_pte = qw * (1.0 - CAND[qnm][1][:, :NL].max(1)) + (1 - qw) * CAND[qnm][1][:, -1]
qthr_cv = fold_thr(Ytr[:, -1], q_oof, Ftr)
fq2 = f1_score(Ytr[:, -1], q_oof > qthr_cv, zero_division=0)
if fq2 > fq:
    print('quality заменён на %s (F1 %.4f -> %.4f)' % (qnm, fq, fq2))
    fq, THR['quality'] = fq2, qthr_cv
else:
    q_oof, q_pte = oof_b[:, -1], pte_b[:, -1]

print('\nИТОГ OOF: f1_final = %.3f  (quality %.4f | artefacts %.4f)  [%s]' %
      (10 * fq + 10 * fa, fq, fa, BEST))

# %%
# --- adversarial validation: применимы ли пороги train к тесту ---
Xav = np.concatenate([XEMB, TABX], 1)
yav = ite.astype(int)
av = lgb.LGBMClassifier(objective='binary', n_estimators=200, learning_rate=0.05,
                        num_leaves=31, feature_fraction=0.3, verbosity=-1,
                        n_jobs=NJOB, seed=SEED)
from sklearn.model_selection import cross_val_predict
pav = cross_val_predict(av, Xav, yav, cv=5, method='predict_proba')[:, 1]
print('adversarial validation AUC train vs test = %.4f  (0.5 = распределения совпадают)'
      % roc_auc_score(yav, pav))

# %%
# --- submission ---
sub = pd.DataFrame({'item_id': full.loc[ite, 'item_id'].values})
for i, l in enumerate(LABELS):
    sub[l] = (pte_b[:, i] > THR[l]).astype(int)
sub['quality'] = (q_pte > THR['quality']).astype(int)
sub = sub[['item_id'] + LABELS + ['quality']]
sub.to_csv(WORK / 'submission.csv', index=False)

np.save(WORK / 'oof_stack.npy', oof_b)
np.save(WORK / 'pte_stack.npy', pte_b)
json.dump({'best': BEST, 'thr': THR, 'quality_src': qnm, 'quality_w': qw,
           'oof_final': float(10 * fq + 10 * fa), 'oof_quality': float(fq),
           'oof_artefacts': float(fa)},
          open(WORK / 'stack_config.json', 'w'), indent=2, default=float)

print('\nдоля положительных в submission:')
print(sub[LABELS + ['quality']].mean().round(4).to_string())
print('в train:')
print(train[LABELS + ['quality']].mean().round(4).to_string())
print('\nвсего:', round((time.time() - T0) / 60, 1), 'мин')