%matplotlib inline
from IPython.html.widgets import interact
from scipy import  stats
import seaborn as sns
import pandas as pd

n0=stats.norm(0,1)
n1=stats.norm(0,10)
xi = linspace(-5,5,100)

fig,ax=subplots()
ax.plot(xi,n0.pdf(xi))
ax.plot(xi,n1.pdf(xi))

def bias_coin(phead = .5):
    while True:
        yield int( np.random.rand() < phead ) 

pct_mixed  = 0.1
bias_coin_gen = bias_coin(pct_mixed)        
dual_set = [n0,n1]
samples  = [ dual_set[bias_coin_gen.next()].rvs() for i in range(500) ]

hist(samples,bins=20)
title('average = %3.3f, median=%3.3f pct_mixed=%3.3f'%(mean(samples),np.median(samples),pct_mixed))

import sympy.stats
from sympy.abc import x
eps = sympy.symbols('epsilon')

mixed_cdf = sympy.stats.cdf(sympy.stats.Normal('x',0,1),'x')(x)*(1-eps) + eps*sympy.stats.cdf(sympy.stats.Normal('x',1,2),'x')(x)
mixed_pdf = sympy.diff(mixed_cdf,x)


def plot_mixed_dist(epsilon=.1):
    n1 = stats.norm(1,2)
    xi = linspace(-5,5,100)
    fig,ax = subplots()
    ax.plot(xi,[sympy.lambdify(x,mixed_pdf.subs(eps,epsilon))(i) for i in xi],label='mixed',lw=2)
    ax.plot(xi,n0.pdf(xi),label='g(x)',linestyle='--')
    ax.plot(xi,n1.pdf(xi),label='h(x)',linestyle='--')
    ax.legend(loc=0)
    ax.set_title('epsilon = %2.2f'%(epsilon))
    ax.vlines(0,0,.4,linestyle='-',color='g')
    ax.vlines(epsilon,0,.4,linestyle='-',color='b')

interact(plot_mixed_dist,epsilon=(0,1,.05))

fig,ax=subplots()
colors=['b','r']
for k in [1,2]:
    ax.plot(xi,np.ma.masked_array(xi,abs(xi)>k),color=colors[k-1])
    ax.plot(xi,np.ma.masked_array(np.sign(xi)*k,abs(xi)<k),color=colors[k-1],label='k=%d'%k)
ax.axis(ymax=2.3,ymin=-2.3)
ax.set_ylabel(r'$\psi(x)$',fontsize=28)
ax.set_xlabel(r'$x$',fontsize=24)
ax.legend(loc='best')
ax.set_title('Huber functions')
ax.grid()

fig,ax=subplots()
ax.plot(xi,np.vstack([np.ones(xi.shape),2/abs(xi)]).min(axis=0),label='k=2')
ax.plot(xi,np.vstack([np.ones(xi.shape),1/abs(xi)]).min(axis=0),label='k=1')
ax.axis(ymax=1.1)
ax.legend(loc=0)
ax.set_title("Huber's weight function")

from sympy import mpmath, symbols, diff, Piecewise, sign, lambdify
from sympy.stats import density, cdf, Normal
from sympy.abc import k,x

eps = symbols('epsilon')
lpdf=diff(cdf(Normal('x',0,1))(x)*(1-eps)+ eps*cdf(Normal('x',0,10))(x),x)
p = Piecewise((x,abs(x)<k),(k*sign(x),True))

def closure_on_asymptotic_variance(mn=(0,0),std=(1,10)):
    from sympy.abc import k,x
    eps = symbols('epsilon')
    lpdf=diff(cdf(Normal('x',mn[0],std[0]))(x)*(1-eps)+ eps*cdf(Normal('x',mn[1],std[1]))(x),x)
    def asymptotic_variance(kval,epsval):
        p = Piecewise((x,abs(x)<kval),(kval*sign(x),True))
        denom=mpmath.quad(lambdify(x,lpdf.subs(eps,epsval)),[-kval,kval])**2
        numer=mpmath.quad(lambdify(x,(p.subs(k,kval))**2*lpdf.subs(eps,epsval)),[-np.inf,np.inf])
        return float(numer/denom)
    return asymptotic_variance

case2=closure_on_asymptotic_variance()

kvals= [.001,0.3,0.5,.7,1.00001,1.4,1.7,2,3,4,5]

fig2,ax=subplots()
ax.plot(kvals,[case2(k,0) for k in kvals],'-o',label='eps=0')
ax.plot(kvals,[case2(k,.05) for k in kvals],'-o',label='eps=.05')
ax.plot(kvals,[case2(k,.1) for k in kvals],'-o',label='eps=0.1')
ax.set_xlabel("k")
ax.set_ylabel("asymptotic variance")
ax.legend(loc=0)
ax.set_title(r"$\mathcal{N}(0,1) , \mathcal{N}(0,10)$ mixed",fontsize=18)

nsamples = 500
ncols = 300
xs = np.array([dual_set[bias_coin_gen.next()].rvs() for i in range(ncols)*nsamples ]).reshape(nsamples,-1)

fig,ax=subplots()
ax.hist(np.mean(xs,0),20,alpha=0.8,label = 'mean')
ax.hist(np.median(xs,0),20,alpha=0.3,label ='median')
ax.legend()

fig,ax=subplots()
sns.violinplot(np.vstack([np.median(xs,axis=0),np.mean(xs,axis=0)]).T,ax=ax,names=['median','mean']);

def huber_estim(x,k=1,tol=1e-6):
    if x.ndim==2: # loop over columns
        out = []        
        for i in range(x.shape[1]):
            out.append(huber_estim(x[:,i],k=k)) # recurse        
        return np.array(out)
    else:
        mu = median(x)
        mad = median(abs(x-mu))*1.4826 # follow MADN convention
        while True:
            mu_i=mean(minimum(maximum(mu-k*mad,x),mu+k*mad))
            if abs(mu-mu_i) < tol*mad: break
            mu = mu_i
        return mu_i

huber_est={k:huber_estim(xs,k) for k in [1,1.5,2,3]}
huber_est[0] = np.median(xs,axis=0)
huber_est[4] = np.mean(xs,axis=0)
fig,ax=subplots()
sns.violinplot(pd.DataFrame(huber_est),ax=ax)
ax.set_xticklabels(['median','1','1.5','2','3','mean']);

huber_est={k:huber_estim(xs,k) for k in kvals}
huber_est_df = pd.DataFrame(huber_est)

fig,ax=subplots()
sns.violinplot(huber_est_df,ax=ax)

fig2,ax=subplots()
ax.plot(kvals,[case2(k,pct_mixed) for k in kvals],'-o',label='eps=0.1')
(huber_est_df.var()*nsamples).plot(marker='o',label='est eps=0.1',ax=ax)
ax.set_xlabel("k")
ax.set_ylabel("asymptotic variance")
ax.legend(loc=0)