import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import os
import pandas as pd

def create_indices(labels, mapping):
    return [mapping[label] for label in labels]

def write_to_csv(predicted, actual, probs, write_path, header):

    label_names = ["Non-Damage", "Earthquake", "Fire", "Flood"]

    if header:
        with open(write_path, "w") as file:
            file.write("Predicted,True,Non_Damage_Score,Earthquake_Score,Fire_Score,Flood_Score\n")

    with open(write_path, "a") as file:
        for i in range(len(actual)):
            file.write(
                f"{label_names[actual[i].item()]},"
                f"{label_names[predicted[i].item()]},"
                f"{probs[i, 0].item()},"
                f"{probs[i, 1].item()},"
                f"{probs[i, 2].item()},"
                f"{probs[i, 3].item()}\n"
            )


class ResNet50():

    def __init__(self, num_classes, lr=0.01, momentum=0.9, mapping=None):
        self.model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        self.num_classes = num_classes
        self.lr = lr
        self.momentum = momentum
        self.num_features = self.model.fc.in_features
        self.model.fc = nn.Linear(self.num_features, self.num_classes)

        self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
        self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr, momentum=self.momentum)

        self.mapping = mapping

    def train(self, epochs, train_loader):
        loss_over_time = []
        num_epochs = list(range(1, epochs + 1))
        for epoch in range(epochs):
            self.model.train()
            current_loss = 0.0
            for i, data in enumerate(train_loader, 0):
                inputs, labels = data
                self.optimizer.zero_grad()
                outputs = self.model(data[inputs].float())
                indices = create_indices(data[labels], self.mapping)
                target = torch.tensor(indices)
                loss = self.criterion(outputs, target)
                loss.backward()
                self.optimizer.step()
                current_loss += loss.item()
            loss_over_time.append(current_loss / len(train_loader))
            print(f"Epoch: {epoch + 1} \t Loss: {current_loss / len(train_loader)}")

        torch.save({
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "epochs": num_epochs,
            "loss": loss_over_time
        }, "model_weights.pth")

    def eval(self, test_loader, write_path=None):
        self.model.eval()
        header = True

        with torch.no_grad():
            correct = 0
            total = 0
            for data in test_loader:
                images, labels = data
                images = data[images].float()
                labels = data[labels]
                # indices = create_indices(labels)
                indices = create_indices(labels, self.mapping)
                labels = torch.tensor(indices)

                outputs = self.model(images)
                _, predicted = torch.max(outputs.data, 1)
                probs = torch.softmax(outputs, dim=1)

                total += len(labels)
                correct += (predicted == labels).sum().item()
                if write_path:
                    write_to_csv(predicted, labels, probs, write_path=write_path, header=header)
                header = False
        
        print(f'Accuracy of the network on the test images: {round(100 * correct / total, 3)}%')