import json
import pandas as pd
import os
from os.path import join
import numpy as np
import sys

cwd = os.getcwd()
data_path = join(cwd, '..', 'Data storage')
idx = pd.IndexSlice

"""
Explanation: National generation and fuel consumption
The data in this notebook is generation and consumption by fuel type for the entire US. These values are larger than what would be calculated by summing facility-level data. Note that the fuel types are somewhat aggregated (coal rather than BIT, SUB, LIG, etc). So when we multiply the fuel consumption by an emissions factor there will be some level of error.
The code assumes that you have already downloaded an ELEC.txt file from EIA's bulk download website.
End of explanation
"""


file_date = '2018-03-06'

%load_ext watermark

%watermark -iv -v

# Load the "autoreload" extension
%load_ext autoreload

# always reload modules marked with "%aimport"
%autoreload 1

# add the 'src' directory as one where we can import modules
src_dir = join(os.getcwd(), os.pardir, 'src')
sys.path.append(src_dir)

%aimport Analysis.index
from Analysis.index import add_datetime, add_quarter

"""
Explanation: Date string for filenames
This will be inserted into all filenames (reading and writing)
End of explanation
"""


cwd = os.getcwd()
path = join(data_path, 'Raw EIA bulk', '{} ELEC.txt'.format(file_date))
with open(path, 'r') as f:
    raw_txt = f.readlines()

"""
Explanation: Read ELECT.txt file
End of explanation
"""


def line_to_df(line):
    """
    Takes in a line (dictionary), returns a dataframe
    """
    for key in ['latlon', 'source', 'copyright', 'description', 
                'geoset_id', 'iso3166', 'name', 'state']:
        line.pop(key, None)

    # Split the series_id up to extract information
    # Example: ELEC.PLANT.GEN.388-WAT-ALL.M
    series_id = line['series_id']
    series_id_list = series_id.split('.')
    # Use the second to last item in list rather than third
    plant_fuel_mover = series_id_list[-2].split('-')
    line['type'] = plant_fuel_mover[0]
#     line['state'] = plant_fuel_mover[1]
    line['sector'] = plant_fuel_mover[2]
    temp_df = pd.DataFrame(line)

    try:
        temp_df['year'] = temp_df.apply(lambda x: x['data'][0][:4], axis=1).astype(int)
        temp_df['month'] = temp_df.apply(lambda x: x['data'][0][-2:], axis=1).astype(int)
        temp_df['value'] = temp_df.apply(lambda x: x['data'][1], axis=1)
        temp_df.drop('data', axis=1, inplace=True)
        return temp_df
    except:
        exception_list.append(line)
        pass

states = ["AL", "AK", "AZ", "AR", "CA", "CO", "CT", "DE",
          "FL", "GA", "HI", "ID", "IL", "IN", "IA", "KS",
          "KY", "LA", "ME", "MD", "MA", "MI", "MN", "MS",
          "MO", "MT", "NE", "NV", "NH", "NJ", "NM", "NY",
          "NC", "ND", "OH", "OK", "OR", "PA", "RI", "SC",
          "SD", "TN", "TX", "UT", "VT", "VA", "WA", "WV", "WI", "WY"]

state_geos = ['USA-{}'.format(state) for state in states]

type(json.loads(raw_txt[0]))

json.loads(raw_txt[0])['geography']

gen_rows = [row for row in raw_txt if 'ELEC.GEN' in row 
            and 'series_id' in row 
            and '-99.M' in row 
            and 'ALL' not in row]

total_fuel_rows = [row for row in raw_txt if 'ELEC.CONS_TOT_BTU' in row 
                   and 'series_id' in row 
                   and '-99.M' in row 
                   and 'ALL' not in row
                   and 'US-99.m' not in row]

eg_fuel_rows = [row for row in raw_txt if 'ELEC.CONS_EG_BTU' in row 
                and 'series_id' in row 
                and '-99.M' in row 
                and 'ALL' not in row
                and 'US-99.m' not in row]

"""
Explanation: Filter lines to only include total generation and fuel use
Only want monthly US data for all sectors
- US-99.M
- ELEC.GEN, ELEC.CONS_TOT_BTU, ELEC.CONS_EG_BTU
- not ALL
Fuel codes:
- WWW, wood and wood derived fuels
- WND, wind
- STH, solar thermal
- WAS, other biomass
- TSN, all solar
- SUN, utility-scale solar
- NUC, nuclear
- NG, natural gas
- PEL, petroleum liquids
- SPV, utility-scale solar photovoltaic
- PC, petroluem coke
- OTH, other
- COW, coal,
- DPV, distributed photovoltaic
- OOG, other gases
- HPS, hydro pumped storage
- HYC, conventional hydroelectric
- GEO, geothermal
- AOR, other renewables (total)
End of explanation
"""


gen_dicts = [json.loads(row) for row in gen_rows]

gen_df = pd.concat([line_to_df(x) for x in gen_dicts
                    if x['geography'] in state_geos])

#drop
gen_df.head()

gen_df['geography'].unique()

"""
Explanation: All generation by fuel
End of explanation
"""


gen_df.loc[:,'value'] *= 1000
gen_df.loc[:,'units'] = 'megawatthours'

gen_df.rename_axis({'value':'generation (MWh)'}, axis=1, inplace=True)

gen_df.loc[gen_df.isnull().any(axis=1)]

gen_df.dropna(inplace=True)

gen_df.set_index(['type', 'year', 'month', 'geography'], inplace=True)

"""
Explanation: Multiply generation values by 1000 and change the units to MWh
End of explanation
"""


total_fuel_dict = [json.loads(row) for row in total_fuel_rows]

total_fuel_df = pd.concat([line_to_df(x) for x in total_fuel_dict
                           if x['geography'] in state_geos])

"""
Explanation: Total fuel consumption
End of explanation
"""


total_fuel_df.loc[:,'value'] *= 1E6
total_fuel_df.loc[:,'units'] = 'mmbtu'

total_fuel_df.rename_axis({'value':'total fuel (mmbtu)'}, axis=1, inplace=True)

total_fuel_df.set_index(['type', 'year', 'month', 'geography'], inplace=True)

"""
Explanation: Multiply generation values by 1,000,000 and change the units to MMBtu
End of explanation
"""


total_fuel_df.loc[total_fuel_df.isnull().any(axis=1)]

total_fuel_df = total_fuel_df.dropna()

"""
Explanation: Drop nans
End of explanation
"""


eg_fuel_dict = [json.loads(row) for row in eg_fuel_rows]

eg_fuel_df = pd.concat([line_to_df(x) for x in eg_fuel_dict
                        if x['geography'] in state_geos])

"""
Explanation: Electric generation fuel consumption
End of explanation
"""


eg_fuel_df.loc[:,'value'] *= 1E6
eg_fuel_df.loc[:,'units'] = 'mmbtu'

eg_fuel_df.rename_axis({'value':'elec fuel (mmbtu)'}, axis=1, inplace=True)

eg_fuel_df.set_index(['type', 'year', 'month', 'geography'], inplace=True)

#drop
eg_fuel_df.head()

"""
Explanation: Multiply generation values by 1,000,000 and change the units to MMBtu
End of explanation
"""


#drop
eg_fuel_df.loc[~(eg_fuel_df['elec fuel (mmbtu)'] >= 0) &
                  ~(eg_fuel_df['elec fuel (mmbtu)'].isnull())]

eg_fuel_df.dropna(inplace=True)

"""
Explanation: I verified on EIA's website that the values below are correct.
End of explanation
"""


fuel_df = pd.concat([total_fuel_df, eg_fuel_df['elec fuel (mmbtu)']], axis=1)

"""
Explanation: Combine three datasets
Need to estimate fuel use for OOG, because EIA doesn't include any (this is only ~2% of OOG fuel for electricity in 2015).
End of explanation
"""


#drop
fuel_df.loc[~(fuel_df['elec fuel (mmbtu)']>=0)]

#drop
fuel_df.loc[~(fuel_df['total fuel (mmbtu)']>=0)]

"""
Explanation: Not sure how this happens in EIA's data, but we do see the negative fuel consumption for electricity generation.
End of explanation
"""


path = join(data_path, 'Final emission factors.csv')
ef = pd.read_csv(path, index_col=0)

"""
Explanation: Add CO<sub>2</sub> emissions
The difficulty here is that EIA combines all types of coal fuel consumption together in the bulk download and API. Fortunately the emission factors for different coal types aren't too far off on an energy basis (BIT is 93.3 kg/mmbtu, SUB is 97.2 kg/mmbtu). I'm going to average the BIT and SUB factors rather than trying to do something more complicated. In 2015 BIT represented 45% of coal energy for electricity and SUB represented 48%.
Same issue with petroleum liquids. Using the average of DFO and RFO, which were the two largest share of petroleum liquids.
End of explanation
"""


fuel_factors = pd.Series({'NG' : ef.loc['NG', 'Fossil Factor'],
                   'PEL': ef.loc[['DFO', 'RFO'], 'Fossil Factor'].mean(),
                   'PC' : ef.loc['PC', 'Fossil Factor'], 
                   'COW' : ef.loc[['BIT', 'SUB'], 'Fossil Factor'].mean(),
                   'OOG' : ef.loc['OG', 'Fossil Factor']}, name='type')

#drop
fuel_factors

fuel_df['all fuel CO2 (kg)'] = (fuel_df['total fuel (mmbtu)']
                                .multiply(fuel_factors, level='type',
                                          fill_value=0))
fuel_df['elec fuel CO2 (kg)'] = (fuel_df['elec fuel (mmbtu)']
                                .multiply(fuel_factors, level='type',
                                          fill_value=0))

fuel_df.head()

fuel_cols = ['total fuel (mmbtu)', 'elec fuel (mmbtu)',
              'all fuel CO2 (kg)', 'elec fuel CO2 (kg)']
gen_fuel_df = pd.concat([gen_df, fuel_df[fuel_cols]], axis=1)

"""
Explanation: Match general types with specific fuel codes
Fuel codes:
- WWW, wood and wood derived fuels
- WND, wind
- STH, solar thermal
- WAS, other biomass
- TSN, all solar
- SUN, utility-scale solar
- NUC, nuclear
- NG, natural gas
- PEL, petroleum liquids
- SPV, utility-scale solar photovoltaic
- PC, petroluem coke
- OTH, other
- COW, coal,
- DPV, distributed photovoltaic
- OOG, other gases
- HPS, hydro pumped storage
- HYC, conventional hydroelectric
- GEO, geothermal
- AOR, other renewables (total)
End of explanation
"""


add_quarter(gen_fuel_df)

gen_fuel_df.head()

"""
Explanation: Add datetime and quarter columns
End of explanation
"""


gen_fuel_df['generation (MWh)'].fillna(value=0, inplace=True)

gen_fuel_df.loc['COW',:].head()

"""
Explanation: No records with positive fuel use but no generation
End of explanation
"""


path = join(data_path, 'EIA state-level gen fuel CO2 2018-03-06.csv')
gen_fuel_df.to_csv(path)

"""
Explanation: Export data
State-level
End of explanation
"""


nat_gen_fuel = gen_fuel_df.groupby(['type', 'year', 'month']).sum()
add_quarter(nat_gen_fuel)

nat_gen_fuel.tail()

path = join(data_path, 'Derived data',
            'EIA country-wide gen fuel CO2 {}.csv'.format(file_date))
gen_fuel_df.to_csv(path)

"""
Explanation: National totals
End of explanation
"""


