from typing import NamedTuple, Sequence, Literal, Union
import subprocess
import vh3


@vh3.decorator.external_operation("https://nirvana.yandex-team.ru/operation/9ff37ad1-f024-4079-9e90-a7b62743449e")
@vh3.decorator.nirvana_output_names("code")
@vh3.decorator.nirvana_names_transformer(vh3.name_transformers.snake_to_dash, options=True)
def arc_export(
    *,
    arc_token: vh3.Secret,
    path: vh3.String,
    ttl: vh3.Integer = 360,
    max_ram: vh3.Integer = 100,
    max_disk: vh3.Integer = 1024,
    retries_on_job_failure: vh3.Integer = 0,
    debug_timeout: vh3.Integer = 0,
    reference: vh3.String = "trunk"
) -> vh3.Binary:
    """
    Arc export

    Checkouts an arc branch to an archive

    :param reference: Reference
      [[Arc commit or branch name or revision]]
    """
    raise NotImplementedError("Write your local execution stub here")


@vh3.decorator.external_operation("https://nirvana.yandex-team.ru/operation/c6d54a7d-e64d-47b7-b83a-c70d58a1e50e")
@vh3.decorator.nirvana_names(tar_mode="tar_mode")
@vh3.decorator.nirvana_output_names("archive")
@vh3.decorator.nirvana_names_transformer(vh3.name_transformers.snake_to_dash, options=True)
def create_tar_archive(
    *,
    ttl: vh3.Integer = 360,
    max_ram: vh3.Integer = 1024,
    max_disk: vh3.Integer = 1024,
    tar_mode: vh3.Enum[Literal["none", "extract", "subdir", "plus"]] = "extract",
    compress: vh3.Enum[Literal["--no-auto-compress", "--gzip", "--bzip2"]] = "--no-auto-compress",
    files: Sequence[Union[vh3.Binary, vh3.Executable, vh3.HTML, vh3.Image, vh3.JSON, vh3.TSV, vh3.Text, vh3.XML]] = ()
) -> vh3.Binary:
    """
    Create TAR archive
    Create TAR archive from given binary files. If the files are archives themselves, unpack them.

    :param tar_mode: Tar processing mode
      [[Action for input tar archives]]
      Action for input tar archives
    :param compress: Tar compress
      [[Tar compress mode]]
      Tar compress mode
    :param files:
      File or archive
    """
    raise NotImplementedError("Write your local execution stub here")


class BuildArcadiaProjectOutput(NamedTuple):
    arcadia_project: vh3.Executable
    sandbox_task_id: vh3.Text


@vh3.decorator.external_operation("https://nirvana.yandex-team.ru/operation/dd4b5735-1ee7-497d-91fc-b81ba8b510fc")
@vh3.decorator.nirvana_output_names(arcadia_project="ARCADIA_PROJECT", sandbox_task_id="SANDBOX_TASK_ID")
def build_arcadia_project(
    *,
    targets: vh3.String,
    arts: vh3.String,
    arcadia_url: vh3.String = "arcadia:/arc/trunk/arcadia",
    arcadia_revision: vh3.Integer = None,
    checkout_arcadia_from_url: vh3.String = None,
    build_type: vh3.Enum[
        Literal["release", "debug", "profile", "coverage", "relwithdebinfo", "valgrind", "valgrind-release"]
    ] = "release",
    arts_source: vh3.String = None,
    result_single_file: vh3.Boolean = False,
    definition_flags: vh3.String = None,
    sandbox_oauth_token: vh3.Secret = None,
    arcadia_patch: vh3.String = None,
    owner: vh3.String = None,
    use_aapi_fuse: vh3.Boolean = True,
    use_arc_instead_of_aapi: vh3.Boolean = True,
    aapi_fallback: vh3.Boolean = False,
    kill_timeout: vh3.Integer = None,
    sandbox_requirements_disk: vh3.Integer = None,
    sandbox_requirements_ram: vh3.Integer = None,
    sandbox_requirements_platform: vh3.Enum[
        Literal[
            "Any",
            "darwin-20.4.0-x86_64-i386-64bit",
            "linux",
            "linux_ubuntu_10.04_lucid",
            "linux_ubuntu_12.04_precise",
            "linux_ubuntu_14.04_trusty",
            "linux_ubuntu_16.04_xenial",
            "linux_ubuntu_18.04_bionic",
            "osx",
            "osx_10.12_sierra",
            "osx_10.13_high_sierra",
            "osx_10.14_mojave",
            "osx_10.15_catalina",
            "osx_10.16_big_sur",
        ]
    ] = None,
    checkout: vh3.Boolean = False,
    clear_build: vh3.Boolean = True,
    strip_binaries: vh3.Boolean = False,
    lto: vh3.Boolean = False,
    thinlto: vh3.Boolean = False,
    musl: vh3.Boolean = False,
    use_system_python: vh3.Boolean = False,
    target_platform_flags: vh3.String = None,
    javac_options: vh3.String = None,
    ya_yt_proxy: vh3.String = None,
    ya_yt_dir: vh3.String = None,
    ya_yt_token_vault_owner: vh3.String = None,
    ya_yt_token_vault_name: vh3.String = None,
    result_rt: vh3.String = None,
    timestamp: vh3.Date = None,
    build_system: vh3.Enum[Literal["ya", "ya_force", "semi_distbuild", "distbuild"]] = "ya"
) -> BuildArcadiaProjectOutput:
    """
    Build Arcadia Project

    Launches YA_MAKE task in Sandbox for provided target and downloads requested artifact.

    :param targets: Target
      [[Multiple targets with ";" are not allowed]]
    :param arts: Build artifact
      [[Multiple artifacts with ";" and custom destination directory with "=" are not allowed]]
    :param arcadia_url: Svn url for arcadia
      [[Should not contain revision]]
    :param arcadia_revision: Arcadia Revision
    :param checkout_arcadia_from_url: Full SVN url for arcadia (Overwrites base URL and revision, use @revision to fix revision)
    :param build_type: Build type
    :param arts_source: Source artifacts (semicolon separated pairs path[=destdir])
      [[Какие файлы из Аркадии поместить в отдельный ресурс (формат тот же, что и у build artifacts)]]
    :param result_single_file: Result is a single file
    :param definition_flags: Definition flags
      [[For example "-Dkey1=val1 ... -DkeyN=valN"]]
    :param sandbox_oauth_token: Sandbox OAuth token
      [[To run task on behalf of specific user]]
      Name of the secret containing oauth token of user the sandbox task should be launched from
    :param arcadia_patch: Apply patch
      [[Diff file rbtorrent, paste.y-t.ru link or plain text. Doc: https://nda.ya.ru/3QTTV4]]
    :param owner: Custom sandbox task owner (should be used only with OAuth token)
      [[OAuth token owner should be a member of sandbox group]]
    :param use_aapi_fuse: Use arcadia-api fuse
    :param use_arc_instead_of_aapi: Use arc fuse instead of aapi
    :param aapi_fallback: Fallback to svn/hg if AAPI services are temporary unavailable
    :param kill_timeout: Kill Timeout (seconds)
    :param sandbox_requirements_disk: Disk requirements in Mb
    :param sandbox_requirements_ram: RAM requirements in Mb
    :param sandbox_requirements_platform: Platform
    :param checkout: Run ya make with --checkout
    :param clear_build: Clear build
    :param strip_binaries: Strip result binaries
    :param lto: Build with LTO
    :param thinlto: Build with ThinLTO
    :param musl: Build with musl-libc
    :param use_system_python: Use system Python to build python libraries
    :param target_platform_flags: Target platform flags (only for cross-compilation)
    :param javac_options: Javac options (semicolon separated)
    :param ya_yt_proxy: YT store proxy
    :param ya_yt_dir: YT store cypress path
    :param ya_yt_token_vault_owner: YT token vault owner
    :param ya_yt_token_vault_name: YT token vault name
    :param result_rt: Result resource type
    :param timestamp: Timestamp
    :param build_system: Build System
    """
    raise NotImplementedError("Write your local execution stub here")


class Python3DeepLearningOutput(NamedTuple):
    data: vh3.Binary
    state: vh3.Binary
    logs: vh3.Binary
    json_output: vh3.JSON


@vh3.decorator.external_operation("https://nirvana.yandex-team.ru/operation/4c2e2df0-4dff-4cf8-b3ce-1f89bd62c9e8")
@vh3.decorator.nirvana_names(
    base_layer="base_layer",
    environment="Environment",
    install_pydl_package="install_pydl_package",
    run_command="run_command",
    openmpi_runner="openmpi_runner",
    openmpi_params="openmpi_params",
    mpi_processes_count="mpi_processes_count",
    before_run_command="before_run_command",
    after_run_command="after_run_command",
    ssh_key="ssh_key",
    auto_snapshot="auto_snapshot",
    nodes_count="nodes_count",
    additional_layers="additional_layers",
    user_requested_secret="user_requested_secret",
)
@vh3.decorator.nirvana_names_transformer(vh3.name_transformers.snake_to_dash, options=True)
def python_3_deep_learning(
    *,
    base_layer: vh3.Enum[
        Literal["PYDL_V3", "PYDL_V4", "PYDL_V5", "PYDL_V5_GPU", "PYDL_V5_GPU_NVIDIA", "NONE"]
    ] = "PYDL_V5_GPU_NVIDIA",
    environment: vh3.MultipleStrings = (),
    pip: vh3.MultipleStrings = (),
    install_pydl_package: vh3.Boolean = True,
    run_command: vh3.String = "python3 $SOURCE_CODE_PATH/__main__.py",
    cpu_cores_usage: vh3.Integer = 1600,
    gpu_count: vh3.Integer = 0,
    openmpi_runner: vh3.Boolean = False,
    openmpi_params: vh3.String = None,
    mpi_processes_count: vh3.Integer = -1,
    before_run_command: vh3.String = None,
    after_run_command: vh3.String = None,
    ttl: vh3.Integer = 1440,
    max_ram: vh3.Integer = 10000,
    max_disk: vh3.Integer = 10000,
    force_tmpfs_disk: vh3.Boolean = False,
    gpu_max_ram: vh3.Integer = 1024,
    gpu_type: vh3.Enum[
        Literal["NONE", "ANY", "CUDA_ANY", "CUDA_3_5", "CUDA_5_2", "CUDA_6_1", "CUDA_7_0", "CUDA_8_0"]
    ] = "NONE",
    strict_gpu_type: vh3.Boolean = False,
    ssh_key: vh3.Secret = None,
    yt_token: vh3.Secret = None,
    mr_account: vh3.String = None,
    auto_snapshot: vh3.Integer = 0,
    nodes_count: vh3.Integer = 1,
    additional_layers: vh3.MultipleStrings = (),
    retries_on_job_failure: vh3.Integer = 0,
    timestamp: vh3.Date = "2017-08-25T16:40:26+0300",
    debug_timeout: vh3.Integer = 0,
    user_requested_secret: vh3.Secret = None,
    pre_hook_enable: vh3.Boolean = True,
    post_hook_enable: vh3.Boolean = True,
    pre_hook_commands_json: vh3.MultipleStrings = (),
    post_hook_commands_json: vh3.MultipleStrings = (),
    job_host_tags: vh3.MultipleStrings = (),
    job_scheduler_instance: vh3.String = None,
    job_scheduler_yt_pool_tree: vh3.String = None,
    job_scheduler_yt_pool: vh3.String = None,
    job_scheduler_yt_token: vh3.Secret = None,
    job_mtn_enabled: vh3.Boolean = True,
    job_scheduler_yt_custom_spec: vh3.String = "{}",
    job_layer_yt_path: vh3.MultipleStrings = (),
    job_variables: vh3.MultipleStrings = (
        "LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/lib64/stubs:/usr/local/cuda/extras/CUPTI/lib64",
    ),
    script: vh3.Binary = None,
    data: Sequence[vh3.Binary] = (),
    state: vh3.Binary = None,
    volume: Sequence[vh3.Binary] = (),
    params: vh3.JSON = None
) -> Python3DeepLearningOutput:
    """
    Python 3 Deep Learning

    Python code runner with support of tensorflow, numpy, theano, torch, keras and nirvana_dl library. See https://wiki.yandex-team.ru/computervision/projects/deeplearning/nirvanadl/ for user manual.

    :param base_layer: Base porto layer
    :param pip: Libraries to Install
      [[List of libraries to be installed using pip install. Some libraries are already available,
      see https://wiki.yandex-team.ru/computervision/projects/deeplearning/nirvanadl/ for details. WARNING: use only with explicit version for reproducibility.]]
    :param install_pydl_package: Install nirvana-dl package
    :param run_command: Run command
      [[Custom bash code]]
    :param openmpi_runner: use openmpi runner
    :param openmpi_params: openmpi runner extra args
    :param before_run_command: Before run command
      [[Command which will be executed before run_command]]
    :param after_run_command: After run command
      [[Command which will be executed after run command]]
    :param strict_gpu_type: strict GPU type
    :param ssh_key: SSH Key
      [[Secret with ssh private key to sync logs with remote server]]
    :param yt_token: YT Token
      [[ID of Nirvana Secret with YT access token (https://nda.ya.ru/3RSzVU). Guide to Nirvana Secrets: https://nda.ya.ru/3RSzWZ]]
    :param mr_account: MR Account:
      [[MR Account Name. By default, output tables and directories will be created in some subdirectory of home/<MR Account>/<workflow owner>/nirvana]]
    :param auto_snapshot: Auto Snapshot
      [[Time interval (minutes) to dump snapshots automatically, without explicit python method call. May cause race condition. If 0, option will be disabled]]
    :param nodes_count: Full number of nodes:
      [[Number of nodes. Should be >= 1]]
      Number of nodes. Should be >= 1
    :param additional_layers: Additional porto layers
      [[IDs of porto layers to put on top of base layer]]
    :param timestamp: Timestamp
    :param job_host_tags: master-job-host-tags
    :param job_mtn_enabled:
      https://st.yandex-team.ru/NIRVANA-12358#5eda1eeb6213d14744d9d791
    :param script:
      Tar-archive with source code inside. Unpacked archive will be available by SOURCE_CODE_PATH environment variable.
    :param data:
      Tar archive(s) with various data. Unpacked data will be available by INPUT_PATH environment variable.
      Multiple archives will be concatenated.
    :param state:
      Saved state of the process, will be used as "state" output snapshot.
    :param volume:
      Additional porto volume(s).
    :param params:
      JSON with additional parameters that can be used in your program.
    """
    raise NotImplementedError("Write your local execution stub here")


class TrafficSignsDatasetPrepareOutput(NamedTuple):
    train_tfrecord: vh3.Binary
    test_tfrecord: vh3.OptionalOutput[vh3.Binary]
    label_map: vh3.Binary


@vh3.decorator.operation(vh3.mr_run_base, deterministic=True, owner='quoter')
@vh3.decorator.autorelease_to_nirvana_on_trunk_commit(
    version='https://nirvana.yandex-team.ru/alias/operation/maps_mrc_trafic_signs_prepare_dataset/0.0.4',
    script_method='ops.traffic_signs_prepare_dataset',
    nirvana_quota='geo-common',
)
@vh3.decorator.nirvana_names(yt_token="yt-token", max_ram="max-ram", max_disk="max-disk")
def traffic_signs_prepare_dataset(
    *,
    executable: vh3.Executable,
    yt_token: vh3.Secret,
    input_table: vh3.String,
    input_table_cache_date: vh3.String,
    arcadia_revision: vh3.Integer = None,
    arcadia_patch: vh3.String = None,
    input_test_table: vh3.String = None,
    min_size: vh3.Integer = None,
    min_amount: vh3.Integer = None,
    train_part: vh3.Number = None,
    join_speed_limits: vh3.Boolean = False,
    augm_flip_horz: vh3.Boolean = False,
    blur_fliped: vh3.Boolean = False,
    filter_class_table: vh3.String = None,
    rename_class_table: vh3.String = None,
    max_ram: vh3.Integer = 1000,
    max_disk: vh3.Integer = 10000,
) -> TrafficSignsDatasetPrepareOutput:
    """
    TrafficSigns dataset prepare

    из таблицы YT с разметкой дорожных знаков на изображениях готовит два файла в формате tfrecord (тренировочный и тестовый), а также файл соответствия индексов и названий классов

    :param yt_token: YT Token:
      [[ID of Nirvana Secret with YT access token (https://nda.ya.ru/3RSzVU). Guide to Nirvana Secrets: https://nda.ya.ru/3RSzWZ]]
    :param input_table: Input table:
      [[Input YT table name]]
    :param arcadia_revision: Arcadia Revision
    :param arcadia_patch: Apply patch
      [[Diff file rbtorrent, paste.y-t.ru link or plain text. Doc: https://nda.ya.ru/3QTTV4]]
    :param min_size: Min signs size:
      [[Minimal value of size (width or height) of signs]]
    :param min_amount: Min amount:
      [[Minimal amount of insances of class]]
    :param train_part: Train part:
      [[Part of dataset for building train.tfrecord (from 0. to 1.0) other part will use for test.tfrecord]]
    """

    out_train = vh3.runtime.get_output_path("train_tfrecord")
    out_test = vh3.runtime.get_output_path("test_tfrecord")
    label_map = vh3.runtime.get_output_path("label_map")

    cmd = [
        executable,
        "--input", input_table,
        "--outtrain", str(out_train),
        "--outtest", str(out_test),
        "--labelmap", str(label_map),
    ]

    if min_size:
        cmd.extend(["--minsize", str(min_size)])

    if input_test_table:
        cmd.extend(["--input-test", input_test_table])

    if min_amount:
        cmd.extend(["--minamount", str(min_amount)])

    if train_part:
        cmd.extend(["--trainpart", str(train_part)])

    if augm_flip_horz:
        cmd.append("--horzflip")

    if blur_fliped:
        cmd.append("--blurfliped")

    if join_speed_limits:
        cmd.append("--join-speed-limits")

    if filter_class_table:
        cmd.extend(["--filter_class_table", filter_class_table])

    if rename_class_table:
        cmd.extend(["--rename_class_table", rename_class_table])

    subprocess.check_call(cmd)
    return TrafficSignsDatasetPrepareOutput(
        train_tfrecord=out_train,
        test_tfrecord=out_test,
        label_map=label_map)


class TrainFasterRcnnOutput(NamedTuple):
    gdef_tgz: vh3.Binary
    last_ckpt_state: vh3.Binary
    all_ckpt_state: vh3.Binary


@vh3.decorator.external_operation("https://nirvana.yandex-team.ru/operation/ff521b71-941d-46dc-befe-10246730e953")
@vh3.decorator.nirvana_names(gpu_count="gpu-count", max_ram="max-ram", gpu_max_ram="gpu-max-ram")
def train_faster_rcnn(
    *,
    tf_research_url: vh3.String,
    num_epoches: vh3.Integer,
    optimizer: vh3.String,
    lr_init: vh3.Number,
    lr_decay_epoches: vh3.Number,
    lr_decay_factor: vh3.Number,
    train_data: Sequence[vh3.Binary],
    label_map: vh3.Binary,
    config_url: vh3.String = "",
    detection_ckpt_url: vh3.String = "",
    save_after_steps: vh3.Integer = 100000,
    save_every_steps: vh3.Integer = 10000,
    gpu_count: vh3.Integer = 1,
    ttl: vh3.Integer = 48000,
    max_ram: vh3.Integer = 32000,
    gpu_max_ram: vh3.Integer = 10500,
    last_ckpt_state: vh3.Binary = None,
    all_ckpt_state: vh3.Binary = None
) -> TrainFasterRcnnOutput:
    """
    Train FasterRCNN

    :param train_data:
      train data in TFRecord format
    :param label_map:
      file with classes
    :param last_ckpt_state:
      tar file with last checkpoint of model
    :param all_ckpt_state:
      checkpoints for save it in state by restart
    """
    raise NotImplementedError("Write your local execution stub here")


@vh3.decorator.external_operation("https://nirvana.yandex-team.ru/operation/34a37d5e-4544-4f06-80ae-ae37b4f871af")
@vh3.decorator.nirvana_names(gpu_count="gpu-count", max_ram="max-ram")
@vh3.decorator.nirvana_output_names("statistic")
def test_faster_rcnn(
    *,
    tf_research_url: vh3.String,
    score_thr: vh3.Number,
    bbox_cmp_thr: vh3.Number,
    bbox_min_size: vh3.Integer,
    nms_thr: vh3.Number,
    gdef_tgz: vh3.Binary,
    test_data: Sequence[vh3.Binary],
    label_map: vh3.Binary,
    bbox_cmp_func: vh3.String = None,
    gpu_count: vh3.Integer = 1,
    ttl: vh3.Integer = 2000,
    max_ram: vh3.Integer = 32000
) -> vh3.Binary:
    """
    Test FasterRCNN

    :param gdef_tgz:
      tar archive with gdef of trained graphs
    :param test_data:
      test data in TFRecord format
    :param label_map:
      file with classes
    :param bbox_cmp_func:
      функция вычисления величины совпадения двух bbox
      сейчас либо iou либо iom
    """
    raise NotImplementedError("Write your local execution stub here")
