import torch
import torchvision
import sys
import os
import random
import zxingcpp
import datetime
import time
from barcode_reader.unet import UNet, L1_DICE, VGGPerceptualLoss
import argparse


class DiffDataset(torch.utils.data.Dataset):
    def __init__(self, root='/data/doc_detect/barcodes/barcode_diffusion/',
                 augment_flag=True,
                 limit=None,
                 mode='train'):
        self.targets, self.features, self.filenames = [], [], []
        self.augment_flag = augment_flag
        self.mode = mode
        self.root = root
        self.image_size = [256, 1024]
        self._load_dataset(mode=mode, limit=limit)

    def _get_files(self, mode):
        files = []
        if mode is None:
            files = os.listdir(self.root + 'features/')
        elif mode == 'train':
            test = self._get_files(mode='test')
            for file in os.listdir(self.root + 'features/'):
                if file not in test:
                    files.append(file)
            del test
        else:
            with open(self.root + mode, 'r') as f:
                lines = f.readlines()
                for line in lines:
                    files.append(line[:line.find("'\'")])
        return files

    def _load_dataset(self, mode, limit):
        files = self._get_files(mode=mode)
        for file in os.listdir(self.root + 'features/'):
            if file in files:
                self.filenames.append(file)
            if limit is not None and len(self.filenames) >= limit:
                break

    def _load_one(self, file):
        target = torch.load(self.root + 'targets/' + file)
        feature = torch.load(self.root + 'features/' + file)
        return feature, target

    def _augment(self, x, y):
        x, y = torchvision.transforms.Resize(self.image_size)(x), torchvision.transforms.Resize(self.image_size)(y)
        pad_y_amount, pad_x_amount = random.randint(0, 30), random.randint(0, 60)
        pad_top, pad_left = random.randint(0, pad_y_amount), random.randint(0, pad_x_amount)
        pad_bottom, pad_right = pad_y_amount - pad_top, pad_x_amount - pad_left
        if self.augment_flag:
            t = torchvision.transforms.Compose([
                torchvision.transforms.Resize([self.image_size[0]-(pad_top + pad_bottom),
                                               self.image_size[1]-(pad_left + pad_right)]),
                torchvision.transforms.Pad(padding=[pad_left, pad_top, pad_right, pad_bottom], padding_mode='constant', fill=1),
            ])
        else:
            t = torchvision.transforms.Compose([
                torchvision.transforms.Resize(self.image_size)
            ])
        return t(x), t(y)

    def __getitem__(self, index):
        feature, target = self._load_one(self.filenames[index])
        feature, target = self._augment(feature, target)
        return feature, target, self.filenames[index], index

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


class UNetDecoder():
    def __init__(self, model_obj, device, num_nets=5):
        self.model = model_obj
        self.device = device
        self.model.eval().to(self.device)
        self.untransform = torchvision.transforms.ToPILImage()
        self.num_nets = num_nets

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

    def decode_count(self, loader, lf=None, zxing_first=False, resize=False, limit=None):
        success, total = 0, 0
        rotation = []
        ave_loss, iterations = 0, 0
        if resize:
            untransform = torchvision.transforms.Compose([
                torchvision.transforms.Resize([256, 1024]),
                torchvision.transforms.ToPILImage()
            ])
        else:
            untransform = torchvision.transforms.ToPILImage()
        count = 0
        for i, (image_batch, target_batch, filename_batch, index_batch) in enumerate(loader):
            batch_size = len(filename_batch)
            count += batch_size
            if limit is not None:
                if count > limit:
                    break
            image_batch = torch.stack(image_batch).to(self.device)
            rotated = torchvision.transforms.functional.rotate(image_batch, 180)
            with torch.no_grad():
                out1 = self.forward(image_batch)
                if lf is not None:
                    target_batch = torch.stack(target_batch).to(self.device)
                    loss = torch.mean(torch.stack(lf(out1, target_batch)))
                    iterations += 1
                    ave_loss += loss.item()
            with torch.no_grad():
                out2 = self.forward(rotated)
            for i in range(len(out1)):
                norm1 = out1[i]
                norm2 = out2[i]
                norm1[norm1 > 1] = 1
                norm1[norm1 < 0] = 0
                norm2[norm2 > 1] = 1
                norm2[norm2 < 0] = 0
                norm1 = untransform(norm1)
                norm2 = untransform(norm2)
                try:
                    result = zxingcpp.read_barcode(norm1)
                    rotated_result = zxingcpp.read_barcode(norm2)
                except Exception:
                    result = None
                    rotated_result = None
                # norm1.save('temp.jpg')
                # result = self.zx.decode('temp.jpg', try_harder=True)
                # norm2.save('temp.jpg')
                # rotated_result = self.zx.decode('temp.jpg', try_harder=True)
                if result is not None:
                    success += 1
                elif rotated_result is not None:
                    rotation.append(filename_batch[i])
                total += 1
        if iterations > 0 and lf is not None:
            return success, total, success/total, rotation, ave_loss/iterations
        return success, total, success/total, rotation


class AdderDiffusion(torch.nn.Module):
    def __init__(self, num_nets=5, magnitude=16, model_path=None, mode='train', cuda=False):
        super().__init__()
        if cuda:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device('cpu')
        self.mode = mode
        self.num_nets = num_nets
        self.regularization_lambda = 1e-6
        layers = []
        for _ in range(num_nets):
            layers.append(UNet(in_channels=3, n_class=3, magnitude=magnitude, dropout_rate=0.0))
        self.model = torch.nn.Sequential(*layers)
        del layers
        if model_path is not None:
            self.load_model(model_path)

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

    def collate_fn(self, batch):
        return tuple(zip(*batch))

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

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

    def l2(self, loss):
        regularization_loss = torch.tensor(0.0).to(self.device)
        for net in self.model:
            for name, block in net.named_children():
                if 'encoder' in name:
                    regularization_loss += torch.norm(block.conv1.weight)
                    regularization_loss += torch.norm(block.conv2.weight)
                elif 'bridge' in name:
                    regularization_loss += torch.norm(block.weight)
                elif 'decoder' in name:
                    regularization_loss += torch.norm(block.conv1.weight)
                    regularization_loss += torch.norm(block.conv2.weight)
                    try:
                        regularization_loss += torch.norm(block.convT1.weight)
                    except Exception:
                        regularization_loss += torch.norm(block.conv3.weight)
        regularization_loss *= self.regularization_lambda
        loss += regularization_loss
        return loss, regularization_loss

    def train_one_epoch(self, epoch):
        header = f"Epoch: [{epoch}]"
        sum_loss, sum_l1, sum_dice, sum_vgg = 0, 0, 0, 0
        test_reg, sum_reg = 0, 0
        num_iterations = 0
        S = time.time()
        test_iterations, test_loss, test_l1, test_dice, test_vgg = 1, 0, 0, 0, 0
        self.model.eval()
        for ind, (image_batch, target_batch, _, index_batch) in enumerate(self.test_loader):
            image_batch = torch.stack(image_batch).to(self.device)
            target_batch = torch.stack(target_batch).to(self.device)
            with torch.no_grad():
                result = self.forward(image_batch)
                loss, dice, l1 = self.criterion(result, target_batch)
            vgg_loss = torch.tensor(0.0)
            reg_loss = torch.tensor(0.0)
            test_reg += reg_loss.item()
            test_iterations += 1
            test_loss += loss.item()
            test_l1 += l1.item()
            test_dice += dice.item()
            test_vgg += vgg_loss.item()
        self.model.train()
        for ind, (image_batch, target_batch, _, index_batch) in enumerate(self.train_loader):
            image_batch = torch.stack(image_batch).to(self.device)
            target_batch = torch.stack(target_batch).to(self.device)
            self.optimizer.zero_grad()
            result = self.forward(image_batch)
            loss, dice, l1 = self.criterion(result, target_batch)
            # vgg_loss, _, _ = self.vgg_loss(result, target_batch)
            vgg_loss = torch.tensor(0.0)
            # l1_cst = 1
            # loss = (vgg_cst) * vgg_loss + (l1_cst) * loss
            # loss, reg_loss = self.l2(loss)
            reg_loss = torch.tensor(0.0)
            sum_reg += reg_loss.item()
            num_iterations += 1
            sum_loss += loss.item()
            sum_l1 += l1.item()
            sum_dice += dice.item()
            sum_vgg += vgg_loss.item()
            loss.backward()
            # torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()
        print(header, '[' + str(self.optimizer.param_groups[0]['lr']) + ']',
              'loss/l1/vgg/reg {:.4f}/{:.4f}/{:.4f}/{:.4f}'.format(sum_loss/num_iterations,
                                                                   sum_l1/num_iterations,
                                                                   sum_vgg/num_iterations,
                                                                   sum_reg/num_iterations),
              '----------------------------- val loss/l1/vgg/reg {:.4f}/{:.4f}/{:.4f}/{:.4f}'.format(test_loss/test_iterations,
                                                                                                     test_l1/test_iterations,
                                                                                                     test_vgg/test_iterations,
                                                                                                     test_reg/test_iterations),
              str(time.time()-S)[:7])
        return sum_loss/num_iterations, test_loss/test_iterations

    def _init_dataloaders(self, batch_size, train, test):
        self.train_loader = torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True,
                                                        collate_fn=self.collate_fn, pin_memory=True,
                                                        prefetch_factor=3, num_workers=4)
        self.test_loader = torch.utils.data.DataLoader(test, batch_size=batch_size, shuffle=True,
                                                       collate_fn=self.collate_fn, pin_memory=True,
                                                       prefetch_factor=3, num_workers=4)

    def train(self, train, test, num_epochs=5000, batch_size=13, learning_rate=0.001, betas=(0.7, 0.99), version=18):
        print('Training', num_epochs, 'Epochs', 'Version:', version, 'Dataset:', len(train))
        self._init_dataloaders(batch_size=batch_size, train=train, test=test)
        params = [p for p in self.model.parameters()]
        self.optimizer = torch.optim.Adam(params, lr=learning_rate, betas=betas, eps=1e-08, weight_decay=0.0)
        self.criterion = L1_DICE(alpha=1.0)
        self.vgg_loss = VGGPerceptualLoss()
        for epoch in range(num_epochs):
            self.model.train().to(self.device)
            train_loss, test_loss = self.train_one_epoch(epoch)
            if epoch % 5 == 0:
                self.model.eval().to(self.device)
                self.log(version, train_loss, test_loss, epoch)

    def log(self, version, train_loss, test_loss, epoch, sample_size=None):
        s = time.time()
        self.model.eval().to(self.device)
        decoder = UNetDecoder(self.model, self.device, num_nets=self.num_nets)
        count, total, acc, rotation_candidates = decoder.decode_count(self.train_loader, lf=None)
        test = decoder.decode_count(self.test_loader, lf=None, limit=sample_size)
        model_name = datetime.datetime.today().strftime('%m%d-%H:%M')
        rotation_candidates += test[3]
        self.save_model('/home/nathan/barcode-reader/barcode_reader/models/adder/' + str(version) + '/' + model_name)
        print('     Model File:                            ', model_name)
        print('     Decoded train accuracy                  [' + str(count) + ', ' + str(total) + '] -', str(acc*100)[:5] + '%')
        print('     Decoded test accuracy                   [' + str(test[0]) + ', ' + str(test[1]) + '] -', str(test[2]*100)[:5] + '%')
        print('     Rotation candidates:                    [' + str(len(rotation_candidates)) + ']')
        print('Completed in:', time.time()-s)
        out_str = 'Epoch [' + str(epoch) + '] -'
        out_str += model_name + '- Train: ' + str(acc)[:7]
        out_str += ' - Test: ' + str(test[2])[:7] + ' - Loss: [' + str(train_loss)[:7] + ', ' + str(test_loss)[:7] + ']'
        out_str += ' Rotation Candidates: ' + str(len(rotation_candidates)) + '\n'
        with open('/home/nathan/barcode-reader/barcode_reader/models/adder/' + str(version) + '/model.log', 'a') as f:
            f.write(out_str)
        if epoch > 200 and acc > 0.5:
            with open('/home/nathan/barcode-reader/barcode_reader/models/adder/' + str(version) + '/rotation.log', 'a') as f:
                for filename in rotation_candidates:
                    f.write(filename + '\n')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', default=13, type=int)
    parser.add_argument('--num_epochs', default=500, type=int)
    parser.add_argument('--limit', default=None, type=int)
    parser.add_argument('--version', default=35, type=int)
    parser.add_argument('--lr', default=0.0001, type=float)
    parser.add_argument('--num_nets', default=5, type=int)
    parser.add_argument('--size', default=512, type=int)
    parser.add_argument('--magnitude', default=16, type=int)
    args = parser.parse_args()
    S = time.time()
    print('Mag', args.magnitude, 'Num nets', args.num_nets)
    if args.limit is None:
        train_lim, test_lim = None, None
    else:
        train_lim, test_lim = int(0.9 * args.limit), int(0.1 * args.limit)
        if train_lim == 0:
            train_lim = 1
        if test_lim == 0:
            test_lim = 1
    train = DiffDataset(limit=train_lim,
                        mode='train',
                        augment_flag=True)
    test = DiffDataset(limit=test_lim,
                       mode='test',
                       augment_flag=False)
    print('Dataset loaded:', len(train), len(test), time.time()-S)
    diffusion = AdderDiffusion(num_nets=args.num_nets,
                               magnitude=args.magnitude,
                               cuda=True,
                               model_path='/home/nathan/barcode-reader/barcode_reader/models/adder/35/1120-12:44'
                               )
    diffusion.train(train=train,
                    test=test,
                    learning_rate=args.lr,
                    batch_size=args.batch_size,
                    num_epochs=args.num_epochs,
                    version=args.version,
                    )


# model 3 -> single UNet, 500 epochs
# model 4 -> double UNet, 170 - epochs, 76% train, 70% test, ~20% raw dataset
# model 5 -> quad UNet, up to 85% train, 80% test, then overfit down to 55% test
# model 6 -> quad UNet, stopped early, around 76% each
# model 7 -> quad UNet, stopped early (only 60 epochs), 65% each
# model 8 -> 5 Unets, multi gpu enabled, more than doubled dataset (L1 Loss)
    # maxed out around 77/67 train/test, collapsed after around 100 epochs
# model 9 -> 5 Unets, try out L2 loss, added dropout, adjusted betas
    # for model 1114-19:50
    # starts to drop off around 250 epochs at 80/75 train/test
    # gets close to 22% on raw dataset (+5%)
    # 1150 new from more_barcoces, 885 new from barcodes/train/features - 2035 total ()
# model 10 -> extension of model 9, changed back to l1 and removed negative padding, added ~2000 more images
    # ~ 700 new barcodes after 60 epochs of training (on top of model 9)
    # added slightly more dropout
# model 11 -> lowered padding amount, working on custom loss - l1_ssim
    # extended from model 10 1115-16:47
    # changed my mind on the padding, reset
    # adjusted alpha in loss function
# model 12 -> fixing stuff I broke
    # change data augmentation to happen within __getitem__ so that every epoch is a different variation
    # broke something in the process - i think it was my model definition using custom modules
    # ~3k new datapoints added
    # using ssim_l1 with low alpha (0.2) again since previous test was invalid
    # total data 11k
# model 13 -> testing out independent graph again
    # second changed alpha
    # third switched to l2
    # fourth turned up beta1 and learning rate, back to l1ssim
    # 5, l1 loss with sgd
# model 14 -> not started yet, just planning
    # once most of rotation is removed I'll restart training, use SGD and higher learning rate with decay
    # probably try just l1 loss again
    # using l1_dice

# baseline zxing test on full train dataset:
    # 3759 11146 0.33725103176027277
# baseline zxing test on docv2 barcodes test dataset:
    # 27.2%
# baseline zxing test on more_barcodes test dataset:
    # 48 %

    # model 19 ->
    # test 75%
    # docv2 36%

# model 35
# test 65 %