# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02-collation.ipynb (unless otherwise specified).

__all__ = ['collate_cleaned_incidents_info', 'get_combined_outages_df', 'check_postcode_is_in_df',
           'get_postcode_location_df', 'get_postcode_url', 'get_postcode_location_api', 'PostCodeLocator',
           'add_outage_longitudes_latitudes', 'construct_outages_gdf', 'generate_map_js', 'clean_outages_df',
           'generate_index_md']

# Cell
import json
import numpy as np
import pandas as pd
import geopandas as gpd

from outage import retrieval

import os
import requests
from warnings import warn
from jinja2 import Template

from ipypb import track

# Cell
def collate_cleaned_incidents_info(dnos=['ukpn', 'ssen', 'wpd', 'sp', 'np', 'enw']):
    cleaned_incidents_info = dict()

    for dno in track(dnos, label='DNOs'):
        raw_dno_incidents_info, cleaned_dno_incidents_info = getattr(retrieval, f'get_{dno}_incidents_info')()
        cleaned_incidents_info[dno] = cleaned_dno_incidents_info
        retrieval.save_json_data(raw_dno_incidents_info, f'{dno}_incidents_info')

    return cleaned_incidents_info

# Cell
def get_combined_outages_df(cleaned_incidents_info):
    df_outages = pd.DataFrame()

    for dno, dno_outages in cleaned_incidents_info.items():
        df_dno_outages = pd.DataFrame(dno_outages).T.query('incident_active==True')

        df_dno_outages['dno'] = dno
        df_dno_outages = df_dno_outages[['received_time', 'estimated_restored_time', 'postcodes_impacted', 'description', 'dno', 'incident_url']]

        df_dno_outages.index.name = 'incident_id'
        df_dno_outages = df_dno_outages.reset_index()

        df_outages = df_outages.append(df_dno_outages)

    df_outages = df_outages.reset_index(drop=True)

    return df_outages

# Cell
check_postcode_is_in_df = lambda postcode, df: postcode in df.index

# Cell
def get_postcode_location_df(df_postcodes, postcode):
    postcode_is_in_df = check_postcode_is_in_df(postcode, df_postcodes)

    if postcode_is_in_df:
        location = df_postcodes.loc[postcode]

        longitude = location['longitude']
        latitude = location['latitude']

    else:
        longitude, latitude = None, None

    return longitude, latitude

# Cell
get_postcode_url = lambda postcode: f'http://api.getthedata.com/postcode/{postcode}'

# Cell
def get_postcode_location_api(postcode):
    longitude, latitude = None, None

    postcode_url = get_postcode_url(postcode)
    r_json = requests.get(postcode_url).json()

    if r_json['status'] == 'match':
        data_keys = r_json['data'].keys()

        if 'longitude' in data_keys and 'latitude' in data_keys:
            longitude = float(r_json['data']['longitude'])
            latitude = float(r_json['data']['latitude'])

    return longitude, latitude

# Cell
class PostCodeLocator:
    def __init__(self, postcodes_fp: str='../data/postcodes.csv'):
        # Loading Post Codes DataFrame
        df_postcodes = pd.read_csv(postcodes_fp, index_col='postcode')
        df_postcodes.index = df_postcodes.index.str.replace(' ', '')

        self.df_postcodes = df_postcodes

    def get_postcode_location(self, postcode):
        postcode_is_in_df = check_postcode_is_in_df(postcode, self.df_postcodes)

        if postcode_is_in_df:
            longitude, latitude = get_postcode_location_df(self.df_postcodes, postcode)

        else:
            longitude, latitude = get_postcode_location_api(postcode)

        if longitude is None or latitude is None:
            warn(f'The longitude and latitude of {postcode} could not be determined')

        return longitude, latitude

# Cell
def add_outage_longitudes_latitudes(df_outages, postcodes_fp: str='../data/postcodes.csv', null_threshold=1):
    pcl = PostCodeLocator(postcodes_fp)

    s_incident_longitudes = pd.Series(index=df_outages.index, dtype='float64')
    s_incident_latitudes = pd.Series(index=df_outages.index, dtype='float64')

    for outage_idx, s_outage in track(df_outages.iterrows(), total=df_outages.shape[0]):
        incident_longitudes = []
        incident_latitudes = []

        for postcode_impacted in s_outage['postcodes_impacted']:
            longitude, latitude = pcl.get_postcode_location(postcode_impacted)

            incident_longitudes += [longitude]
            incident_latitudes += [latitude]

        s_incident_longitudes.loc[outage_idx] = pd.Series(incident_longitudes).mean()
        s_incident_latitudes.loc[outage_idx] = pd.Series(incident_latitudes).mean()

    df_outages = (df_outages
                  .assign(longitude=s_incident_longitudes)
                  .assign(latitude=s_incident_latitudes)
                 )

    idxs_to_keep = (df_outages.isnull().sum(axis=1)<=null_threshold).replace(False, np.nan).dropna().index
    df_outages = df_outages.loc[idxs_to_keep]

    return df_outages

# Cell
construct_outages_gdf = lambda df_outages: gpd.GeoDataFrame(
    data=df_outages.drop(columns=['longitude', 'latitude']),
    geometry=gpd.points_from_xy(df_outages['longitude'], df_outages['latitude']),
    crs='EPSG:4326'
)

# Cell
def generate_map_js(gdf_outages, js_template_fp='templates/map_js.j2', js_docs_fp='docs/map.js'):
    outages_geojson = json.loads(gdf_outages.to_json(drop_id=True))

    rendered_map_js = Template(open(js_template_fp).read()).render(
        zoom=5,
        center=[53.96, -3.22],
        geojson_features=str(outages_geojson).replace('None', 'null')
    )

    with open(js_docs_fp, 'w', encoding='utf8') as fp:
        fp.write(rendered_map_js)

    return

# Cell
def clean_outages_df(df_outages):
    df_outages_clean = pd.DataFrame()

    df_outages_clean['Incident'] = df_outages.agg(lambda x: f"<a href=\"{x['incident_url']}\" target=\"_blank\">{x['incident_id']}</a>", axis=1)
    df_outages_clean['Received Time'] = pd.to_datetime(df_outages['received_time']).dt.strftime('%Y-%m-%d %H:%M').fillna('Unknown')
    df_outages_clean['Estimated Restored Time'] = pd.to_datetime(df_outages['estimated_restored_time']).dt.strftime('%Y-%m-%d %H:%M').fillna('Unknown')
    df_outages_clean['Postcodes Impacted'] = df_outages['postcodes_impacted'].str.join(', ')
    df_outages_clean['Operator'] = df_outages['dno'].str.upper()
    df_outages_clean['Description'] = df_outages['description'].str.replace('\n', '<br>')

    df_outages_clean = df_outages_clean.sort_values('Received Time', ascending=False)

    return df_outages_clean

# Cell
def generate_index_md(df_outages_clean, md_template_fp='templates/index_md.j2', docs_md_fp='docs/index.md'):
    rendered_index_md = Template(open(md_template_fp).read()).render(
        map_width='600px',
        map_height='500px',
        md_table=df_outages_clean.to_markdown(index=False)
    )

    with open(docs_md_fp, 'w', encoding='utf8') as fp:
        fp.write(rendered_index_md)

    return