import os
import sys
import requests
import logging
import urllib
import datetime as dt
import json
import csv
import zipfile
import time
import pandas as pd
from pathlib import Path
from datetime import date
from datetime import timedelta
import shutil

# copy ghost standards in current directory for easier import
shutil.copy('/esarchive/scratch/rgrodofz/software/ghost/GHOST/processing_modules/GHOST_standards.py', '/esarchive/scratch/rgrodofz/software/ghost/GHOST/download_scripts/GHOST_standards.py')
from GHOST_standards import standard_parameters

#from convert_json_csv import convert_json_csv


class aqs_epa_api(object):
    """defines and configures API to query the aqms database"""

    def __init__(self, ):
        self.logger = logging.getLogger(__file__)
        self.url_api = "https://aqs.epa.gov/data/api/"
        self.headers = {'content-type': 'application/json', 'accept': 'application/json'}

        self.get_states = 'list/states'
        self.get_counties = 'list/countiesByState'
        self.get_sites = 'list/sitesByCounty'

        self.get_parameters = "list/parametersByClass"
        self.get_data_by_site = 'sampleData/bySite'
        self.get_Dailydata_by_site = 'dailyData/bySite'

        return
    
    def get_parameters_in_list(self, ):
        """build a query to return data by state"""
        query = urllib.parse.urljoin(self.url_api, self.get_parameters)
        params = {'email': 'raphael.grodofzig@bsc.es', 'key': 'bluewolf34', 'pc': 'ALL'}
        response = requests.get(url=query, data='', params=params)
        return response.text

    def states(self, ):
        """build a query to return states in list"""
        query = urllib.parse.urljoin(self.url_api, self.get_states)
        params = {'email': 'raphael.grodofzig@bsc.es', 'key': 'bluewolf34'}
        response = requests.get(url=query, data='', params=params)
        if response.status_code == 200:
            return response.text
        else:
            return print('Response error')
    
    def counties(self, code):
        """build a query to return list of counties by state"""
        query = urllib.parse.urljoin(self.url_api, self.get_counties)
        params = {'email': 'raphael.grodofzig@bsc.es', 'key': 'bluewolf34', 'state': str(code)}
        response = requests.get(url=query, data='', params=params)
        if response.status_code == 200:
            return response.text
        else:
            return 'Response error'

    def sites(self, code, ccode):
        """build a query to return list of sites per county"""
        query = urllib.parse.urljoin(self.url_api, self.get_sites)
        params = {'email': 'raphael.grodofzig@bsc.es', 'key': 'bluewolf34', 'state': str(code), 'county': str(ccode)}
        response = requests.get(url=query, data='', params=params)
        if response.status_code == 200:
            return response.text
        else:
            return 'Response error'

    def request_data(self, n_max_tries, max_time_per_dl, query, params):
        n_tries = 0
        errcode = 999

        while (n_tries < n_max_tries) and (errcode != 200):
            response = requests.get(url=query, data='', params=params, timeout=max_time_per_dl)
                    
            if response.status_code == 200:
                return response.text
                errcode = response.status_code
                
            else:
                # try again
                print('Response error {}, attempt {}'.format(response.status_code, n_tries))
                errcode = response.status_code
                n_tries += 1

        if n_tries == n_max_tries:
            print('Failed downloading {}_{}_{}_{}_{} {} times in {} seconds, error code {}'.format(code, ccode, scode, pcode, bbdate[:4], n_tries, max_time_per_dl, errcode))
            return None


    def get_data_site(self, code, ccode, scode, pcode, bbdate, eedate):
        """build a query to return data by site"""
        query = urllib.parse.urljoin(self.url_api, self.get_data_by_site)
        params = {'email': 'raphael.grodofzig@bsc.es', 'key': 'bluewolf34', 'param': pcode, 'bdate': bbdate, 'edate': eedate, 'state': str(code), 'county': str(ccode), 'site': str(scode)}
        return query, params


    def get_Dailydata_site(self, code, ccode, scode, pcode, bbdate, eedate):
        """build a query to return data by site"""
        query = urllib.parse.urljoin(self.url_api, self.get_Dailydata_by_site)
        params = {'email': 'raphael.grodofzig@bsc.es', 'key': 'bluewolf34', 'param': pcode, 'bdate': bbdate, 'edate': eedate, 'state': str(code), 'county': str(ccode), 'site': str(scode)}
        return query, params

    


def download_data(mode, version, n_max_tries, max_time_per_dl, start, stop, dl_root):

    if mode == 'all':
        bdate = date(2022, 1, 1) #date(1960, 1, 1) # date before record starts
        edate = date.today() + timedelta(days = 365)

        os.makedirs('{}/US_EPA_AQS/original_files/{}/'.format(dl_root, version), exist_ok=True)

    elif mode == 'nrt':
        version = mode
        bdate = start
        edate = stop 

        os.makedirs('{}/US_EPA_AQS/original_files/nrt/'.format(dl_root), exist_ok=True)
        download_location = '{}/US_EPA_AQS/original_files/nrt/'.format(dl_root)

    else:
        print('time mode inapplicable')


    time_resolutions = ['hourly', 'daily']
    years = pd.date_range(bdate, edate, freq='Y').strftime('%Y').tolist()
    print(years)

    components = {}
    for standard_parameter in standard_parameters:
        aqs_code = standard_parameters[standard_parameter]['aqs_code']
        if aqs_code: # check if aqs_code is available
            components[standard_parameter] = standard_parameters[standard_parameter]['aqs_code']


    if mode == 'all': # use US EPA AQS API

        # params_list = json.loads(aqs_epa_api().get_parameters_in_list())['Data']
        # with open('US_EPA_AQS_parameters.json', 'w') as f:
        #     f.write(json.dumps(params_list, indent=4))
        # params_list = pd.DataFrame(params_list)
        # params_list.to_csv('/esarchive/scratch/rgrodofz/software/ghost/GHOST/download_scripts/US_EPA_AQS_parameterlist.csv', index=False)

        # find all states
        states = json.loads(aqs_epa_api().states())['Data']
        for state in states:
            code = state.get('code')
            state_name = state.get('value_represented')
            print('Downloading... '+state_name)

            # find all counties
            counties = json.loads(aqs_epa_api().counties(code))['Data']
            for county in counties:
                ccode = county.get('code')
                county_name = state.get('value_represented')

                # find all sites
                sites = json.loads(aqs_epa_api().sites(code, ccode))['Data']
                for site in sites:
                    scode = site.get('code')
                    site_name = site.get('value_represented')

                    if ((scode is not None) and (site_name is not None)): # maybe too strict ???

                        for time_resolution in time_resolutions:

                            for component in components:
                                pcode = components[component][0] # select first aqs_code if there are multiple

                                for year in years:
                                    # yearly files
                                    bbdate = date(int(year), 1, 1).strftime('%Y%m%d')
                                    eedate = date(int(year), 12, 31).strftime('%Y%m%d')

                                    if time_resolution == 'hourly':
                                        query, params = aqs_epa_api().get_data_site(code, ccode, scode, pcode, bbdate, eedate)
                                    elif time_resolution == 'daily':
                                        query, params = aqs_epa_api().get_Dailydata_site(code, ccode, scode, pcode, bbdate, eedate)
                                    else:
                                        print('Time resolution problem')

                                    data_per_site = aqs_epa_api().request_data(n_max_tries, max_time_per_dl, query, params)
                                    if data_per_site == None:
                                        continue # skip if download failed

                                    # check if data is available for that site
                                    if json.loads(data_per_site)['Header'][0]['status'] == 'Success':
                                        # save files
                                        filename = '{}_{}_{}_{}_{}_{}.json'.format(time_resolution, code, ccode, scode, pcode, year)

                                        # folder per year
                                        os.makedirs('{}/US_EPA_AQS/original_files/{}/{}/{}/'.format(dl_root, version, component, year), exist_ok=True)
                                        download_location = '{}/US_EPA_AQS/original_files/{}/{}/{}/'.format(dl_root, version, component, year)

                                        with open(download_location+filename, 'w') as outfile:
                                            outfile.write(data_per_site)
                                        print("Downloaded {}".format(filename))
                                    else:
                                        print('No data found for {}_{}_{}_{}_{}_{}'.format(time_resolution, code, ccode, scode, pcode, year))
                                        continue


    elif mode == 'nrt': # use AirNow

        base_url = 'https://s3-us-west-1.amazonaws.com//files.airnowtech.org/airnow/{}/{}/HourlyData_{}.dat'

        # create hourly array
        hours = pd.date_range(bdate, edate, freq='H').strftime('%Y%m%d%H').tolist()[:-1]

        # download
        for hour in hours:
            url = base_url.format(hour[:4], hour[:8], hour)
            n_tries = 0
            errcode = 999

            while (n_tries < n_max_tries) and (errcode != 200):
                r = requests.get(url, timeout=max_time_per_dl)
                if r.status_code == 200:
                    with open(download_location+'HourlyData_'+hour+'.dat', 'wb') as file:
                        file.write(r.content)
                    print('Downloaded {}'.format(url))
                    errcode = r.status_code

                elif r.status_code == 404:
                    print("No data found, error 404")
                    errcode = 200

                elif r.status_code == 403:
                    print("Permission denied for {}".format(url))
                    errcode = 200

                else:
                    # try again
                    print('Response error {}, attempt {}'.format(r.status_code, n_tries))
                    errcode = r.status_code
                    n_tries += 1
                    max_time_per_dl = max_time_per_dl*2
                    time.sleep(n_tries ** 2) # wait a lil more every time 

            if n_tries == n_max_tries:
                print('Failed downloading {} {} times in {} seconds, error code {}'.format(url, n_tries, max_time_per_dl, errcode))




def download_metadata(n_max_tries, max_time_per_dl, version, dl_root):

    urls = ['https://aqs.epa.gov/aqsweb/airdata/aqs_monitors.zip', 'https://aqs.epa.gov/aqsweb/airdata/aqs_sites.zip']
    filenames = ['aqs_monitors.zip', 'aqs_sites.zip']
    today = date.today()

    os.makedirs('{}/US_EPA_AQS/metadata/{}/network_provided/temp/'.format(dl_root, version), exist_ok=True)
    download_location = '/{}/US_EPA_AQS/metadata/{}/network_provided/temp/'.format(dl_root, version)

    i=0
    for url in urls:
        
        n_tries = 0
        errcode = 999
        while (n_tries < n_max_tries) and (errcode != 200):
            r = requests.get(url, timeout=max_time_per_dl)
            if r.status_code == 200:
                urllib.request.urlretrieve(url, download_location+filenames[i])
                print('Downloaded {}'.format(filenames[i]))

                # unzip
                with zipfile.ZipFile(download_location+filenames[i], 'r') as zip_ref:
                    zip_ref.extractall(download_location)
        
                os.remove(download_location+filenames[i])
                errcode = r.status_code

            elif r.status_code == 404:
                print("No data found, error 404")
                errcode = 200

            elif r.status_code == 403:
                print("Permission denied for {}".format(url))
                errcode = 200

            else:
                # try again
                print('Response error {}, attempt {}'.format(r.status_code, n_tries))
                errcode = r.status_code
                n_tries += 1
                max_time_per_dl = max_time_per_dl*2
                time.sleep(n_tries ** 2) # wait a lil more every time 

        if n_tries == n_max_tries:
            print('Failed downloading {} {} times in {} seconds, error code {}'.format(url, n_tries, max_time_per_dl, errcode))


        
    # JSON ===============================================================================================
        """
        # create json from original metadata file =====================================================================================
        with open('{}/US_EPA_AQS/metadata/{}/network_provided/{}.csv'.format(dl_root, version, filenames[i][:-4]), 'r') as file:
            csv_filedata = csv.DictReader(file)

            state_codes = []
            [state_codes.append(row['State Code']) for row in csv_filedata if row['State Code'] not in state_codes]

            if filenames[i] == 'aqs_sites.zip': #sites
                json_metadata = {}
                file.seek(0) # set iterator to first line
                for row in csv_filedata: 
                    update_date = today.strftime('%Y-%m-%d')
                    key = row['State Code']+'_'+row['County Code']+'_'+row['Site Number']
                    for parameter in row:
                        row[parameter] = {'values': [row[parameter]], 'update_time': [update_date]} # create inner dictionary for every parameter
                        json_metadata[key] = row

                    with open('{}/US_EPA_AQS/metadata/{}/processed/US_EPA_AQS_META_{}.json'.format(dl_root, version, filenames[i][4:-4]), 'w', encoding='utf-8') as f:
                        f.write(json.dumps(json_metadata, indent=4, ensure_ascii=False))

            elif filenames[i] == 'aqs_monitors.zip': #monitors
                os.makedirs('{}/US_EPA_AQS/metadata/{}/processed/monitors_by_state/'.format(dl_root, version), exist_ok=True)

                for state_code in state_codes:
                    json_metadata = {}
                    file.seek(0)
                    for row in csv_filedata: 
                        if row['State Code'] == state_code: # only include rows from one state
                            update_date = today.strftime('%Y-%m-%d')
                            key = row['State Code']+'_'+row['County Code']+'_'+row['Site Number']+'_'+row['Parameter Code']
                            for parameter in row:
                                row[parameter] = {'values': [row[parameter]], 'update_time': [update_date]} # create inner dictionary for every parameter
                                json_metadata[key] = row

                    with open('{}/US_EPA_AQS/metadata/{}/processed/monitors_by_state/US_EPA_AQS_META_monitors_{}.json'.format(dl_root, version, state_code), 'w', encoding='utf-8') as f:
                        f.write(json.dumps(json_metadata, indent=4, ensure_ascii=False))

            else:
                print("Key error")"""
        
        
        # create json in desired shape from current metadata file
        with open('{}/US_EPA_AQS/metadata/{}/network_provided/temp/{}.csv'.format(dl_root, version, filenames[i][:-4]), 'r') as file:
            csv_filedata = csv.DictReader(file)
            state_codes = []
            [state_codes.append(row['State Code']) for row in csv_filedata if row['State Code'] not in state_codes]

            if filenames[i] == 'aqs_sites.zip': #sites
                json_metadata_now = {}
                file.seek(0)
                for row in csv_filedata: 
                    update_date = today.strftime('%Y-%m-%d')
                    key = row['State Code']+'_'+row['County Code']+'_'+row['Site Number']
                    for parameter in row:
                        row[parameter] = {'values': [row[parameter]], 'update_time': [update_date]} # create inner dictionary for every parameter
                        json_metadata_now[key] = row

                    with open('{}/US_EPA_AQS/metadata/{}/network_provided/US_EPA_AQS_META_{}_{}.json'.format(dl_root, version, filenames[i][4:-4], today.strftime('%Y%m%d')), 'w', encoding='utf-8') as f:
                        f.write(json.dumps(json_metadata_now, indent=4, ensure_ascii=False))

            elif filenames[i] == 'aqs_monitors.zip': #monitors
                # create directory for monitor meta files per state
                os.makedirs('{}/US_EPA_AQS/metadata/{}/network_provided/monitors_by_state_{}/'.format(dl_root, version, today.strftime('%Y%m%d')), exist_ok=True)

                for state_code in state_codes:
                    json_metadata_now = {}
                    file.seek(0)
                    for row in csv_filedata: 
                        if row['State Code'] == state_code: # only include rows from one state
                            update_date = today.strftime('%Y-%m-%d')
                            key = row['State Code']+'_'+row['County Code']+'_'+row['Site Number']+'_'+row['Parameter Code']
                            for parameter in row:
                                row[parameter] = {'values': [row[parameter]], 'update_time': [update_date]} # create inner dictionary for every parameter
                                json_metadata_now[key] = row

                    with open('{}/US_EPA_AQS/metadata/{}/network_provided/monitors_by_state_{}/US_EPA_AQS_META_monitors_{}_{}.json'.format(dl_root, version, today.strftime('%Y%m%d'), state_code, today.strftime('%Y%m%d')), 'w', encoding='utf-8') as f:
                        f.write(json.dumps(json_metadata_now, indent=4, ensure_ascii=False))

            else:
                print("Key error")



        # print("Creating and comparing json of {}".format(filenames[i]))
        # if filenames[i] == 'aqs_sites.zip': #sites
        #     json_metadata_now = {}  
        #     with open(download_location+'{}.csv'.format(filenames[i][:-4]), encoding='utf-8') as file:
        #         csv_filedata = csv.DictReader(file)

        #         for row in csv_filedata:  
        #             if filenames[i] == 'aqs_sites.zip': #sites
        #                 key = row['State Code']+'_'+row['County Code']+'_'+row['Site Number']
        #             elif filenames[i] == 'aqs_monitors.zip': #monitors
        #                 key = row['State Code']+'_'+row['County Code']+'_'+row['Site Number']+'_'+row['Parameter Code']
        #             else:
        #                 print("Key error")
                    
        #             update_date = today.strftime('%Y-%m-%d')
        #             for parameter in row:
        #                 row[parameter] = {'values': [row[parameter]], 'update_time': [update_date]} # create inner dictionary for every parameter
        #             json_metadata_now[key] = row
            
        #     with open('{}/US_EPA_AQS/metadata/{}/network_provided/US_EPA_AQS_META_{}_{}.json'.format(fl_root, version, filenames[i][4:-4], today.strftime('%Y%m%d')), 'w', encoding='utf-8') as f:
        #         f.write(json.dumps(json_metadata_now, indent=4, ensure_ascii=False))

        # # alternative reading huge csv file with limited ram
        # elif filenames[i] == 'aqs_monitors.zip': #monitors
        #     json_metadata_now = {}
        #     meta_monitors = pd.read_csv(download_location+'{}.csv'.format(filenames[i][:-4]), encoding='utf-8')
        #     print(meta_monitors)

        #     for j in meta_monitors.index: # loop over rows
        #         if filenames[i] == 'aqs_sites.zip': #sites
        #             key = str(meta_monitors.iloc[j,:]['State Code'])+'_'+str(meta_monitors.iloc[j,:]['County Code'])+'_'+str(meta_monitors.iloc[j,:]['Site Number'])
        #         elif filenames[i] == 'aqs_monitors.zip': #monitors
        #             key = str(meta_monitors.iloc[j,:]['State Code'])+'_'+str(meta_monitors.iloc[j,:]['County Code'])+'_'+str(meta_monitors.iloc[j,:]['Site Number'])+'_'+str(meta_monitors.iloc[j,:]['Parameter Code'])
        #         else:
        #             print("Key error")
        #         print(key)
                
        #         update_date = today.strftime('%Y-%m-%d')
        #         json_metadata_now[key] = {}
        #         for k in range(len(meta_monitors.iloc[j,:])): # loop over columns
        #             parameter = {'values': [str(meta_monitors.iloc[j,k])], 'update_time': [update_date]} # create inner dictionary for every parameter
        #             json_metadata_now[key].update({meta_monitors.columns[k]: parameter})
            
        #     with open('{}/US_EPA_AQS/metadata/{}/network_provided/US_EPA_AQS_META_{}_{}.json'.format(dl_root, version, filenames[i][4:-4], today.strftime('%Y%m%d')), 'w', encoding='utf-8') as f:
        #         #f.write(json.dumps(json_metadata_now, indent=4, ensure_ascii=False))
        #         f.write('[')
        #         for item in json_metadata_now:
        #             print(item)
        #             f.write('{')
        #             f.write(json.dumps(item, indent=4, ensure_ascii=False))
        #             f.write(': ')
        #             f.write(json.dumps(json_metadata_now[item], indent=4, ensure_ascii=False))
        #             f.write('},\n')
        #         f.write(']')
        # else:
        #     pass
        
  
        i=i+1


    # SITES ================================================================================================================================
    #read newly scraped metadata
    with open('{}/US_EPA_AQS/metadata/{}/network_provided/US_EPA_AQS_META_sites_{}.json'.format(dl_root, version, today.strftime('%Y%m%d')), 'r', encoding='utf-8') as f:
        json_metadata_now = json.loads(f.read())
        
    # read standardised file to compare!
    with open('{}/US_EPA_AQS/metadata/{}/processed/US_EPA_AQS_META_sites.json'.format(dl_root, version), 'r', encoding='utf-8') as f:
        json_metadata = json.loads(f.read())
    
    for station in json_metadata: # loop through all the old stations
        if station in json_metadata_now.keys(): # if station is in current meta data, go on
            for parameter in json_metadata[station]:
                if json_metadata[station][parameter]['values'][-1] != json_metadata_now[station][parameter]['values'][0]: # compare last entry in standardised file to value in new file
                    # if different value, append the standardised metadeta file
                    print("old {} --- new {}".format(json_metadata[station][parameter]['values'][-1], json_metadata_now[station][parameter]['values'][0]))
                    json_metadata[station][parameter]['values'].append(json_metadata_now[station][parameter]['values'][0])
                    json_metadata[station][parameter]['update_time'].append(json_metadata_now[station][parameter]['update_time'][0])
                else:
                    pass
        else:
            print('Station {} was abolished'.format(station))

    for station in json_metadata_now: # loop through all the new stations
        if station in json_metadata.keys(): # if station is in old meta data
            pass # comparison was done before
        else: # new station appeared!
            print('New station {}'.format(station))
            json_metadata.update({station: json_metadata_now[station]})

    # safe
    with open('{}/US_EPA_AQS/metadata/{}/processed/US_EPA_AQS_META_sites.json'.format(dl_root, version), 'w', encoding='utf-8') as f:
        f.write(json.dumps(json_metadata, indent=4, ensure_ascii=False))


    # MONITORS ============================================================================================
    files_now = os.listdir('{}/US_EPA_AQS/metadata/{}/network_provided/monitors_by_state_{}/'.format(dl_root, version, today.strftime('%Y%m%d')))
    files_old = os.listdir('{}/US_EPA_AQS/metadata/{}/processed/monitors_by_state/'.format(dl_root, version))

    for file_now in files_now:
        for file_old in files_old:
            if file_now.split('_')[5] == file_old.split('_')[5][:-5]: # match the files by state!

                with open('{}/US_EPA_AQS/metadata/{}/network_provided/monitors_by_state_{}/{}'.format(dl_root, version, today.strftime('%Y%m%d'), file_now), 'r', encoding='utf-8') as f:
                    json_metadata_now = json.loads(f.read())
                    
                # read standardised file to compare!
                with open('{}/US_EPA_AQS/metadata/{}/processed/monitors_by_state/{}'.format(dl_root, version, file_old), 'r', encoding='utf-8') as f:
                    json_metadata = json.loads(f.read())

                for station in json_metadata: # loop through all the old stations
                    if station in json_metadata_now.keys(): # if station is in current meta data, go on
                        for parameter in json_metadata[station]:
                            if json_metadata[station][parameter]['values'][-1] != json_metadata_now[station][parameter]['values'][0]: # compare last entry in standardised file to value in new file
                                # if different value, append the standardised metadeta file
                                print("old {} --- new {}".format(json_metadata[station][parameter]['values'][-1], json_metadata_now[station][parameter]['values'][0]))
                                json_metadata[station][parameter]['values'].append(json_metadata_now[station][parameter]['values'][0])
                                json_metadata[station][parameter]['update_time'].append(json_metadata_now[station][parameter]['update_time'][0])
                            else:
                                pass
                    else:
                        print('Station {} was abolished'.format(station))

                for station in json_metadata_now: # loop through all the new stations
                    if station in json_metadata.keys(): # if station is in old meta data
                        pass # comparison was done before
                    else: # new station appeared!
                        print('New station {}'.format(station))
                        json_metadata.update({station: json_metadata_now[station]})

                # safe
                with open('{}/US_EPA_AQS/metadata/{}/processed/monitors_by_state/US_EPA_AQS_META_monitors_{}.json'.format(dl_root, version, file_old.split('_')[5][:-5]), 'w', encoding='utf-8') as f:
                    f.write(json.dumps(json_metadata, indent=4, ensure_ascii=False))