from astroplan import Observer
from astropy import units as u
from astropy.time import Time, TimeDelta
from numpy import nan, where
from pandas import read_csv
from urllib.request import urlretrieve
from zipfile import ZipFile

if __name__ == "__main__":
    urlretrieve('http://bossa.pl/pub/futures/mstock/mstfut.zip', 'mstfut.zip')
    with ZipFile('mstfut.zip', 'r') as fd:
        fd.extract('FW20.mst', '.')
    fw20 = read_csv('FW20.mst', parse_dates=['<DTYYYYMMDD>'])
    fw20 = fw20.rename(columns={'<DTYYYYMMDD>': 'date', '<CLOSE>': 'close'})

    gpw = Observer(longitude=21.023616 * u.deg, latitude=52.230268 * u.deg,
                   elevation=100 * u.m, name="gpw", timezone="Europe/Warsaw")

    fw20XXI = fw20[fw20.date.dt.year > 2000]
    times = Time(fw20XXI.date, format='datetime64')
    times_hour_ago = times - TimeDelta(3600, format='sec')
    change = fw20XXI.close - fw20XXI.close.shift()
    moon = gpw.moon_illumination(times)
    waxing = gpw.moon_illumination(times_hour_ago) < gpw.moon_illumination(times)
    flag = ((moon > 0.3) & waxing) | (moon > 0.95)
    fw20XXI = fw20XXI.assign(lycan=where(flag, change, nan), vampire=where(flag, nan, change))
    stats = fw20XXI[['vampire', 'lycan']].groupby(fw20XXI.date.dt.year).sum()
    stats = stats.assign(vampire_win=stats.vampire > stats.lycan)
    print(stats)
    print(fw20XXI[['vampire', 'lycan']].sum())