#! /usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pylab as plt
import matplotlib.ticker as ticker
import urllib3
import ftplib
import shutil

plt.rcdefaults()
plt.style.use('classic')

def saveurl(url, fname):
    if url.startswith("ftp"):
        ftp = ftplib.FTP('/'.join(url.split('/')[2:3]))
        ftp.login('', '')
        with open(fname, 'wb') as f:
            res = ftp.retrbinary('RETR %s' % '/'.join(url.split('/')[3:]), f.write)
    else:
        c = urllib3.PoolManager()
        with c.request('GET', url, preload_content=False) as resp, open(fname, 'wb') as out_file:
            shutil.copyfileobj(resp, out_file)
    print('saved', fname, 'from', url)

def smooth_convolution(signal, window_size):
    window_size = 1 + 2 * (window_size // 2) # should be odd
    w = np.hamming(window_size)
    y = (np.convolve(w, signal) / np.convolve(w, np.ones_like(signal)))
    return y[(len(w)-1)//2:-(len(w)//2)]

def smooth_boxcar(signal, window_size):
    w = np.ones(window_size)
    return np.convolve(w, signal, mode="valid") / window_size

def smooth_lowess(signal, window_size):
    s = window_size // 2
    window = np.hamming(1+2*s) / np.hamming(1+2*s).sum()
    smoothed = np.empty_like(signal)
    for i in range(len(signal)):
        i0, i1 = max(i-s, 0), min(i+s, len(signal)-1) + 1
        # linear LOWESS smoothing
        smoothed[i] = np.poly1d(np.polyfit(np.arange(i0, i1), signal[i0:i1], deg=1,
            w=np.sqrt(window[i0-i+s:i1-i+s])))(i)
    return smoothed

saveurl("https://crudata.uea.ac.uk/cru/data/temperature/HadCRUT5.1Analysis_gl.txt", "hadcrut5gl.txt")
with open("hadcrut5gl.txt") as temp_file:
    temp_years, temps = [], []
    last_year = None
    for l in temp_file.readlines():
        yT = l.split()
        if (not l.startswith('#')) and len(yT) > 1:
            year = float(yT[0])
            if year != last_year:
                temp_years += [year + (0.5 + i) / 12. for i in range(len(yT) - 2)]
                temps += [float(T) for T in yT[1:13]]
                last_year = year
    temp_years = smooth_boxcar(temp_years, 12)
    temps = smooth_boxcar(temps, 12)
    #print('T', list(zip(temp_years, temps)))

saveurl("ftp://aftp.cmdl.noaa.gov/products/trends/co2/co2_mm_mlo.txt", "co2_mm_mlo.dat")
with open("co2_mm_mlo.dat") as co2_file:
    data_rows = [x.split() for x in co2_file if not x.startswith("#") and x.find("-99.99") == -1]
    co2_years = [float(x[2]) for x in data_rows]
    co2concs = [float(x[3]) for x in data_rows]
    #print('co2', list(zip(co2_years, co2concs)))

saveurl("ftp://ftp.ncdc.noaa.gov/pub/data/paleo/icecore/antarctica/law/law_co2.txt", "law_co2.txt")
with open("law_co2.txt", encoding="ISO-8859-1") as lawco2_file:
    data_rows = [x.split() for x in lawco2_file if x.startswith("     1")]
    del data_rows[ [float(x[0]) for x in data_rows].index(1010.):] 
    lawco2_years = [float(x[0]) for x in data_rows]
    lawco2concs = [float(x[-1]) for x in data_rows]
    #print('lawco2', list(zip(lawco2_years, lawco2concs)))

saveurl("http://www.sidc.be/silso/DATA/SN_y_tot_V2.0.txt", "yearssn.dat")
with open("yearssn.dat") as sunspot_file:
    data_rows = [x.split() for x in sunspot_file if "*" not in x]
    sun_years = [float(x[0]) for x in data_rows]
    sunspots = [float(x[1]) for x in data_rows]
    #print('sunspots', list(zip(sun_years, sunspots)))

t_min, t_max = 1850, max(temp_years[-1], co2_years[-1], sun_years[-1])
smoothed_temps = smooth_lowess(temps, 25 * 12)
smoothed_sunspots = smooth_convolution(sunspots, 25)

rect = [0.125, 0.1, 0.775, 0.8]
base_ax = plt.axes(rect)
base_ax.yaxis.tick_left()
plt.yticks([])
plt.xlim(t_min, t_max)
plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(20))
plt.xlabel("Year", size=16)
plt.title(r"Temperature, CO$_2$, and Sunspots", size=22)

temp_ax = plt.axes([rect[0],rect[1]+.42*rect[3],rect[2],.58*rect[3]], frameon=False)
temp_ax.yaxis.tick_left()
plt.plot(temp_years, temps, '#FF2200')
tline = plt.plot(temp_years, smoothed_temps, '#AA0000', lw=3)
plt.xlim(t_min, t_max)
plt.ylim(min(temps)*1.03 - 0.03*max(temps), max(temps)*1.03 - 0.03*min(temps))
plt.yticks(np.arange(-0.6, 1.3, 0.2))
plt.ylabel(u'Temperature anomaly (°C)', size=14, color='#AA0000')
plt.xticks([])

co2_ax = plt.axes([rect[0],rect[1]+.25*rect[3],rect[2],.5*rect[3]], frameon=False)
co2_ax.yaxis.tick_right()
co2_ax.yaxis.set_label_position("right")
co2_ax.xaxis.tick_bottom()
plt.plot(co2_years, co2concs, '#44AAFF')
cline = plt.plot(lawco2_years, lawco2concs, '#2288EE', lw=2)
plt.xlim(t_min, t_max)
plt.ylabel(r'CO$_2$ (ppm)', size=14, color='#2288EE')
plt.xticks([])

sun_ax = plt.axes([rect[0],rect[1],rect[2],.5*rect[3]], frameon=False)
sun_ax.yaxis.tick_left()
plt.plot(sun_years, sunspots, "#FFDD00")
sline = plt.plot(sun_years, smoothed_sunspots, "#FF9900", lw=3)
plt.xlim(t_min, t_max)
plt.yticks(np.arange(0,250,50))
plt.ylabel("Sunspot number", size=14, color='#FF9900')
plt.xticks([])

plt.sca(base_ax)
t_proxy = plt.Line2D([0], [0], c='#AA0000', lw=3)
c_proxy = plt.Line2D([0], [0], c='#2288EE', lw=3)
s_proxy = plt.Line2D([0], [0], c='#FF9900', lw=3)
plt.legend((t_proxy, c_proxy, s_proxy), ("Temperature", r"CO$_2$", "Sunspots"), loc="upper left")

plt.savefig("Temp-co2-sunspot.svg")
plt.show()