# -------------------------------------------------------------------- #
# ------------------------  space weather  --------------------------- #
# -------------------------------------------------------------------- #

'''
pyatmos spaceweather

This submodule defines the following functions:

download_sw - Download or update the space weather file from www.celestrak.com

read_sw     - Read the space weather file

get_sw      - Extract the space weather data    
'''

import numpy as np
from datetime import datetime,timedelta
from os import path,makedirs,remove
from pathlib import Path
from urllib.request import urlretrieve

# =================== download and update sw data  =================== #

def download_sw(direc=None):
    '''
    Download or update the space weather file from www.celestrak.com

    Usage: 
    swfile = download_sw([direc])

    Inputs: 
    direc -> [str, optionanl, default = $HOME+'/src/sw-data/'] Directory for storing sw file
    
    Outputs: 
    swfile -> [str] Path of sw file

    Examples:
    >>> swfile = download_sw()
    Downloading the latest space weather data ... Finished
    >>> print(swfile)
    /Users/lichunxiao/src/sw-data/SW-All.txt
    >>> swfile = download_sw('sw-data/')
    Downloading the latest space weather data ... Finished
    >>> swfile = download_sw('sw-data/')
    The existing space weather data is already up to date
    >>> print(swfile)
    sw-data/SW-All.txt
    '''
    
    if direc is None:
        home = str(Path.home())
        direc = home + '/src/sw-data/'
    
    swfile = direc + 'SW-All.txt'
    url = 'https://www.celestrak.com/SpaceData/SW-All.txt'

    if not path.exists(direc): makedirs(direc)

    if not path.exists(swfile):
        print('Downloading the latest space weather data from celestrak',end=' ... ')
        urlretrieve(url, swfile)
        print('Finished')

    else:
        modified_time = datetime.fromtimestamp(path.getmtime(swfile))
        if datetime.now() > modified_time + timedelta(days=1):
            remove(swfile)
            print('Updating the space weather data from celestrak',end=' ... ')
            urlretrieve(url, swfile)    
            print('Finished')
        else:
            print('The space weather data in {:s} is already the latest.').format(direc)   
    return swfile

# =========================== read sw file ========================== #

def read_sw(swfile):
    '''
    Parse and read the space weather file

    Usage: 
    sw_obs_pre = read_sw(swfile)

    Inputs: 
    swfile -> [str] Path of sw file
    
    Outputs: 
    sw_obs_pre -> [2d str array] sw data

    Examples:
    >>> swfile = 'sw-data/SW-All.txt'
    >>> sw_obs_pre = read_sw(swfile)
    >>> print(sw_obs_pre)
    [['2020' '01' '07' ... '72.4' '68.0' '71.0']
    ['2020' '01' '06' ... '72.4' '68.1' '70.9']
    ['2020' '01' '05' ... '72.4' '68.2' '70.9']
    ...
    ['1957' '10' '03' ... '266.3' '268.1' '232.7']
    ['1957' '10' '02' ... '253.3' '267.4' '231.7']
    ['1957' '10' '01' ... '269.3' '266.6' '230.9']]
    '''
    sw_data = open(swfile,'r').readlines()
    SW_OBS,SW_PRE = [],[]
    flag1 = flag2 = 0
    for line in sw_data:
        if line.startswith('BEGIN OBSERVED'): 
            flag1 = 1
            continue
        if line.startswith('END OBSERVED'): flag1 = 0 
        if flag1 == 1: 
            sw_p = line.split()
            if len(sw_p) == 30:
                del sw_p[24]
            elif len(sw_p) == 31: 
                sw_p = np.delete(sw_p,[23,25]) 
            else: 
                sw_p = np.delete(sw_p,[23,24,25,27])
            SW_OBS.append(sw_p)
            
        if line.startswith('BEGIN DAILY_PREDICTED'): 
            flag2 = 1
            continue    
        if line.startswith('END DAILY_PREDICTED'): break 
        if flag2 == 1: SW_PRE.append(line.split())    
    SW_OBS_PRE = np.vstack((np.array(SW_OBS),np.array(SW_PRE)))   
    # inverse sort
    SW_OBS_PRE = np.flip(SW_OBS_PRE,0)
    return SW_OBS_PRE

# ========================== extract sw data ========================== #   

def get_sw(SW_OBS_PRE,t_ymd,hour):
    j = 0
    for ymd in SW_OBS_PRE[:,:3]:
        if np.array_equal(t_ymd,ymd): break
        j+=1 
    f107A,f107,ap = float(SW_OBS_PRE[j,27]),float(SW_OBS_PRE[j+1,26]),int(SW_OBS_PRE[j,22])
    aph_tmp_b0 = SW_OBS_PRE[j,14:22]   
    i = int(np.floor_divide(hour,3))
    ap_c = aph_tmp_b0[i]
    aph_tmp_b1 = SW_OBS_PRE[j+1,14:22]
    aph_tmp_b2 = SW_OBS_PRE[j+2,14:22]
    aph_tmp_b3 = SW_OBS_PRE[j+3,14:22]
    aph_tmp = np.hstack((aph_tmp_b3,aph_tmp_b2,aph_tmp_b1,aph_tmp_b0))[::-1].astype(np.float)
    apc_index = 7-i
    aph_c369 = aph_tmp[apc_index:apc_index+4]
    aph_1233 = np.average(aph_tmp[apc_index+4:apc_index+12])
    aph_3657 = np.average(aph_tmp[apc_index+12:apc_index+20])
    aph = np.hstack((ap,aph_c369,aph_1233,aph_3657))
    return f107A,f107,ap,aph