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

print(torch.cuda.is_available())
torch.cuda.get_device_name(0)

torch.cuda.memory_allocated()

torch.cuda.memory_reserved()

df = pd.read_csv('PYTORCH_NOTEBOOKS/Data/NYCTaxiFares.csv')

df.head()

df.fare_amount.describe()

def haversine_distance(df, lat1, long1, lat2, long2):
    """
    Calculates the haversine distance between 2 sets of GPS coordinates in df
    """
    r = 6371  # average radius of Earth in kilometers
       
    phi1 = np.radians(df[lat1])
    phi2 = np.radians(df[lat2])
    
    delta_phi = np.radians(df[lat2]-df[lat1])
    delta_lambda = np.radians(df[long2]-df[long1])
     
    a = np.sin(delta_phi/2)**2 + np.cos(phi1) * np.cos(phi2) * np.sin(delta_lambda/2)**2
    c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1-a))
    d = (r * c) # in kilometers

    return d

df['dist_km'] = haversine_distance(df, 'pickup_latitude', 'pickup_longitude','dropoff_latitude','dropoff_longitude')

df.info()

df['pickup_datetime'] = df['pickup_datetime'].astype('datetime64')

df['EDTdate'] = df.pickup_datetime - pd.Timedelta(hours = 4)

df['Hour'] = df.EDTdate.dt.hour
df['AMorPM'] = np.where(df.Hour>=12, 'PM','AM')

df['Weekday'] = df['EDTdate'].dt.strftime('%a') # day of the week

cat_cols = ['Hour','AMorPM','Weekday']
cont_cols = ['pickup_longitude',
       'pickup_latitude', 'dropoff_longitude', 'dropoff_latitude',
       'passenger_count', 'dist_km']

y_col = ['fare_amount']

df[cat_cols] = df[cat_cols].astype('category')

cats = df.select_dtypes('category').apply(lambda x: x.cat.codes).values
cont = df[cont_cols].values
cont

cats = torch.tensor(cats, dtype = torch.int64)
cont = torch.tensor(cont, dtype = torch.float)

y = torch.tensor(df[y_col].values,dtype = torch.float)

print(cats.shape)
print(cont.shape)
print(y.shape)

cat_size = [len(df[col].cat.categories) for col in cat_cols]
cat_size

emb_size = [(size, min(50,size+1)//2) for size in cat_size]
emb_size # corresponde al numero de categorías y el número al que se va a reducir

selfembeds = nn.ModuleList([nn.Embedding(ni,nf) for ni,nf in emb_size])
selfembeds

class TabularModel(nn.Module):
    def __init__(self,emb_szs, n_cont, out_sz, layers, p=0.5):
        #layers = [200,100,50] unidades en cada capa
        super().__init__()
        self.embeds = nn.ModuleList([nn.Embedding(ni,nf) for ni,nf in emb_size])
        self.emb_drop = nn.Dropout(p)
        self.bn_cont = nn.BatchNorm1d(n_cont)
        
        layerlist = []
        n_emb = sum([nf for ni,nf in emb_szs])
        n_in = n_emb + n_cont
        
        for i in layers:
            layerlist.append(nn.Linear(n_in,i))
            layerlist.append(nn.ReLU(inplace = True))
            layerlist.append(nn.BatchNorm1d(i))
            layerlist.append(nn.Dropout(p))
            n_in = i
        
        layerlist.append(nn.Linear(layers[-1],out_sz))
        
        self.layers = nn.Sequential(*layerlist)
    def forward(self, x_cat, x_cont):
        embeddings = []
        
        for i,e in enumerate(self.embeds):
            embeddings.append(e(x_cat[:,i]))
        
        x = torch.cat(embeddings,1)
        x = self.emb_drop(x)
        x_cont = self.bn_cont(x_cont)
        x = torch.cat([x,x_cont],1)
        x = self.layers(x)
        return x
    
      

torch.manual_seed(33)
model = TabularModel(emb_size, n_cont = cont.shape[1],out_sz = 1,layers = [200,100],p = 0.4)
model

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

gpumodel = model.to(device)

torch.cuda.memory_allocated()

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr = 0.001)

batch_size = 60000
test_size = int(0.2*batch_size)

cat_train = cats[:batch_size-test_size].to(device)
cat_test = cats[batch_size-test_size:batch_size].to(device)
con_train = cont[:batch_size-test_size].to(device)
con_test = cont[batch_size-test_size:batch_size].to(device)

y_train = y[:batch_size-test_size].to(device)
y_test = y[batch_size-test_size:batch_size].to(device)

import time
start_time = time.time()
epochs = 300
losses = []

for i in range(epochs):
    i +=1
    y_pred = gpumodel(cat_train, con_train)
    loss = torch.sqrt(criterion(y_pred, y_train))
    losses.append(loss)
    
    if i%10 == 1:
        print(f'epoch {i} loss is {loss}')
        
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

duration = time.time()- start_time
print(f'Training took {duration/60} minutes')

plt.plot(range(epochs), losses)

model.eval()
with torch.no_grad():
    y_val = model(cat_test, con_test)
    loss = torch.sqrt(criterion(y_val, y_test))

loss

for i in range(10):
    print(f'{i+1}.) PREDICTED: {y_val[i].item():8.2f} TRUE: {y_test[i].item():.2f}')

torch.save(model.state_dict(),'TaxiModel.pt')