from mrz_decoder.config import Config
import numpy as np
import torch
import torchvision
import cv2
import itertools
import PIL.Image


class MRZProcessingPipeline:
    def __init__(self):
        super().__init__()

    def _filter_character_list(self, bounds, peaks=44):
        scored = []
        for (x1, x2) in bounds:
            region_width = x2 - x1
            scored.append(((x1, x2), region_width))
        median = np.median([score for _, score in scored])
        scored_sorted = sorted(scored, key=lambda x: abs(x[1] - median))
        top_bounds = [b[0] for b in scored_sorted[:peaks]]
        return sorted(top_bounds, key=lambda x: x[0])
    
    def _find_character_bounds(self, projection, min_width=2, threshold=5):
        bounds = []
        in_char = False
        start = 0
        for i, val in enumerate(projection):
            if not in_char and val > threshold:
                in_char = True
                start = i
            elif in_char and val <= threshold:
                in_char = False
                if i - start >= min_width:
                    bounds.append((start, i))
        bounds = self._filter_character_list(bounds)
        return bounds
    
    def grid_image(self, image, crop=False):
        if crop:
            image = self.process_mrz(image)
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        binary_inv = 255 - binary
        projection = np.sum(binary_inv > 0, axis=0)
        bounds = self._find_character_bounds(projection, min_width=5, threshold=5)
        line1 = []
        line2 = []
        padding = 5
        for x1, x2 in bounds:
            x_min = max(x1-padding, 0)
            x_max = min(x2+padding, image.shape[1])
            col = image[:, x_min:x_max]
            h, _, _ = col.shape
            line1.append(col[:h//2, :])
            line2.append(col[h//2:, :])
        return line1 + line2, None


class MRZDecoder:
    def __init__(self, model, device='cpu'):
        self.processing_pipeline = MRZProcessingPipeline()
        self.transform = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Resize([28, 28])
        ])
        self.device = torch.device(device)
        self.model = model.eval().to(self.device)
        self.model_config = Config
    
    def _process_mrz(self, image):
        image = image.convert('RGB')
        image = np.array(image)
        cells, image = self.processing_pipeline.grid_image(image, crop=False)
        cells = torch.stack([self.transform(c) for c in cells]).to(self.device)
        return cells
    
    def _label_correction(self, label):
        label = list(label.replace(' ', ''))
        changing = True
        while changing:
            changing = False
            for i in range(len(label)):
                try:
                    if (i > 21 or i > 65) and (label[i] != '<') and (label[i-1] == '<' or label[i-1] == 'K') and (label[i+1] == '<' or label[i+1] == 'K'):
                        label[i] = '<'
                        changing = True
                    if (label[i] != '<') and (label[i-1] == '<') and (label[i+1] == '<'):
                        label[i] = '<'
                        changing = True
                except Exception as e:
                    print(e)
                    continue
        if label[43] == 'K':
            label[43] = '<'
        return ''.join(label)

    def fuzzy_decode(self, predictions, thresh=0.9):
        predictions = torch.softmax(predictions, dim=1)
        topk_probs, topk_classes = torch.topk(predictions, 2, dim=1) # this takes the 2 most likely classes per image
        class_variations = []
        for i, (probs, classes) in enumerate(zip(topk_probs, topk_classes)):
            if probs[0] < thresh and i > 44:
                val = [list(self.model_config.character_lookup.keys())[p.item()].upper() for p in classes]
                class_variations.append(val)
            else:
                val = [list(self.model_config.character_lookup.keys())[classes[0].item()].upper()]
                class_variations.append(val)
        variants = 2 ** len([x for x in class_variations if len(x) > 1])
        combinations = []
        if variants < 5000:
            combinations = list(itertools.product(*class_variations))
            combinations = [''.join(comb) for comb in combinations]
            for raw_decoding in combinations:
                fields, checks = self.parse_mrz(raw_decoding)
                validation = self.validate(raw_decoding, fields, checks)
                if validation:
                    return validation, raw_decoding, fields, checks
        predictions = torch.argmax(predictions, dim=1)
        predictions = [list(self.model_config.character_lookup.keys())[p].upper() for p in predictions]
        raw_decoding = ''.join(predictions)
        fields, checks = self.parse_mrz(raw_decoding)
        validation = self.validate(raw_decoding, fields, checks)
        return validation, raw_decoding, fields, checks
        

    def decode_mrz(self, image: PIL.Image.Image, fuzzy=False):
        output = {'Valid': False, 'Raw': '', 'Parsed': {}}
        try:
            cells = self._process_mrz(image)
            self.cells = cells
            assert len(cells) == 88
        except Exception as e:
            output['Error'] = e
            return output
        predictions = self.model(cells)
        if fuzzy:
            validation, raw_decoding, fields, checks = self.fuzzy_decode(predictions)
        else:
            predictions = torch.argmax(predictions, dim=1)
            predictions = [list(self.model_config.character_lookup.keys())[p].upper() for p in predictions]
            raw_decoding = ''.join(predictions)
            fields, checks = self.parse_mrz(raw_decoding)
            try:
                validation = self.validate(raw_decoding, fields, checks)
            except:
                validation = False
        return {'Valid': validation, 'Raw': raw_decoding, 'Parsed': fields, 'Checks': checks}
    
    def parse_mrz(self, mrz: str):
        parsed_fields = {}
        parsed_checks = {}
        if len(mrz) < 88:
            return parsed_fields, parsed_checks
        line1 = mrz[:44]
        line2 = mrz[44:]
        parsed_fields['Document Type'] = line1[0]
        parsed_fields['Issuer'] = line1[2:5]
        names_raw = line1[5:].split('<<', 1)
        parsed_fields['Surname'] = names_raw[0].replace('<', '')
        parsed_fields['Given Names'] = names_raw[1].replace('<', ' ').strip() if len(names_raw) > 1 else ''
        parsed_fields['Document Number'] = line2[:9].replace('<', '')
        parsed_checks['Document Number'] = line2[9]
        parsed_fields['Nationality'] = line2[10:13].replace('<', '')
        parsed_fields['Date of Birth'] = line2[13:19].replace('<', '')
        parsed_checks['Date of Birth'] = line2[19]
        parsed_fields['Sex'] = line2[20]
        parsed_fields['Expiration Date'] = line2[21:27].replace('<', '')
        parsed_checks['Expiration Date'] = line2[27]
        parsed_fields['Personal Number'] = line2[28:42].replace('<', '')
        if len(parsed_fields['Personal Number']) > 0:
            parsed_checks['Personal Number'] = line2[42]
        parsed_checks['Overall'] = line2[43]
        return parsed_fields, parsed_checks
    
    def validate(self, raw_decoding, fields, checks):
        if len(set(raw_decoding)) < 3:
            return False
        for item in fields:
            if item in checks:
                if checks[item] != self.check_digit(fields[item]):
                    return False
        overall_check = raw_decoding[44:][:10] + raw_decoding[44:][13:20] + raw_decoding[44:][21:43]
        if self.check_digit(overall_check) != checks['Overall']:
            return False
        return True
    
    def check_digit(self, check):
        weights = [7, 3, 1]
        value = 0
        for i, c in enumerate(check):
            digit = int(c) if c.isdigit() else ord(c) - 55
            if c == '<':
                digit = 0
            value += digit * weights[i % 3]
        return str(value % 10)