import logging
import time

import cv2
import numpy as np
from sdsvkie.utils.word_formation import merge_boxes
import torch
from easydict import EasyDict
from PIL import Image

from sdsvkie.cfg import load_cfg
# from sdsvkie.models.layoutlm import LayoutLM
from sdsvkie.models.layoutlmv2 import LayoutLMv2
from sdsvkie.models.ocr import OCREngine
from sdsvkie.utils import invoice_postprocess  # invoice
from sdsvkie.utils import receipt_postprocess  # receipt
from sdsvkie.utils import POSTPROCESS_FUNC
from sdsvkie.utils import (Word, construct_word_groups_to_kie_label,
                           normalize_box, sliding_windows, sort_words,
                           unnormalize_box, words_to_lines,
                           pdf_to_image)


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Predictor:
    def __init__(self, cfg, **kwargs) -> None:
        """

        Args:
            cfg (dict / str): config
            **kwargs: device=..., weights=..., text_det=..., text_reg=...
        """

        if isinstance(cfg, str):
            cfg = load_cfg(cfg, kwargs)
        self.cfg = EasyDict(cfg)
        print(self.cfg)

        self.model = None
        self.processor = None
        self.ocr_engine = None
        self.classes = self.cfg["classes"]

        self.max_num_words = self.cfg.max_num_words
        self.model = None
        self.processor = None 
        self.ocr_engine = None

    def _init_predictor(self, model=None, proccessor=None, ocr_engine=None):

        if self.cfg.device != "cpu" and not torch.cuda.is_available():
            logger.info("Can not found cuda, training with CPU!!!")
            self.cfg.device = "cpu"

        if self.cfg["inference_weights"] is None:
            logger.info(
                "Not yet set value for inference weights, use weights instead!!!"
            )
        if model is None:
            self.model = LayoutLMv2._load_model(self.cfg)
        if proccessor is None:
            self.processor = LayoutLMv2._load_processor(self.cfg)

        if ocr_engine is None:
            self.ocr_engine = OCREngine(
                text_det=self.cfg["text_det"],
                text_recog=self.cfg["text_reg"],
                device=self.cfg["device"],
            )
        
        
    def __call__(self, input , ocr_output=None, return_raw=False):
        """Inference KIE -
        Pipeline:  Img -> OCR -> box + text (word-level) -> sort by x, y-axis -> LayoutLM -> Word formation -> result

        Args:
            input (np.ndarray / str / pdf path): BGR image (cv2)

        Returns:
            (dict): {
                'kie_raw_output': 
                'kie_post_output':
                'end2end_results': (dict) : {kie_label : value}
            }
        """
        if self.model is None:
            self._init_predictor()
        #check single or multi images
        if isinstance(input, np.ndarray):
            final_out= self.predict_single_image(input, ocr_output=ocr_output, return_raw=return_raw)
        elif isinstance(input, list):
            items = [self.predict_single_image(im, ocr_output=ocr_output, return_raw=return_raw) for im in input]
            final_out = self.aggregate_outputs(items)
        else:  #pdf
            import time 
            t1 = time.time()
            imgs = pdf_to_image(input)
            print("1. pdf2img: ", round(time.time() - t1, 4))
            t2 = time.time()
            items = [self.predict_single_image(im, ocr_output=ocr_output, return_raw=return_raw) for im in imgs]
            final_out = self.aggregate_outputs(items)
            print("2. kie: ", round(time.time() - t2, 4))
            print(f"3. full pipeline for {len(imgs)} pages: {round(time.time() - t1, 4)}")
        return final_out

    def predict_single_image(self, img, ocr_output=None, return_raw=False):

        if ocr_output is None:
            ocr_output = self.ocr_engine(img, extend_ratio=[0.1, 0.3], ratio_thr=5)   # solve long box
            # ocr_output = self.ocr_engine(img)
        kie_input = self.prepare_inputs(img, ocr_output)
        kie_output = self.predict(kie_input)
        kie_post_output = self.postprocessing(kie_output)
        formated_output = self.format_output(kie_post_output)
        output = {
            "kie_raw_output": kie_output if return_raw else None,  # raw output from layoutlm model
            "kie_post_output": kie_post_output if return_raw else None,  # wordgroup
            "end2end_results": formated_output,                         # field_key + field_value
        }
        return output

    def predict(self, inputs: dict):
        """predict

        Args:
            inputs (dict): format
                {
                    'image": PIL RGB,
                    'boxes': list[],
                    'texts'
                }

        Returns:
            list[Word]: list of Word object
        """
        window_size = self.cfg.max_num_words
        slice_interval = self.cfg.slice_interval
        
        image, batch_boxes, batch_words = (
            inputs["img"],
            inputs["boxes"],
            inputs["texts"],
        )
        results = []
        non_norm_boxes = inputs['non_norm_boxes']
        if len(batch_boxes) == 0:
            logger.info("Not found any words in image!!! Continue...")
            return results
        
        text_windows = sliding_windows(batch_words, window_size, slice_interval)
        box_windows = sliding_windows(batch_boxes, window_size, slice_interval)
        # print([len(t) for t in text_windows])

        out_boxes_windows = []
        # out_labels_windows = []
        out_logits_windows = []

        for i in range(len(text_windows)):
            words = text_windows[i]  # len: MAX_N_WORDS
            boxes = box_windows[i]

            # Preprocess
            dummy_word_labels = [0] * len(words)
            encoding = self.processor(
                image,
                text=words,
                boxes=boxes,
                word_labels=dummy_word_labels,
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=self.cfg.max_seq_length,
            )

            label_ = encoding.pop('labels')
            # Run model
            for k, v in encoding.items():
                encoding[k] = v.to(self.cfg.device)
            with torch.no_grad():
                output = self.model(**encoding)


            logits = output.logits.squeeze()  # seq_len * classes
            predictions = output.logits.argmax(-1).squeeze().tolist()
            token_boxes = encoding.bbox.squeeze().tolist()

            # Postprocess
            # is_subword = (encoding["labels"] == -100).detach().cpu()[0]
            is_subword = (label_ == -100).detach().cpu()[0]
            logit_predictions = logits[torch.logical_not(is_subword), :]
            true_boxes = torch.Tensor(
                [
                    unnormalize_box(box, image.size[0], image.size[1])
                    for idx, box in enumerate(token_boxes)
                    if not is_subword[idx]
                ]
            )

            # print("logit shape: ", logit_predictions.shape)
            # print("box shape: ", true_boxes.shape)
            # true_boxes = torch.Tensor(
            #     [
            #         unnormalize_box(box, image.size[0], image.size[1])
            #         for idx, box in enumerate(boxes)
            #     ]
            # )
            out_boxes_windows.append(true_boxes)
            out_logits_windows.append(logit_predictions)

        # merge output by average logits between overlap window
        merged_out_boxes = out_boxes_windows[0]
        merged_out_logit = out_logits_windows[0]
        overlap = window_size - slice_interval 
        for i in range(1, len(out_boxes_windows)):
            if overlap != 0:
                prev_overlap_logits = merged_out_logit[-overlap:, :]  
                curr_overlap_logits = out_logits_windows[i][:overlap, :]  
                avg_overlap_logits = (
                    prev_overlap_logits + curr_overlap_logits
                ) / 2  
                curr_logits = torch.cat(
                    [avg_overlap_logits, out_logits_windows[i][overlap:, :]], dim=0
                )
                merged_out_logit = torch.cat(
                    [merged_out_logit[:-overlap, :], curr_logits], dim=0
                )
            else:
                merged_out_logit = torch.cat(
                    [merged_out_logit, out_logits_windows[i]], dim=0
                )

            merged_out_boxes = torch.cat(
                [merged_out_boxes, out_boxes_windows[i][overlap:, :]], dim=0
            )

        # print(f"merged_out_logit={len(merged_out_logit)} - merged_out_boxes={len(merged_out_boxes)}")
        # from IPython import embed; embed()
        assert len(merged_out_logit) == len(
            merged_out_boxes
        ), f"{len(merged_out_logit)} # {len(merged_out_boxes)}"
        predictions = merged_out_logit.argmax(-1).squeeze().tolist()
        if not isinstance(predictions, list):
            predictions = [predictions]


        assert len(predictions) == len(batch_words), f"{len(predictions)} # {len(batch_words)}"
        # for word_index in range(len(batch_words)):
        for word_index, word in enumerate(batch_words):
            word = batch_words[word_index]
            # bndbox = [int(coord) for coord in merged_out_boxes[word_index]]
            bndbox = non_norm_boxes[word_index]
            kie_label = self.cfg.classes[predictions[word_index]]
            results.append(
                Word(
                    text=word,
                    bndbox=bndbox,
                    kie_label=kie_label,
                    conf_cls=inputs["recog_confs"][word_index]
                )
            )

        return results

    def postprocessing(self, words):
        """Post processing for output of kie
        - Merge wordgroup list for each field

        Args:
            items (dict): _description_
        """
        list_lines, _ = words_to_lines(words)
        list_word_group = []
        for line in list_lines:
            for word_group in line.list_word_groups:
                word_group.update_kie_label()
                word_group.update_conf()
                list_word_group.append(word_group)

        kie_dict = construct_word_groups_to_kie_label(list_word_group)
        #receipt postprocess
        if 'postprocess_type' in self.cfg and self.cfg.postprocess_type == "receipt_postprocess":
            kie_dict = receipt_postprocess(kie_dict, words) 
        else:  #invoice_postprocess
            # kie_dict = invoice_postprocess(kie_dict)
            kie_dict = self._postprocess_kie_wordgroups(kie_dict, doc_type=self.cfg.postprocess_type)
        return kie_dict

    def format_output(self, kie_dict):
        """
        Args:
            kie_dict (dict) : format 
                {
                    'field_name': list[Wordgroup]
                }

        Returns:
            (dict): format 
                {
                    'field_name': {
                        'box': list,
                        'value': str,
                        'conf': float
                    }
                }
        """
        end2end_results = {}
        filtered_dict = {k:v for k,v in kie_dict.items() if "key" not in k}
        for field_name, wg_list in filtered_dict.items():
            wg_list = [wg_list] if not isinstance(wg_list, list) else wg_list

            if len(wg_list) == 0:
                text, conf, box = "", 0.0, []
            else:
                text = " ".join([wg.text for wg in wg_list]).strip().replace("✪", " ")
                conf = sum(wg.conf for wg in wg_list) / len(wg_list)
                box = merge_boxes([wg.boundingbox for wg in wg_list])

            end2end_results[field_name] = {
                "box": box,
                "value": text, 
                "conf": conf
            }
        
        # add empty values for missing fields
        for class_name in self.classes:
            if "key" not in class_name and class_name not in end2end_results and class_name.lower() not in ['other', 'others']:
                end2end_results[class_name] = {
                    "box": [],
                    "value": "", 
                    "conf": 0.0
                }

        end2end_results = self._postprocess_recognized_text(end2end_results, doc_type=self.cfg.postprocess_type)
        # sort by keys
        end2end_results = dict(sorted(end2end_results.items()))
        return end2end_results
    
    
    def _postprocess_kie_wordgroups(self, result, doc_type, metadata=None):
        """post process for wordgroup outputs

        Args:
            result (dict): {'field_name': list[Wordgroup]}
            doc_type (str): invoice / receipt
            metadata (_type_, optional): _description_. Defaults to None.

        Returns:
            _type_: _description_
        """
        for field_name in result.keys():
            if field_name not in POSTPROCESS_FUNC[doc_type]:
                continue
            postprocess_func = POSTPROCESS_FUNC[doc_type][field_name].get("wordgroup", None)
            if postprocess_func is None:
                continue
            result[field_name] = postprocess_func(result[field_name], metadata={"field_name": field_name, "wg_res": result})
        return result
    
    def _postprocess_recognized_text(self, result, doc_type, metadata=None):
        for field_name in result.keys():
            if field_name not in POSTPROCESS_FUNC[doc_type]:
                continue
            postprocess_func = POSTPROCESS_FUNC[doc_type][field_name].get("text", None)
            if postprocess_func is None:
                continue
            result[field_name]["value"] = postprocess_func(result[field_name]['value'], metadata)
        return result

    def prepare_inputs(self, img, ocr_output):
        """Prepare input for KIE  model

        Args:
            img (np.ndarray): BGR image
            ocr_output (dict): format
                {
                    "img_path": img_path,
                    "img": image,
                    "boxes": boxes,
                    "texts": words,
                    "kie_labels": word_labels
                }
        """
        assert "boxes" in ocr_output, "boxes not exists in ocr_output"
        assert "texts" in ocr_output, "texts not exists in ocr_output"

        # cv2 to PIL (RGB)
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        w, h = img.size
        texts = ocr_output["texts"]
        boxes = ocr_output["boxes"]
        recog_confs = ocr_output['recog_confs']
        det_confs = ocr_output['det_confs']
        texts = [text.replace(" ", "✪") for text in texts]   # layoutlm will throw an error if the input has space characters
        word_items = {
            'boxes': boxes,
            'texts': texts,
            'det_confs': det_confs,
            'recog_confs': recog_confs
        }
        sorted_word_items = sort_words(word_items)
        (boxes, texts, det_confs, recog_confs) = (
            sorted_word_items['boxes'], 
            sorted_word_items['texts'], 
            sorted_word_items['det_confs'], 
            sorted_word_items['recog_confs']
        )
        non_norm_boxes = sorted_word_items['boxes']
        boxes = [normalize_box(box, width=w, height=h) for box in boxes]
        out_item = {
            "img": img, 
            "boxes": np.array(boxes), 
            "texts": texts,
            "det_confs": det_confs,
            "recog_confs": recog_confs,
            "non_norm_boxes": non_norm_boxes
        }
        return out_item


    
    def aggregate_outputs(self, outs):
        f"""Postprocess the outputs of the muliple pages

        Args:
            outs (_type_): _description_
            
        Returns:
            output: (dict): final output 
            
        """
        combine_out = {
            'kie_raw_output': [],
            'kie_post_output': [],
            'end2end_result_each_page': [],
            'end2end_results': None
        }
        
        for idx, out in enumerate(outs):
            combine_out['kie_raw_output'].append(out['kie_raw_output'])
            combine_out['kie_post_output'].append(out['kie_post_output'])
            combine_out['end2end_result_each_page'].append(out['end2end_results'])
            
        #merge end2end result
        end2end_results = combine_out['end2end_result_each_page'][0]
        
        for page_id, end2end_results_page in enumerate(combine_out['end2end_result_each_page'][1:]):
            for field_key, field_value in end2end_results_page.items():
                if "value" in end2end_results[field_key] \
                  and (end2end_results[field_key]['value'] == "" or end2end_results[field_key]['value'] == "0"):
                    end2end_results[field_key] = field_value
        combine_out['end2end_results'] = end2end_results
        
        return combine_out
        