import pandas as pd 
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Assuming that we are on a CUDA machine, this should print a CUDA device:
print(device)

#Import data 

x=np.array(x)
xt=np.array(xt)
y=np.array(y)
yt=np.array(yt)

sx=np.std(x,axis=0)
x=x[:,sx!=0]
xt=xt[:,sx!=0]
print(x.shape)

mx=np.mean(x,axis=0)
sx=np.std(x,axis=0)
print(mx.shape)
x=torch.tensor((x-mx)/sx)
xt=torch.tensor((xt-mx)/sx)

n=x.shape[0]
y=torch.tensor(y.reshape(n,1))
y0=y.clone()
y0[y0<1]=0

n1=xt.shape[0]
yt=torch.tensor(yt.reshape(n1,1))
yt0=yt.clone()
yt0[yt0<1]=0

x=x.to(device)
y=y.to(device)
xt=xt.to(device)
yt=yt.to(device)


#loss function for classification

def loss(u):
    # x is nxp
    # w is px1
    # u=y*(x@w+w0)
    l2=torch.log(1+torch.exp(-u)) 
    #l2[u<0]=u[u<0]+torch.log(1+torch.exp(u[u<0]))#log(exp(-u)(exp(u)+1))=-u+log(1+exp(u))
    return torch.mean(l2)

# error function for classification 
def err(xw,y):
    py=torch.sign(xw)
    #print(py.shape,y.shape)
    return torch.mean((py!=y).float())


#loss function for regression 
def loss(u,y):
    l1=torch.mean((u-y)**2)
    return l1

#%% error function for regression 
def err(xw,y):
    py=(xw)
    #print(py.shape,y.shape)
    return torch.mean(torch.abs(py-y)).float()
    
    
import torch.optim as optim

k=10
Niter=300
mu=100
p=x.shape[1]
print(n,p)
s=0.1
errs=[]
errst=[]
idx=torch.tensor(np.arange(0,p)).long().to(device)
losses=[]
w=torch.zeros((p,1))
w0=torch.zeros(1)
w=w.to(device)
x=x.to(device)
y=y.to(device)
w0=w0.to(device)
w0.requires_grad=True
w.requires_grad=True
optimizer = optim.Adam([w,w0], lr=0.001)
for i in range(Niter):
    optimizer.zero_grad()   # zero the gradient buffers
    xw=x[:,idx].float()@w.view(-1,1)+w0
    yxw=y*xw
    loss1 = loss(yxw)+s*torch.sum(w**2)+s*w0**2
    loss1.backward()
    optimizer.step()
    m=int(k+(p-k)*max(0,(Niter-2*i)/(2*i*mu+Niter)))
    if m<w.shape[0]:
        sw=-torch.sort(-torch.abs(w.view(-1)))[0]
        thr=sw[m-1].item()
        j=torch.where(torch.abs(w.view(-1))>=thr)[0]
        idx=idx[j]
        w=w[torch.abs(w)>=thr].detach().clone()
        w.requires_grad=True
        optimizer = optim.Adam([w,w0], lr=0.001)
    print(i,loss1.item(),m,idx.shape) # r2_loss(xw,y))#, r2_loss(xw,y))
    losses.append(loss1.item())
plt.plot(losses)
plt.show()
xw=x[:,idx].float()@w.view(-1,1)+w0
er=err(xw,y)
xw=xt[:,idx].float().to(device)@w.view(-1,1)+w0
errt=err(xw,yt.to(device))
errs.append(er.item())
errst.append(errt.item())
print(torch.sum(w!=0).item(),er.item(),errt.item()) #,r2_loss(xw, yt))#, r2_loss(xw, yt))
print(loss(yt*xw)) #yt*xw for class (xw,yt) for reg











