"""Plots the evolution of the simulated annealing algorithm from a log file"""

import matplotlib.pyplot as plt
import numpy as np
import pickle
import pandas
from scipy.signal import savgol_filter

from analysis import utils

def deroll(arr, limits, start=0):
    """Derolls a log array. It returns a likely guess of what an array would
    have been before applying a mod operator to bring it into the limits
    region. Example, for limits = [0, 1] the array [0.5, 0.7, 0.9, 0.1] would
    return [0.5, 0.7, 0.9, 1.1].

    Parameters:
        arr (array): array to deroll
        limit (tuple): (lower limit, upper limit)
        start (int): the first start values of the array will not be derolled

    Returns:
        The derolled array.
    """
    for i in range(start,len(arr)-1): # Do not deroll before start
        if np.abs(arr[i+1]-arr[i]) > (limits[1]-limits[0])/2:
            # Continue the array in the closest possible way
            if arr[i+1]>arr[i]: arr[i+1:] -= (limits[1]-limits[0])
            else: arr[i+1:] += (limits[1]-limits[0])
    return arr

def returnBars(arr, n):
    """Calculates the 10th-90th percentile running confidence interval

    Parameters:
        arr (arr): The array to calculate error bars for
        n (int): The smoothing of the confidence interval

    Returns:
        The smoothed running confidence interval
    """
    r = pandas.Series(arr).rolling(window = n, center = False)
    s1, s2 =  r.quantile(.90), r.quantile(.1)
    return savgol_filter(s1[n:], 101, 3), savgol_filter(s2[n:], 101, 3)

# Load the data saved by the simulating annealing algorithm
log = pickle.load(open('data/logs/log.pickle', "rb" ) )[:1400]

# Define the parameters we wish to plot
mask = [0,1,2,3,5,7,8] #Non-fixed parameters
limits = np.array([[0, np.pi], [-np.pi, np.pi], 
    [0, np.pi], [-np.pi, np.pi], 
    [.5,1.0], [.55,.8], [.55,.8]])
labels = [r'$i_1$ / rad', r'$\omega_1$ / rad', 
    r'$i_2$ / rad', r'$\omega_2$ / rad', 
    'e', r'$R_1$', r'$R_2$', r'$\mu$']
ticks = [[0, np.pi], [-np.pi,0, np.pi],
    [0, np.pi],[-np.pi, 0, np.pi],
    [.5, 1.0], [.55, .8], [.55, .8]]
ticklabels = [[0,r'$\pi$'],[r'$-\pi$',0,r'$\pi$'],
    [0,r'$\pi$'],[r'$-\pi$',0,r'$\pi$'],
    ['.5','1.0'], ['.55','.8'], ['.55','.8']]

# Mask away the parameters we don't want to plot
scores = np.array([l[0] for l in log])
paramss = np.array([l[1] for l in log])[:,mask]
# Fix conventions for inclination
paramss[:,0] = paramss[:,0] - np.pi
paramss[:,2] = np.pi - paramss[:,2]

# Start plotting
i = np.arange(len(log))
f, axs = plt.subplots(1+len(paramss[0]), 1,  figsize=(10, 10), 
    sharex=True, gridspec_kw = {'height_ratios':[2., 1., 1, 1, 1, 1, 1, 1]})
plt.tight_layout()
utils.stylizePlot(axs)

# Plot metric
axs[0].scatter(i, scores, marker='x', c='black', s=5, linewidth=.5)
axs[0].fill_between(i[20:], *returnBars(scores, 20), color='r', alpha=.2)
utils.setSize(axs[0], x=(0, None), y=(0.8, None))
utils.setAxes(axs[0], y='Metric')
# Get twin axis to mark temperature in it
ax2 = axs[0].twiny()
ax2.set_xscale('log')
ax2.set_xlabel('Temperature', fontsize=14)
ax2.invert_xaxis()
ax2.set_xlim((.25, 0.015129))
ax2.set_xticks([.2, .1, .09, .08, .07, .06, .05, .04, .03, .02, .01])
ax2.set_xticklabels([str(i) 
    for i in [.2, .1, .09, .08, .07, .06, .05, .04, .03, .02, .01]])

# Plot parameters one by one
for j in range(0, len(paramss[0])):
    utils.setSize(axs[j+1], x=(0, len(log)), y=limits[j])
    axs[j+1].set_ylabel(labels[j], fontsize=14)
    axs[j+1].set_yticks(ticks[j])
    axs[j+1].set_yticklabels(ticklabels[j])
    # Be careful plotting cyclic parameters
    derolled = deroll(paramss[:,j], limits[j], start=500)
    bar1, bar2 = returnBars(paramss[:,j], 20)
    for k in range(-3, 3): #
        # Plot the confidence intervals an data multiple times
        # to deal with cyclic parameters
        axs[j+1].scatter(i, derolled + k*(limits[j][1]-limits[j][0]), 
            marker='x', c='black', s=5, linewidth=.5)
        axs[j+1].fill_between(i[20:], bar1 + k*(limits[j][1]-limits[j][0]),
            bar2 + k*(limits[j][1]-limits[j][0]), color='r', alpha=.2)

f.align_ylabels(axs[:])
axs[-1].set_xlabel('Iteration', fontsize=14)
plt.show()