import torch
import torch.nn as nn 
import torch.optim as optim  
import torchvision.transforms as transforms  
import torchvision
import os
import numpy as np
import pandas as pd
from skimage import io
from torch.utils.data import (
    Dataset,
    DataLoader,
)  
import torch.nn.functional as F
import cv2


class LaserDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        self.annotations = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform
        
        bad = []
        for idx in range(len(self.annotations)):
            print(f'Filtering: {idx}\r', end='')
            try:
                self.__getitem__(idx)
            except ValueError:
                bad.append(idx)
        for b in bad:
            self.annotations.drop(b, inplace=True)
        print(f'Filtering complete. Dropped {len(bad)} images')

    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, index):
        img_path = os.path.join(self.root_dir, self.annotations.iloc[index, 0])
        image = io.imread(img_path)
        # dsize
        dsize = (128, 128)
        # resize image
        image = cv2.resize(image, dsize)
        y_label = torch.tensor(int(self.annotations.iloc[index, 4]))

        if self.transform:
            image = self.transform(image)

        return (image, y_label)



class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # convolutional layer (sees 128x128x3 image tensor)
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        # convolutional layer (sees 64x64x16 tensor)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        # convolutional layer (sees 32x32x32 tensor)
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
        # max pooling layer
        self.pool = nn.MaxPool2d(2, 2)
        # linear layer (64 * 16 * 16 -> 8000)
        self.fc1 = nn.Linear(64 * 16 * 16, 8000)
        # linear layer (8000 -> 500)
        self.fc2 = nn.Linear(8000, 500)
        # linear layer (500->2)
        self.fc3 = nn.Linear(500, 2)
        # dropout layer (p=0.25)
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        # add sequence of convolutional and max pooling layers
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        # flatten image input
        x = x.view(-1, 64 * 16 * 16)
        # add dropout layer=1
        x = self.dropout(x)
        # add 1st hidden layer, with relu activation function
        x = F.relu(self.fc1(x))
        # add dropout layer
        x = self.dropout(x)
        # add 2nd hidden layer, with relu activation function
        x = F.relu(self.fc2(x))
        # add dropout layer 
        x = self.dropout(x)
        # add 3rd hidden layer, with sigmoid activation function
        x = torch.sigmoid(self.fc3(x))

        return x

def main():
    # Set device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


    # Hyperparameters
    in_channel = 3
    num_classes = 2
    learning_rate = 0.01
    batch_size = 36
    num_epochs = 30


    # Load Data
    dataset = LaserDataset(
        csv_file="labels.csv",
        root_dir="data",
        transform=transforms.ToTensor(),
    )

    train_size = int(len(dataset) * 0.8)
    train_set, test_set = torch.utils.data.random_split(dataset, [train_size,len(dataset) - train_size])
    train_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=True)

    # create a complete CNN
    model = Net()
    #print(model)

    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()

    # specify optimizer
    optimizer = optim.SGD(model.parameters(), learning_rate)


    # Train Network
    for epoch in range(num_epochs):
        losses = []

        for batch_idx, (data, targets) in enumerate(train_loader):
            # Get data to cuda if possible
            data = data.to(device=device)
            targets = targets.to(device=device)

            # forward
            scores = model(data)
            loss = criterion(scores, targets)

            losses.append(loss.item())

            # backward
            optimizer.zero_grad()
            loss.backward()

            # gradient descent or adam step
            optimizer.step()

        print(f"Cost at epoch {epoch} is {sum(losses)/len(losses)}")

    # Check accuracy on training to see how good our model is
    def check_accuracy(loader, model):
        num_correct = 0
        num_samples = 0
        model.eval()

        with torch.no_grad():
            for x, y in loader:
                x = x.to(device=device)
                y = y.to(device=device)

                scores = model(x)
                _, predictions = scores.max(1)
                num_correct += (predictions == y).sum()
                num_samples += predictions.size(0)

            print(
                f"Got {num_correct} / {num_samples} with accuracy {float(num_correct)/float(num_samples)*100:.2f}"
            )

        model.train()


    print("Checking accuracy on Training Set")
    check_accuracy(train_loader, model)

    print("Checking accuracy on Test Set")
    check_accuracy(test_loader, model)


if __name__ == '__main__':
    main()