# AUTOGENERATED! DO NOT EDIT! File to edit: ../00_core.ipynb.

# %% auto 0
__all__ = ['extract_spd_list', 'get_spd', 'main_spd']

# %% ../00_core.ipynb 3
import pandas as pd
import numpy as np
import httpx
import json
import calendar
import datetime
import lxml
import asyncio
import re
from zipfile import ZipFile
from io import BytesIO
from io import StringIO

# %% ../00_core.ipynb 5
def _date(x):
    ''' Parse the filename to return file date '''
    tmp = re.search(r'(\d{1,2})111(\d{4})(\d{2})', x).groups()
    return f'{tmp[0]}/{tmp[2]}/{tmp[1]}'    

def extract_spd_list(year): 
    ''' Get list of available downloads from the casefiles '''
    return (pd.read_html(f'https://www.emi.ea.govt.nz/Wholesale/Datasets/FinalPricing/CaseFiles/{year}')[0] 
            .assign(file_date = lambda df_: pd.to_datetime(df_.Name.apply(_date), dayfirst=True)+ pd.DateOffset(days=1),
                    file_size = lambda df_: (df_['File size'].str.replace(',', '').str[:-3]).astype(int))
            .pipe(lambda df_: df_.loc[df_.file_size > 1000].drop_duplicates(subset=['file_date']))
            .set_index('file_date')
            .sort_index()
            )

# %% ../00_core.ipynb 7
def _branch(file):
    ''' Get raw spd files and return as dataframe for branches '''
    df = pd.concat((pd.DataFrame(_extract(file, fn, 'SOLUTION,BRANCH')[1:], 
                                    columns = _extract(file, fn, 'SOLUTION,BRANCH')[0])
                        for fn in file.namelist() if '.SPDSOLVED' in fn[-10:]), axis=0)
    cols = ['INTERVAL', 'BRANCHNAME', 'FROM_STATION', 'TO_STATION', 'FROM_MW', 'TO_MW', 'BRANCHLOSSES']
    return (df
        [cols]
        .assign(INTERVAL = lambda df_: pd.to_datetime(df_.INTERVAL.replace({';':':'}, regex=True),  dayfirst = True),
                TO_STATION = lambda df_: df_.TO_STATION.str.strip(),
                FROM_STATION = lambda df_: df_.FROM_STATION.str.strip(),
                FROM_MW = lambda df_: df_.FROM_MW.astype('float32'),
                TO_MW = lambda df_: df_.TO_MW.astype('float32'),
                BRANCHLOSSES = lambda df_: df_.BRANCHLOSSES.astype('float32'),
                ))

def _bus(file):
    ''' Get raw spd files and return as dataframe for bus '''
    df = pd.concat((pd.DataFrame(_extract(file, fn, 'SOLUTION,BUS')[1:], 
                                    columns = _extract(file, fn, 'SOLUTION,BUS')[0])
                        for fn in file.namelist() if '.SPDSOLVED' in fn[-10:]), axis=0)
    cols = ['INTERVAL', 'ID_BUS', 'ID_ST', 'LOAD', 'GENERATION']
    return (df
        [cols]
        .assign(INTERVAL = lambda df_: pd.to_datetime(df_.INTERVAL.replace({';':':'}, regex=True),  dayfirst = True),
                ID_BUS = lambda df_: df_.ID_BUS.str.strip(),
                ID_ST = lambda df_: df_.ID_ST.str.strip(),
                LOAD = lambda df_: df_.LOAD.astype('float32'),
                GENERATION = lambda df_: df_.GENERATION.astype('float32')
                ))

def _extract(file, fn, text):
    ''' parse zip file '''
    lines_to_keep = []
    with file.open(fn) as f:
        for line in f.readlines():
            if text in line.decode('utf8'):
                lines_to_keep.append(line.decode('utf8').replace('\n','').replace('\r','').split(','))
    return lines_to_keep

# %% ../00_core.ipynb 11
async def get_spd(client, url):
    res = await client.get(url, follow_redirects=True)
    return _branch(ZipFile(BytesIO(res.content))), _bus(ZipFile(BytesIO(res.content)))

async def main_spd(spd_df, year):
    print(year)
    async with httpx.AsyncClient(timeout=None) as client:
        tasks = []
        for file_name in spd_df.Name.to_list():
            url = f'https://www.emi.ea.govt.nz/Wholesale/Datasets/FinalPricing/CaseFiles/{year}/{file_name}'
            tasks.append(asyncio.create_task(get_spd(client, url)))
        res = await asyncio.gather(*tasks)
        return pd.concat([x[0] for x in res]), pd.concat([x[1] for x in res])
