#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Apr 23 08:00:26 2021

@author: doug
"""
import numpy as np
from scipy.spatial.distance import cdist


def OLS(p,t,eg,sp,mn):
    nr1,q=np.shape(p)
    nr2,q=np.shape(t)
    b=np.sqrt(-np.log(0.5))/sp
    
    #Radial Basis Layer Outputs
    P=radbas(dist(p.T,p)*b)
    PP=np.sum(P**2,axis=0).T
    #print('Shape of t',t.shape)
    d=t.T
    #print('Shape of d',d.shape)
    dd=np.sum(d**2,axis=0)
    
    #Calculate "errors" associated with vectors
    #print(P.shape,d.shape,dd.shape,PP.shape)
    e=(((P.T @ d).T)**2)/( dd[:,np.newaxis] @ PP[:,np.newaxis].T )
    
    #Pick vector with most error
    pick=findLargeColumn(e)
    left=np.arange(0,q)
    W=P[:,pick]
    W=W[:,np.newaxis]
    P=np.delete(P,pick,1)
    PP=np.delete(PP,pick,0)
    e=np.delete(e,pick,1)
    used=[left[pick]]
    left=np.delete(left,pick,0)
    
    #Calculate actual error
    w1=p[:,used].T
    a1=radbas(dist(w1,p)*b)
    w2,b2=solvelin2(a1,t)
    a2=w2 @ a1 + b2[:,np.newaxis]
    MSE=mse(t,a2)
    
    # START
    iterations=min(mn,q)
    tr=np.zeros((iterations))
    for k in range(0,iterations):
        # Calculate errors
        wj = W[:,k];
        a = (wj.T @ P) / (wj.T @ wj)
        P = P - ( wj[:,np.newaxis] @ a[np.newaxis,:] )
        PP = np.sum(P**2,axis=0).T
        e = (((P.T @ d).T)**2) / (dd[:,np.newaxis] @ PP[:,np.newaxis].T)
        #Pick vector with most error
        pick=findLargeColumn(e)
        u2=P[:,pick]
        W=np.concatenate((W,u2[:,np.newaxis]),axis=1)
        uu=np.delete(P,pick,1)
        P=uu
        PP=np.delete(PP,pick,0)
        e=np.delete(e,pick,1)
        used=np.append(used,left[pick])
        left=np.delete(left,pick,0)
        #Calculate actual error
        w1=p[:,used].T
        a1=radbas(dist(w1,p)*b)
        w2,b2=solvelin2(a1,t)
        a2=w2 @ a1 + b2[:,np.newaxis]
        MSE=mse(t,a2)
        #Track error
        tr[k]=MSE
        if MSE<eg:
            break
        
    S1,R=w1.shape
    C=w1
    W=w2
    b1=b
    b=b2
    return C,b1,W,b,tr
    


def findLargeColumn(m):
    m[np.isnan(m)]=0    #Replace NaN's with 0's
    m = np.sum(m**2,axis=0)
    i = np.argmax(m)
    return i

def solvelin2(p,t):
    pr,pc=np.shape(p)
    temp=np.ones((1,pc))
    p=np.vstack((p,temp))
    xt=np.linalg.lstsq(p.T,t.T,rcond=None)
    xt=xt[0].T
    w=xt[:,:pr]
    b=xt[:,pr]
    return w,b
    
def mse(t,y):
    temp=np.sum((t-y)**2,axis=1)
    return np.mean(temp)

def radbas(n):
    a=np.exp(-(n**2))
    return a


def dist(w,p):
   p=p.T
   temp=cdist(w,p)
   return temp

   
def rbfOLS(X,T,Goal,MaxNeurons):
    r1,c1=np.shape(X)
    r2,c2=np.shape(T)
    assert c1==c2, "Dimension mismatch in rbfOLS"
    C,b1,W,b,tr = OLS(X,T,Goal,1,MaxNeurons)
    NetParams={'C':C, 'W':W.T, 'b':b.T, 'spr':b1, 'tr':tr}
    return NetParams
    
def rbfPredict(NetParams,X):
    C=NetParams.get('C')
    W=NetParams.get('W')
    b=NetParams.get('b')
    spr=NetParams.get('spr')
    A=radbas(dist(X,C.T)*spr)
    Y= A @ W + b
    Y=Y.T
    return Y