import statsmodels.api as sm
import pandas as pd
data_loader = sm.datasets.sunspots.load_pandas()
df = data_loader.data

df

df.head()

df.tail()

fractional_nums = df['SUNACTIVITY'].apply(lambda x: x % 1) #Take the modulo of each value with 1 to get the fractional part
fractional_nums[fractional_nums > 0].head()

print df['SUNACTIVITY'].describe()

df.plot(x='YEAR', y='SUNACTIVITY', xlim=(1700,2008))

pd.tools.plotting.autocorrelation_plot(df['SUNACTIVITY'])

import numpy as np
N2 = df.shape[0] / 2
freqs = np.linspace(0, 0.5, num=N2, endpoint=False)[1:] #Nyquist range

import scipy as sp
periodogram = sp.signal.lombscargle(df['YEAR'], df['SUNACTIVITY'], freqs * 2 * np.pi)
plt.plot(freqs, periodogram, color='green')

freq_index_at_max_power = np.argmax(periodogram)
print 'Frequency and corresponding time in years at max power: %.2f, %.1f' % (freqs[freq_index_at_max_power], 1 / freqs[freq_index_at_max_power])

v_data = df['SUNACTIVITY'].max() - df['SUNACTIVITY'].min()
h_data = df['YEAR'].max() - df['YEAR'].min()
v_data_diffs = df['SUNACTIVITY'].diff().apply(np.abs)
vbar_data_diffs = v_data_diffs / v_data    
h_data_diffs = df['YEAR'].diff().apply(np.abs)
hbar_data_diffs = h_data_diffs / h_data

def objective_fcn(width_height, target):
    dev = setup_device_coords(figsize=width_height)
    lengths = segment_lengths(dev['aspect ratio'], dev['horizontal_device'])
    weighted_avg_banking =  np.sum(segment_orientations(dev['aspect ratio']) * lengths) / np.sum(lengths)  
    return np.abs(weighted_avg_banking - target)

def setup_device_coords(figsize=(8,6)):
    h_device, v_device = figsize
    fig, ax = plot_sunspots(figsize)
    device_coords = [ax.transData.transform(data_coords) for data_coords in df.values]
    df_device = pd.DataFrame(device_coords, columns=['YEAR', 'SUNACTIVITY'])    
    v_device = df_device['SUNACTIVITY'].max() - df_device['SUNACTIVITY'].min()
    h_device = df_device['YEAR'].max() - df_device['YEAR'].min()
    aspect_ratio = v_device / h_device
    v_conversion = v_device / v_data  
    h_conversion = h_device / h_data
    fig.clear()
    return {'aspect ratio': aspect_ratio,
            'vertical_device': v_device,
            'horizontal_device': h_device,
            'vertical conversion': v_conversion,
            'horizontal conversion': h_conversion}

def plot_sunspots(figsize, color='blue'):
    fig = plt.figure(figsize=figsize)
    fig.canvas.set_window_title('%.1f by %.1f inch window' % (figsize[0], figsize[1]))
    ax1 = fig.add_subplot(111)
    df.plot(x='YEAR', y='SUNACTIVITY', ax=ax1, linewidth=2, color=color)
    fig.tight_layout()
    ax1.set_xlim(right=df['YEAR'].max())
    ax1.set_ylim(top=df['SUNACTIVITY'].max())
    ax1.set_ylabel('Observed Sunspots')
    ax1.set_title('Sunspot Activity Over Time')
    return (fig, ax1)

def segment_lengths(aspect_ratio, h_device):
    return h_device * np.sqrt(hbar_data_diffs.dropna()**2 + aspect_ratio**2 * vbar_data_diffs.dropna()**2)

def segment_orientations(aspect_ratio):
    return np.arctan(aspect_ratio * vbar_data_diffs / hbar_data_diffs)

import scipy.optimize as spo
target = np.radians(45)   
slice_obj = np.s_[20:26:0.5, # widths
                  1:4:0.5]   # heights              
results = spo.brute(objective_fcn, slice_obj, args=[target], full_output=True, finish=None)
optimal_dims, objective_val, search_grid, objective_grid = results

plt.close('all')
ax1 = plot_sunspots((8,6))
print '\n\nWeighted-average banking using default 8 x 6 inch plot window: 87.3 degrees (goal is 45 degrees)' 

ax2 = plot_sunspots(optimal_dims, color='red')
print '\n\nOptimal width and height found to be %.1f by %.1f inches' % (optimal_dims[0], optimal_dims[1])    
banking = [np.degrees(target - objective_val),
           np.degrees(target + objective_val)]    
print 'Average banking interval at optimized aspect ratio: (%.2f, %.2f)' % (banking[0], banking[1])