# -*- coding: utf-8 -*-
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point
import json
import time
import tqdm

start_time = time.time()

# same region dictionary
R = {"Алтайский край": "22", "Амурская область": "28", "Архангельская область": "29", "Астраханская область": "30",
     "Белгородская область": "31", "Брянская область": "32", "Владимирская область": "33",
     "Волгоградская область": "34", "Вологодская область": "35", "Воронежская область": "36",
     "Еврейская автономная область": "79", "Забайкальский край": "75", "Ивановская область": "37",
     "Иркутская область": "38", "Кабардино-Балкария": "7", "Калининградская область": "39", "Калужская область": "40",
     "Камчатский край": "41", "Карачаево-Черкесия": "9", "Кемеровская область": "42", "Кировская область": "43",
     "Костромская область": "44", "Краснодарский край": "23", "Красноярский край": "24", "Курганская область": "45",
     "Курская область": "46", "Ленинградская область": "47", "Липецкая область": "48", "Магаданская область": "49",
     "Москва": "77", "Московская область": "50", "Мурманская область": "51", "Ненецкий автономный округ": "83",
     "Нижегородская область": "52", "Новгородская область": "53", "Новосибирская область": "54", "Омская область": "55",
     "Оренбургская область": "56", "Орловская область": "57", "Пензенская область": "58", "Пермский край": "59",
     "Приморский край": "25", "Псковская область": "60", "Адыгея": "1", "Республика Алтай": "4", "Башкортостан": "2",
     "Республика Бурятия": "3", "Дагестан": "5", "Ингушетия": "6", "Калмыкия": "8", "Республика Карелия": "10",
     "Республика Коми": "11", "Автономна Республіка Крим": "90", "Республика Крым": "90", "Марий Эл": "12",
     "Мордовия": "13", "Республика Саха (Якутия)": "14", "Северная Осетия — Алания": "15", "Татарстан": "16",
     "Республика Тыва": "17", "Республика Хакасия": "19", "Ростовская область": "61", "Рязанская область": "62",
     "Самарская область": "63", "Санкт-Петербург": "78", "Саратовская область": "64", "Сахалинская область": "65",
     "Свердловская область": "66", "Севастополь": "91", "Смоленская область": "67", "Ставропольский край": "26",
     "Тамбовская область": "68", "Тверская область": "69", "Томская область": "70", "Тульская область": "71",
     "Тюменская область": "72", "Удмуртия": "18", "Ульяновская область": "73", "Хабаровский край": "27",
     "Ханты-Мансийский автономный округ — Югра": "86", "Челябинская область": "74", "Чечня": "20", "Чувашия": "21",
     "Чукотский автономный округ": "87", "Ямало-Ненецкий автономный округ": "89", "Ярославская область": "76"}

print("Reading Russian geometry")

df = gpd.read_file('russia.geo.json')[['id', 'geometry', 'tags']]
df['reg'] = df['tags'].map(lambda x: int(R.get(x.get('name', 0), 0)))
df = df[['geometry', 'reg']]
df['reg'] = df['reg'].astype(int)
df = df[df.reg != 0]
df = df.to_crs('EPSG:4326')

print("Reading addresses")

ad = pd.read_csv('addresses_adverts.csv', sep="\t")
# ad = pd.read_csv('addresses_adverts.csv', sep="\t", nrows=10000)

print("Creating GeoDataFrame for all points (this may take some time)")

geometry = [Point(xy) for xy in zip(ad.geo_lon, ad.geo_lat)]
pnts = gpd.GeoDataFrame(ad, geometry=geometry, crs='EPSG:4326')

chunk_size = 100000  # adjust this value to your needs
chunks = [pnts[i:i + chunk_size] for i in range(0, pnts.shape[0], chunk_size)]

print("Mapping addresses to regions")

joined_chunks = []

for chunk in tqdm.tqdm(chunks):
    joined_chunk = gpd.sjoin(chunk, df, how='left', predicate='within')
    joined_chunks.append(joined_chunk)

# Concatenate all chunks at once
joined = pd.concat(joined_chunks)

print("Cleaning the result")

# Keep only the first match for each point
joined = joined.sort_index()
joined = joined.loc[~joined.index.duplicated(keep='first')]

joined['reg'] = joined['reg'].fillna(-1).astype(int)  # fill not matched with -1

# Create a new DataFrame with the desired columns
result = pd.DataFrame({
    'id': joined['id'],
    'geo_lat': joined['geo_lat'],
    'geo_lon': joined['geo_lon'],
    'region_id': joined['reg']
})

print("Writing into file")
result.to_json('result_new.json', orient='records', double_precision=13)
result.to_csv('result_new.csv', sep="\t")

print(f"Finished in: {time.time() - start_time} seconds")