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

__all__ = ['get_accounts_raw_search', 'extract_search_df', 'get_num_operating_accounts_pages', 'get_full_search_df',
           'get_search_df', 'account_id_to_url', 'retry_request', 'extract_key_table_soups', 'try_convert',
           'extract_compliance_df', 'filter_for_year_indexes', 'extract_single_row_table_info',
           'extract_master_general_info', 'extract_child_general_info', 'extract_contact_info', 'extract_address_info',
           'clean_dict_2nd_level_nulls', 'extract_page_info', 'collate_owner_info', 'collate_unit_info',
           'construct_ets_unit_dfs', 'constuct_da_ts_from_ts_dfs', 'ts_dfs_to_separate_vars',
           'construct_installation_dfs', 'get_installation_dfs', 'construct_aircraft_dfs', 'get_aircraft_dfs',
           'retrieve_all_data']

# Cell
import pandas as pd
import numpy as np
import xarray as xr

import re
import os
import copy
import requests
from bs4 import BeautifulSoup as bs

import matplotlib.pyplot as plt

from ipypb import track
import FEAutils as hlp

# Cell
get_accounts_raw_search = lambda page_num=0: requests.get(f'https://ec.europa.eu/clima/ets/oha.do?form=oha&languageCode=en&accountHolder=&installationIdentifier=&installationName=&permitIdentifier=&mainActivityType=-1&searchType=oha&currentSortSettings=&backList=%3CBack&resultList.currentPageNumber={page_num}')

# Cell
def extract_search_df(r):
    soup = bs(r.text, features='lxml')
    results_table = soup.find('table', attrs={'id': 'tblAccountSearchResult'})

    df_search = (pd
                 .read_html(str(results_table))
                 [0]
                 .iloc[1:, :-2]
                 .reset_index(drop=True)
                 .T
                 .set_index(0)
                 .T
                 .reset_index(drop=True)
                 .rename(columns={
                     'National Administrator': 'country',
                     'Account Type': 'account_type',
                     'Account Holder Name': 'account_holder_name',
                     'Installation/Aircraft ID': 'installation_or_aircraft_id',
                     'Installation Name/Aircraft Operator Code*': 'operator_code',
                     'Company Registration No': 'company_registration_number',
                     'Permit/Plan ID': 'permit_or_plan_id',
                     'Permit/Plan Date': 'permit_or_plan_date',
                     'Main Activity Type': 'main_activity_type',
                     'Latest Compliance Code': 'latest_compliance_code'
                 })
                )

    df_search['account_id'] = [a['href'].split('accountID=')[-1].split('&')[0] for a in results_table.findAll('a', text=re.compile('Details - All Phases'))]

    return df_search

# Cell
def get_num_operating_accounts_pages():
    r = get_accounts_raw_search()
    soup = bs(r.text, features='lxml')

    soup_pn = soup.find('input', attrs={'name': 'resultList.lastPageNumber'})
    num_pages = int(soup_pn['value'])

    return num_pages

# Cell
def get_full_search_df(num_pages):
    df_search = pd.DataFrame()

    for page_num in track(range(num_pages), label='Accounts'):
        r = get_accounts_raw_search(page_num=page_num)
        df_search_page = extract_search_df(r)
        df_search = df_search.append(df_search_page)

    df_search = (df_search
                 .reset_index(drop=True)
                 .drop_duplicates()
                )

    return df_search

def get_search_df(data_dir='data', num_pages=None, redownload=False):
    if num_pages is None:
        num_pages = get_num_operating_accounts_pages()

    if not os.path.exists(data_dir):
        os.makedirs(data_dir)

    if redownload == True:
        df_search = get_full_search_df(num_pages)
        df_search.to_csv(f'{data_dir}/account_search.csv', index=False)
    else:
        df_search = pd.read_csv(f'{data_dir}/account_search.csv')

    return df_search

# Cell
account_id_to_url = lambda account_id: f'https://ec.europa.eu/clima/ets/ohaDetails.do?accountID={account_id}&action=all'

# Cell
def retry_request(root_url, params={}, n_retries=10, **kwargs):
    for i in range(n_retries):
        try:
            r = requests.get(root_url, params=params, **kwargs)
            return r
        except Exception as e:
            err = e
            continue

    raise err

def extract_key_table_soups(account_url):
    r = retry_request(account_url)
    soup = bs(r.text, features='lxml')

    operator_master_table = soup.find('table', attrs={'summary': 'Master account details'})
    operator_child_table, compliance_table = soup.findAll('table', attrs={'summary': 'Child account details'})

    return operator_master_table, operator_child_table, compliance_table

# Cell
def try_convert(value, default, type_):
    try:
        return type_(value)
    except:
        pass

    return default

filter_for_year_indexes = lambda df: df.loc[pd.Series(df.index).apply(try_convert, args=(np.nan, float)).dropna().astype(int).astype(str).values]

def extract_compliance_df(compliance_table):
    df_compliance = (pd
                     .read_html(str(compliance_table))[1].iloc[1:, :-2]
                     .reset_index(drop=True)
                     .T
                     .set_index(0)
                     .T
                     .reset_index(drop=True)
                     .drop(columns=['EU ETS Phase', 'Cumulative Surrendered Units**', 'Cumulative Verified Emissions***'])
                     .set_index('Year')
                     .pipe(filter_for_year_indexes)
                     .rename(columns={
                         'Allowances in Allocation': 'allocated_allowances',
                         'Verified Emissions': 'verified_emissions',
                         'Units Surrendered': 'units_surrendered',
                         'Compliance Code': 'compliance_code'
                     })
                    )

    return df_compliance

# Cell
extract_single_row_table_info = lambda df_info, num_excess_start_cols, last_end_col: df_info.iloc[num_excess_start_cols:, :last_end_col].reset_index(drop=True).T.set_index(0)[1].to_dict()

extract_master_general_info = lambda df_master_general_info: extract_single_row_table_info(df_master_general_info, 1, 6)
extract_child_general_info = lambda df_child_general_info: extract_single_row_table_info(df_child_general_info, 1, 10)
extract_contact_info = lambda df_contact_info: extract_single_row_table_info(df_contact_info, 1, 11)
extract_address_info = lambda df_address_info: extract_single_row_table_info(df_address_info, 1, 8)

# Cell
def clean_dict_2nd_level_nulls(dict_):
    dict_ = {
        k1: {
            k2: (
                v2
                if v2
                not in [np.nan, 'nan', '-']
                else None
               )
            for k2, v2
            in v1.items()
        }
        for k1, v1
        in dict_.items()
    }

    return dict_

def extract_page_info(
    account_id,
    master_general_info_func=extract_master_general_info,
    child_general_info_func=extract_child_general_info,
):
    # Retrieving table html
    account_url = account_id_to_url(account_id)
    operator_master_table, operator_child_table, compliance_table = extract_key_table_soups(account_url)

    # Extracting raw dataframes
    _ , df_master_general_info, _, df_contact_info = pd.read_html(str(operator_master_table))
    _ , df_child_general_info, df_address_info = pd.read_html(str(operator_child_table))

    # Parsing to clean dictionaries
    master_general_info = master_general_info_func(df_master_general_info)
    child_general_info = child_general_info_func(df_child_general_info)
    contact_info = extract_contact_info(df_contact_info)
    address_info = extract_address_info(df_address_info)

    # Collating data
    page_info = {
        'master_general_info': master_general_info,
        'child_general_info': child_general_info,
        'contact_info': contact_info,
        'address_info': address_info
    }

    # Cleaning null values
    page_info = clean_dict_2nd_level_nulls(page_info)

    # Extracting time-series data
    df_ts = extract_compliance_df(compliance_table)

    return page_info, df_ts

# Cell
def collate_owner_info(installation_page_info):
    general_info = installation_page_info['master_general_info']
    contact_info = installation_page_info['contact_info']

    owner_info = copy.deepcopy(general_info)
    owner_info.update(copy.deepcopy(contact_info))

    assert len(general_info) + len(contact_info) - len(owner_info) == 0, 'There are duplicate entries in the dictionary keys'

    return owner_info

def collate_unit_info(installation_page_info):
    general_info = installation_page_info['child_general_info']
    address_info = installation_page_info['address_info']

    installation_info = copy.deepcopy(general_info)
    installation_info.update(copy.deepcopy(address_info))

    assert len(general_info) + len(address_info) - len(installation_info) == 0, 'There are duplicate entries in the dictionary keys'

    return installation_info

# Cell
def construct_ets_unit_dfs(account_ids, owners_col_rename_map, units_col_rename_map, label=None):
    df_owners = pd.DataFrame(index=account_ids, columns=owners_col_rename_map.keys())
    df_units = pd.DataFrame(index=account_ids, columns=units_col_rename_map.keys())
    ts_dfs = {}

    for account_id in track(account_ids, label=label):
        page_info, df_ts = extract_page_info(account_id)

        df_owners.loc[account_id] = pd.Series(collate_owner_info(page_info))
        df_units.loc[account_id] = pd.Series(collate_unit_info(page_info))
        ts_dfs[account_id] = df_ts

    df_owners = df_owners.rename(columns=owners_col_rename_map)
    df_units = df_units.rename(columns=units_col_rename_map)

    return df_owners, df_units, ts_dfs

def constuct_da_ts_from_ts_dfs(ts_dfs):
    arr = np.stack([df.values for df in ts_dfs.values()])

    coords = {
        'account_id': list(ts_dfs.keys()),
        'year': list(ts_dfs.values())[0].index.values,
        'variable': list(ts_dfs.values())[0].columns.values
    }

    da_ts = xr.DataArray(arr, coords=coords, dims=coords.keys())

    return da_ts

def ts_dfs_to_separate_vars(ts_dfs):
    da_ts = constuct_da_ts_from_ts_dfs(ts_dfs)
    ts_var_dfs = {}

    for variable in da_ts['variable'].values:
        ts_var_dfs[variable] = (da_ts
                                .sel(variable=variable)
                                .to_dataframe(name=variable)
                                [variable]
                                .reset_index()
                                .pivot('account_id', 'year', variable)
                               )

    return ts_var_dfs

def construct_installation_dfs(account_ids):
    installation_owners_col_rename_map = {
        'National Administrator': 'national_administrator',
        'Account Type': 'account_type',
        'Account Holder Name': 'account_holder_name',
        'Installation ID': 'installation_id',
        'Company Registration No': 'company_registration_number',
        'Account Status': 'account_status',
        'Type': 'type',
        'Name': 'name',
        'Legal Entity Identifier': 'legal_entity_identifier',
        'Main Address Line': 'first_address_line',
        'Secondary Address Line': 'second_address_line',
        'Postal Code': 'postcode',
        'City': 'city',
        'Country': 'country',
        'Telephone 1': 'telephone_1',
        'Telephone 2': 'telephone_2',
        'E-Mail Address': 'email'
    }

    installations_col_rename_map = {
        'Installation ID': 'installation_id',
        'Installation Name': 'installation_name',
        'Permit ID': 'permit_id',
        'Permit Entry Date': 'permit_entry_date',
        'Permit Expiry/Revocation Date': 'permit_expiration_Date',
        'Name of Subsidiary undertaking': 'subsidiary_undertaking_name',
        'Name of Parent undertaking': 'parent_undertaking_name',
        'E-PRTR identification': 'EPRTR_id',
        'First Year of Emissions': 'initial_emissions_year',
        'Last Year of Emissions': 'final_emissions_year',
        'Main Address Line': 'first_address_line',
        'Secondary Address Line': 'second_address_line',
        'Postal Code': 'postcode',
        'City': 'city',
        'Country': 'country',
        'Latitude': 'lat',
        'Longitude': 'lon',
        'Main Activity': 'main_activity'
    }

    df_owners, df_installations, ts_dfs = construct_ets_unit_dfs(account_ids, installation_owners_col_rename_map, installations_col_rename_map, label='Installations')
    installation_dfs = ts_dfs_to_separate_vars(ts_dfs)

    installation_dfs.update({
        'owners': df_owners,
        'installations': df_installations
    })

    return installation_dfs

def get_installation_dfs(df_search, data_dir='data/installations', redownload=False):
    df_search_installations = df_search.query("account_type=='Operator Holding Account'")
    account_ids = df_search_installations['account_id']

    if not os.path.exists(data_dir):
        os.makedirs(data_dir)

    if redownload == True:
        installation_dfs = construct_installation_dfs(account_ids)

        for filename, df_installation in installation_dfs.items():
            df_installation.to_csv(f'{data_dir}/{filename}.csv')

    else:
        installation_dfs = dict()
        filenames = [f[:-4] for f in os.listdir(data_dir) if '.csv' in f]

        for filename in filenames:
            installation_dfs[filename] = pd.read_csv(f'{data_dir}/{filename}.csv')

    return installation_dfs

# Cell
def construct_aircraft_dfs(account_ids):
    aircraft_owners_col_rename_map = {
        'National Administrator': 'national_administrator',
        'Account Type': 'account_type',
        'Account Holder Name': 'account_holder_name',
        'Aircraft Operator ID': 'aircraft_operator_id',
        'Company Registration No': 'company_registration_number',
        'Account Status': 'account_status',
        'Type': 'type',
        'Name': 'name',
        'Legal Entity Identifier': 'legal_entity_identifier',
        'Main Address Line': 'first_address_line',
        'Secondary Address Line': 'second_address_line',
        'Postal Code': 'postcode',
        'City': 'city',
        'Country': 'country',
        'Telephone 1': 'telephone_1',
        'Telephone 2': 'telephone_2',
        'E-Mail Address': 'email'
    }

    aircraft_col_rename_map = {
        'Aircraft Operator ID': 'aircraft_operator_id',
        'Unique Code under Commission Regulation (EC) No 748/2009': '',
        'Monitoring Plan ID': 'monitoring_plan_id',
        'Monitoring plan — first year of applicability': 'monitoring_plan_start_date',
        'Monitoring plan — year of expiry': 'monitoring_plan_expiration_Date',
        'Name of Subsidiary undertaking': 'subsidiary_undertaking_name',
        'Name of Parent undertaking': 'parent_undertaking_name',
        'E-PRTR identification': 'EPRTR_id',
        'Call Sign (ICAO designator)': 'call_sign',
        'First Year of Emissions': 'initial_emissions_year',
        'Main Address Line': 'first_address_line',
        'Secondary Address Line': 'second_address_line',
        'Postal Code': 'postcode',
        'City': 'city',
        'Country': 'country',
        'Latitude': 'lat',
        'Longitude': 'lon',
        'Main Activity': 'main_activity'
    }

    df_owners, df_aircraft, ts_dfs = construct_ets_unit_dfs(account_ids, aircraft_owners_col_rename_map, aircraft_col_rename_map, label='Aircraft')
    aircraft_dfs = ts_dfs_to_separate_vars(ts_dfs)

    aircraft_dfs.update({
        'owners': df_owners,
        'aircraft': df_aircraft
    })

    return aircraft_dfs

def get_aircraft_dfs(df_search, data_dir='data/aircraft', redownload=False):
    df_search_aircraft = df_search.query("account_type=='Aircraft Operator Account'")
    account_ids = df_search_aircraft['account_id']

    if not os.path.exists(data_dir):
        os.makedirs(data_dir)

    if redownload == True:
        aircraft_dfs = construct_aircraft_dfs(account_ids)

        for filename, df_aircraft in aircraft_dfs.items():
            df_aircraft.to_csv(f'{data_dir}/{filename}.csv')

    else:
        aircraft_dfs = dict()
        filenames = [f[:-4] for f in os.listdir(data_dir) if '.csv' in f]

        for filename in filenames:
            aircraft_dfs[filename] = pd.read_csv(f'{data_dir}/{filename}.csv')

    return aircraft_dfs

# Cell
def retrieve_all_data(data_dir='data', redownload_search=True, redownload_installations=True, redownload_aircraft=True):
    all_dfs = dict()
    all_dfs['account_search'] = get_search_df(data_dir=data_dir, redownload=(redownload_installations|redownload_installations|redownload_aircraft))

    installation_dfs = get_installation_dfs(all_dfs['account_search'], data_dir=f'{data_dir}/installations', redownload=redownload_installations)
    aircraft_dfs = get_aircraft_dfs(all_dfs['account_search'], data_dir=f'{data_dir}/aircraft', redownload=redownload_aircraft)

    all_dfs.update(installation_dfs)
    all_dfs.update(aircraft_dfs)

    return all_dfs