#%%

from collections import defaultdict

import numpy as np
from sklearn.model_selection import cross_val_score 
from sklearn.ensemble import RandomForestRegressor as RFR 
from sklearn.datasets import load_boston
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  

X, y = load_boston(return_X_y=True) 

def rfrcv(target, train_data, param1, param2):
    val = cross_val_score(
        RFR(min_samples_split=param1, 
            min_samples_leaf=param2,
            criterion='mse',
            random_state=2,
        ),
        train_data, target, 
        cv=6, n_jobs=7
    ).mean()
    return val
 
def lerp(a, b, f):
    return a + f * (b - a)

params = {'min_samples_leaf': (0.01, 0.35),
         'min_samples_split': (0.01, 0.65), 
        }

#%% exploring variation in hyperparameter over sample size
size_grid = list(range(100, len(X), 100)) 
# size_grid = [100, 200, 300, 400, 500]
steps = 10

grid_search = dict()
for size in size_grid:
    print(size, len(X))
    grid_search[size] = list() 
    
    idx = np.random.randint(len(X), size=size)
    train_data = X[idx,:]
    target = y[idx] 
     
    for n1 in range(steps): 
        print(n1)
        param1 = round(lerp(params['min_samples_split'][0], params['min_samples_split'][1], n1/(steps-1)), 3) 
        for n2 in range(steps): 
            param2 = round(lerp(params['min_samples_leaf'][0], params['min_samples_leaf'][1], n2/(steps-1)), 3)         
            val = rfrcv(target=target, train_data=train_data, param1=param1, param2=param2)   
            grid_search[size].append([val, param1, param2]) 

# #%% marginal plots

# for size in size_grid:
#     d = grid_search[size]
#     score, p1, p2 = list(map(list, zip(*d))) 
    
#     plt.scatter(p1, p2, c=score)
#     plt.title(str(size))
#     plt.show()

#%% 3d view is more appealing for the parameter space dimension
import matplotlib as mpl

colormap = mpl.cm.autumn 
colorst = [colormap(40*i) for i in range(len(size_grid))]
title_font = {'fontname':'Arial', 'size':'26', 'color':'black', 'weight':'normal',
              'verticalalignment':'bottom'}

for i, size in enumerate(size_grid):
    d = grid_search[size]  
    x, y, z = list(map(list, zip(*d))) 

    x = np.array(x) 
    y = np.array(y) 
    z = np.array(z) 
          
    for degree in range(60, 420, 360):
        print(i) 
        ax = Axes3D(plt.figure(figsize=(15, 15)))
        ax.plot_trisurf(y, z, x, alpha=0.6, color=colorst[i])
        # ax.scatter(y, z, x, s=10, marker='o', c='red')
        ax.view_init(25, degree)   
        # ax._axis3don = False
        # ax.w_zaxis.line.set_lw(0.)
        # ax.set_zticks([])
        ax.set_ylabel('min_samples_leaf')
        ax.set_xlabel('min_samples_split')
        ax.text2D(0.05, 0.45, "2D Text", transform=ax.transAxes)

        plt.title('Sample size: ' + str(round(size/5))+ '%', **title_font) 
        plt.show()
        
'''
As in the one parameter case, the surface generated by the performance of the 
model by the parameter tuple is pretty similar as we vary the sample size.
'''

#%% exploring variance in results from sample sizes
 
size_grid = list(range(100, len(X), 200)) 
# size_grid = [len(X)]
repeats = 10 # amount of models to estimate variation caused by different samples
steps = 10

variance = defaultdict(list) 
for it in range(repeats):
    print(f'iteration: {it}')
    for size in size_grid: 
        idx = np.random.randint(len(X), size=size)
        train_data = X[idx,:]
        target = y[idx] 
        
        '''if we don't sample when the pct is 100, then there is no variance in 
        the final result. Variation in the order of the data leads to different
        results even with the seed fixed on the algo side.'''
        # train_data = X 
        # target = y 
        for n1 in range(steps): 
            
            param1 = round(lerp(params['min_samples_split'][0], params['min_samples_split'][1], n1/(steps-1)), 3) 
            for n2 in range(steps): 
                # print(round((n1 + n2)/steps**2,2))
                param2 = round(lerp(params['min_samples_leaf'][0], params['min_samples_leaf'][1], n2/(steps-1)), 3)         
                val = rfrcv(target=target, train_data=train_data, param1=param1, param2=param2)   
                variance[size].append([val, param1, param2]) 

 
for size in size_grid:
    d = variance[size]
    x, y, z = list(map(list, zip(*d))) 

    x = np.array(x) 
    y = np.array(y) 
    z = np.array(z) 
          
    for i in range(60, 420, 360):
        # print(i) 
        ax = Axes3D(plt.figure(figsize=(15, 15)))
        ax.plot_trisurf(y, z, x, alpha=0.9)
        ax.scatter(y, z, x, s=10, marker='o', c='red', alpha=0.2)
        ax.view_init(25, i)   
        # ax._axis3don = False
        ax.w_zaxis.line.set_lw(0.)
        ax.set_zticks([])
        plt.title(str(size))
        plt.show()

'''
There is more variance for smaller samples, so everything observed in the single
parameter case generalizes to a N-parameter case.
'''
#%% comparing mean/median of scores from smaller and biggest size. Do they match?
 
param_median = defaultdict(lambda: defaultdict(list))
param_mean = defaultdict(lambda: defaultdict(list))

for size in size_grid:
    p1_set = set(sorted([p1 for s, p1, p2 in variance[size]]))
    p2_set = set(sorted([p2 for s, p1, p2 in variance[size]]))
    for p1_ in p1_set:
        for p2_ in p2_set:
            s_lst = [s for s, p1, p2 in variance[size] if p1 == p1_ and p2 == p2_]
            param_median[p1_][p2_].append(round(np.median(s_lst),3))
            param_mean[p1_][p2_].append(round(np.mean(s_lst),3))
        # print(f'{size}: {p1}   {np.median(s_lst)}')

'''
It seems to hold
'''

#%% focusing on small samples distributions for all parameter values: gaussian or not?
 
'''
 
'''

 