import configparser

import psycopg2
from psycopg2 import sql

cfg = configparser.ConfigParser()
cfg.read('config.ini')
assert "POSTGRES" in cfg, "missing POSTGRES in config.ini"
param_postgres = cfg["POSTGRES"]


def get_year_columns(cursor):
    columns = []
    query = sql.SQL("SELECT column_name FROM information_schema.columns WHERE table_schema = 'public' AND table_name = 'stations';")
    cursor.execute(query)
    results = cursor.fetchall()
    for result in results:
        try:
            columns.append(int(result[0]))
        except ValueError:
            pass
    return columns


def get_neighbours(cursor, lat, lon, columns):
    values = ''  # Used in second parameter of cursor.execute() (Avoids SQL injection)
    for n in [lat, lon]:
        values = (*values, n)  # adding n to existing tuple

    query = sql.SQL("""
                SELECT array_to_json(array_agg(row_to_json(t))) from (
                    SELECT {columns} FROM stations 
                    ORDER BY ST_Distance(ST_MakePoint(lon, lat), ST_MakePoint({lon}, {lat})) 
                    LIMIT 10
                ) t;
                    """).format(columns=columns, lon=sql.Placeholder(), lat=sql.Placeholder())
    cursor.execute(query, values)
    neighbours = cursor.fetchall()[0][0]
    return neighbours


def calc_averages(neighbours, years):
    averages = {}
    for year in years:
        values = []
        # print(neighbours)
        for neighbour in neighbours:

            # print(neighbour[str(year)])
            if not neighbour[str(year)] == 'NaN': values.append(neighbour[str(year)])
        avg = round(sum(values) / len(values), 3)
        averages[year] = avg
    return averages


def get_average_data_for_point(lat, lon, columns):
    with psycopg2.connect(database=param_postgres["dbName"], user=param_postgres["user"], password=param_postgres["password"], host=param_postgres["host"], port=param_postgres["port"]) as connection:
        with connection.cursor() as cursor:
            print(columns)
            if '*' in str(columns):
                year_columns = get_year_columns(cursor)
            else:
                year_columns = (str(columns).replace("""SQL('""", "").replace('"', '').replace("')", "")).split(',')
                print(year_columns)
            neighbours = get_neighbours(cursor, lat, lon, columns)
            avg_data = calc_averages(neighbours, year_columns)
            print(avg_data)
            return avg_data


# get_average_data_for_point(52.5, 13.4)