import geojson
import numpy as np
import pandas as pd
import sqlalchemy
from django.conf import settings
from django.contrib.gis.geos import GEOSGeometry
from geojson import MultiPolygon
from tqdm import tqdm

from service import models


def raschet_real(df, koefs):
    koef_summ = sum(koefs.values())
    df['rate'] = df[list(koefs.keys())].apply(lambda x: sum([x[k] * koefs[k] for k in list(koefs.keys())]), axis=1)
    df['rate'] = df['rate'] / koef_summ
    return df


def raschet(tables, filters, koefs, method):
    conn = sqlalchemy.create_engine(settings.DB_URL, connect_args={'options': '-csearch_path=public'})
    msk_ao = filters.get('msk_ao')
    msk_rayon = filters.get('msk_rayon')
    rate_from = filters.get('rate_from', None)
    rate_to = filters.get('rate_to', None)
    category = filters.get('category', [])
    categories = ','.join([f"'{c}'" for c in category])
    points_df = None
    nets_df = None
    for table in tables:
        if 'point' in table:
            if msk_ao is not None:
                query = f"select * from {table} where msk_ao={msk_ao} and category in ({categories});" if len(
                    categories) > 0 else f"select * from {table} where msk_ao={msk_ao};"
            elif msk_rayon is not None:
                query = f"select * from {table} where msk_rayon={msk_rayon} and category in ({categories});" if len(
                    categories) > 0 else f"select * from {table} where msk_rayon={msk_rayon};"
            else:
                query = f"select * from {table} where category in ({categories});" if len(
                    categories) > 0 else f"select * from {table};"
            points = pd.read_sql(query, conn)
            if method == 'rate':
                points_df = raschet_real(points, koefs)
                if rate_from is not None and rate_to is not None:
                    points_df = points_df[(points_df['rate'] >= rate_from) & (points_df['rate'] <= rate_to)]
            else:
                points_df = points[(points['model'] >= rate_from) & (points['model'] <= rate_to)]

        if 'net' in table:
            if msk_ao is not None:
                query = f"select * from {table} where msk_ao={msk_ao};"
            elif msk_rayon is not None:
                query = f"select * from {table} where msk_rayon={msk_rayon};"
            else:
                query = f"select * from {table};"
            nets = pd.read_sql(query, conn)
            if method == 'rate':
                nets_df = raschet_real(nets, koefs)
                if rate_from is not None and rate_to is not None:
                    nets_df = nets_df[(nets_df['rate'] >= rate_from) & (nets_df['rate'] <= rate_to)]
            else:
                nets_df = nets[(nets['model'] >= rate_from) & (nets['model'] <= rate_to)]

    return points_df, nets_df


# ==============

def load_data(filepath: str):
    models.PlacementPoint.objects.all().delete()
    df = pd.read_csv(filepath)
    df = df.replace(np.nan, None)
    df = df.replace('NaT', None)
    for row in tqdm(df.to_dict('records'), desc='Loading data...'):
        data = {k: row[k] for k in row.keys() if k not in ['id', 'okrug', 'rayon', 'age_month']}
        data['location_id'] = row['id']
        data['okrug'] = models.AO.objects.get(name=row['okrug'])
        data['rayon'] = models.Rayon.objects.get(name=row['rayon'])
        models.PlacementPoint.objects.create(**data)


def load_ao_and_rayons(ao_filepath: str,
                       rayons_filepath: str):
    models.AO.objects.all().delete()
    models.Rayon.objects.all().delete()
    gj = geojson.load(ao_filepath)
    objs = gj['features']
    for obj in tqdm(objs, desc='Loading AOs...'):
        name = obj['properties']['okrug']
        coords = obj['geometry']['coordinates']
        models.AO.objects.create(**{'name': name, 'polygon': GEOSGeometry(str(MultiPolygon(coords)))})
    gj = geojson.load(rayons_filepath)
    objs = gj['features']
    for obj in tqdm(objs, desc='Loading Rayons...'):
        name = obj['properties']['rayon']
        coords = obj['geometry']['coordinates']
        okr = obj['properties']['okrug']
        ao = models.AO.objects.get(name=okr)
        models.Rayon.objects.create(**{'name': name, 'polygon': GEOSGeometry(str(MultiPolygon(coords))), 'AO': ao})
    print('AO and Rayons loaded')