from maps.wikimap.stat.assessment.tasks_dataset.lib.tree_dimensions import (
    action_by_to_action_by_tree,
    action_to_action_tree,
    region_id_to_region_name_tree,
    task_id_to_task_name_tree,
)
from maps.wikimap.stat.libs import nile_ut
from nile.api.v1 import Record
from cyson import UInt


def test_should_convert_action_to_action_tree():
    result = nile_ut.yql_run(
        action_to_action_tree,
        stream=nile_ut.Table([
            Record(action='action1', payload_column=91),
            Record(action='action2', payload_column=92),
        ]),
    )
    assert sorted([
        Record(action_tree='\tall\t',          payload_column=91),
        Record(action_tree='\tall\taction1\t', payload_column=91),

        Record(action_tree='\tall\t',          payload_column=92),
        Record(action_tree='\tall\taction2\t', payload_column=92),
    ]) == sorted(result)


def test_should_convert_region_id_to_region_name_tree():
    result = nile_ut.yql_run(
        region_id_to_region_name_tree,
        stream=nile_ut.Table([
            Record(region_id=1, payload_column=91),
            Record(region_id=3, payload_column=93),
        ]),
        major_regions_map=nile_ut.Table([
            Record(region_id="1", region_name='\tЗемля\t',                  test_column=911),
            Record(region_id="1", region_name='\tЗемля\tСтрана1\t',         test_column=912),
            Record(region_id="1", region_name='\tЗемля\tСтрана1\tГород1\t', test_column=913),

            Record(region_id="2", region_tree='\tЗемля\t',                  test_column=921),
            Record(region_id="2", region_tree='\tЗемля\tСтрана2\t',         test_column=922),
            Record(region_id="2", region_tree='\tЗемля\tСтрана2\tГород2\t', test_column=923),
        ]),
    )
    assert sorted([
        Record(region_name_tree='\tЗемля\t',                  payload_column=91),
        Record(region_name_tree='\tЗемля\tСтрана1\t',         payload_column=91),
        Record(region_name_tree='\tЗемля\tСтрана1\tГород1\t', payload_column=91),
        Record(region_name_tree='\tЗемля\t',                  payload_column=93),
    ]) == sorted(result)


def test_should_convert_action_by_to_action_by_tree():
    result = nile_ut.yql_run(
        action_by_to_action_by_tree,
        stream=nile_ut.Table([
            Record(action_by=1, payload_column=91),
            Record(action_by=3, payload_column=92),
        ]),
        puid_map=nile_ut.Table([
            Record(puid=UInt(1), puid_tree=b'\tall\t',                   test_column=991),
            Record(puid=UInt(1), puid_tree=b'\tall\tcategory1\t',        test_column=992),
            Record(puid=UInt(1), puid_tree=b'\tall\tcategory1\tuser1\t', test_column=993),

            Record(puid=UInt(2), puid_tree=b'\tall\t',                   test_column=994),
            Record(puid=UInt(2), puid_tree=b'\tall\tcategory2\t',        test_column=995),
            Record(puid=UInt(2), puid_tree=b'\tall\tcategory2\tuser2\t', test_column=996),
        ]),
    )
    assert sorted([
        Record(action_by_tree='\tall\t',                   payload_column=91),
        Record(action_by_tree='\tall\tcategory1\t',        payload_column=91),
        Record(action_by_tree='\tall\tcategory1\tuser1\t', payload_column=91),

        Record(action_by_tree='\tall\t',                   payload_column=92),
        Record(action_by_tree='\tall\tunknown\t',          payload_column=92),
        Record(action_by_tree='\tall\tunknown\t3\t',       payload_column=92),
    ]) == sorted(result)


def test_should_convert_task_id_to_task_name_tree():
    result = nile_ut.yql_run(
        task_id_to_task_name_tree,
        stream=nile_ut.Table([
            Record(action_date='date1', task_id='task_id1', payload_column=91),
            Record(action_date='date1', task_id='task_id2', payload_column=92),
            Record(action_date='date1', task_id='task_id3', payload_column=93),

            Record(action_date='date2', task_id='task_id1', payload_column=94),
            Record(action_date='date2', task_id='task_id2', payload_column=95),
        ]),
        tariffs_with_dates=nile_ut.Table([
            Record(tariff_date='date1', task_id='task_id1', task_name_tree='\tall\t'),
            Record(tariff_date='date1', task_id='task_id1', task_name_tree='\tall\ttask_category11\t'),
            Record(tariff_date='date1', task_id='task_id1', task_name_tree='\tall\ttask_category11\ttask11\t'),

            Record(tariff_date='date1', task_id='task_id2', task_name_tree='\tall\t'),
            Record(tariff_date='date1', task_id='task_id2', task_name_tree='\tall\ttask_category12\t'),
            Record(tariff_date='date1', task_id='task_id2', task_name_tree='\tall\ttask_category12\ttask12\t'),

            Record(tariff_date='date2', task_id='task_id1', task_name_tree='\tall\t'),
            Record(tariff_date='date2', task_id='task_id1', task_name_tree='\tall\ttask_category21\t'),
            Record(tariff_date='date2', task_id='task_id1', task_name_tree='\tall\ttask_category21\ttask21\t'),

            Record(tariff_date='date2', task_id='task_id3', task_name_tree='\tall\t'),
            Record(tariff_date='date2', task_id='task_id3', task_name_tree='\tall\ttask_category23\t'),
            Record(tariff_date='date2', task_id='task_id3', task_name_tree='\tall\ttask_category23\ttask23\t'),
        ]),
    )
    assert sorted([
        Record(action_date='date1', task_name_tree='\tall\t',                          payload_column=91),
        Record(action_date='date1', task_name_tree='\tall\ttask_category11\t',         payload_column=91),
        Record(action_date='date1', task_name_tree='\tall\ttask_category11\ttask11\t', payload_column=91),

        Record(action_date='date1', task_name_tree='\tall\t',                          payload_column=92),
        Record(action_date='date1', task_name_tree='\tall\ttask_category12\t',         payload_column=92),
        Record(action_date='date1', task_name_tree='\tall\ttask_category12\ttask12\t', payload_column=92),

        Record(action_date='date1', task_name_tree='\tall\t',                          payload_column=93),
        Record(action_date='date1', task_name_tree='\tall\tunknown\t',                 payload_column=93),
        Record(action_date='date1', task_name_tree='\tall\tunknown\ttask_id3\t',       payload_column=93),

        Record(action_date='date2', task_name_tree='\tall\t',                          payload_column=94),
        Record(action_date='date2', task_name_tree='\tall\ttask_category21\t',         payload_column=94),
        Record(action_date='date2', task_name_tree='\tall\ttask_category21\ttask21\t', payload_column=94),

        Record(action_date='date2', task_name_tree='\tall\t',                          payload_column=95),
        Record(action_date='date2', task_name_tree='\tall\tunknown\t',                 payload_column=95),
        Record(action_date='date2', task_name_tree='\tall\tunknown\ttask_id2\t',       payload_column=95),
    ]) == sorted(result)
