---

# ============================================================
# МТС "Умная полка" Track A
# Быстрый демонстрационный режим обучения ~ до 10 минут
# Скачивание zip из облака -> обучение -> pred_a_fast.csv
# ============================================================

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

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

packages = [
    "timm",
    "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

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 = "resnet18"

# Меньший размер изображения
IMG_SIZE = 160

NUM_CLASSES = 5

# Большой batch для ускорения
BATCH_SIZE = 128

# Только 1 фолд и 1 эпоха
FOLDS_TO_RUN = [0]
EPOCHS = 1

# Берём не весь train, а часть для демонстрации
MAX_TRAIN_ROWS = 4000

LR = 1e-3
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=(10, 120)) 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("Исходное распределение классов:")
print(train_df["class"].value_counts().sort_index())

# ----------------------------
# 9. Уменьшаем train для быстрого обучения
# ----------------------------

sample_parts = []
per_class = MAX_TRAIN_ROWS // NUM_CLASSES

for cls in sorted(train_df["class"].unique()):
    cls_df = train_df[train_df["class"] == cls]

    take_n = min(len(cls_df), per_class)

    sample_parts.append(
        cls_df.sample(
            n=take_n,
            random_state=SEED
        )
    )

train_df = pd.concat(sample_parts, axis=0).sample(
    frac=1,
    random_state=SEED
).reset_index(drop=True)

print("Быстрый train shape:", train_df.shape)
print("Распределение классов после сэмплирования:")
print(train_df["class"].value_counts().sort_index())

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

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-строки без картинки."

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

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())

# ----------------------------
# 12. Быстрая предобработка без albumentations
# ----------------------------

MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

def letterbox_resize(img, size=160):
    h, w = img.shape[:2]

    if h == 0 or w == 0:
        img = np.zeros((size, size, 3), dtype=np.uint8)
        return img

    scale = size / max(h, w)
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))

    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)

    canvas = np.zeros((size, size, 3), dtype=np.uint8)

    y0 = (size - new_h) // 2
    x0 = (size - new_w) // 2

    canvas[y0:y0 + new_h, x0:x0 + new_w] = resized

    return canvas

def preprocess_image(img, train=True):
    img = letterbox_resize(img, IMG_SIZE)

    if train:
        if random.random() < 0.5:
            img = cv2.flip(img, 1)

        if random.random() < 0.25:
            alpha = random.uniform(0.85, 1.15)
            beta = random.uniform(-10, 10)
            img = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)

    img = img.astype(np.float32) / 255.0
    img = (img - MEAN) / STD

    img = np.transpose(img, (2, 0, 1))
    img = torch.tensor(img, dtype=torch.float32)

    return img

# ----------------------------
# 13. Dataset
# ----------------------------

class ShelfDataset(Dataset):
    def __init__(self, df, is_test=False, train_mode=False, pad=0.04):
        self.df = df.reset_index(drop=True)
        self.is_test = is_test
        self.train_mode = train_mode
        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)
        img = preprocess_image(img, train=self.train_mode)

        if self.is_test:
            return img

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

# ----------------------------
# 14. Модель ResNet18 с заморозкой backbone
# ----------------------------

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

    # Замораживаем все параметры
    for name, param in model.named_parameters():
        param.requires_grad = False

    # Размораживаем только последний классификатор
    for name, param in model.named_parameters():
        if (
            "fc" in name
            or "classifier" in name
            or "head" in name
        ):
            param.requires_grad = True

    return model

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

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

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,
        is_test=False,
        train_mode=True,
        pad=0.04
    )

    valid_ds = ShelfDataset(
        va_df,
        is_test=False,
        train_mode=False,
        pad=0.04
    )

    train_loader = DataLoader(
        train_ds,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=2,
        pin_memory=(DEVICE == "cuda"),
        drop_last=False
    )

    valid_loader = DataLoader(
        valid_ds,
        batch_size=BATCH_SIZE * 2,
        shuffle=False,
        num_workers=2,
        pin_memory=(DEVICE == "cuda")
    )

    model = make_model().to(DEVICE)

    trainable_params = [p for p in model.parameters() if p.requires_grad]
    print("Обучаемых параметров:", sum(p.numel() for p in trainable_params))

    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)

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

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

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

    best_acc = 0.0
    best_path = f"/content/best_fast_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

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

best_models = []

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

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

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

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

    test_loader = DataLoader(
        test_ds,
        batch_size=BATCH_SIZE * 2,
        shuffle=False,
        num_workers=2,
        pin_memory=(DEVICE == "cuda")
    )

    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, pad=0.04)
    )

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())

# ----------------------------
# 18. Создание 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_fast.csv"
sub.to_csv(SUB_PATH, index=False)

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

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

# ----------------------------
# 19. Проверка формата submission
# ----------------------------

validate_script = DATA_DIR / "validate_submission.py"

if validate_script.exists():
    print("\nПроверяю формат submission...")
    subprocess.run(
        [
            sys.executable,
            str(validate_script),
            "--track", "a",
            "--pred", SUB_PATH
        ],
        cwd=str(ROOT_DIR),
        check=False
    )
else:
    print("validate_submission.py не найден, пропускаю проверку.")

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

files.download(SUB_PATH)