import numpy as np

from det.mlt_det import MLT_Det
from det.utils import get_word_pic_acd_box,draw_text
from rec.sp_rec import SP_Rec
import cv2
import mlt_cfg as cfg
from rec.sp_rec import SP_Rec
from rec.language_cls import LangCls
class MLT_Rec:
    def __init__(self):
        self.Latin_model=SP_Rec(cfg.latin_config,cfg.latin_checkpoint)
        self.Korean_model=SP_Rec(cfg.korean_config,cfg.korean_checkpoint)
        self.Hindi_model=SP_Rec(cfg.hindi_config,cfg.hindi_checkpoint)
        self.Chinese_model = SP_Rec(cfg.chinese_config, cfg.chinese_checkpoint)
        self.LangCls=LangCls(cfg.lang_checkpoint)
    def __call__(self, ims):
        if cfg.lang is not None:
            langs=[cfg.lang]*len(ims)
        else:
            langs=self.LangCls(ims)
        index_dict=dict(latin=[],hindi=[],korean=[],chinese=[])
        latin_ims=[]
        korean_ims=[]
        hindi_ims=[]
        chinese_ims=[]
        for i in range(len(langs)):
            if langs[i]=='Latin':
                latin_ims.append(ims[i])
                index_dict['latin'].append(i)
            elif langs[i]=='Hindi':
                hindi_ims.append(ims[i])
                index_dict['hindi'].append(i)
            elif langs[i]=='Korean':
                korean_ims.append(ims[i])
                index_dict['korean'].append(i)
            elif langs[i]=='Chinese':
                chinese_ims.append(ims[i])
                index_dict['chinese'].append(i)
        latin_res=self.Latin_model.batch_predict(latin_ims)
        hindi_res = self.Hindi_model.batch_predict(hindi_ims)
        korean_res = self.Korean_model.batch_predict(korean_ims)
        chinese_res = self.Chinese_model.batch_predict(chinese_ims)
        all_res=['']*len(ims)
        for i in range(len(index_dict['latin'])):
            all_res[index_dict['latin'][i]]=latin_res[i]
        for i in range(len(index_dict['hindi'])):
            all_res[index_dict['hindi'][i]]=hindi_res[i]
        for i in range(len(index_dict['korean'])):
            all_res[index_dict['korean'][i]]=korean_res[i]
        for i in range(len(index_dict['chinese'])):
            all_res[index_dict['chinese'][i]]=chinese_res[i]
        return all_res,langs


class E2E:
    def __init__(self):
        self.det_model=MLT_Det(cfg.det_config,cfg.det_checkpoint)
        self.rec_model=MLT_Rec()

    def __call__(self, im):
        polys = self.det_model.predict_polys(im)
        crops = []
        for poly in polys:
            crop_img = get_word_pic_acd_box(im, poly)
            crops.append(crop_img)
        rec_res,langs=self.rec_model(crops)
        return polys,rec_res,langs
    def viz(self,im):
        polys,rec_res,langs=self(im)
        polys=np.array(polys,dtype=np.int32)
        polys=polys.reshape(-1,4,2)
        im=cv2.polylines(im,polys,True,(255,0,0),thickness=2)
        text_im=draw_text(im,polys,rec_res,langs)
        return im,text_im

if __name__=='__main__':
    model=E2E()
    im = cv2.imread(
        '/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/PublicDataSet/OCR/hccdoc/image/2000/002018.jpg')
    im,text_im=model.viz(im)
    # im=cv2.resize(im,(512,512))
    # text_im=cv2.resize(text_im,(512,512))
    cv2.imshow('a',im)
    cv2.imshow('b',text_im)
    cv2.waitKey(0)