import logging
import os
import pickle
import random
from copy import deepcopy
from pathlib import Path

import cv2
import pandas as pd
import torch
from datasets import (Array2D, Array3D, ClassLabel, Dataset, Features,
                      Sequence, Value, concatenate_datasets)
from easydict import EasyDict
from PIL import Image
from tqdm import tqdm
from sdsvkie.utils.io_file import read_txt

from sdsvkie.utils import normalize_box, visualize_kie
from sdsvkie.utils.augmentation import perturbate_character, sampling_data
from sdsvkie.utils.word_formation import sliding_windows, sort_words
import glob 

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


IMG_EXTENSION = [".jpg", ".jpeg", ".png"]

class BaseDataset:
    def __init__(self, cfg):

        self.cfg = cfg
        self.feature_format = Features(
            {
                "image": Array3D(
                    dtype="int64", shape=(3, self.cfg.img_size, self.cfg.img_size)
                ),
                "input_ids": Sequence(feature=Value(dtype="int64")),
                "attention_mask": Sequence(Value(dtype="int64")),
                "bbox": Array2D(dtype="int64", shape=(self.cfg.max_seq_length, 4)),
                "labels": Sequence(ClassLabel(names=self.cfg.classes)),
            }
        ) if not self.cfg.v3 else \
            Features(
                {
                    "pixel_values": Array3D(
                        dtype="float32", shape=(3, self.cfg.img_size, self.cfg.img_size)
                    ),
                    "input_ids": Sequence(feature=Value(dtype="int64")),
                    "attention_mask": Sequence(Value(dtype="int64")),
                    "bbox": Array2D(dtype="int64", shape=(self.cfg.max_seq_length, 4)),
                    "labels": Sequence(ClassLabel(names=self.cfg.classes)),
                }
            )
        
        logger.info("Feature format: {}".format(self.feature_format.keys()))


    def _build_df(self, data_dir):
        """Build dataframe from data directory

        Args:
            data_dir (str): structure data folder
                - data_dir
                    - img1.jpg
                    - img1.txt
                    - ...
        """
        data_dir = Path(data_dir)

        # img_paths = glob.glob("*") + glob.glob("*/*")

        img_paths = [
            path for path in list(data_dir.rglob("*")) 
            
            if ".txt" not in str(path) and path.with_suffix(".txt").exists() and path.suffix.lower() in IMG_EXTENSION
        ]


        label_paths = [str(path.with_suffix(".txt")) for path in img_paths]


        img_paths = [str(path) for path in img_paths]
        

        
        assert len(label_paths) == len(img_paths)
        # remove empty txt
        ids = [id for id in range(len(label_paths)) if len(read_txt(label_paths[id])) > 0]
        label_paths = [label_paths[id] for id in ids]
        img_paths = [img_paths[id] for id in ids]

        dataframe = pd.DataFrame.from_dict(
            {"image_path": img_paths, "label": label_paths}
        )
        return dataframe

    def build_dataloader_from_dataset(
        self,
        dataset,
        processor,
        device,
        batch_size,
        shuffle=True,
        num_workers=4,
        cache_file="./cache.pkl",
        use_sampling=False,
    ):

        if not os.path.exists(cache_file):
            self._build_cache(
                dataset,
                processor,
                cache_file=cache_file,
                max_seq_length=self.cfg.max_seq_length,
            )
        cache = self._load_cache(cache_file)
        dataset = dataset.map(
            self._prepare_data,
            fn_kwargs={"cache": cache, "sampling": use_sampling},
            remove_columns=dataset.column_names,
            features=self.feature_format,
            batched=False,
            batch_size=self.cfg.batch_size,
        )

        dataset.set_format(type="torch", device=device)
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers
        )
        return dataloader

    def build_dataloader_from_dir(
        self,
        data_dir,
        processor,
        device,
        batch_size,
        shuffle=True,
        num_workers=4,
        cache_file="./cache.pkl",
        use_sampling=False,
    ):

        dataset = self._build_dataset(data_dir)
        dataloader = self.build_dataloader_from_dataset(
            dataset,
            processor,
            device,
            batch_size,
            shuffle,
            num_workers,
            cache_file,
            use_sampling,
        )
        return dataloader

    def _build_dataset(self, data_dir):
        df = self._build_df(data_dir)
        dataset = Dataset.from_pandas(df)
        logger.info(f"Load example for {data_dir}")
        dataset = dataset.map(lambda example: self._load_example_info(example))
        return dataset

    def _build_cache(self, dataset, processor, max_seq_length, cache_file=""):
        logger.info(f"Caching {cache_file}...")
        cache = {}
        for examples in tqdm(dataset):
            encoding_inputs = self._cache_feature(examples, processor, max_seq_length)
            cache[examples["image_path"]] = encoding_inputs

        with open(cache_file, "wb") as f:
            pickle.dump(cache, f)

    def _load_cache(self, cache_file):
        with open(cache_file, "rb") as f:
            cache = pickle.load(f)
        return cache

    def _prepare_data(self, example, cache, sampling):
        encoded_inputs_windows = cache[example["image_path"]]
        if len(encoded_inputs_windows) == 0:
            raise Exception("Empty encoded_inputs_windows")
        if sampling:
            if random.random() < 0.6:
                encoded_inputs = encoded_inputs_windows[-1]
            else:
                encoded_inputs = random.choice(encoded_inputs_windows)
        else:
            encoded_inputs = encoded_inputs_windows[-1]
        for k, v in encoded_inputs.items():
            if k in ["image", "pixel_values"]:
                encoded_inputs[k] = encoded_inputs[k][0]
        return encoded_inputs

    def _cache_feature(self, example, processor, max_seq_length=512):
        """Sampling (optional) + apply LayoutLMProcessor

        Args:
            examples (dict): dict {'image_path', 'words', 'bbox', 'word_labels'}
            max_seq_length (int, optional): _description_. Defaults to 512.

        Returns:
            list[dict]: list encoding inputs
        """

        image = Image.open(example["image_path"]).convert("RGB")
        batch_words = example["words"]
        batch_boxes = example["bbox"]
        batch_labels = example["word_labels"]

        window_size = self.cfg.max_num_words
        slice_interval = self.cfg.slice_interval

        word_windows = sliding_windows(batch_words, window_size, slice_interval)
        box_windows = sliding_windows(batch_boxes, window_size, slice_interval)
        label_windows = sliding_windows(batch_labels, window_size, slice_interval)

        encoded_inputs = []
        for words, boxes, labels in zip(word_windows, box_windows, label_windows):
            # Process examples
            encoded_input = processor(
                image,
                padding="max_length",
                truncation=True,
                text=words,
                boxes=boxes,
                word_labels=labels,
                max_length=max_seq_length,
            )

            encoded_inputs.append(encoded_input)

        # full page
        encoded_input = processor(
            image,
            padding="max_length",
            truncation=True,
            text=batch_words,
            boxes=batch_boxes,
            word_labels=batch_labels,
            max_length=max_seq_length,
        )
        encoded_inputs.append(encoded_input)

        return encoded_inputs

    def _load_example_info(self, example, aug_prob=0.0):
        """_summary_

        Args:
            example (_type_): _description_
            aug_prob (float, optional): _description_. Defaults to 0.0.
        """

        image_path = example["image_path"]
        label_path = example["label"]
        assert os.path.exists(image_path)
        assert os.path.exists(label_path)
        # try:
        image = cv2.imread(image_path)
        h, w, _ = image.shape
        with open(label_path) as f:
            lines = [line.replace("\n", "").replace("\r", "") for line in f.readlines()]

        words, boxes, labels = [], [], []
        # print(label_path)
        for i, line in enumerate(lines):
            x1, y1, x2, y2, text, label = line.split("\t")
            box = [int(x1), int(y1), int(x2), int(y2)]
            if text != " ":
                words.append(text)
                boxes.append(box)
                labels.append(label)

        if aug_prob > 0:
            p_words = perturbate_character(words, aug_prob)
            logging.info("{} - {}".format(len(p_words), len(words)))

        # custom for sort boxes
        items = {
            'boxes': boxes,
            'texts': words,
            'labels': labels
        }
        # boxes, words, labels = sort_words(boxes, words, labels)
        sorted_items = sort_words(items)
        boxes, words, labels = sorted_items['boxes'], sorted_items['texts'], sorted_items['labels']

        # print(image_path)
        # print(image_path)
        labels = [self.cfg.classes.index(label) for label in labels]
        if self.cfg.debug:
            visualize_kie(
                img=image,
                boxes=boxes,
                pred_labels=labels,
                outdir="wordirs/debug_{}".format(
                    "val" if "train" not in image_path else "train"
                ),
                image_name=os.path.basename(image_path),
            )

        boxes = [normalize_box(box, width=w, height=h) for box in boxes]

        example["words"] = words
        example["bbox"] = boxes  # TODO: Check this
        example["word_labels"] = labels
        # except Exception as err:
        #     logger.info(f"Exception: {err} at image path: {example['image_path']}")
        #     example["words"] = []
        #     example["bbox"] = []  # TODO: Check this
        #     example["word_labels"] = []
        return example
