---

# ============================================================
# МТС "Умная полка" Track A
# Скачивание zip из облака -> обучение -> pred_a_improved.csv
# ============================================================

import os
import sys
import subprocess
import zipfile
import random
from pathlib import Path

# ----------------------------
# 1. Установка библиотек
# ----------------------------

packages = [
    "timm",
    "albumentations",
    "opencv-python-headless",
    "scikit-learn",
    "tqdm",
    "requests"
]

subprocess.check_call(
    [sys.executable, "-m", "pip", "install", "-q"] + packages
)

# ----------------------------
# 2. Импорты
# ----------------------------

import requests
import numpy as np
import pandas as pd
import cv2
from tqdm.auto import tqdm

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import timm
import albumentations as A
from albumentations.pytorch import ToTensorV2

from sklearn.model_selection import StratifiedGroupKFold
from google.colab import files

# ----------------------------
# 3. Настройки
# ----------------------------

DATA_URL = "https://storage.yandexcloud.net/codenrock-mts-production-platform-public/competitions/umnaya-polka/umnaya-polka.zip"

ZIP_PATH = Path("/content/umnaya-polka.zip")
OUT_DIR = Path("/content/data")

SEED = 42
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

MODEL_NAME = "convnext_tiny.fb_in22k_ft_in1k"
IMG_SIZE = 320
NUM_CLASSES = 5

BATCH_SIZE = 32

# Быстрый вариант для первой проверки:
FOLDS_TO_RUN = [0]
EPOCHS = 5

# Для результата лучше, но дольше:
# FOLDS_TO_RUN = [0, 1, 2]
# EPOCHS = 9

LR = 2e-4
WEIGHT_DECAY = 1e-4
N_FOLDS = 5

print("DEVICE:", DEVICE)
if DEVICE == "cpu":
    print("ВНИМАНИЕ: GPU не включён. Обучение будет очень долгим.")
    print("Включи Runtime -> Change runtime type -> T4 GPU и перезапусти ячейку.")

# ----------------------------
# 4. Seed
# ----------------------------

def seed_everything(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

seed_everything(SEED)

# ----------------------------
# 5. Надёжное скачивание zip из облака
# ----------------------------

def is_good_zip(path: Path):
    if not path.exists():
        return False
    if path.stat().st_size < 10_000_000:
        return False
    return zipfile.is_zipfile(path)

def download_with_requests(url, path):
    headers = {
        "User-Agent": "Mozilla/5.0"
    }
    with requests.get(url, stream=True, headers=headers, timeout=60) as r:
        r.raise_for_status()
        total = int(r.headers.get("content-length", 0))
        with open(path, "wb") as f, tqdm(
            total=total,
            unit="B",
            unit_scale=True,
            desc="Downloading zip"
        ) as pbar:
            for chunk in r.iter_content(chunk_size=1024 * 1024):
                if chunk:
                    f.write(chunk)
                    pbar.update(len(chunk))

OUT_DIR.mkdir(parents=True, exist_ok=True)

if is_good_zip(ZIP_PATH):
    print("Zip уже скачан и корректен:", ZIP_PATH)
else:
    if ZIP_PATH.exists():
        print("Найден битый или неполный zip. Удаляю его.")
        ZIP_PATH.unlink()

    print("Скачиваю архив из облака...")

    success = False

    try:
        subprocess.run(
            [
                "curl",
                "-L",
                "--fail",
                "--retry", "10",
                "--retry-delay", "5",
                "-o", str(ZIP_PATH),
                DATA_URL
            ],
            check=True
        )
        success = is_good_zip(ZIP_PATH)
    except Exception as e:
        print("curl не смог скачать файл:", e)

    if not success:
        try:
            if ZIP_PATH.exists():
                ZIP_PATH.unlink()
            subprocess.run(
                [
                    "wget",
                    "--tries=10",
                    "--timeout=60",
                    "-O", str(ZIP_PATH),
                    DATA_URL
                ],
                check=True
            )
            success = is_good_zip(ZIP_PATH)
        except Exception as e:
            print("wget не смог скачать файл:", e)

    if not success:
        print("Пробую скачать через Python requests...")
        if ZIP_PATH.exists():
            ZIP_PATH.unlink()
        download_with_requests(DATA_URL, ZIP_PATH)
        success = is_good_zip(ZIP_PATH)

    assert success, "Не удалось скачать корректный zip-файл."

print("Размер zip-файла:")
subprocess.run(["ls", "-lh", str(ZIP_PATH)])

# ----------------------------
# 6. Распаковка
# ----------------------------

print("Распаковываю архив...")

if OUT_DIR.exists():
    subprocess.run(["rm", "-rf", str(OUT_DIR)])

OUT_DIR.mkdir(parents=True, exist_ok=True)

with zipfile.ZipFile(ZIP_PATH, "r") as z:
    z.extractall(OUT_DIR)

print("Архив распакован в:", OUT_DIR)
print("Содержимое /content/data:")
print(os.listdir(OUT_DIR))

# ----------------------------
# 7. Правильные пути
# ----------------------------

ROOT_DIR = OUT_DIR / "umnaya-polka"

if not ROOT_DIR.exists():
    if (OUT_DIR / "images").exists() and (OUT_DIR / "track_a").exists():
        ROOT_DIR = OUT_DIR
    else:
        candidates = list(OUT_DIR.rglob("track_a"))
        assert len(candidates) > 0, "Не найдена папка track_a после распаковки."
        ROOT_DIR = candidates[0].parent

DATA_DIR = ROOT_DIR / "track_a"
IMG_DIR = ROOT_DIR / "images"

print("ROOT_DIR:", ROOT_DIR)
print("DATA_DIR:", DATA_DIR)
print("IMG_DIR:", IMG_DIR)

assert DATA_DIR.exists(), f"Не найдена папка track_a: {DATA_DIR}"
assert IMG_DIR.exists(), f"Не найдена папка images: {IMG_DIR}"

assert (DATA_DIR / "windows_train_a.csv").exists(), "Не найден windows_train_a.csv"
assert (DATA_DIR / "windows_test_a.csv").exists(), "Не найден windows_test_a.csv"

print("Файлы в track_a:")
print(os.listdir(DATA_DIR))

print("Примеры файлов в images:")
print(os.listdir(IMG_DIR)[:10])

# ----------------------------
# 8. Чтение CSV
# ----------------------------

train_df = pd.read_csv(DATA_DIR / "windows_train_a.csv")
test_df = pd.read_csv(DATA_DIR / "windows_test_a.csv")

train_df.columns = train_df.columns.str.strip()
test_df.columns = test_df.columns.str.strip()

print("train shape:", train_df.shape)
print("test shape:", test_df.shape)
print("train columns:", train_df.columns.tolist())
print("test columns:", test_df.columns.tolist())
print("Распределение классов:")
print(train_df["class"].value_counts().sort_index())

# ----------------------------
# 9. Поиск путей к изображениям
# ----------------------------

image_files = []
for ext in ["*.jpg", "*.jpeg", "*.png", "*.webp"]:
    image_files.extend(list(IMG_DIR.rglob(ext)))

print("Найдено изображений:", len(image_files))
assert len(image_files) > 0, "В папке images не найдены изображения."

by_name = {p.name: p for p in image_files}
by_stem = {p.stem: p for p in image_files}

def resolve_image_path(image_id):
    s = str(image_id)

    if s in by_name:
        return str(by_name[s])

    if s in by_stem:
        return str(by_stem[s])

    for ext in [".jpg", ".jpeg", ".png", ".webp"]:
        if s + ext in by_name:
            return str(by_name[s + ext])

    return None

train_df["image_path"] = train_df["image_id"].apply(resolve_image_path)
test_df["image_path"] = test_df["image_id"].apply(resolve_image_path)

print("Пропуски картинок train:", train_df["image_path"].isna().sum())
print("Пропуски картинок test:", test_df["image_path"].isna().sum())

if train_df["image_path"].isna().sum() > 0:
    print("Примеры train image_id без картинки:")
    print(train_df.loc[train_df["image_path"].isna(), "image_id"].head(20).tolist())

if test_df["image_path"].isna().sum() > 0:
    print("Примеры test image_id без картинки:")
    print(test_df.loc[test_df["image_path"].isna(), "image_id"].head(20).tolist())

assert train_df["image_path"].notna().all(), "Есть train-строки без картинки."
assert test_df["image_path"].notna().all(), "Есть test-строки без картинки."

# ----------------------------
# 10. Фолды
# ----------------------------

train_df["fold"] = -1

sgkf = StratifiedGroupKFold(
    n_splits=N_FOLDS,
    shuffle=True,
    random_state=SEED
)

for fold, (tr_idx, val_idx) in enumerate(
    sgkf.split(train_df, train_df["class"], groups=train_df["image_id"])
):
    train_df.loc[val_idx, "fold"] = fold

print("Фолды:")
print(train_df["fold"].value_counts().sort_index())

# ----------------------------
# 11. Dataset
# ----------------------------

class ShelfDataset(Dataset):
    def __init__(self, df, transform=None, is_test=False, pad=0.04):
        self.df = df.reset_index(drop=True)
        self.transform = transform
        self.is_test = is_test
        self.pad = pad

    def __len__(self):
        return len(self.df)

    def crop_window(self, row):
        img = cv2.imread(row["image_path"])

        if img is None:
            raise RuntimeError(f"Не удалось прочитать изображение: {row['image_path']}")

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        H, W = img.shape[:2]

        x = float(row["x"])
        y = float(row["y"])
        w = float(row["w"])
        h = float(row["h"])

        px = w * self.pad
        py = h * self.pad

        x1 = max(0, int(round(x - px)))
        y1 = max(0, int(round(y - py)))
        x2 = min(W, int(round(x + w + px)))
        y2 = min(H, int(round(y + h + py)))

        crop = img[y1:y2, x1:x2]

        if crop.size == 0:
            crop = img

        return crop

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img = self.crop_window(row)

        if self.transform:
            img = self.transform(image=img)["image"]

        if self.is_test:
            return img

        label = int(row["class"])
        return img, torch.tensor(label, dtype=torch.long)

# ----------------------------
# 12. Аугментации
# ----------------------------

train_transform = A.Compose([
    A.LongestMaxSize(max_size=IMG_SIZE),
    A.PadIfNeeded(
        min_height=IMG_SIZE,
        min_width=IMG_SIZE,
        border_mode=cv2.BORDER_CONSTANT,
        fill=(0, 0, 0)
    ),
    A.RandomCrop(height=IMG_SIZE, width=IMG_SIZE),

    A.HorizontalFlip(p=0.5),
    A.Affine(
        translate_percent=(-0.02, 0.02),
        scale=(0.92, 1.08),
        rotate=(-3, 3),
        border_mode=cv2.BORDER_CONSTANT,
        fill=(0, 0, 0),
        p=0.5
    ),
    A.RandomBrightnessContrast(p=0.35),
    A.HueSaturationValue(p=0.20),
    A.GaussianBlur(p=0.05),

    A.Normalize(
        mean=(0.485, 0.456, 0.406),
        std=(0.229, 0.224, 0.225)
    ),
    ToTensorV2()
])

valid_transform = A.Compose([
    A.LongestMaxSize(max_size=IMG_SIZE),
    A.PadIfNeeded(
        min_height=IMG_SIZE,
        min_width=IMG_SIZE,
        border_mode=cv2.BORDER_CONSTANT,
        fill=(0, 0, 0)
    ),
    A.CenterCrop(height=IMG_SIZE, width=IMG_SIZE),
    A.Normalize(
        mean=(0.485, 0.456, 0.406),
        std=(0.229, 0.224, 0.225)
    ),
    ToTensorV2()
])

valid_transform_flip = A.Compose([
    A.LongestMaxSize(max_size=IMG_SIZE),
    A.PadIfNeeded(
        min_height=IMG_SIZE,
        min_width=IMG_SIZE,
        border_mode=cv2.BORDER_CONSTANT,
        fill=(0, 0, 0)
    ),
    A.CenterCrop(height=IMG_SIZE, width=IMG_SIZE),
    A.HorizontalFlip(p=1.0),
    A.Normalize(
        mean=(0.485, 0.456, 0.406),
        std=(0.229, 0.224, 0.225)
    ),
    ToTensorV2()
])

# ----------------------------
# 13. Модель
# ----------------------------

def make_model():
    model = timm.create_model(
        MODEL_NAME,
        pretrained=True,
        num_classes=NUM_CLASSES,
        drop_rate=0.15
    )
    return model

def accuracy_from_logits(logits, targets):
    preds = logits.argmax(dim=1)
    return (preds == targets).float().mean().item()

# ----------------------------
# 14. Обучение одного фолда
# ----------------------------

def train_one_fold(fold):
    print(f"\n========== FOLD {fold} ==========")

    tr_df = train_df[train_df["fold"] != fold].copy()
    va_df = train_df[train_df["fold"] == fold].copy()

    train_ds = ShelfDataset(
        tr_df,
        transform=train_transform,
        is_test=False,
        pad=0.04
    )

    valid_ds = ShelfDataset(
        va_df,
        transform=valid_transform,
        is_test=False,
        pad=0.04
    )

    train_loader = DataLoader(
        train_ds,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=2,
        pin_memory=True,
        drop_last=True
    )

    valid_loader = DataLoader(
        valid_ds,
        batch_size=BATCH_SIZE * 2,
        shuffle=False,
        num_workers=2,
        pin_memory=True
    )

    model = make_model().to(DEVICE)

    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=LR,
        weight_decay=WEIGHT_DECAY
    )

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=EPOCHS * len(train_loader),
        eta_min=1e-6
    )

    scaler = torch.cuda.amp.GradScaler(enabled=(DEVICE == "cuda"))

    best_acc = 0.0
    best_path = f"/content/best_fold_{fold}.pth"

    for epoch in range(1, EPOCHS + 1):
        model.train()

        train_loss = 0.0
        train_acc = 0.0

        pbar = tqdm(
            train_loader,
            desc=f"Fold {fold} | Epoch {epoch}/{EPOCHS} | train"
        )

        for imgs, labels in pbar:
            imgs = imgs.to(DEVICE, non_blocking=True)
            labels = labels.to(DEVICE, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)

            with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
                logits = model(imgs)
                loss = criterion(logits, labels)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()

            acc = accuracy_from_logits(logits.detach(), labels)

            train_loss += loss.item()
            train_acc += acc

            pbar.set_postfix({
                "loss": f"{loss.item():.4f}",
                "acc": f"{acc:.4f}"
            })

        train_loss /= len(train_loader)
        train_acc /= len(train_loader)

        model.eval()

        valid_loss = 0.0
        valid_acc = 0.0

        with torch.no_grad():
            for imgs, labels in tqdm(valid_loader, desc=f"Fold {fold} | valid"):
                imgs = imgs.to(DEVICE, non_blocking=True)
                labels = labels.to(DEVICE, non_blocking=True)

                with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
                    logits = model(imgs)
                    loss = criterion(logits, labels)

                acc = accuracy_from_logits(logits, labels)

                valid_loss += loss.item()
                valid_acc += acc

        valid_loss /= len(valid_loader)
        valid_acc /= len(valid_loader)

        print(
            f"Epoch {epoch}: "
            f"train_loss={train_loss:.4f}, "
            f"train_acc={train_acc:.4f}, "
            f"valid_loss={valid_loss:.4f}, "
            f"valid_acc={valid_acc:.4f}"
        )

        if valid_acc > best_acc:
            best_acc = valid_acc
            torch.save(model.state_dict(), best_path)
            print(f"Сохранена лучшая модель: {best_path}, acc={best_acc:.5f}")

    return best_path, best_acc

# ----------------------------
# 15. Запуск обучения
# ----------------------------

best_models = []

for fold in FOLDS_TO_RUN:
    path, acc = train_one_fold(fold)
    best_models.append(path)

print("\nГотовые модели:")
print(best_models)

# ----------------------------
# 16. Предсказание
# ----------------------------

def predict_with_model(model_path, transform, pad=0.04):
    test_ds = ShelfDataset(
        test_df,
        transform=transform,
        is_test=True,
        pad=pad
    )

    test_loader = DataLoader(
        test_ds,
        batch_size=BATCH_SIZE * 2,
        shuffle=False,
        num_workers=2,
        pin_memory=True
    )

    model = make_model().to(DEVICE)
    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
    model.eval()

    all_probs = []

    with torch.no_grad():
        for imgs in tqdm(test_loader, desc=f"Predict {model_path}"):
            imgs = imgs.to(DEVICE, non_blocking=True)

            with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
                logits = model(imgs)
                probs = torch.softmax(logits, dim=1)

            all_probs.append(probs.cpu().numpy())

    return np.concatenate(all_probs, axis=0)

all_predictions = []

for model_path in best_models:
    all_predictions.append(
        predict_with_model(model_path, valid_transform, pad=0.04)
    )

    all_predictions.append(
        predict_with_model(model_path, valid_transform_flip, pad=0.04)
    )

    all_predictions.append(
        predict_with_model(model_path, valid_transform, pad=0.00)
    )

    all_predictions.append(
        predict_with_model(model_path, valid_transform, pad=0.08)
    )

mean_probs = np.mean(all_predictions, axis=0)
test_pred = mean_probs.argmax(axis=1)

print("\nРаспределение предсказанных классов:")
print(pd.Series(test_pred).value_counts().sort_index())

# ----------------------------
# 17. Создание submission
# ----------------------------

sample_path = DATA_DIR / "sample_submission_a.csv"

if sample_path.exists():
    sub = pd.read_csv(sample_path)
else:
    sub = test_df[["window_id"]].copy()
    sub["class"] = 0

sub["class"] = test_pred.astype(int)

SUB_PATH = "/content/pred_a_improved.csv"
sub.to_csv(SUB_PATH, index=False)

print("\nПервые строки submission:")
print(sub.head())

print("\nФайл сохранён:", SUB_PATH)

# ----------------------------
# 18. Скачивание файла
# ----------------------------

files.download(SUB_PATH)