import numpy as np
from PIL import Image, ImageEnhance
import torchvision
import torch
import zxingcpp
from barcode_reader import unet, adder
import importlib.resources
from autocapture_model import detector


class AdderDecoder:
    def __init__(self, model_path=None, num_nets=5, magnitude=16, cuda=False, cropped=False):
        self.untransform = torchvision.transforms.ToPILImage()
        self.image_size = (256, 1024)
        if cuda:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device('cpu')
        if model_path is not None:
            add = adder.AdderDiffusion(num_nets=num_nets, model_path=model_path, pretrained=False, magnitude=magnitude)
        else:
            path = importlib.resources.open_binary('barcode_reader.models', 'adder.pytorch.model')
            add = adder.AdderDiffusion(num_nets=num_nets, model_path=path, magnitude=magnitude, cuda=cuda)
        self.model = add.model
        self.model.eval().to(self.device)
        self.cropped = cropped
        if not self.cropped:
            self.document_detector = detector.DocumentDetector()

    def extract_barcode(self, im, box=None):
        if type(im) is torch.Tensor:
            im = torchvision.transforms.ToPILImage()(im)
        im_arr = np.array(im)
        if self.cropped:
            return im
        if box is None:
            result = self.document_detector.detect_document(name='test', image=im)
            if result[0] is True:
                if 'barcode' in result[1]['raw_result'].keys():
                    box = result[1]['raw_result']['barcode']['original_box']
                else:
                    return None
        padding = int((0.15) * (box[3] - box[1]))
        pad_x, pad_y = padding, padding
        return im_arr[box[1]-pad_y:box[3]+pad_y, box[0]-pad_x:box[2]+pad_x]

    def _input_transform(self, im):
        pad_top, pad_left, pad_bottom, pad_right = 10, 15, 10, 15
        pad_top, pad_left, pad_bottom, pad_right = 0, 0, 0, 0
        transform = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            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)
        ])
        return transform(im)

    def decode(self, im, box=None):
        im = self.extract_barcode(im, box)
        if im is None:
            return None
        im = self._input_transform(im)
        self.input_image = im
        im = im.unsqueeze(0).to(self.device)
        with torch.no_grad():
            out = self.model(im)
        out = out[0]
        out = torch.clamp(out, min=0, max=1)
        out = self.untransform(out)
        self.output_image = out
        result = zxingcpp.read_barcode(out)
        if result is not None:
            return result.text
        return None
