import os
import sys
import wget

from pathlib import Path
import json
import numpy as np
import torch
import cv2
from paddleocr import PaddleOCR

FILE = Path(__file__).resolve()
FILE_DIR = os.path.join(os.path.dirname(__file__))
ROOT = FILE.parents[0]
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
	sys.path.append(str(ROOT))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))


from ecos_core.base_model import BaseModel

class Model(BaseModel):
    NAME = "paddle_ocr"

    def __init__(self, opt):
        self.opt = opt
        self.opt_path = "./opt.json"
        self.device = "gpu:0" if torch.cuda.is_available() else "cpu"


    def __call__(self,x):
        return None

    def get_transform(self):
        """Get image data transform function for preprocessing

        Returns:
            transform (func): transform function.
                This function will take input as image path and output 
                - raw_image: image numpy mat
                - image_transform: image tensor (pytorch)
                Example: torchvision.transforms.transforms: data transforms function
        """
        def transform(image_path):
            return image_path, torch.tensor(0)
        return transform

    def build_model(self,custom_weight=None):
        """Build model in PaddleOCR model

        Args:
            custom_weight (str, optional): custom weight path. Defaults to None.

        Returns:
            model (Model): PaddleOCR model
        """
        text_detection_model_dir = self.opt.text_detection_model_dir or None
        text_recognition_model_dir = self.opt.text_recognition_model_dir or None
        self.model = PaddleOCR(
            device=self.device,
            text_detection_model_dir=text_detection_model_dir,
            text_recognition_model_dir=text_recognition_model_dir,
            use_doc_orientation_classify=self.opt.use_doc_orientation_classify,
            use_doc_unwarping=self.opt.use_doc_unwarping,
            use_textline_orientation=self.opt.use_textline_orientation)
        

        return self.model


    def update_opt(self, option, output_path="./opt.json"):
        """Update config

        Args:
            option (Namespace/Dict): option to update
            output_path (str, optional): path to save opt json. Defaults to "./opt.json".
        """
        for key, value in vars(option).items():
            if hasattr(self.opt, key):
                setattr(self.opt, key, value)

        # write to json
        with open(output_path, "w") as f:
            json.dump(vars(self.opt), f)
        
        print("Configuration updated successfully!")

    def reload_param(self, opt_path="./opt.json"):
        """Reload parameters

        Args:
            opt_path (str, optional): Opt path. Defaults to "./opt.json".
        """
        if not os.path.isfile(opt_path):
            opt_path = self.opt_path
        
        with open(opt_path, 'r') as f:
            data = json.load(f)
            for key, value in data.items():
                setattr(self.opt, key, value)

    def fit(self):
        """
        Train the model

        Returns:
            results: training results
        """

        return None
    
    def predict(self, source=None):
        """
        Run inference on images or folders.

        Args:
            source (str, optional): image path or folder path. Defaults to None.
            classes (list, optional): list of class ids to filter. Defaults to None.

        Returns:
            results: prediction results
        """
        text_det_thresh = self.opt.text_det_thresh or None
        text_rec_score_thresh = self.opt.text_rec_score_thresh or None
        return self.model.predict(
                input=source,
                text_det_thresh=text_det_thresh,
                text_rec_score_thresh=text_rec_score_thresh)

    def process_predictions(self, net_output, raw_image_path, raw_image_mat, image_transform, save_image_path):
        """Process predictions and generate annotated images.

        Args:
            net_output: output of detection net
            raw_image_path (str): raw image path
            raw_image_mat (numpy.Array): raw image numpy mat type
            image_transform (func): image_transform function
            save_image_path (str): save image path
            
        Returns:
            save_image_path:
                inform that the annotation image has been written successfully.
        """
        output = cv2.imread(raw_image_path)
        predict_index = 0
        boxes = []
        confs = []
        text = ""
        try:
            height, width = output.shape[:2]
            image = output
            resized_image =  cv2.resize(image, (0, 0), fx=0.5, fy=0.5)
            result = self.model.predict(
                input=resized_image)

            for res in result:
                res.print()

                pil_img = res.img['ocr_res_img']
                output = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)


                if text == "" and len(res['rec_texts']) > 0:
                    # get text with longest length
                    text = max(res['rec_texts'], key=len)


            os.makedirs(os.path.dirname(os.path.abspath(save_image_path)), exist_ok=True)
            

        except Exception as e:
            print(e)
            predict_index = 1

        cv2.imwrite(save_image_path, output)
        predict_class = self.opt.class_names[predict_index]
        save_txt_path = "{}.json".format(os.path.splitext(save_image_path)[0])

        with open(save_txt_path, 'w') as result_file:
            json.dump({
                "save_image_path": save_image_path,
                "score": {
                    "image_score": 0.25,
                },
                "predict_index": predict_index,
                "class_name": predict_class,
                "bbox": boxes,
                "confidence": confs,
                "content": text
            }, result_file)

        return save_image_path
