import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

## Explore Turkish stock exchange dataset

tse = pd.read_excel('../../data/data_akbilgic.xlsx',skiprows=1)

tse = tse.rename(columns={'ISE':'TLISE','ISE.1':'USDISE'})

def const_wave(T,a,b):
    wave = np.zeros(T)
    s1 = (b-a) // 2
    s2 = (b-a) - s1
    norm_C = (s1*s2 / (s1+s2))**0.5
    wave[a:a+s1] = norm_C / s1
    wave[a+s1:b] = -norm_C / s2
    return wave

def _const_wave_basis(T,a,b):
    if b-a < 2:
        return []
    wave_basis = []
    wave_basis.append(const_wave(T,a,b))
    mid_pt = a + (b-a)//2
    wave_basis += _const_wave_basis(T,a,mid_pt)
    wave_basis += _const_wave_basis(T,mid_pt,b)
    return wave_basis

def const_wave_basis(T,a,b):
    father = np.ones(T) / T**0.5
    return [father] + _const_wave_basis(T,a,b)

# Construct discrete Haar wavelet basis
T,p = tse.shape
wave_basis = const_wave_basis(T,0,T)
W = np.array(wave_basis).T

_ = plt.plot(W[:,:3])

def soft(y,lamb):
    pos_part = (y - lamb) * (y > lamb)
    neg_part = (y + lamb) * (y < -lamb)
    return pos_part + neg_part

## Volatility seems most interesting
## will construct local measure of volatility
## remove rolling window estimate (local centering)
## square the residuals

#tse = tse.set_index('date')
tse_trem = tse - tse.rolling("7D").mean()
tse_vol = tse_trem**2.

## Make wavelet transformation and soft threshold

tse_wave = W.T @ tse_vol.values
lamb = .001
tse_soft = soft(tse_wave,lamb)
tse_rec = W @ tse_soft
tse_den = tse_vol.copy()
tse_den.iloc[:,:] = tse_rec

_ = tse_vol.plot(subplots=True,figsize=(10,10))

_ = tse_den.plot(subplots=True,figsize=(10,10))

"""
Explanation: The Lasso
Stats 208: Lecture 5
Prof. Sharpnack

Lecture slides at course github page
Some content of these slides are from STA 251 notes and STA 141B lectures. 
Some content is from Elements of Statistical Learning

Recall Convex Optimization
Def A function $f : \mathbb R^p \to \mathbb R$ is convex if for any $0 \le \alpha \le 1$, $x_0, x_1 \in \mathbb R^p$,
$$
f(\alpha x_0 + (1 - \alpha) x_1) \le \alpha f(x_0) + (1 - \alpha) f(x_1).
$$

For convex functions, local minima are global minima

Recall 1st Order Condition. If f is differentiable then it is convex if 
$$
f(x) \ge f(x_0) + \nabla f(x_0)^\top (x - x_0), \forall x,x_0
$$
and when $\nabla f(x_0) = 0$ then 
$$
f(x) \ge f(x_0), \forall x
$$
so any fixed point of gradient descent is a global min (for convex, differentiable f)
Subdifferential
Def. $g(x_0) \in \mathbb R^p$ is a subgradient of $f$ at $x_0$ if
$$
f(x) \ge f(x_0) + g(x_0)^\top (x - x_0), \forall x.
$$
The set of all subgradients at $x_0$ is call the subdifferential, denoted $\partial f(x_0)$.

For any global optima, $0 \in \partial f(x_0)$.

Wavelet denoising
Soft thresholding is commonly used for orthonormal bases.
- Suppose that we have a vector $y_1,\ldots, y_T$ (like a time series).
- And we want to reconstruct $y$ with $W \beta$ where $\beta$ has a small sum of absolute values $\sum_i |\beta_i|$ 
- $W$ is $T \times T$ and $W W^\top = W^\top W = I$ (orthonormal full rank design)
Want to minimize 
$$
\frac 12 \sum_{i=1}^T (y - W \beta)i^2 + \lambda \sum{i=1}^T |\beta_i|.
$$
End of explanation
"""


plt.plot(tse_soft[:,4])
high_idx = np.where(np.abs(tse_soft[:,5]) > .0001)[0]
print(high_idx)

fig, axs = plt.subplots(len(high_idx) + 1,1)
for i, idx in enumerate(high_idx):
    axs[i].plot(W[:,idx])
plt.plot(tse_den['FTSE'],c='r')

"""
Explanation: Wavelet reconstruction
Can reconstruct the sequence by
$$
\hat y = W \hat \beta.
$$
The objective is likelihood term + L1 penalty term,
$$
\frac 12 \sum_{i=1}^T (y - W \beta)i^2 + \lambda \sum{i=1}^T |\beta_i|.
$$

The L1 penalty "forces" some $\beta_i = 0$, inducing sparsity
End of explanation
"""


# %load ../standard_import.txt
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn import preprocessing, model_selection, linear_model

%matplotlib inline

## Modified from the github repo: https://github.com/JWarmenhoven/ISLR-python 
## which is based on the book by James et al. Intro to Statistical Learning.

df = pd.read_csv('../../data/Hitters.csv', index_col=0).dropna()
df.index.name = 'Player'
df.info()

## Simulate a dataset for lasso

n=100
p=1000
X = np.random.randn(n,p)
X = preprocessing.scale(X)

## Subselect true active set

sprob = 0.02
Sbool = np.random.rand(p) < sprob
s = np.sum(Sbool)
print("Number of non-zero's: {}".format(s))

## Construct beta and y

mu = 100.
beta = np.zeros(p)
beta[Sbool] = mu * np.random.randn(s)

eps = np.random.randn(n)
y = X.dot(beta) + eps

## Run lars with lasso mod, find active set

larper = linear_model.lars_path(X,y,method="lasso")
S = set(np.where(Sbool)[0])

def plot_it():
    for j in S:
        _ = plt.plot(larper[0],larper[2][j,:],'r')
    for j in set(range(p)) - S:
        _ = plt.plot(larper[0],larper[2][j,:],'k',linewidth=.75)
    _ = plt.title('Lasso path for simulated data')
    _ = plt.xlabel('lambda')
    _ = plt.ylabel('Coef')

plot_it()

## Hitters dataset

df = pd.read_csv('../../data/Hitters.csv', index_col=0).dropna()
df.index.name = 'Player'
df.info()

df.head()

dummies = pd.get_dummies(df[['League', 'Division', 'NewLeague']])
dummies.info()
print(dummies.head())

y = df.Salary

# Drop the column with the independent variable (Salary), and columns for which we created dummy variables
X_ = df.drop(['Salary', 'League', 'Division', 'NewLeague'], axis=1).astype('float64')
# Define the feature set X.
X = pd.concat([X_, dummies[['League_N', 'Division_W', 'NewLeague_N']]], axis=1)
X.info()

X.head(5)

loo = model_selection.LeaveOneOut()
looiter = loo.split(X)
hitlasso = linear_model.LassoCV(cv=looiter)
hitlasso.fit(X,y)

print("The selected lambda value is {:.2f}".format(hitlasso.alpha_))

hitlasso.coef_

"""
Explanation: Non-orthogonal design
The objective is likelihood term + L1 penalty term,
$$
\frac 12 \sum_{i=1}^T (y - X \beta)i^2 + \lambda \sum{i=1}^T |\beta_i|.
$$
does not have closed form for $X$ that is non-orthogonal.

it is convex
it is non-smooth (recall $|x|$)
has tuning parameter $\lambda$

Compare to best subset selection (NP-hard):
$$
\min \frac 12 \sum_{i=1}^T (y - X \beta)_i^2.
$$
for
$$
\| \beta \|_0 = |{\rm supp}(\beta)| < s.
$$
Image of Lasso solution
<img src="lasso_soln.PNG" width=100%>
Solving the Lasso
The lasso can be written in regularized form,
$$
\min \frac 12 \sum_{i=1}^T (y - X \beta)i^2 + \lambda \sum{i=1}^T |\beta_i|,
$$
or in constrained form,
$$
\min \frac 12 \sum_{i=1}^T (y - X \beta)i^2, \quad \textrm{s.t.} \sum{i=1}^T |\beta_i| \le C,
$$

For every $\lambda$ there is a $C$ such that the regularized form and constrained form have the same argmin
This correspondence is data dependent

Solving Lasso
A quadratic program (QP) is a convex optimization of the form 
$$
\min \beta^\top Q \beta + \beta^\top a \quad \textrm{ s.t. } A\beta \le c
$$
where $Q$ is positive semi-definite.
claim: The lasso (constrained form) is a QP.
$$
\sum_{i=1}^T (y - X \beta)_i^2 = \frac 12 \beta^\top (X^\top X) \beta + \beta^\top (X^\top y) + C
$$
but what about $\| \beta \|_1$?

Solving the lasso
For a single $\lambda$ (or $C$ in constrained form) can solve the lasso with many specialized methods
- quadratic program solver
- proximal gradient
- alternating direction method of multipliers
but $\lambda$ is a tuning parameter.  Options
1. Construct a grid of $\lambda$ and solve each lasso
2. Solve for all $\lambda$ values - path algorithm
Active sets and why lasso works better

Let $\hat \beta_\lambda$ be the $\hat \beta$ at tuning parameter $\lambda$.
Define $\mathcal A_\lambda = {\rm supp}(\hat \beta_\lambda)$ the non-zero elements of $\hat \beta_\lambda$.
For large $\lambda = \infty$, $|\mathcal A_\lambda| = 0$ (penalty dominates)
For small $\lambda = 0$, $|\mathcal A_\lambda| = p$ (loss dominates)

Forward greedy selection only adds elements to the active set, does not remove elements.

Lasso Path

Start at $\lambda = +\infty, \hat \beta = 0$.
Decrease $\lambda$ until $\hat \beta_{j_1} \ne 0$, $\mathcal A \gets {j_1}$. (Hitting event)

Continue decreasing $\lambda$ updating $\mathcal A$ with hitting and leaving events


$x_{j_1}$ is the predictor variable most correlated with $y$

Hitting events are when element is added to $\mathcal A$
Leaving events are when element is removed from $\mathcal A$
$\hat \beta_{\lambda,j}$ is piecewise linear, continuous, as a function of $\lambda$
knots are at "hitting" and "leaving" events


from sklearn.org
Least Angle Regression (LAR)

Standardize predictors and start with residual $r = y - \bar y$, $\hat beta = 0$
Find $x_j$ most correlated with $r$
Move $\beta_j$ in the direction of $x_j^\top r$ until the residual is more correlated with another $x_k$
Move $\beta_j,\beta_k$ in the direction of their joint OLS coefficients of $r$ on $(x_j,x_k)$ until some other competitor $x_l$ has as much correlation with the current residual
Continue until all predictors have been entered.

Lasso modification
4.5 If a non-zero coefficient drops to 0 then remove it from the active set and recompute the restricted OLS.

from ESL
End of explanation
"""


bforw = [-0.21830515,  0.38154135,  0.        ,  0.        ,  0.        ,
        0.16139123,  0.        ,  0.        ,  0.        ,  0.        ,
        0.09994524,  0.56696569, -0.16872682,  0.16924078,  0.        ,
        0.        ,  0.        , -0.19429699,  0.        ]

print(", ".join(X.columns[(hitlasso.coef_ != 0.) != (bforw != 0.)]))

"""
Explanation: We can also compare this to the selected model from forward stagewise regression:
[-0.21830515,  0.38154135,  0.        ,  0.        ,  0.        ,
        0.16139123,  0.        ,  0.        ,  0.        ,  0.        ,
        0.09994524,  0.56696569, -0.16872682,  0.16924078,  0.        ,
        0.        ,  0.        , -0.19429699,  0.        ]
This is not exactly the same model with differences in the inclusion or exclusion of AtBat, HmRun, Runs, RBI, Years, CHmRun, Errors, League_N, Division_W, NewLeague_N
End of explanation
"""


