import cv2
import glob
import os
import fitz
import json
import numpy as np
import pandas as pd
from time import time
import yaml
from tqdm.auto import tqdm
import xml.etree.ElementTree as ET
from sklearn.model_selection import StratifiedKFold

# from textdetection.src.serve_model import Predictor as TextDetector

FOLDER = "/mnt/ssd1T/tuanlv/06.KVUCombineStage/preprocess/data/invoices-receipts/SBT/nttmai_renamed/"
TXT_DIR = "/mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/pseudo_ocr/invoice_receipt_sbt"
OUT_FOLDER = "/mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/pseudo_ocr/sbt_batches"
ANN_OUT = "processed_textdet_batch%d.json"
N_BATCHES = 3

os.makedirs(OUT_FOLDER, exist_ok=True)


if __name__ == "__main__":
    all_files = []
    doc_types = []
    for file_name in tqdm(sorted(os.listdir(FOLDER))):
        try:
            file_name_no_ext = file_name.split('.')[0]
            all_files.append(file_name)
            doc_type = "1"
            doc_types.append(doc_type)
            file_path = os.path.join(FOLDER, file_name)

        except Exception as ex:
            print('image:', file_name, '. Error:', ex)

    df = pd.DataFrame({'file_name': all_files, 'doc_type': doc_types})

    df = df[df.doc_type.isin(['1'])].reset_index(drop=True)

    kfold = StratifiedKFold(n_splits=N_BATCHES)
    fold_num = 0
    for train_inds, val_inds in kfold.split(df, df['doc_type']):
        df.loc[val_inds, 'fold'] = fold_num
        fold_num+=1

    df.to_csv(f'{OUT_FOLDER}/tmp.csv', index=False)

    text_detector = TextDetector(setting['text_detection']['setting'], "textdetection")
    
    for batch in range(N_BATCHES):
        print(f"================== Batch {batch} ================")
        fold_df = df.loc[df.fold == batch]
        # tree = ET.parse(XML_IN)
        # root = tree.getroot()

        # for image in root.findall('image'):
        #     root.remove(image)
        with open('/home/sds/namnt/FWD_Data/coco_template.json', 'r') as f:
            coco_annotations = json.load(f)
        
        count = 2
        img_id = 1
        ann_id = 1

        all_images = []
        all_annotations = []

        for row_num, row in tqdm(fold_df.iterrows(), total=len(fold_df)):
            # try:
            file_name = row['file_name']
            file_name_no_ext = file_name.split('.')[0]
            doc_type = row['doc_type']
            file_path = os.path.join(FOLDER, file_name)
            images = pdf2np_fitz(file_path, _type='fname')
            
            images, batch_boxes = text_detector(images)
            for page_num, (img, boxes) in enumerate(zip(images, batch_boxes)):
                os.makedirs(os.path.join(OUT_FOLDER, f"batch{batch}"), exist_ok=True)
                out_img_path = os.path.join(OUT_FOLDER, f"batch{batch}", f"batch{batch}_{img_id:04d}_{file_name_no_ext}_{page_num}.jpg")
                cv2.imwrite(out_img_path, img[:,:,::-1])
                H, W = img.shape[:2]
                c_img = {
                    "id": int(img_id),
                    "width": W,
                    "height": H,
                    "file_name": os.path.join(f"batch{batch}", f"batch{batch}_{img_id:04d}_{file_name_no_ext}_{page_num}.jpg"),
                    "license": 0,
                    "flickr_url": "",
                    "coco_url": "",
                    "date_captured": 0
                }
                all_images.append(c_img)

                for box in boxes:
                    x1,y1,x2,y2 = box
                    w, h = x2-x1, y2-y1
                    c_ann = {
                        "id": int(ann_id),
                        "image_id": int(img_id),
                        "category_id": 1,
                        "segmentation": [],
                        "area": w*h,
                        "bbox": [x1,y1,w,h],
                        "iscrowd": 0,
                        "attributes": {
                            "occluded": False,
                            "rotation": 0.0
                        }
                    }
                    all_annotations.append(c_ann)
                    ann_id += 1
                
                img_id += 1
            # if count == 1:
            #     break
            # else:
            #     count -= 1
            # except Exception as ex:
            #     print('image:', file_name, '. Error:', ex)

        coco_annotations['categories'] = [{
            "id": 1,
            "name": "text",
            "supercategory": ""
        }]
        coco_annotations['images'] = all_images
        coco_annotations['annotations'] = all_annotations
        with open(os.path.join(OUT_FOLDER, ANN_OUT%(batch)), 'w') as f:
            json.dump(coco_annotations, f)

        # break
        

    