import copy
import pandas as pd

# TODO: Remove pandas and write direct to csv


def download_olap(yt_client, table, filename):

    data = []
    interest_id = {}

    expected_segments = [
        'shanson', 'rock', 'rap', 'folk', 'child', 'bicycle', 'ready-food', 'cat', 'dog',
        'higher-education', 'restaurants', 'football', 'japan-food', 'alcohol', 'autoservice',
        'want-to-buy-car', 'gamers', 'cinema', 'android', 'ios', 'dachniki',
    ]
    interest_id = {item: key for key, item in enumerate(expected_segments)}

    for row in yt_client.read_table(table):
        age, sex, income = row['cat'].split('#')
        counter = row['counters']['count']

        if counter < 200:
            continue

        poligon_id = int(row['geo_id'], 16)
        base_new_rec = {
            'polygon_id': poligon_id,
            'age': age.replace('-', '').replace('_', ''),
            'sex': sex,
            'incomelevel': income,
            'socdem_counts': counter,
            'city_id': row['city_id'],
            'hash_9': poligon_id,
        }
        for key in expected_segments:
            if key in row['counters']:
                new_rec = copy.copy(base_new_rec)
                new_rec['interest_id'] = interest_id[key]
                new_rec['interest_name'] = key.replace('-', '').replace('_', '')
                new_rec['interest_counts'] = row['counters'][key]

                data.append(new_rec)
    order = [
        'polygon_id', 'age', 'sex', 'incomelevel', 'city_id', 'hash_9',
        'socdem_counts', 'interest_id', 'interest_counts', 'interest_name',
    ]
    pd.DataFrame.from_dict(data)[order].to_csv(filename, sep=',', index=False, header=False)


def download_polygons(yt_client, table, filename):
    data = []
    for row in yt_client.read_table(table):
        new_rec = {
            'city': row['city'],
            'city_ru': row['city_ru'],
            'city_id': row['city_id'],
            'country': row['country'],
            'country_ru': row['country_ru'],
            'country_id': row['country_id'],
            'district': row['district'],
            'district_ru': row['district_ru'],
            'district_id': row['district_id'],
            'hash_6': int(row['hash_6'], 16),
            'hash_6_coords': '[' + row['hash_6_coords'] + ']',
            'hash_7': int(row['hash_7'], 16),
            'hash_7_coords': '[' + row['hash_7_coords'] + ']',
            'hash_8': int(row['hash_8'], 16),
            'hash_8_coords': '[' + row['hash_8_coords'] + ']',
            'hash_9': int(row['hash_9'], 16),
            'hash_9_coords': '[' + row['hash_9_coords'] + ']',
            'lat': row['lat'],
            'lon': row['lon'],
        }
        data.append(new_rec)
    order = [
        'lat', 'lon', 'country_id', 'country', 'city_id', 'city', 'district_id', 'district',
        'hash_6', 'hash_6_coords', 'hash_7', 'hash_7_coords', 'hash_8', 'hash_8_coords', 'hash_9', 'hash_9_coords',
        'country_ru', 'city_ru', 'district_ru',
    ]
    pd.DataFrame.from_dict(data)[order].to_csv(filename, sep=',', index=False, header=False)
