"""
Module containing all lark internal_transformer classes
"""
from __future__ import annotations

from collections import defaultdict
import re
from typing import (
    TYPE_CHECKING,
    DefaultDict,
    Dict,
    List,
    Optional,
    Pattern,
    Set,
    Tuple,
    Union,
)

import ibis
from ibis.expr.operations import CrossJoin
from ibis.expr.types import AnyColumn, TableExpr
from lark import Token, Tree, v_args

from sql_to_ibis.exceptions.sql_exception import (
    InvalidQueryException,
    NeedsAggOrGroupQueryException,
    TableExprDoesNotExist,
)
from sql_to_ibis.parsing.transformers import InternalTransformer, TransformerBaseClass
from sql_to_ibis.query_info import OrderByInfo, QueryInfo
from sql_to_ibis.sql.sql_clause_objects import LimitExpression, WhereExpression
from sql_to_ibis.sql.sql_objects import AliasRegistry, AmbiguousColumn
from sql_to_ibis.sql.sql_value_objects import (
    Aggregate,
    Column,
    CountStar,
    DerivedColumn,
    GroupByColumn,
    Literal,
    NestedCrossJoin,
    NestedJoin,
    NestedJoinBase,
    Subquery,
    Table,
    TableOrJoinbase,
    Value,
)

if TYPE_CHECKING:
    from lark.tree import Meta

TableWithColumn = Tuple[Table, AnyColumn]
TableWithColumnCollection = DefaultDict[str, List[TableWithColumn]]

GET_TABLE_REGEX: Pattern[str] = re.compile(
    r"^(?P<table>[a-z_]\w*)\.(?P<column>[a-z_]\w*)$", re.IGNORECASE
)
PANDAS_TYPE_PYTHON_TYPE_FUNCTION = {
    "object": str,
    "string": str,
    "int16": int,
    "int32": int,
    "int64": int,
    "float16": float,
    "float32": float,
    "float64": float,
    "bool": bool,
}

TYPE_TO_PANDAS_TYPE = {
    "varchar": "string",
    "integer": "int32",
    "int": "int32",
    "bigint": "int64",
    "float": "float64",
    "timestamp": "datetime64",
    "datetime64": "datetime64",
    "timedelta[ns]": "timedelta[ns]",
    "category": "category",
}

class TupleTree(Tree[Subquery]):
    def __init__(
        self,
        data: str,
        children: Tuple[Tree, Subquery],
        meta: Optional[Meta] = None,
    ) -> None:
        super().__init__(data, list(children), meta)

for TYPE in PANDAS_TYPE_PYTHON_TYPE_FUNCTION:
    TYPE_TO_PANDAS_TYPE[TYPE] = TYPE

@v_args(inline=True)
class SQLTransformer(TransformerBaseClass):
    """
    Transformer for the lark sql_to_ibis parser
    """

    def __init__(
        self,
        table_name_map: dict,
        table_map: dict,
        column_name_map: dict,
        column_to_table_name: dict,
    ) -> None:
        super().__init__(
            table_name_map,
            table_map,
            column_name_map,
            column_to_table_name,
            _temp_dataframes_dict={},
        )
        self._alias_registry = AliasRegistry()

    def add_column_to_column_to_table_name_map(self, column: str, table: str) -> None:
        """
        Adds a column to the _column_to_table_name map
        :param column:
        :param table:
        :return:
        """
        if self._column_to_table_name.get(column) is None:
            self._column_to_table_name[column] = table
            return
        table_name = self._column_to_table_name[column]
        if isinstance(table_name, AmbiguousColumn):
            table_name.add_table(table)
        else:
            original_table = table_name
            self._column_to_table_name[column] = AmbiguousColumn(
                {original_table, table}
            )

    def table(self, table_name: str, alias: str = "") -> Table:
        """
        Check for existence of table with same name
        If not exists raise TableExprDoesNotExist
        Otherwise return the name of the actual TableExpr
        :return:
        """
        table_name = table_name.lower()
        if table_name not in self._table_name_map:
            raise TableExprDoesNotExist(table_name)
        true_name = self._table_name_map[table_name]
        if isinstance(alias, Tree) and alias.data == "alias_string":
            alias = str(alias.children[0])
        table_value_from_map = self._table_map[true_name]
        assert isinstance(table_value_from_map, Table)
        table = Table(
            value=table_value_from_map.get_table_expr(),
            name=true_name,
            alias=alias,
        )
        if alias:
            self._alias_registry.add_to_registry(alias, table)
        return table

    def order_by_expression(self, rank_tree: Tree[Tree[str]]) -> OrderByInfo:
        """
        Returns the column name for the order sql_object
        :param rank_tree: Tree containing order info
        :return:
        """
        order_type = rank_tree.data
        ascending = order_type == "order_asc"
        order_by_column_tree = rank_tree.children[0]
        order_by_column_name = order_by_column_tree.children[0]
        assert isinstance(order_by_column_name, str)
        return OrderByInfo(order_by_column_name, ascending)

    def integer_(self, integer_token: Token) -> int:
        """
        Returns the integer value
        :param integer_token:
        :return:
        """
        integer_value = int(integer_token.value)
        return integer_value

    def limit_count(self, limit_count_value: int) -> LimitExpression:
        """
        Returns a limit token_or_tree
        :param limit_count_value:
        :return:
        """
        return LimitExpression(limit_count_value)

    def query_expr(
        self, query_info: QueryInfo, *args: List[Union[OrderByInfo, LimitExpression]]
    ) -> QueryInfo:
        """
        Handles the full query, including order and set operations such as union
        :param query_info: Map of all query information
        :param args: Additional arguments aside from query info
        :return: Query info
        """
        for token in args:
            if isinstance(token, OrderByInfo):
                query_info.add_order_by_info(token)
            if isinstance(token, LimitExpression):
                query_info.limit = token.limit
        return query_info

    def _handle_join_subqueries(self, join: NestedJoinBase) -> QueryInfo:
        info = QueryInfo(
            InternalTransformer(
                join.get_tables(),
                join.get_table_map(),
                self._column_name_map,
                self._column_to_table_name,
                self._table_name_map,
                self._alias_registry,
            ),
        )
        info.add_table(join)
        info.add_column(Column(name="*"))
        return info

    def subquery(
        self, query_object: Union[QueryInfo, NestedJoinBase], alias: Tree
    ) -> Subquery:
        """
        Handle subqueries amd return a subquery object
        :param query_object:
        :param alias:
        :return:
        """
        alias_name = str(alias.children[0])
        if isinstance(query_object, NestedJoinBase):
            query_info = self._handle_join_subqueries(query_object)
        else:
            query_info = query_object
        subquery_value = self._to_ibis_table(query_info)
        subquery = Subquery(name=alias_name, value=subquery_value)
        self._table_map[alias_name] = subquery
        self._column_name_map[alias_name] = {}
        for column in subquery.column_names:
            self.add_column_to_column_to_table_name_map(column.lower(), alias_name)
            self._column_name_map[alias_name][column.lower()] = column
        return subquery

    def column_name(self, *names: Optional[str]) -> Tree:
        full_name = ".".join([str(name) for name in names if name is not None])
        return Tree("column_name", [full_name])

    def join(self, join_expression: NestedJoin) -> NestedJoin:
        """
        Handle join tree
        :param join_expression:
        :return:
        """
        return join_expression

    def comparison_type(self, comparison: Tree) -> Tree:
        """
        Return the comparison expression
        :param comparison:
        :return:
        """
        return comparison

    def join_expression(self, *args: Union[NestedJoin, Table, str, Tree]) -> NestedJoin:
        # There will only ever be four args
        # args_with_join: Tuple[Table, Optional[str], Table, Tree] = args
        table1 = args[0]
        assert isinstance(table1, (Table, NestedJoin))
        join_type = args[1] if args[1] is not None else "inner"
        assert isinstance(join_type, str)
        table2 = args[2]
        assert isinstance(table2, Table)
        join_condition = args[3]
        assert isinstance(join_condition, Tree)
        if "outer" in join_type:
            match = re.match(r"(?P<type>.*)\souter", join_type)
            if match:
                join_type = match.group("type")
        if join_type in {"full", "cross"}:
            join_type = "outer"
        return NestedJoin(
            left_table=table1,
            right_table=table2,
            join_type=join_type,
            join_condition=join_condition,
        )

    def select(self, *select_expressions: Tree[NestedJoinBase]) -> QueryInfo:
        """
        Forms the final sequence of methods that will be executed
        :param select_expressions:
        :return:
        """
        tables: List[TableOrJoinbase] = []
        having_expr = None
        where_expr = None
        for select_expression in select_expressions:
            if isinstance(select_expression, Tree):
                if select_expression.data == "from_expression":
                    table_object = select_expression.children[0]
                    if isinstance(table_object, NestedJoinBase):
                        tables += [
                            table_object.right_table,
                            table_object.left_table,
                        ]
                    elif (
                        isinstance(table_object, Tree)
                        and table_object.data == "cross_join_expression"
                    ):
                        cross_join = table_object.children[0]
                        assert isinstance(cross_join, NestedCrossJoin)
                        tables += [
                            cross_join.right_table,
                            cross_join.left_table,
                        ]
                    elif isinstance(table_object, Table):
                        tables.append(table_object)
                    else:
                        raise Exception(f"Invalid type found: {type(table_object)}")
                elif select_expression.data == "having_expr":
                    having_expr = select_expression
                elif select_expression.data == "where_expr":
                    where_expr = select_expression

        internal_transformer = InternalTransformer(
            tables,
            self._table_map,
            self._column_name_map,
            self._column_to_table_name,
            self._table_name_map,
            self._alias_registry,
        )

        select_expressions_no_boolean_clauses: List[Union[str, Tree]] = []
        distinct = False
        for select_expression in select_expressions:
            if isinstance(select_expression, Tree) and select_expression.data not in (
                "having_expr",
                "where_expr",
            ):
                select_expressions_no_boolean_clauses.append(select_expression)
            if (
                isinstance(select_expression, Token)
                and select_expression.value.lower() == "distinct"
            ):
                distinct = True

        return QueryInfo(
            internal_transformer,
            select_expressions_no_boolean_clauses,
            having_expr=having_expr,
            where_expr=where_expr,
            distinct=distinct,
        )

    def cross_join(self, table1: Table, table2: Table) -> NestedCrossJoin:
        """
        Returns the crossjoin between two dataframes
        :param table1: TableExpr1
        :param table2: TableExpr2
        :return: Crossjoined dataframe
        """
        return NestedCrossJoin(
            left_table=table1,
            right_table=table2,
        )

    def from_item(self, item: Subquery) -> Subquery:
        return item

    def _handle_count_star(
        self, aggregate: Aggregate, relation: TableExpr
    ) -> Aggregate:
        if isinstance(aggregate.value, CountStar):
            aggregate.value = relation.count()
        return aggregate

    def _get_aggregate_ibis_columns(
        self, aggregates: Dict[str, Aggregate], relation: TableExpr
    ) -> List[AnyColumn]:
        aggregate_ibis_columns: List[AnyColumn] = []
        for aggregate_column_name in aggregates:
            aggregate_column = aggregates[aggregate_column_name]
            self._handle_count_star(aggregate_column, relation)
            column = aggregate_column.value.name(aggregate_column_name)
            aggregate_ibis_columns.append(column)
        return aggregate_ibis_columns

    def _handle_having_expressions(
        self,
        having_expr: Optional[Tree],
        internal_transformer: InternalTransformer,
        table: TableExpr,
        aggregates: Dict[str, Aggregate],
        group_column_names: List[str],
    ) -> Optional[AnyColumn]:
        having = None
        if having_expr:
            transform_result = internal_transformer.transform(having_expr.children[0])
            assert isinstance(transform_result, Value)
            having = transform_result.value
        if having is not None and not aggregates:
            for column in table.columns:
                if column not in group_column_names:
                    raise NeedsAggOrGroupQueryException(column)
        return having

    def handle_aggregation(
        self,
        aggregates: Dict[str, Aggregate],
        group_columns: List[GroupByColumn],
        table: TableExpr,
        having_expr: Optional[Tree],
        internal_transformer: InternalTransformer,
        selected_columns: List[Value],
    ) -> TableExpr:
        """
        Handles all aggregation operations when translating from dictionary info
        to dataframe
        """
        selected_column_names = {
            column.get_name().lower() for column in selected_columns
        }
        aggregate_ibis_columns = self._get_aggregate_ibis_columns(aggregates, table)
        having = self._handle_having_expressions(
            having_expr,
            internal_transformer,
            table,
            aggregates,
            [group_column.get_name() for group_column in group_columns],
        )

        if group_columns and not selected_column_names:
            for group_column in group_columns:
                group_column.set_ibis_name_to_name()

        if group_columns and having is not None and not aggregates:
            raise NotImplementedError(
                "Group by, having, without aggregation not yet implemented"
            )
        if group_columns and not aggregates:
            for column in [
                selected_column.get_name() for selected_column in selected_columns
            ]:
                if column not in group_columns:
                    raise NeedsAggOrGroupQueryException(column)
            table = table.distinct()
        elif aggregates and not group_columns:
            table = table.aggregate(aggregate_ibis_columns, having=having)
        elif aggregates and group_columns:
            table = table.group_by(
                [group_column.value for group_column in group_columns]
            )
            if having is not None:
                table = table.having(having)
            table = table.aggregate(aggregate_ibis_columns)

        non_selected_columns = []
        if group_columns and aggregates:
            for group_column in group_columns:
                if group_column.get_name().lower() not in selected_column_names:
                    non_selected_columns.append(group_column.group_by_name)
            if non_selected_columns:
                table = table.drop(non_selected_columns)

        return table

    def handle_filtering(
        self,
        ibis_table: TableExpr,
        where_expr: Optional[Tree],
        internal_transformer: InternalTransformer,
    ) -> TableExpr:
        """
        Returns frame with appropriately selected and named columns
        :param ibis_table: Ibis expression table to manipulate
        :param where_expr: Syntax tree containing where clause
        :param internal_transformer: Transformer to transform the where clauses
        :return: Filtered TableExpr
        """
        if where_expr is not None:
            where_expression = internal_transformer.transform(where_expr)
            assert isinstance(where_expression, WhereExpression)
            return ibis_table.filter(where_expression.value.get_value())
        return ibis_table

    def subquery_in(self, column: Tree, subquery: Subquery) -> TupleTree:
        return TupleTree("subquery_in", (column, subquery))

    def handle_selection(
        self, ibis_table: TableExpr, columns: List[Value]
    ) -> TableExpr:
        column_mutation = []
        for column in columns:
            if column.get_name() == "*":
                return ibis_table
            column_value = column.get_value().name(column.get_name())
            column_mutation.append(column_value)
        if column_mutation:
            return ibis_table.projection(column_mutation)
        return ibis_table

    @staticmethod
    def _columns_have_select_star(columns: List[Value]) -> bool:
        for column in columns:
            if column.get_name() == "*":
                return True
        return False

    def _get_all_columns_rename_duplicates(
        self, tables: List[Table]
    ) -> List[AnyColumn]:
        columns = {table: table.get_ibis_columns() for table in tables}

        def set_dict_column_name(table: Table, col_name: str) -> None:
            index = table.column_names.index(col_name)
            columns[table][index] = columns[table][index].name(
                f"{table.get_alias_else_name()}.{col_name}"
            )

        for i, table1 in enumerate(tables):
            for table2 in tables[i + 1 :]:
                table1_column_names = set(table1.column_names)
                table2_column_names = set(table2.column_names)
                intersect = table1_column_names.intersection(table2_column_names)
                for column_name in intersect:
                    for table in [table1, table2]:
                        set_dict_column_name(table, column_name)

        all_columns = []
        for table in columns:
            all_columns += columns[table]
        return all_columns

    @staticmethod
    def _get_overlapping_column_names(
        columns_by_name: TableWithColumnCollection,
    ) -> Set[AnyColumn]:
        overlapping = set()
        for column_name, column_list in columns_by_name.items():
            if len(column_list) > 1:
                for table, column in column_list:
                    overlapping.add(column.name(table.name + "." + column_name))
        return overlapping

    def handle_join(
        self,
        join: NestedJoinBase,
        columns: List[Value],
        internal_transformer: InternalTransformer,
    ) -> TableExpr:
        """
        Return the table expr resulting from the join
        :param join:
        :param columns: List of all column values
        :return:
        """

        def get_join_table(
            join: NestedJoin, left_table: TableExpr, right_table: TableExpr
        ) -> TableExpr:
            compiled_condition = internal_transformer.transform(join.join_condition)
            assert isinstance(compiled_condition, Value)
            return left_table.join(
                right_table,
                predicates=compiled_condition.get_value(),
                how=join.join_type,
            )

        def get_cross_table(
            join: CrossJoin, left_table: TableExpr, right_table: TableExpr
        ) -> TableExpr:
            return ibis.cross_join(left_table, right_table)

        def rename_or_add_columns(
            table: Table, columns: TableWithColumnCollection
        ) -> None:
            for column in table.column_names:
                columns[column].append(
                    (table, table.get_table_expr().get_column(column))
                )

        def resolve_table(
            table: TableOrJoinbase,
            columns: TableWithColumnCollection,
        ) -> TableExpr:
            if isinstance(table, NestedJoinBase):
                left = resolve_table(table.left_table, columns)
                right = resolve_table(table.right_table, columns)
                join_actions = {
                    NestedJoin: get_join_table,
                    NestedCrossJoin: get_cross_table,
                }
                return join_actions[type(table)](table, left, right)
            rename_or_add_columns(table, columns)
            return table.get_table_expr()

        result: TableExpr = None
        columns_by_name: TableWithColumnCollection = defaultdict(lambda: [])
        left_table = Table(resolve_table(join.left_table, columns_by_name), "left")
        right_table = Table(resolve_table(join.right_table, columns_by_name), "right")

        all_columns: List[Value] = []
        if self._columns_have_select_star(columns):
            all_columns = join.get_all_join_columns_handle_duplicates(
                left_table, right_table
            )

        left_ibis_table = left_table.get_table_expr()
        right_ibis_table = right_table.get_table_expr()
        if isinstance(join, NestedJoin):
            compiled_condition = internal_transformer.transform(join.join_condition)
            assert isinstance(compiled_condition, Value)
            result = left_ibis_table.join(
                right_ibis_table,
                predicates=compiled_condition.get_value(),
                how=join.join_type,
            )
        if isinstance(join, NestedCrossJoin):
            result = ibis.cross_join(left_ibis_table, right_ibis_table)

        def get_renamed_columns(
            column_collection: TableWithColumnCollection,
        ) -> List[AnyColumn]:
            renamed_columns: List[AnyColumn] = []
            for column_name in column_collection:
                if len(column_collection[column_name]) > 1:
                    for table, column in column_collection[column_name]:
                        renamed_columns.append(
                            column.name(f"{table.get_alias_else_name()}.{column_name}")
                        )
                else:
                    original_column = column_collection[column_name][0][1]
                    renamed_columns.append(original_column)
            return renamed_columns

        overlapping = self._get_overlapping_column_names(columns_by_name)
        if overlapping and not all_columns:
            renamed_columns = get_renamed_columns(columns_by_name)
            return result.projection(renamed_columns)

        if all_columns:
            return result[all_columns]
        return result

    def _set_casing_for_groupby_names(self, query_info: QueryInfo) -> None:
        lower_case_to_true_column_name = {
            column.get_name().lower(): column.get_name()
            for column in query_info.columns
        }
        for groupby_column in query_info.group_columns:
            lower_case_group_name = groupby_column.get_name().lower()
            if lower_case_group_name in lower_case_to_true_column_name:
                selection_statement_name = lower_case_to_true_column_name[
                    lower_case_group_name
                ]
                groupby_column.group_by_name = selection_statement_name
                groupby_column.value = groupby_column.value.name(
                    selection_statement_name
                )

    def get_table_value(
        self, table: Union[Table, NestedJoinBase, Subquery]
    ) -> Optional[Union[TableExpr, NestedJoinBase]]:
        if isinstance(table, Table):
            return table.get_table_expr()
        if isinstance(table, NestedJoinBase):
            return table
        return None

    def _get_relation(self, query_info: QueryInfo) -> TableExpr:
        tables = query_info.tables
        if not tables:
            raise Exception("No table specified")
        first_table = self.get_table_value(tables[0])
        if isinstance(first_table, NestedJoinBase):
            first_table = self.handle_join(
                join=first_table,
                columns=query_info.columns,
                internal_transformer=query_info.internal_transformer,
            )
        for table in tables[1:]:
            next_table = self.get_table_value(table)
            first_table = first_table.cross_join(next_table)
        if len(tables) > 1 and self._columns_have_select_star(query_info.columns):
            all_columns = self._get_all_columns_rename_duplicates(
                [table for table in tables if isinstance(table, Table)]
            )
            first_table = first_table[all_columns]
        return first_table

    def _to_ibis_table(self, query_info: QueryInfo) -> TableExpr:
        """
        Returns the dataframe resulting from the SQL query
        :return:
        """
        query_info.perform_transformation()

        relation = self._get_relation(query_info)

        self._set_casing_for_groupby_names(query_info)

        selected_group_columns = []
        group_column_names = {
            group_column.group_by_name for group_column in query_info.group_columns
        }
        if query_info.group_columns and query_info.aggregates:
            selection_column_names = [
                column.get_name() for column in query_info.columns
            ]
            remaining_columns = set(selection_column_names) - group_column_names
            if remaining_columns:
                raise InvalidQueryException(
                    f"Must have grouping for columns in {remaining_columns}"
                )
            selected_group_columns = [
                column
                for column in query_info.columns
                if column.get_name() in set(group_column_names)
            ]

            query_info.columns = [
                column
                for column in query_info.columns
                if column.get_name() in remaining_columns
            ]
        new_table = self.handle_filtering(
            relation,
            query_info.where_expr,
            query_info.internal_transformer,
        )
        new_table = self.handle_selection(new_table, query_info.columns)
        new_table = self.handle_aggregation(
            query_info.aggregates,
            query_info.group_columns,
            new_table,
            query_info.having_expr,
            query_info.internal_transformer,
            selected_group_columns,
        )

        if query_info.distinct:
            new_table = new_table.distinct()

        order_by = query_info.order_by
        if order_by:
            new_table = new_table.sort_by(order_by)

        if query_info.limit is not None:
            new_table = new_table.head(query_info.limit)

        return new_table

    def set_expr(self, query_info: QueryInfo) -> TableExpr:
        """
        Return different sql_object with set relational operations performed
        :param query_info:
        :return:
        """
        return self._to_ibis_table(query_info)

    def union_all(
        self,
        expr1: TableExpr,
        expr2: TableExpr,
    ) -> TableExpr:
        """
        Return union distinct of two TableExpr
        :param expr1: Left TableExpr
        :param expr2: Right TableExpr
        :return:
        """
        return expr1.union(expr2)

    def union_distinct(
        self,
        expr1: TableExpr,
        expr2: TableExpr,
    ) -> TableExpr:
        """
        Return union distinct of two TableExpr
        :param expr1: Left TableExpr
        :param expr2: Right TableExpr
        :return:
        """
        return expr1.union(expr2, distinct=True)

    def intersect_distinct(self, expr1: TableExpr, expr2: TableExpr) -> TableExpr:
        """
        Return distinct intersection of two TableExpr
        :param expr1: Left TableExpr
        :param expr2: Right TableExpr
        :return:
        """
        return expr1.intersect(expr2)

    def except_distinct(self, expr1: TableExpr, expr2: TableExpr) -> TableExpr:
        """
        Return distinct set difference of two TableExpr
        :param expr1: Left TableExpr
        :param expr2: Right TableExpr
        :return:
        """
        return expr1.difference(expr2).distinct()

    def except_all(self, expr1: TableExpr, expr2: TableExpr) -> TableExpr:
        """
        Return set difference of two TableExpr
        :param expr1: Left TableExpr
        :param expr2: Right TableExpr
        :return:
        """
        return expr1.difference(expr2)

    def final(self, table: TableExpr) -> TableExpr:
        """
        Returns the final dataframe
        :param table:
        :return:
        """
        DerivedColumn.reset_expression_count()
        Literal.reset_literal_count()
        return table


"""DOCUMENTDISTANCE"""
import math
FILENAME = "stopwords.txt"
def similarity(dictionary_1, dictionary_2):
    """DISTANCE"""
    list_1 = ''
    for i in dictionary_1:
        for j in i:
            if j not in '!@#$%^&*()_+-=,.?1234567890':
                if j not in "'":
                    list_1 = list_1 + j
    list_2 = ''
    for i in dictionary_2:
        for j in i:
            if j not in '!@#$%^&*()_+-=",.?1234567890':
                if j not in "'":
                    list_2 = list_2 + j

    list_1 = list_1.split()
    list_2 = list_2.split()
    list_3 = list_1 + list_2
    dictionary3 = {}
    for word in list_3:
        if word not in load_stopwords(FILENAME).keys():
            dictionary3[word] = (list_1.count(word), list_2.count(word))
    numerator = 0
    denominator = 0
    sum1 = 0
    sum2 = 0
    for i in dictionary3:
        numerator = numerator + (dict3[i][0] * dict3[i][1])
        sum1 = sum1 + dictionary3[i][0] ** 2
        sum2 = sum2 + dictionary3[i][1] ** 2
    denominator = math.sqrt(sum1) * math.sqrt(sum2)
    similarity1 = (numerator / denominator)
    return similarity1
def load_stopwords(filename):
    """DISTABCE"""
    stopwords = {}
    with open(filename, 'r') as filename:
        for line in filename:
            stopwords[line.strip()] = 0
    return stopwords

def main():
    """DISTANCE"""
    input1 = input().lower()
    input2 = input().lower()

    print(similarity(input1, input2))

if __name__ == '__main__':
    main()


#!/usr/bin/env python3

import re
import sys
import requests

from datetime import datetime
from html.parser import HTMLParser
from typing import List, Optional, Tuple

class WrongPageDateException(Exception):
    pass

def get_href_value(attrs: List[Tuple[str, str]]) -> str:
    href_value = ""
    for attr in attrs:
        if attr[0] == "href":
            href_value = attr[1]
            break
    return href_value

class ParseInfo250(HTMLParser):
    def __init__(self, year: int, month: int, day: Optional[int] = None):
        super().__init__()
        self.titles = []
        self.year = year
        self.month = month
        self.day = day

        self.searching_td = False
        self.searching_movie_link = False
        self.searching_date_link = False
        self.searching_span = False
        self.movie_id = ""
        self.next_date_candidate = ""
        self.next_date = None

    def handle_starttag(self, tag, attrs):
        if tag == "td":
            self.searching_td = True

        date_href_pattern = re.compile(r"/charts/\?(\d{4}/\d{2}/\d{2})")
        movie_href_pattern = re.compile(r"/movie/\?(\w+)")
        href_value = get_href_value(attrs)
        movie_href_match = movie_href_pattern.match(href_value)
        date_href_match = date_href_pattern.match(href_value)
        if tag == "a":
            if movie_href_match:
                self.movie_id = movie_href_match.group(1)
                self.searching_movie_link = True
            if date_href_match:
                self.next_date_candidate = date_href_match.group(1)
                self.searching_date_link = True
        if tag == "span":
            self.searching_span = True

    def handle_endtag(self, tag):
        if tag == "td":
            self.searching_td = False
        if tag == "a":
            self.searching_movie_link = False
            self.searching_date_link = False
        if tag == "span":
            self.searching_span = False

    def handle_data(self, data):
        content = str(data).strip()
        if self.searching_td:
            if self.searching_date_link:
                next_page_pattern = re.compile(r"→")
                if next_page_pattern.match(content):
                    date_parts = [
                        int(part) for part in self.next_date_candidate.split("/")
                    ]
                    self.next_date = datetime(
                        date_parts[0], date_parts[1], date_parts[2]
                    )
            if self.searching_movie_link and self.searching_span:
                title_pattern = re.compile(r"(.*)(\([\d]{4}\))")
                title_match = title_pattern.match(content)
                if title_match:
                    title = title_match.groups()[0].strip()
                    if title:
                        date = f"{str(self.year)}-{str(self.month)}-{str(self.day)}"
                        self.titles.append(
                            [
                                f"tt{self.movie_id}",
                                date,
                                date,
                                "",
                                title,
                                f"https://www.imdb.com/title/{self.movie_id}",
                                "movie",
                                "",
                                "",
                                "",
                                "",
                                "",
                                "",
                                "",
                            ]
                        )

def get_list_page_html(year: int, month: int, day: Optional[int] = None):
    url = "http://top250.info/charts/?"
    url = f"{url}{year}"
    month_string = f"0{month}" if month < 10 else month
    url = f"{url}/{month_string}"
    if day:
        day_string = f"0{day}" if day < 10 else day
        url = f"{url}/{day_string}"

    resp = requests.get(url)
    response = resp.content.decode("utf8")
    clean_html = re.sub(r"[\n\r\t]", "", response)
    return clean_html

try:
    date_pieces = sys.argv[1].split("-")
    year = int(date_pieces[0])
    month = int(date_pieces[1])
    day = int(date_pieces[2])
except:
    print(
        "WARNING: no date or wrong date format provided, should be in the form: 'YYYY-MM-DD'. Will parse from 1996-04-26."
    )
    year = 1996
    month = 4
    day = 26

start_date = datetime(year=year, month=month, day=day)
today = datetime.now()
end_year = today.year
end_month = today.month
end_day = today.day
max_days = (today - start_date).days
date = start_date
date_counter = 0
start_row = 0
while True:
    list_page_html = get_list_page_html(year=date.year, month=date.month, day=date.day)
    parser = ParseInfo250(year=date.year, month=date.month, day=date.day)
    parser.feed(list_page_html)
    titles = parser.titles
    title_count = len(titles)
    if title_count > 0:
        if start_row == 0:
            print(
                "Position,Const,Created,Modified,Description,Title,URL,"
                + "Title Type,IMDb Rating,Runtime (mins),Year,Genres,Num Votes,"
                + "Release Date,Directors"
            )
            start_row += 1
        print(
            *[
                f"{start_row+index}," + ",".join(titles[index])
                for index in range(title_count)
            ],
            sep="\n",
        )
        start_row += title_count
    if not parser.next_date:
        break
    else:
        date = parser.next_date

    date_counter += 1
    if date_counter > max_days:
        break


import os
from netstat import *
from tkinter import *
from tkinter.messagebox import *

temp_netstat_list = []
netstat_object_list = []
pid_set = set()

win_width = 400
win_height = 150
title = "Kill Process"

root = Tk()
root.title(title)

screen_width = root.winfo_screenwidth()
screen_height = root.winfo_screenheight()

x = int((screen_width - win_width) / 2)
y = int((screen_height - win_height) / 2)

root.geometry("%sx%s+%s+%s" % (win_width, win_height, x, y))

root.resizable(False, False)

frame = Frame(root)
frame.place(rely=.5, relx=.5, x=-80, y=-40, width=245, height=100)

Label(frame, text="Port").grid(row=0)

port_var = StringVar()
Entry(frame, textvariable=port_var).grid(row=0, column=1, columnspan=2, padx=10)

def kill_process():
    if port_var.get() == "" or port_var.get() is None:
        showerror("Error", "Please input port number!")
        return

    netstat_list = os.popen("netstat -ano | findstr \"" + port_var.get() + "\"").read().split(" ")

    for i in netstat_list:
        if '\n' in i:
            netstat_object_list.append(netstat(temp_netstat_list[0], temp_netstat_list[1], temp_netstat_list[2],
                                               temp_netstat_list[3], str(i).replace("\n", "")))
            pid_set.add(str(i).replace("\n", ""))
        elif i != '':
            temp_netstat_list.append(i)

    for i in netstat_object_list:
        print(i)

    showinfo("Info", os.popen("taskkill /f /pid " + pid_set.pop()).read())

def reset():
    port_var.set("")

Button(frame, text="Confirm", command=kill_process, width=8).grid(row=1, column=1, pady=8)
Button(frame, text="Reset", command=reset, width=8).grid(row=1, column=2, padx=2, pady=8)

root.mainloop()


#Header
import sys
import matrix
import stats
import matplotlib.pyplot as plt
import numpy as np
#import pyqtgraph-0.10.0/pyqtgraph as pg
#Classes

#Functions
def mean(x):
  sum = 0
  for i in range(len(x)):
    sum += i
  sum = sum/i
  return sum

#Main
def main():
  x = np.linspace(-np.pi, np.pi, 256, endpoint=True)
  C,S = np.cos(x),np.sin(x)
  plt.plot(x,C)
  plt.plot(x,S)

  plt.show()
#init
if __name__ == "__main__":
  main()


from yt_env_setup import (
    YTEnvSetup,
    Restarter,
    SCHEDULERS_SERVICE,
    CONTROLLER_AGENTS_SERVICE,
    NODES_SERVICE,
    MASTERS_SERVICE,
)

from yt_commands import (
    authors, map_reduce, print_debug, update_op_parameters,
    wait, wait_no_assert, wait_breakpoint, release_breakpoint, with_breakpoint, events_on_fs,
    create, ls, get, set, remove, exists, create_user, create_pool, create_pool_tree, abort_transaction, read_table,
    write_table, map, vanilla, run_test_vanilla, suspend_op,
    get_operation, get_operation_cypress_path, PrepareTables, sorted_dicts,
    update_scheduler_config, update_controller_agent_config)

from yt_helpers import get_current_time, parse_yt_time

from yt.common import YtError

import pytest
from flaky import flaky

import pprint
import random
import time
from datetime import timedelta
from io import BytesIO

##################################################################

class TestSchedulerRandomMasterDisconnections(YTEnvSetup):
    NUM_TEST_PARTITIONS = 2
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1
    NUM_CONTROLLER_AGENTS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
            "testing_options": {
                "enable_random_master_disconnection": False,
                "random_master_disconnection_max_backoff": 10000,
                "finish_operation_transition_delay": {
                    "duration": 1000,
                },
            },
            "controller_agent_tracker": {
                "tag_to_alive_controller_agent_thresholds": {
                    "default": {
                        "absolute": 0,
                        "relative": 1.0,
                    }
                },
                "min_agent_count": 0,
            },
        }
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "operation_time_limit_check_period": 100,
            "snapshot_period": 3000,
        }
    }

    OP_COUNT = 8

    def _create_table(self, table):
        create("table", table)
        set(table + "/@replication_factor", 1)

    def _prepare_tables(self):
        self._create_table("//tmp/t_in")
        write_table("//tmp/t_in", {"foo": "bar"})

        for index in range(self.OP_COUNT):
            self._create_table("//tmp/t_out" + str(index))
            self._create_table("//tmp/t_err" + str(index))

    @authors("ignat")
    def test_many_operations(self):
        self._prepare_tables()

        ops = []
        for index in range(self.OP_COUNT):
            op = map(
                track=False,
                command="sleep 1; echo 'AAA' >&2; cat",
                in_="//tmp/t_in",
                out="//tmp/t_out" + str(index),
                spec={
                    "stderr_table_path": "//tmp/t_err" + str(index),
                },
            )
            ops.append(op)

        try:
            update_scheduler_config("testing_options/enable_random_master_disconnection", True)
            for index, op in enumerate(ops):
                try:
                    op.track()
                    assert read_table("//tmp/t_out" + str(index)) == [{"foo": "bar"}]
                except YtError:
                    assert op.get_state() == "failed"
        finally:
            update_scheduler_config("testing_options/enable_random_master_disconnection", False)
            time.sleep(2)

    @authors("ignat")
    @pytest.mark.timeout(150)
    def test_many_operations_hard(self):
        self._prepare_tables()

        ops = []
        for index in range(self.OP_COUNT):
            op = map(
                track=False,
                command="sleep 15; echo 'AAA' >&2; cat",
                in_="//tmp/t_in",
                out="//tmp/t_out" + str(index),
                spec={
                    "stderr_table_path": "//tmp/t_err" + str(index),
                    "testing": {
                        "delay_inside_revive": 2000,
                    },
                },
            )
            ops.append(op)

        try:
            update_scheduler_config("testing_options/enable_random_master_disconnection", True)
            for index, op in enumerate(ops):
                try:
                    op.track()
                    assert read_table("//tmp/t_out" + str(index)) == [{"foo": "bar"}]
                except YtError:
                    assert op.get_state() == "failed"
        finally:
            update_scheduler_config("testing_options/enable_random_master_disconnection", False)
            time.sleep(2)

    @authors("ignat")
    @pytest.mark.timeout(300)
    def test_many_operations_controller_disconnections(self):
        self._prepare_tables()

        ops = []
        for index in range(self.OP_COUNT):
            op = map(
                track=False,
                command="sleep 15; echo 'AAA' >&2; cat",
                in_="//tmp/t_in",
                out="//tmp/t_out" + str(index),
                spec={
                    "stderr_table_path": "//tmp/t_err" + str(index),
                    "testing": {
                        "delay_inside_revive": 2000,
                    },
                },
            )
            ops.append(op)

        ok = False
        for iter in range(100):
            time.sleep(random.randint(5, 15) * 0.5)
            with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
                pass

            completed_count = 0
            for index, op in enumerate(ops):
                assert op.get_state() not in ("aborted", "failed")
                if op.get_state() == "completed":
                    completed_count += 1
            if completed_count == len(ops):
                ok = True
                break
        assert ok

class TestSchedulerRestart(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
        }
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "operation_time_limit_check_period": 100,
            "snapshot_period": 3000,
        }
    }

    @authors("ignat")
    def test_live_preview(self):
        create_user("u")

        data = [{"foo": i} for i in range(3)]

        create("table", "//tmp/t1")
        write_table("//tmp/t1", data)

        create("table", "//tmp/t2")

        op = map(
            wait_for_jobs=True,
            track=False,
            command=with_breakpoint("BREAKPOINT ; cat"),
            in_="//tmp/t1",
            out="//tmp/t2",
            spec={"data_size_per_job": 1},
        )

        jobs = wait_breakpoint(job_count=2)

        async_transaction_id = get(op.get_path() + "/@async_scheduler_transaction_id")
        assert exists(op.get_path() + "/output_0", tx=async_transaction_id)

        release_breakpoint(job_id=jobs[0])
        release_breakpoint(job_id=jobs[1])
        wait(lambda: op.get_job_count("completed") == 2)

        wait(lambda: len(read_table(op.get_path() + "/output_0", tx=async_transaction_id)) == 2)
        live_preview_data = read_table(op.get_path() + "/output_0", tx=async_transaction_id)
        assert all(record in data for record in live_preview_data)

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            abort_transaction(async_transaction_id)

        op.wait_for_fresh_snapshot()

        wait(lambda: op.get_state() == "running")

        new_async_transaction_id = get(op.get_path() + "/@async_scheduler_transaction_id")
        assert new_async_transaction_id != async_transaction_id

        async_transaction_id = new_async_transaction_id
        assert exists(op.get_path() + "/output_0", tx=async_transaction_id)
        live_preview_data = read_table(op.get_path() + "/output_0", tx=async_transaction_id)
        assert all(record in data for record in live_preview_data)

        release_breakpoint()
        op.track()
        assert sorted_dicts(read_table("//tmp/t2")) == sorted_dicts(data)

    @authors("max42", "ignat")
    def test_brief_spec(self):
        create("table", "//tmp/t1")
        write_table("//tmp/t1", [{"foo": 0}])

        create("table", "//tmp/t2")

        op = map(
            wait_for_jobs=True,
            track=False,
            command=with_breakpoint("BREAKPOINT ; cat"),
            in_="//tmp/t1",
            out="//tmp/t2",
            spec={"data_size_per_job": 1},
        )

        wait_breakpoint()

        op.wait_for_fresh_snapshot()

        brief_spec = get(op.get_path() + "/@brief_spec")

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            pass

        release_breakpoint()

        wait(lambda: op.get_state() == "completed")

        assert brief_spec == get(op.get_path() + "/@brief_spec")

    @authors("gritukan", "gepardo")
    def test_description_after_revival(self):
        create("table", "//tmp/t1")
        write_table("//tmp/t1", [{"foo": 0}])

        create("table", "//tmp/t2")

        op = map(
            wait_for_jobs=True,
            track=False,
            command=with_breakpoint("BREAKPOINT ; cat"),
            in_="//tmp/t1",
            out="//tmp/t2",
            spec={"annotations": {"abc": "def"}, "description": {"foo": "bar"}},
        )

        wait_breakpoint()

        op.wait_for_fresh_snapshot()

        annotations_path = op.get_path() + "/@runtime_parameters/annotations"

        required_annotations = {
            "abc": "def",
            "description": {
                "foo": "bar",
            },
        }

        def check_cypress():
            assert exists(annotations_path) and get(annotations_path) == required_annotations

        def check_get_operation():
            result = get_operation(op.id, attributes=["runtime_parameters"])["runtime_parameters"]
            assert result.get("annotations", None) == required_annotations

        wait_no_assert(check_cypress)
        wait_no_assert(check_get_operation)

        update_op_parameters(op.id, parameters={
            "annotations": {
                "abc": "ghi",
            }
        })
        required_annotations["abc"] = "ghi"

        wait_no_assert(check_cypress)
        wait_no_assert(check_get_operation)

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            pass

        wait_no_assert(check_cypress)
        wait_no_assert(check_get_operation)

        release_breakpoint()
        op.track()

    # COMPAT(levysotsky): Remove when commit with this line is on every cluster.
    @authors("levysotsky")
    def test_bad_acl_during_revival(self):
        create("table", "//tmp/t1")
        write_table("//tmp/t1", [{"foo": 0}])

        create("table", "//tmp/t2")

        op = map(
            wait_for_jobs=True,
            track=False,
            command=with_breakpoint("BREAKPOINT ; cat"),
            in_="//tmp/t1",
            out="//tmp/t2",
        )
        wait_breakpoint()
        op.wait_for_fresh_snapshot()

        bad_acl = [
            {
                "permission": ["read", "manage"],
                "subject": ["harold_finch"],
                "action": "allow",
            },
        ]

        acl_path = op.get_path() + "/@spec/acl"

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            set(acl_path, bad_acl)

        release_breakpoint()

        op.track()

    @authors("eshcherbin")
    def test_restart_during_materialization_and_revive_from_snapshot(self):
        update_scheduler_config("operations_update_period", 10000)
        update_controller_agent_config("snapshot_period", 500)

        op = run_test_vanilla(
            "sleep 1",
            job_count=2,
            spec={
                "testing": {"delay_inside_materialize_scheduler": {"duration": 10000, "type": "async"}},
            },
        )
        op.wait_for_fresh_snapshot()

        time.sleep(3.0)

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            pass

        op.wait_for_state("completed")

##################################################################

class TestControllerAgentReconnection(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
            "testing_options": {
                "finish_operation_transition_delay": {
                    "duration": 2000,
                },
            },
        }
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "snapshot_period": 500,
            "operation_time_limit_check_period": 100,
            "operation_build_progress_period": 100,
        }
    }

    def _create_table(self, table):
        create("table", table, attributes={"replication_factor": 1})

    def _wait_for_state(self, op, state):
        wait(lambda: op.get_state() == state)

    @authors("ignat")
    @flaky(max_runs=3)
    def test_connection_time(self):
        def get_connection_time():
            controller_agents = ls("//sys/controller_agents/instances")
            assert len(controller_agents) == 1
            return parse_yt_time(
                get("//sys/controller_agents/instances/{}/@connection_time".format(controller_agents[0]))
            )

        time.sleep(3)

        assert get_current_time() - get_connection_time() > timedelta(seconds=3)

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            pass

        assert get_current_time() - get_connection_time() < timedelta(seconds=3)

    @authors("ignat")
    @pytest.mark.parametrize("wait_transition_state", [False, True])
    def test_abort_operation_without_controller_agent(self, wait_transition_state):
        def get_agent_states():
            return [agent_info["state"]
                    for agent_info in get("//sys/scheduler/orchid/scheduler/controller_agents").values()]

        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", {"foo": "bar"})

        for iter in range(2):
            op = map(
                command="sleep 1000",
                in_=["//tmp/t_in"],
                out="//tmp/t_out",
                track=False,
            )

            self._wait_for_state(op, "running")

            with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
                wait(lambda: all(state == "unregistered" for state in get_agent_states()),
                     iter=30, sleep_backoff=1.0)
                if wait_transition_state:
                    self._wait_for_state(op, "waiting_for_agent")
                op.abort()

            self._wait_for_state(op, "aborted")

    @authors("ignat")
    def test_complete_operation_without_controller_agent(self):
        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", {"foo": "bar"})

        op = map(command="sleep 1000", in_=["//tmp/t_in"], out="//tmp/t_out", track=False)
        self._wait_for_state(op, "running")

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            with pytest.raises(YtError):
                op.complete()

        self._wait_for_state(op, "running")
        op.complete()
        self._wait_for_state(op, "completed")

    @authors("ignat")
    def test_complete_operation_on_controller_agent_connection(self):
        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", {"foo": "bar"})

        op = map(
            command="sleep 1000",
            in_=["//tmp/t_in"],
            out="//tmp/t_out",
            spec={
                "testing": {
                    "delay_inside_revive": 10000,
                }
            },
            track=False,
        )
        self._wait_for_state(op, "running")

        snapshot_path = op.get_path() + "/snapshot"
        wait(lambda: exists(snapshot_path))

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            pass

        with pytest.raises(YtError):
            op.complete()

        self._wait_for_state(op, "running")
        op.complete()

        self._wait_for_state(op, "completed")

    @authors("alexkolodezny")
    @flaky(max_runs=3)
    def test_complete_map_reduce_operation_on_controller_agent_connection(self):
        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", {"foo": "bar"})

        op = map_reduce(
            map_command="sleep 1000",
            reduce_command="cat",
            in_="//tmp/t_in",
            out="//tmp/t_out",
            sort_by=["foo"],
            spec={
                "testing": {
                    "delay_inside_revive": 10000,
                }
            },
            track=False,
        )
        self._wait_for_state(op, "running")

        snapshot_path = op.get_path() + "/snapshot"
        wait(lambda: exists(snapshot_path))

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            pass

        with pytest.raises(YtError):
            op.complete()

        self._wait_for_state(op, "running")
        op.complete()

        self._wait_for_state(op, "completed")

    @authors("ignat")
    def test_abort_operation_on_controller_agent_connection(self):
        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", {"foo": "bar"})

        op = map(
            command="sleep 1000",
            in_=["//tmp/t_in"],
            out="//tmp/t_out",
            spec={
                "testing": {
                    "delay_inside_revive": 10000,
                }
            },
            track=False,
        )
        self._wait_for_state(op, "running")

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            pass

        op.abort()
        self._wait_for_state(op, "aborted")

@authors("levysotsky")
class TestControllerAgentZombieOrchids(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "zombie_operation_orchids": {
                "clean_period": 15 * 1000,
            },
        }
    }

    def _create_table(self, table):
        create("table", table, attributes={"replication_factor": 1})

    def _get_operation_orchid_path(self, op):
        controller_agent = get(op.get_path() + "/@controller_agent_address")
        return "//sys/controller_agents/instances/{}/orchid/controller_agent/operations/{}".format(
            controller_agent, op.id
        )

    def test_zombie_operation_orchids(self):
        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", {"foo": "bar"})

        op = map(command="cat", in_=["//tmp/t_in"], out="//tmp/t_out")

        orchid_path = self._get_operation_orchid_path(op)
        wait(lambda: exists(orchid_path))
        assert get(orchid_path + "/state") == "completed"
        wait(lambda: not exists(orchid_path))

    def test_retained_finished_jobs(self):
        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", [{"foo": "bar1"}, {"foo": "bar2"}])

        op = map(
            command='if [[ "$YT_JOB_INDEX" == "0" ]] ; then exit 1; fi; cat',
            in_=["//tmp/t_in"],
            out="//tmp/t_out",
            spec={
                "data_size_per_job": 1,
            },
            fail_fast=False
        )

        orchid_path = self._get_operation_orchid_path(op)
        wait(lambda: exists(orchid_path))
        retained_finished_jobs = get(orchid_path + "/retained_finished_jobs")
        assert len(retained_finished_jobs) == 1
        ((job_id, attributes),) = list(retained_finished_jobs.items())
        assert attributes["job_type"] == "map"
        assert attributes["state"] == "failed"

##################################################################

class TestRaceBetweenShardAndStrategy(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 2  # snapshot upload replication factor is 2; unable to configure
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "fair_share_update_period": 100,
        }
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "snapshot_period": 500,
            "operation_time_limit_check_period": 100,
            "operation_build_progress_period": 100,
        }
    }

    @authors("renadeen")
    def test_race_between_shard_and_strategy(self):
        # Scenario:
        # 1. operation is running
        # 2. controller fails, scheduler disables operation and it disappears from tree snapshot
        # 3. job completed event arrives to node shard but is not processed until job revival
        # 4. controller returns, scheduler revives operation
        # 5. node shard revives job and sets JobsReady=true
        # 6. scheduler should enable operation but a bit delayed
        # 7. node shard manages to process job completed event cause job is revived
        #    but operation still is not present at tree snapshot (due to not being enabled)
        #    and strategy discards job completed event telling node to remove it forever
        # 8. job resource usage is stuck in scheduler and next job won't be scheduled ever

        op = run_test_vanilla(
            with_breakpoint("BREAKPOINT"),
            job_count=2,
            spec={
                "testing": {"delay_inside_materialize_scheduler": {"duration": 1000, "type": "async"}},
                "resource_limits": {"user_slots": 1},
            },
        )
        wait_breakpoint()
        op.wait_for_fresh_snapshot()

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            release_breakpoint()

        wait(lambda: op.get_state() == "completed")

##################################################################

class OperationReviveBase(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
            "testing_options": {
                "finish_operation_transition_delay": {
                    "duration": 2000,
                },
            },
        }
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "snapshot_period": 500,
            "operation_time_limit_check_period": 100,
            "operation_build_progress_period": 100,
        }
    }

    def _create_table(self, table):
        create("table", table, attributes={"replication_factor": 1})

    def _prepare_tables(self):
        self._create_table("//tmp/t_in")
        write_table("//tmp/t_in", {"foo": "bar"})

        self._create_table("//tmp/t_out")

    def _wait_for_state(self, op, state):
        wait(lambda: op.get_state() == state)

    @authors("ignat")
    def test_missing_transactions(self):
        self._prepare_tables()

        op = self._start_op(with_breakpoint("echo '{foo=bar}'; BREAKPOINT"), track=False)

        for iter in range(5):
            self._wait_for_state(op, "running")
            with Restarter(self.Env, SCHEDULERS_SERVICE):
                set(op.get_path() + "/@input_transaction_id", "0-0-0-0")
            time.sleep(1)

        release_breakpoint()
        op.track()

        assert op.get_state() == "completed"

    # NB: we hope that we check aborting state before operation comes to aborted state but
    # we cannot guarantee that this happen.
    @authors("ignat")
    @flaky(max_runs=3)
    def test_aborting(self):
        self._prepare_tables()

        op = self._start_op("echo '{foo=bar}'; sleep 10", track=False)

        self._wait_for_state(op, "running")

        op.abort(ignore_result=True)

        self._wait_for_state(op, "aborting")

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            assert op.get_state() == "aborting"

        with pytest.raises(YtError):
            op.track()

        assert op.get_state() == "aborted"

    # NB: we hope that complete finish first phase before we kill scheduler. But we cannot guarantee that this happen.
    @authors("ignat")
    @flaky(max_runs=3)
    @pytest.mark.parametrize("use_tx_action", [False, True])  # COMPAT(kvk1920)
    def test_completing(self, use_tx_action):
        update_controller_agent_config(
            "set_committed_attribute_via_transaction_action",
            use_tx_action)

        self._prepare_tables()

        op = self._start_op("echo '{foo=bar}'; sleep 15", track=False)

        self._wait_for_state(op, "running")

        op.complete(ignore_result=True)

        self._wait_for_state(op, "completing")

        wait(lambda: exists(op.get_path() + "/@committed"))

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            assert op.get_state() == "completing"

        op.track()

        assert op.get_state() == "completed"

        if self.OP_TYPE == "map":
            assert read_table("//tmp/t_out") == []

    @authors("kvk1920")
    def test_operation_committed_attribute(self):
        self._prepare_tables()

        op = self._start_op("echo '{foo=bar}'; sleep 15", track=False)

        self._wait_for_state(op, "running")

        op.complete(ignore_result=True)

        self._wait_for_state(op, "completing")

        wait(lambda: exists(op.get_path() + "/@committed"))

        with Restarter(self.Env, MASTERS_SERVICE):
            pass

        assert get(op.get_path() + "/@committed")

    # NB: test rely on timings and can flap if we hang at some point.
    @authors("ignat")
    @flaky(max_runs=3)
    @pytest.mark.parametrize("stage", ["stage" + str(index) for index in range(1, 8)])
    def test_completing_with_sleep(self, stage):
        self._create_table("//tmp/t_in")
        write_table("//tmp/t_in", [{"foo": "bar"}] * 2)

        self._create_table("//tmp/t_out")

        op = self._start_op(
            "echo '{foo=bar}'; " + events_on_fs().execute_once("sleep 100"),
            track=False,
            spec={
                "testing": {
                    "delay_inside_operation_commit": 60000,
                    "no_delay_on_second_entrance_to_commit": True,
                    "delay_inside_operation_commit_stage": stage,
                },
                "job_count": 2,
            },
        )

        self._wait_for_state(op, "running")

        wait(lambda: op.get_job_count("completed") == 1 and op.get_job_count("running") == 1)

        op.wait_for_fresh_snapshot()

        # This request will be retried with the new incarnation of the scheduler.
        op.complete(ignore_result=True)

        self._wait_for_state(op, "completing")
        wait(lambda: get(op.get_path() + "/controller_orchid/testing/commit_sleep_started"))

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            assert op.get_state() == "completing"
            set(op.get_path() + "/@testing", {"commit_sleep_started": True})

        # complete_operation retry may come when operation is in reviving state.
        # In this case we should complete operation again.
        wait(lambda: op.get_state() in ("running", "completed"))

        if op.get_state() == "running":
            op.complete()

        op.track()

        events = get(op.get_path() + "/@events")

        events_prefix = [
            "starting",
            "waiting_for_agent",
            "initializing",
            "preparing",
            "pending",
            "materializing",
            "running",
            "completing",
            "orphaned",
        ]
        if stage <= "stage5":
            expected_events = events_prefix + [
                "waiting_for_agent",
                "revive_initializing",
                "reviving",
                "pending",
                "reviving_jobs",
                "running",
                "completing",
                "completed",
            ]
        else:
            expected_events = events_prefix + ["completed"]

        actual_events = [event["state"] for event in events]

        print_debug("Expected: ", expected_events)
        print_debug("Actual:   ", actual_events)
        assert expected_events == actual_events

        assert op.get_state() == "completed"

        if self.OP_TYPE == "map":
            assert read_table("//tmp/t_out") == [{"foo": "bar"}]

    @authors("ignat")
    def test_abort_during_complete(self):
        self._create_table("//tmp/t_in")
        write_table("//tmp/t_in", [{"foo": "bar"}] * 2)

        remove("//tmp/t_out", force=True)
        self._create_table("//tmp/t_out")

        op = self._start_op(
            "echo '{foo=bar}'; " + events_on_fs().execute_once("sleep 100"),
            track=False,
            spec={
                "testing": {
                    "delay_inside_operation_commit": 4000,
                    "delay_inside_operation_commit_stage": "stage4",
                },
                "job_count": 2,
            },
        )

        self._wait_for_state(op, "running")

        op.wait_for_fresh_snapshot()

        op.complete(ignore_result=True)

        self._wait_for_state(op, "completing")

        # Wait to perform complete before sleep.
        time.sleep(2)

        op.abort()
        op.track()

        assert op.get_state() == "completed"

    @authors("ignat")
    @flaky(max_runs=3)
    def test_failing(self):
        self._prepare_tables()

        op = self._start_op("exit 1", track=False, spec={"max_failed_job_count": 1})

        self._wait_for_state(op, "failing")

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            assert op.get_state() == "failing"

        with pytest.raises(YtError):
            op.track()

        assert op.get_state() == "failed"

    @authors("ignat")
    def test_revive_failed_jobs(self):
        self._create_table("//tmp/t_in")
        self._create_table("//tmp/t_out")
        write_table("//tmp/t_in", {"foo": "bar"})

        op = self._start_op("sleep 1; false", spec={"max_failed_job_count": 10000}, track=False)

        self._wait_for_state(op, "running")

        def failed_jobs_exist():
            assert op.get_job_count("failed") >= 3

        wait_no_assert(failed_jobs_exist)

        suspend_op(op.id)

        op.wait_for_fresh_snapshot()

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            pass

        wait_no_assert(failed_jobs_exist)

class TestSchedulerReviveForMap(OperationReviveBase):
    OP_TYPE = "map"

    def _start_op(self, command, **kwargs):
        return map(command=command, in_=["//tmp/t_in"], out="//tmp/t_out", **kwargs)

class TestSchedulerReviveForVanilla(OperationReviveBase):
    OP_TYPE = "vanilla"

    def _start_op(self, command, **kwargs):
        spec = kwargs.pop("spec", {})
        job_count = spec.pop("job_count", 1)
        spec["tasks"] = {"main": {"command": command, "job_count": job_count}}
        return vanilla(spec=spec, **kwargs)

################################################################################

class TestJobRevivalBase(YTEnvSetup):
    def _wait_for_single_job(self, op_id):
        path = get_operation_cypress_path(op_id) + "/controller_orchid"
        for i in range(500):
            time.sleep(0.1)
            if get(path + "/state", default=None) != "running":
                continue

            jobs = None
            try:
                jobs = ls(path + "/running_jobs")
            except YtError as err:
                if err.is_resolve_error():
                    continue
                raise

            if len(jobs) > 0:
                assert len(jobs) == 1
                return jobs[0]

        assert False, "Wait failed"

    def _kill_and_start(self, components):
        with Restarter(self.Env, components):
            pass

################################################################################

class TestJobRevival(TestJobRevivalBase):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
            "operations_update_period": 100,
            "static_orchid_cache_update_period": 100,
        },
        "cluster_connection": {
            "transaction_manager": {
                "default_transaction_timeout": 3000,
                "default_ping_period": 200,
            }
        },
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "operation_time_limit_check_period": 100,
            "snapshot_period": 500,
            "operations_update_period": 100,
        }
    }

    DELTA_NODE_CONFIG = {
        "exec_node": {
            "job_controller": {
                "resource_limits": {"user_slots": 5, "cpu": 5},
            }
        }
    }

    @authors("max42")
    @pytest.mark.parametrize(
        "components_to_kill",
        [["schedulers"], ["controller_agents"], ["schedulers", "controller_agents"]],
    )
    def test_job_revival_simple(self, components_to_kill):
        create("table", "//tmp/t_in")
        create("table", "//tmp/t_out")

        write_table("//tmp/t_in", [{"a": 0}])

        map_cmd = " ; ".join(
            [
                "sleep 2",
                events_on_fs().notify_event_cmd("snapshot_written"),
                events_on_fs().wait_event_cmd("scheduler_reconnected"),
                "echo {a=1}",
            ]
        )
        op = map(track=False, command=map_cmd, in_="//tmp/t_in", out="//tmp/t_out")

        job_id = self._wait_for_single_job(op.id)

        events_on_fs().wait_event("snapshot_written")
        op.wait_for_fresh_snapshot()

        self._kill_and_start(components_to_kill)

        orchid_path = "//sys/scheduler/orchid/scheduler/operations/{0}".format(op.id)

        wait(lambda: exists(orchid_path), "Operation did not re-appear")

        assert self._wait_for_single_job(op.id) == job_id

        events_on_fs().notify_event("scheduler_reconnected")
        op.track()

        assert get("{0}/@progress/jobs/aborted/total".format(op.get_path())) == 0
        assert read_table("//tmp/t_out") == [{"a": 1}]

    @authors("max42")
    @pytest.mark.skipif("True", reason="YT-8635")
    @pytest.mark.timeout(600)
    def test_many_jobs_and_operations(self):
        create("table", "//tmp/t_in")

        row_count = 20
        op_count = 20

        output_tables = []
        for i in range(op_count):
            output_table = "//tmp/t_out{0:02d}".format(i)
            create("table", output_table)
            output_tables.append(output_table)

        for i in range(row_count):
            write_table("<append=%true>//tmp/t_in", [{"a": i}])

        ops = []

        for i in range(op_count):
            ops.append(
                map(
                    track=False,
                    command="sleep 0.$(($RANDOM)); cat",
                    in_="//tmp/t_in",
                    out=output_tables[i],
                    spec={"data_size_per_job": 1},
                )
            )

        def get_total_job_count(category):
            total_job_count = 0
            operations = get("//sys/operations", verbose=False)
            for key in operations:
                if len(key) != 2:
                    continue
                for op_id in operations[key]:
                    total_job_count += get(
                        get_operation_cypress_path(op_id) + "/@progress/jobs/{}".format(category),
                        default=0,
                        verbose=False,
                    )
            return total_job_count

        # We will switch scheduler when there are 40, 80, 120, ..., 400 completed jobs.

        for switch_job_count in range(40, 400, 40):
            while True:
                completed_job_count = get_total_job_count("completed/total")
                aborted_job_count = get_total_job_count("aborted/total")
                aborted_on_revival_job_count = get_total_job_count("aborted/scheduled/revival_confirmation_timeout")
                print_debug("completed_job_count =", completed_job_count)
                print_debug("aborted_job_count =", aborted_job_count)
                print_debug("aborted_on_revival_job_count =", aborted_on_revival_job_count)
                if completed_job_count >= switch_job_count:
                    if (switch_job_count // 40) % 2 == 0:
                        with Restarter(self.Env, SCHEDULERS_SERVICE):
                            pass
                    else:
                        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
                            pass
                    if switch_job_count % 3 == 0:
                        with Restarter(self.Env, NODES_SERVICE):
                            pass
                    break
                time.sleep(1)

        for op in ops:
            op.track()

        if aborted_job_count != aborted_on_revival_job_count:
            print_debug("There were aborted jobs other than during the revival process:")
            for op in ops:
                output = BytesIO()
                pprint.pprint(dict(get(op.get_path() + "/@progress/jobs/aborted")), stream=output)
                print_debug(output.getvalue())

        for output_table in output_tables:
            assert sorted(read_table(output_table, verbose=False)) == [{"a": i} for i in range(op_count)]

    @authors("max42", "ignat")
    @pytest.mark.parametrize(
        "components_to_kill",
        [["schedulers"], ["controller_agents"], ["schedulers", "controller_agents"]],
    )
    def test_user_slots_limit(self, components_to_kill):
        create("table", "//tmp/t_in")
        create("table", "//tmp/t_out")

        row_count = 20
        for i in range(row_count):
            write_table("<append=%true>//tmp/t_in", [{"a": i}])
        user_slots_limit = 10

        map_cmd = " ; ".join(
            [
                "sleep 2",
                "echo '{a=1};'",
                events_on_fs().notify_event_cmd("ready_for_revival_${YT_JOB_INDEX}"),
                events_on_fs().wait_event_cmd("complete_operation"),
                "echo '{a=2};'",
            ]
        )

        op = map(
            track=False,
            command=map_cmd,
            in_="//tmp/t_in",
            out="//tmp/t_out",
            spec={
                "data_size_per_job": 1,
                "resource_limits": {"user_slots": user_slots_limit},
                "auto_merge": {
                    "mode": "manual",
                    "chunk_count_per_merge_job": 3,
                    "max_intermediate_chunk_count": 100,
                },
            },
        )

        # Comment about '+10' - we need some additional room for jobs that can be non-scheduled aborted.
        wait(
            lambda: sum(
                [events_on_fs().check_event("ready_for_revival_" + str(i)) for i in range(user_slots_limit + 10)]
            )
            == user_slots_limit
        )

        self._kill_and_start(components_to_kill)
        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            pass

        orchid_path = "//sys/scheduler/orchid/scheduler/scheduling_info_per_pool_tree/default/fair_share_info"
        wait(lambda: exists(orchid_path + "/operations/" + op.id))

        for i in range(1000):
            user_slots = None
            user_slots_path = orchid_path + "/operations/{0}/resource_usage/user_slots".format(op.id)
            try:
                user_slots = get(user_slots_path, verbose=False)
            except YtError:
                pass

            for j in range(10):
                try:
                    jobs = get(op.get_path() + "/@progress/jobs", verbose=False)
                    break
                except YtError:
                    time.sleep(0.1)
                    continue
            else:
                assert False
            if i == 300:
                events_on_fs().notify_event("complete_operation")
            running = jobs["running"]
            aborted = jobs["aborted"]["total"]
            assert running <= user_slots_limit or user_slots is None or user_slots <= user_slots_limit
            assert aborted == 0

        op.track()

    @authors("ignat", "ni-stoiko")
    def test_revival_of_pending_operation(self):
        def get_job_nodes(op):
            jobs_path = op.get_path() + "/controller_orchid/running_jobs"
            try:
                jobs = ls(jobs_path)
            except YtError:
                return []

            result = []
            for job_id in jobs:
                try:
                    job_node = get("{0}/{1}".format(jobs_path, job_id))["address"]
                except YtError:
                    continue  # The job has already completed, Orchid is lagging.

                result.append((job_id, job_node))
            return result

        set("//sys/pool_trees/default/@config/nodes_filter", "!other")

        nodes = ls("//sys/cluster_nodes")
        set("//sys/cluster_nodes/" + nodes[0] + "/@user_tags/end", "other")

        create_pool_tree("other", config={"nodes_filter": "other"})
        create_pool("my_pool", pool_tree="other")

        create_pool("my_pool", pool_tree="default")
        set("//sys/pool_trees/default/my_pool/@max_running_operation_count", 2)
        set("//sys/pool_trees/default/my_pool/@max_operation_count", 2)

        create("table", "//tmp/t_in")
        create("table", "//tmp/t_out")

        row_count = 20
        for i in range(row_count):
            write_table("<append=%true>//tmp/t_in", [{"a": i}])

        op = map(
            track=False,
            command=with_breakpoint("BREAKPOINT"),
            in_="//tmp/t_in",
            out="//tmp/t_out",
            spec={
                "pool_trees": ["default", "other"],
                "pool": "my_pool",
                "data_size_per_job": 1,
            },
        )

        job_ids = wait_breakpoint(job_count=15)

        assert len(job_ids) == 15

        # There was often an error when checking for equality due to the race.
        # If it falls when checking for more or equal, I suggest you delete this assertion.
        # (ni-stoiko)
        assert op.get_job_count("running") >= 15

        op.wait_for_fresh_snapshot()

        create("table", "//tmp/t_out_concurrent")
        concurrent_op = map(
            track=False,
            command="sleep 1000",
            in_="//tmp/t_in",
            out="//tmp/t_out_concurrent",
            spec={
                "pool_trees": ["default"],
                "pool": "my_pool",
            },
        )

        wait(lambda: concurrent_op.get_state() == "running")

        set(concurrent_op.get_path() + "/@registration_index", 1)
        set(op.get_path() + "/@registration_index", 2)

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            set("//sys/pool_trees/default/my_pool/@max_running_operation_count", 1)

        wait(lambda: concurrent_op.get_state() == "running")

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            pass

        wait(lambda: concurrent_op.get_state() == "running")

        with Restarter(self.Env, CONTROLLER_AGENTS_SERVICE):
            pass

        wait(lambda: concurrent_op.get_state() == "running")

##################################################################

class TestHasFailedJobsOnRevive(YTEnvSetup):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
            "operations_update_period": 100,
            "static_orchid_cache_update_period": 100,
        },
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "operation_time_limit_check_period": 100,
            "snapshot_period": 500,
            "operations_update_period": 100,
        }
    }

    @authors("omgronny")
    @pytest.mark.parametrize("revive_from_snapshot", [True, False])
    def test_has_failed_jobs_after_revive(self, revive_from_snapshot):
        if not revive_from_snapshot:
            update_controller_agent_config("snapshot_period", 1000000000)

        failed_job_cmd = events_on_fs().wait_event_cmd("second_iteration_started", timeout=timedelta(seconds=1))
        op = vanilla(
            track=False,
            spec={
                "max_failed_job_count": 100,
                "tasks": {
                    "failed_task": {
                        "job_count": 1,
                        "command": failed_job_cmd
                    },
                    "task": {
                        "job_count": 1,
                        "command": "sleep 1000"
                    }
                }
            }
        )

        def get_has_failed_jobs(op):
            return get(op.get_path() + "/@has_failed_jobs", default=False)

        wait(lambda: get_has_failed_jobs(op))
        events_on_fs().notify_event("second_iteration_started")

        if revive_from_snapshot:
            op.wait_for_fresh_snapshot()

        assert get_has_failed_jobs(op)

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            pass

        op.wait_for_state("running")
        wait(lambda: get_has_failed_jobs(op) == revive_from_snapshot)

##################################################################

class TestDisabledJobRevival(TestJobRevivalBase):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
            "lock_transaction_timeout": 3000,
            "operations_update_period": 100,
        },
        "cluster_connection": {
            "transaction_manager": {
                "default_transaction_timeout": 3000,
                "default_ping_period": 200,
            }
        },
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "snapshot_period": 500,
            "operations_update_period": 100,
            "operation_time_limit_check_period": 100,
            "enable_job_revival": False,
        }
    }

    @authors("max42", "ignat")
    @pytest.mark.parametrize(
        "components_to_kill",
        [["schedulers"], ["controller_agents"], ["schedulers", "controller_agents"]],
    )
    def test_disabled_job_revival(self, components_to_kill):
        create("table", "//tmp/t_in")
        create("table", "//tmp/t_out")

        write_table("//tmp/t_in", [{"a": 0}])

        map_cmd = " ; ".join(
            [
                events_on_fs().notify_event_cmd("snapshot_written"),
                events_on_fs().wait_event_cmd("scheduler_reconnected"),
                "echo {a=1}",
            ]
        )
        op = map(track=False, command=map_cmd, in_="//tmp/t_in", out="//tmp/t_out")

        orchid_path = "//sys/scheduler/orchid/scheduler/operations/{0}".format(op.id)

        job_id = self._wait_for_single_job(op.id)

        events_on_fs().wait_event("snapshot_written")
        op.wait_for_fresh_snapshot()

        self._kill_and_start(components_to_kill)

        wait(lambda: exists(orchid_path), "Operation did not re-appear")

        # Here is the difference from the test_job_revival_simple.
        assert self._wait_for_single_job(op.id) != job_id

        events_on_fs().notify_event("scheduler_reconnected")
        op.track()

        assert read_table("//tmp/t_out") == [{"a": 1}]

    @authors("gritukan")
    def test_fail_on_job_restart(self):
        create("table", "//tmp/t_in")
        create("table", "//tmp/t_out")

        write_table("//tmp/t_in", [{"a": 0}])

        map_cmd = " ; ".join(
            [
                events_on_fs().notify_event_cmd("snapshot_written"),
                events_on_fs().wait_event_cmd("scheduler_reconnected"),
                "echo {a=1}",
            ]
        )
        op = map(
            track=False,
            command=map_cmd,
            in_="//tmp/t_in",
            out="//tmp/t_out",
            spec={"fail_on_job_restart": True})

        self._wait_for_single_job(op.id)

        events_on_fs().wait_event("snapshot_written")
        op.wait_for_fresh_snapshot()

        self._kill_and_start("controller_agents")

        with pytest.raises(YtError):
            op.track()

##################################################################

class TestPreserveSlotIndexAfterRevive(YTEnvSetup, PrepareTables):
    NUM_MASTERS = 1
    NUM_NODES = 3
    NUM_SCHEDULERS = 1

    DELTA_SCHEDULER_CONFIG = {
        "scheduler": {
            "connect_retry_backoff_time": 100,
            "fair_share_update_period": 100,
            "profiling_update_period": 100,
            "fair_share_profiling_period": 100,
        }
    }

    DELTA_CONTROLLER_AGENT_CONFIG = {
        "controller_agent": {
            "operation_time_limit_check_period": 100,
        }
    }

    @authors("ignat")
    def test_preserve_slot_index_after_revive(self):
        self._create_table("//tmp/t_in")
        write_table("//tmp/t_in", [{"x": "y"}])

        def get_slot_index(op):
            wait(lambda: op.get_runtime_progress("scheduling_info_per_pool_tree/default/slot_index") is not None)
            return op.get_runtime_progress("scheduling_info_per_pool_tree/default/slot_index")

        for i in range(3):
            self._create_table("//tmp/t_out_" + str(i))

        op1 = map(
            command="sleep 1000; cat",
            in_="//tmp/t_in",
            out="//tmp/t_out_0",
            track=False,
        )
        op2 = map(command="sleep 2; cat", in_="//tmp/t_in", out="//tmp/t_out_1", track=False)
        op3 = map(
            command="sleep 1000; cat",
            in_="//tmp/t_in",
            out="//tmp/t_out_2",
            track=False,
        )

        assert get_slot_index(op1) == 0
        assert get_slot_index(op2) == 1
        assert get_slot_index(op3) == 2

        op2.track()  # this makes slot index 1 available again since operation is completed

        with Restarter(self.Env, SCHEDULERS_SERVICE):
            pass

        time.sleep(2.0)

        assert get_slot_index(op1) == 0
        assert get_slot_index(op3) == 2

        op2 = map(
            command="sleep 1000; cat",
            in_="//tmp/t_in",
            out="//tmp/t_out_1",
            track=False,
        )

        assert get_slot_index(op2) == 1

##################################################################