# 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", "alekstop/06-stack", "alekstop/07-finetune", "alekstop/10-visfeat"]}

# %% [markdown]
# # 08 — Финальный стек: таблица + видимость + замороженные эмбеддинги + дообученный ViT
#
# `06-stack` собрал стек из трёх точек зрения (LightGBM на эмбеддингах+таблице,
# логрег на эмбеддингах, табличный LightGBM) и дал OOF **14.088**. С тех пор
# появились два новых источника сигнала:
#
# * **`07-finetune`** — дообученный DINOv3 с multi-view агрегацией, первая в
#   пайплайне модель, которая *видит* картинку, а не её замороженное описание;
# * **`10-visfeat`** — 36 признаков видимости, ставших возможными после
#   восстановления камеры: какая доля поверхности не видна ни с одного ракурса,
#   распределение площади по числу видов, толщина объекта вдоль луча. Это
#   буквальные определения `partial`, `open` и `intersection`.
#
# Ядро измеряет вклад каждой добавки **на одних и тех же фолдах**:
#
# 1. уровень 1 пересчитывается с признаками видимости и сравнивается с прежними
#    числами из `04-tabular` / `06-stack` — сколько дала видимость;
# 2. стек собирается с FT и без FT — сколько дал fine-tune, то есть окупились ли
#    GPU-часы;
# 3. пороги пер-лейбл подбираются усреднением по фолдам, пишется `submission.csv`;
# 4. `quality` (половина метрики) разбирается отдельно: прямая голова стека
#    против производной «1 − max по дефектам» против их бленда.

# %%
import os, 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, NT = len(LABELS), len(LABELS) + 1

train = train.merge(folds, on='item_id', how='left')
Ftr = train.fold.astype(int).values
Ytr = train[TARGETS].astype(int).values
print('train %d | test %d' % (len(train), len(test)))

# %%
# --- признаки и эмбеддинги ---
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'))
IMGF = pd.read_parquet(find('img_features.parquet'))
try:
    VIS = pd.read_parquet(find('vis_features.parquet'))
    print('признаки видимости:', VIS.shape)
except StopIteration:
    VIS = None
    print('vis_features.parquet не найден — считаем без признаков видимости')

full = pd.concat([train.assign(split='train'), test.assign(split='test')], ignore_index=True)
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

TAB = full[['item_id']].merge(GEOM, on='item_id', how='left').merge(IMGF, on='item_id', how='left')
if VIS is not None:
    TAB = TAB.merge(VIS, 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
XEMB = np.concatenate([E[row].astype(np.float32).mean(1) for E in EMB.values()], 1)
XFULL = np.concatenate([XEMB, TABX], 1)
print('emb %s | tab %s | full %s' % (XEMB.shape, TABX.shape, XFULL.shape))

# %%
# --- метрика и пороги (идентичны 04/05/06/07) ---
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):
    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, per


# %%
# --- уровень 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

STACK_DIR = find('oof_embtab.npy').parent
TAB_DIR = find('oof_tabular.npy').parent
FT_DIR = find('oof_ft.npy').parent
FT_META = json.load(open(FT_DIR / 'ft_meta.json'))
print('FT: backbone=%s img=%d folds_done=%s' % (FT_META['backbone'], FT_META['img'],
                                                FT_META['folds_done']))

# прежние числа (без признаков видимости) — база для сравнения
OLD = {'emb_tab_lgb': (np.load(STACK_DIR / 'oof_embtab.npy'), np.load(STACK_DIR / 'pte_embtab.npy')),
       'tab_lgb': (np.load(TAB_DIR / 'oof_tabular.npy'), np.load(TAB_DIR / 'pte_tabular.npy'))}

L1 = {}
# FT сохранён в порядке train.csv/test.csv — тот же, что и здесь; сверяем явно
assert FT_META['ids_train'] == list(train.item_id), 'порядок train не совпал'
assert FT_META['ids_test'] == list(test.item_id), 'порядок test не совпал'
L1['ft'] = (np.load(FT_DIR / 'oof_ft.npy'), np.load(FT_DIR / 'pte_ft.npy'))

FT_DONE = FT_META['folds_done']
if len(FT_DONE) < NFOLD:
    # неполный OOF: строки недосчитанных фолдов остались нулями и испортили бы
    # и стек, и пороги. Заполняем их средним предсказанием по посчитанным фолдам.
    m = np.isin(Ftr, FT_DONE)
    print('FT неполный: %d/%d фолдов, %d строк заполняем средним'
          % (len(FT_DONE), NFOLD, int((~m).sum())))
    o = L1['ft'][0].copy()
    o[~m] = o[m].mean(0)
    L1['ft'] = (o, L1['ft'][1])


def fit_predict(head, A, y, B, C):
    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 + предсказание теста усреднением по фолд-моделям.

    Логрег однопоточный (lbfgs), поэтому 11 голов идут параллельно по ядрам;
    LightGBM сам занимает все ядра — его крутим последовательно.
    """
    Xtr, Xte = X[istr], X[ite]
    oof = np.zeros((len(Xtr), NT), np.float32)
    pte = np.zeros((len(Xte), NT), 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(NT))
        else:
            A, B, C = Xtr[tr], Xtr[va], Xte
            out = [fit_predict(head, A, Ytr[tr, j], B, C) for j in range(NT)]
        for j, (pv, pt) in enumerate(out):
            oof[va, j] = pv
            pte[:, j] += pt / NFOLD
    print('  %s: %.1f мин' % (tag, (time.time() - t) / 60), flush=True)
    return oof, pte


L1['tab_lgb'] = run_level1(TABX, 'lgbm', 'lgbm(tab+vis)')
L1['emb_tab_lgb'] = run_level1(XFULL, 'lgbm', 'lgbm(emb+tab+vis)')
L1['emb_lr'] = run_level1(XEMB, 'logreg', 'logreg(emb)')

print('\nсколько дали признаки видимости (одни фолды, одна модель):')
for nm in ('tab_lgb', 'emb_tab_lgb'):
    a = score_oof(OLD[nm][0])[0]
    b = score_oof(L1[nm][0])[0]
    print('  %-14s без vis %.3f -> с vis %.3f  (%+.3f балла)' % (nm, a, b, b - a))

print()
for nm in ('emb_tab_lgb', 'tab_lgb', 'emb_lr', 'ft'):
    score_oof(L1[nm][0], 'L1 ' + nm)

np.save(WORK / 'oof_embtabvis.npy', L1['emb_tab_lgb'][0])
np.save(WORK / 'pte_embtabvis.npy', L1['emb_tab_lgb'][1])
np.save(WORK / 'oof_tabvis.npy', L1['tab_lgb'][0])
np.save(WORK / 'pte_tabvis.npy', L1['tab_lgb'][1])

# %%
# --- уровень 2: стекинг, с FT и без него ---
def meta_matrix(keys, idx):
    cols = [L1[k][idx] for k in keys]
    der = [1.0 - c[:, :NL].max(1, keepdims=True) for c in cols]
    return np.concatenate(cols + der, 1).astype(np.float32)


def run_meta(keys, head):
    MTR, MTE = meta_matrix(keys, 0), meta_matrix(keys, 1)
    oof = np.zeros((len(MTR), NT), np.float32)
    pte = np.zeros((len(MTE), NT), 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(NT):
            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


BASE = ['emb_tab_lgb', 'emb_lr', 'tab_lgb']
CAND = {}
for tag, keys in (('no_ft', BASE), ('with_ft', BASE + ['ft'])):
    for head in ('logreg', 'lgbm'):
        nm = 'stack_%s_%s' % (head, tag)
        CAND[nm] = run_meta(keys, head)
        score_oof(CAND[nm][0], 'L2 ' + nm)

# простой бленд стека с FT — страховка на случай, если мета-модель переобучится
best_stack = max(CAND, key=lambda n: score_oof(CAND[n][0])[0])
S_oof, S_pte = CAND[best_stack]
F_oof, F_pte = L1['ft']
bw, bs = 0.0, -1.0
for w in np.arange(0, 0.51, 0.05):
    s = score_oof((1 - w) * S_oof + w * F_oof)[0]
    if s > bs:
        bw, bs = float(w), s
print('blend %s + %.2f*ft -> %.3f' % (best_stack, bw, bs))
if bw > 0:
    CAND['blend_ft'] = ((1 - bw) * S_oof + bw * F_oof, (1 - bw) * S_pte + bw * F_pte)

# %%
# --- отдельная работа над quality (половина метрики) ---
print('\nquality по кандидатам:')
QBEST = None
for nm, (o, p) in CAND.items():
    d = o[:, -1]
    v = 1.0 - o[:, :NL].max(1)
    for w in np.arange(0, 1.01, 0.05):
        _, f1 = best_thr(Ytr[:, -1], w * v + (1 - w) * d)
        if QBEST is None or f1 > QBEST[0]:
            QBEST = (f1, nm, float(w))
    print('  %-24s 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)' % QBEST)

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

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

qf1, qnm, qw = 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 = fold_thr(Ytr[:, -1], q_oof, Ftr)
fq2 = f1_score(Ytr[:, -1], q_oof > qthr, zero_division=0)
if fq2 > fq:
    print('quality заменён на %s w=%.2f (F1 %.4f -> %.4f)' % (qnm, qw, fq, fq2))
    fq, THR['quality'] = fq2, qthr
else:
    q_oof, q_pte, qnm, qw = oof_b[:, -1], pte_b[:, -1], BEST, 0.0

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

# %%
# --- сколько дал fine-tune: честное сравнение на одних фолдах ---
a = max(score_oof(CAND[n][0])[0] for n in CAND if n.endswith('no_ft'))
b = max(score_oof(CAND[n][0])[0] for n in CAND if 'with_ft' in n or n == 'blend_ft')
print('стек без FT: %.3f | с FT: %.3f | прирост %+.3f балла' % (a, b, b - a))

# %%
# --- adversarial validation: применимы ли пороги train к тесту ---
from sklearn.model_selection import cross_val_predict

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)
pav = cross_val_predict(av, XFULL, yav, cv=5, method='predict_proba')[:, 1]
print('adversarial validation AUC train vs test = %.4f' % roc_auc_score(yav, pav))

# %%
# --- submission ---
sub = pd.DataFrame({'item_id': test.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_final.npy', oof_b)
np.save(WORK / 'pte_final.npy', pte_b)
np.save(WORK / 'oof_quality.npy', q_oof)
np.save(WORK / 'pte_quality.npy', q_pte)
json.dump({'best': BEST, 'thr': THR, 'quality_src': qnm, 'quality_w': qw,
           'blend_w': float(bw), 'blend_base': best_stack, 'cands': sorted(CAND),
           'oof_final': float(FINAL), 'oof_quality': float(fq), 'oof_artefacts': float(fa),
           'per_label': per, 'gain_from_ft': float(b - a), 'backbone': FT_META['backbone']},
          open(WORK / 'final_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всего: %.1f мин' % ((time.time() - T0) / 60))