#!/usr/bin/env bash
#
# DXA Quality Assessment — единая точка входа.
#
# Использование: ./run.sh <команда> [аргументы]
#   ./run.sh train [доп. аргументы для src.dxa.train]
#   ./run.sh infer <вход> <выход.xlsx|.csv> [доп. аргументы]
#   ./run.sh serve [порт]
#   ./run.sh label
#   ./run.sh split
#   ./run.sh rename [--apply]
#   ./run.sh compare
#   ./run.sh test
#
# Скрипт намеренно не содержит своей реализации обучения и инференса: вся
# логика живёт в Python-модулях, а здесь только вызовы. Дублирование цикла
# обучения (как было раньше) приводило к расхождению правил между режимами.
set -euo pipefail

GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; NC='\033[0m'

PYTHON=${PYTHON:-python3}
DATA_ROOT=${DATA_ROOT:-dataset_hack}
ANNOTATION_PATH=${ANNOTATION_PATH:-"dataset_hack/НД_для_обучения/разметка.xlsx"}
MODEL_PATH=${MODEL_PATH:-models/dxa_model.pth}
LABELS_DIR=${LABELS_DIR:-labels}
LABELS_CSV=${LABELS_CSV:-labels/labels_images.csv}
LABELS_REFERENCE_CSV=${LABELS_REFERENCE_CSV:-labels/labels_images_expert.csv}
SPLIT_FILE=${SPLIT_FILE:-labels/split_expert_seed42.json}
RENAME_MAP=${RENAME_MAP:-labels/rename_map.csv}
PORT=${PORT:-8000}

command=${1:-help}
shift || true

case "$command" in
    train)
        echo -e "${YELLOW}Обучение классификатора качества DXA${NC}"
        echo "  метки: $LABELS_CSV"
        "$PYTHON" -m src.dxa.train \
            --data-root "$DATA_ROOT" \
            --annotation-path "$ANNOTATION_PATH" \
            --labels-csv "$LABELS_CSV" \
            --output-dir "$(dirname "$MODEL_PATH")" \
            "$@"
        echo -e "${GREEN}Готово. Чекпоинт: ${MODEL_PATH}${NC}"
        ;;

    infer)
        input_path=${1:-dataset_hack/Для теста}
        output_path=${2:-results.xlsx}
        shift 2 2>/dev/null || true
        echo -e "${YELLOW}Пакетная обработка${NC}"
        echo "  вход:  $input_path"
        echo "  выход: $output_path"
        "$PYTHON" -m src.dxa.inference \
            --input-path "$input_path" \
            --output-path "$output_path" \
            --model-path "$MODEL_PATH" \
            "$@"
        echo -e "${GREEN}Готово: ${output_path}${NC}"
        ;;

    serve)
        echo -e "${YELLOW}Запуск API на порту ${PORT}${NC}"
        DXA_MODEL_PATH="$MODEL_PATH" \
            "$PYTHON" -m uvicorn src.main:app --host 0.0.0.0 --port "$PORT"
        ;;

    label)
        echo -e "${YELLOW}Разметка датасета на уровне снимков${NC}"
        echo "  таблица: $ANNOTATION_PATH"
        "$PYTHON" -m src.dxa.excel_labels \
            --data-root "$DATA_ROOT" \
            --excel "$ANNOTATION_PATH" \
            --out-dir "$LABELS_DIR" \
            "$@"
        ;;

    split)
        echo -e "${YELLOW}Фиксация разбиения для сравнения правил разметки${NC}"
        echo "  стратификация по эталону: $LABELS_REFERENCE_CSV"
        "$PYTHON" -m src.dxa.train \
            --data-root "$DATA_ROOT" \
            --annotation-path "$ANNOTATION_PATH" \
            --labels-csv "$LABELS_REFERENCE_CSV" \
            --export-split "$SPLIT_FILE" \
            "$@"
        ;;

    rename)
        echo -e "${YELLOW}Приведение имён DICOM к единому виду${NC}"
        if [ "${1:-}" = "--apply" ]; then
            echo -e "${RED}Правка имён. Карта отката: $RENAME_MAP${NC}"
        else
            echo "  режим плана; для правки добавьте --apply"
        fi
        "$PYTHON" -m src.dxa.rename_files \
            --data-root "$DATA_ROOT" \
            --mapping "$RENAME_MAP" \
            "$@"
        ;;

    compare)
        echo -e "${YELLOW}Сравнение источников разметки на одном held-out наборе${NC}"
        echo "  разбиение: $SPLIT_FILE"
        if [ ! -f "$SPLIT_FILE" ]; then
            echo -e "${RED}Нет файла разбиения. Сначала: ./run.sh split${NC}" >&2
            exit 1
        fi
        "$PYTHON" -m src.dxa.compare_labels \
            --data-root "$DATA_ROOT" \
            --annotation-path "$ANNOTATION_PATH" \
            --split-file "$SPLIT_FILE" \
            "$@"
        ;;

    test)
        echo -e "${YELLOW}Запуск тестов${NC}"
        "$PYTHON" -m pytest tests/ -q
        ;;

    help|*)
        cat <<'USAGE'
DXA Quality Assessment — управление запуском

Команды:
  train [аргументы]              Обучить классификатор качества.
                                 Пример: ./run.sh train --epochs 100 --head mlp
  infer <вход> <выход> [арг.]    Пакетная обработка DICOM (файл или каталог).
                                 Пример: ./run.sh infer "dataset_hack/Для теста" results.xlsx
                                 Дополнительно: --zip-out masks.zip
  serve [порт]                   Запустить HTTP API и веб-интерфейс.
  label [аргументы]              Разметить датасет на уровне снимков по Excel.
                                 Результат: labels/labels_images.csv|.xlsx
  split [аргументы]              Зафиксировать разбиение по исследованиям.
                                 Результат: labels/split_expert_seed42.json
  rename [--apply]               Привести имена DICOM к виду область_NN[_метка].
                                 Без --apply — только план; карта отката пишется
                                 до правки (labels/rename_map.csv)
  compare [аргументы]            Сравнить источники разметки на одном held-out
                                 наборе (нужен сначала `./run.sh split`).
  test                           Запустить тесты.

Переменные окружения:
  DATA_ROOT         Каталог датасета (по умолчанию dataset_hack)
  ANNOTATION_PATH   Excel с экспертной разметкой исследований
  MODEL_PATH        Путь к чекпоинту (по умолчанию models/dxa_model.pth)
  LABELS_DIR        Каталог с результатом разметки (по умолчанию labels)
  LABELS_CSV        Разметка снимков для обучения (пустая строка — метки из имён файлов)
  SPLIT_FILE        Файл фиксированного разбиения (labels/split_expert_seed42.json)
  LABELS_REFERENCE_CSV  Эталон для стратификации разбиения (labels/labels_images_expert.csv)
  RENAME_MAP        Карта переименований для отката (по умолчанию labels/rename_map.csv)
  PORT              Порт API (по умолчанию 8000)
  PYTHON            Интерпретатор (по умолчанию python3)

Типовой порядок работы:
  ./run.sh label                       # разметить датасет по Excel
  ./run.sh train                       # обучить модель
  ./run.sh infer dataset_hack results.xlsx
  ./run.sh serve                       # веб-интерфейс на http://localhost:8000
  ./run.sh split && ./run.sh compare   # сравнить варианты разметки
USAGE
        ;;
esac