# the basic imports
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch

# this loads the data, apparently
from torch.utils.data import DataLoader

# load data (only get rid of the Id column)
wine_data = pd.read_csv("./WineQT.csv", delimiter=",").drop("Id", axis=1)

# split the dataset into model and test subsets
wine_model = wine_data.sample(frac=0.7, random_state=123)
wine_test = wine_data.drop(wine_train.index)

# further split the training set into train and validation
wine_train = wine_model.sample(frac=0.7, random_state=123)
wine_validate = wine_model.drop(wine_train.index)

# verify the produced data size matches the source
print("source:", len(wine_data))
print("train:", len(wine_train), "validate:", len(wine_validate), "test:", len(wine_test), "==", len(wine_train + wine_validate + wine_test))

wine_train

wine_validate

wine_test

# isolate results from features
# now this is the hassle most of the time - convert whatever the input is to whatever it's expected to be
# in our case, pandas data frames to pytorch tensors
train_features = torch.tensor(wine_train.drop('quality', axis=1).values.astype(np.float32))
train_target = torch.tensor(wine_train['quality'].values.astype(np.int64))
validate_features = torch.tensor(wine_validate.drop('quality', axis=1).values.astype(np.float32))
validate_target = torch.tensor(wine_validate['quality'].values.astype(np.int64))
test_features = torch.tensor(wine_test.drop('quality', axis=1).values.astype(np.float32))
test_target = torch.tensor(wine_test['quality'].values.astype(np.int64))

train_data = torch.utils.data.TensorDataset(train_features, train_target)
validate_data = torch.utils.data.TensorDataset(validate_features, validate_target)
test_data = torch.utils.data.TensorDataset(test_features, test_target)

train_data[0]

validate_data[0]

test_data[0]

# create data loaders
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
validate_loader = torch.utils.data.DataLoader(validate_data, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

# check what device to use
from torch.accelerator import is_available, current_accelerator
if (is_available()):
    device = current_accelerator().type
    print(f"accelerated by {device}")
else:
    device = "cpu"
    print("no accelerator")

# define the model
print("**** DEFINING Linear/ReLU layer stack ****")

from torch import nn

class SeqNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(11, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.Softmax(),
        )

    def forward(self, x):
        logits = self.linear_relu_stack(x)
        return logits

# show what the model looks like
model = SeqNN() #.to(device)
print(model)

# define the loss function and optimizer
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.90)

def train_epoch(ep_num):
    running_loss = 0.
    last_loss = 0.

    report_every = 5

    for idx, batch in enumerate(train_loader):
        # unpack the next batch
        features, labels = batch
    
        # ensure features are on the accelerator device (if any)
        features.to(device)
        
        # zero the optimizer gradients
        optimizer.zero_grad()
    
        # feed it to model
        outputs = model(features)
    
        # calculate the loss against the expected label
        loss = loss_fn(outputs, labels)
        loss.backward()
    
        # feed the new loss info to optimizer
        optimizer.step()
    
        # calculate whether this is an improvement or a degradation
        running_loss += loss.item()

        # report loss uppon every 20 items
        if idx % report_every == (report_every - 1):
            last_loss = running_loss / report_every
            print("  batch {} loss: {}".format(idx + 1, last_loss))
            tb_x = ep_num * len(train_loader) + idx + 1
            print("  loss/train: {}/{}".format(last_loss, tb_x))
            running_loss = 0.

    return last_loss

# train the model
print("**** TRAINING NN on supplied data ****")

epoch_no = 0
num_epochs = 5
best_vloss = 1000000

for epoch in range(num_epochs):
    print("EPOCH", epoch + 1)

    model.train(True)
    avg_loss = train_epoch(epoch_no)

    running_vloss = 0.0
    model.eval()
    with torch.no_grad():
        for idx, vdata in enumerate(validate_loader):
            vfeatures, vlabels = vdata
            voutputs = model(vfeatures)
            vloss = loss_fn(voutputs, vlabels)
            running_vloss += vloss

    avg_vloss = running_vloss / (idx + 1)
    print("LOSS train {} / valid {}".format(avg_loss, avg_vloss))

    if avg_vloss < best_vloss:
        best_vloss = avg_vloss

    epoch_no += 1

model.eval()