import torchvision
import torch
import os
import time
from mrz_decoder.pipeline import MRZDecoder
from mrz_decoder.config import Config


class CharacterDataset(torch.utils.data.Dataset):
    def __init__(self, features=None, labels=None, root='/data/doc_detect/mrzs/chars/'):
        self.features = []
        self.targets = []
        self.filenames = []
        self.config = Config
        if features is not None and labels is not None:
            self.features = [self.config.input_transform(f) for f in features]
            self.targets = [self.config.character_lookup[target.lower()] for target in labels]
            self.filenames = labels
        else:
            self._load_images(root)
            print('Dataset loaded:', len(self.filenames))
    
    def _load_images(self, root):
        for file in os.listdir(root):
            self.features.append(self.config.input_transform(torch.load(root + file, map_location='cpu')))
            target = file[:2].replace('_', '').lower()
            target = self.config.character_lookup[target]
            self.targets.append(target)
            self.filenames.append(file)

    def __getitem__(self, index):
        return self.features[index], self.targets[index], self.filenames[index]

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


class CharacterModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.config = Config  
        self.model = torchvision.models.resnet18(pretrained=True)
        self.model.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.model.maxpool = torch.nn.Identity()
        self.model.fc = torch.nn.Linear(512, len(self.config.character_lookup))

    def forward(self, x):
        return self.model(x)

    def collate_fn(self, batch):
        return tuple(zip(*batch))
    
    def train_one_epoch(self, epoch):
        header = f"Epoch: [{epoch}]"
        sum_loss, num_iterations = 0, 0
        for i, (image_batch, target_batch, filenames) in enumerate(self.train_loader):
            self.optimizer.zero_grad()
            image_batch = torch.stack(image_batch).to(self.config.device)
            target_batch = torch.tensor(target_batch).to(self.config.device)
            predictions = self.model(image_batch)
            loss = self.criterion(predictions, target_batch)
            loss.backward()
            self.optimizer.step()
            sum_loss += loss.item()
            num_iterations += 1
        train_acc = self.compute_accuracy(self.train_loader)
        test_acc = self.compute_accuracy(self.test_loader)
        print(header, '[Loss/Train/Test] ', "{:.4f}/{:.4f}/{:.4f}".format(sum_loss/num_iterations, train_acc, test_acc))

    def save_model(self, path):
        torch.save(self.model.state_dict(), path)
        # torch.save(self.model, path)

    def load_model(self, path):
        sd = torch.load(path, map_location='cpu')
        self.model.load_state_dict(sd)
        print('model', path, 'loaded')

    def train(self, train, test):
        self.train_loader = torch.utils.data.DataLoader(train, batch_size=512, shuffle=True,
                                                        collate_fn=self.collate_fn)
        self.test_loader = torch.utils.data.DataLoader(test, batch_size=512, shuffle=True,
                                                        collate_fn=self.collate_fn)
        num_epochs = 50
        self.criterion = torch.nn.CrossEntropyLoss()
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.00001)
        self.model.train().to(self.config.device)
        for epoch in range(num_epochs):
            self.train_one_epoch(epoch)
            if epoch % 5 == 0:
                self.save_model('char_model.pytorch.model')
                if epoch > 0:
                    s = time.time()
                    decode_rate = self.decode_rate()
                    print('     Decode Rate:', decode_rate, time.time() - s)


    def compute_accuracy(self, loader):
        self.model.eval().to(self.config.device)
        accuracy, total = 0, 0
        for i, (image_batch, target_batch, filenames) in enumerate(loader):
            image_batch = torch.stack(image_batch).to(self.config.device)
            target_batch = torch.tensor(target_batch).to(self.config.device)
            predictions = self.model(image_batch)
            predictions = torch.argmax(predictions, dim=1)
            accuracy += sum(predictions == target_batch)
            total += len(target_batch)
        self.model.train().to(self.config.device)
        return (accuracy/total).item()
    
    def decode_rate(self):
        self.model.eval()
        decoder = MRZDecoder(self.model)
        root = '/data/doc_detect/mrzs/mlweb2/'
        count = 0
        for i, file in enumerate(os.listdir(root)):
            im = torch.load(root + file)
            im = torchvision.transforms.ToPILImage()(im).convert('RGB')
            decoding = decoder.decode_mrz(im, fuzzy=True)
            if decoding['Valid'] is True:
                count += 1
        return count / i

if __name__ == "__main__":
    root = '/data/doc_detect/mrzs/chars_histogram/'
    dataset = CharacterDataset(root=root)
    train_size = int(0.85 * len(dataset))
    test_size = len(dataset) - train_size
    train, test = torch.utils.data.random_split(dataset, [train_size, test_size])
    del dataset
    print('Datasets loaded:', len(train), len(test))
    characterModel = CharacterModel()
    characterModel.train(train, test)

# linear gets like 60%
# lenet gets like 70%
# resnet does like 80

# 90% on ~4k examples
# 95% on ~8k examples