import torch
import numpy as np
import torchvision.transforms as transforms
import math, yaml
from easydict import EasyDict as edict
from PIL import Image
import cv2
from torch.autograd import Variable
import time
import tensorrt as trt
def trt_version():
   return trt.__version__
 
def torch_device_from_trt(device):
   if device == trt.TensorLocation.DEVICE:
       return torch.device("cuda")
   elif device == trt.TensorLocation.HOST:
       return torch.device("cpu")
   else:
       return TypeError("%s is not supported by torch" % device)
 
 
def torch_dtype_from_trt(dtype):
   if dtype == trt.int8:
       return torch.int8
   elif trt_version() >= '7.0' and dtype == trt.bool:
       return torch.bool
   elif dtype == trt.int32:
       return torch.int32
   elif dtype == trt.float16:
       return torch.float16
   elif dtype == trt.float32:
       return torch.float32
   else:
       raise TypeError("%s is not supported by torch" % dtype)
 
def OcrTrtForward(engine,inputs,contextFlag=False):
    
    t0=time.time()
    #with engine.create_execution_context() as context:
    if not contextFlag: context = engine.create_execution_context() 
    else: context=contextFlag
    
    namess=[ engine.get_tensor_name(index) for index in range(engine.num_bindings) ]
    input_names = [namess[0]];output_names=namess[1:]

    batch_size = inputs[0].shape[0]
    bindings = [None] * (len(input_names) + len(output_names))
    t1=time.time() 
    # 创建输出tensor，并分配内存
    outputs = [None] * len(output_names)
    for i, output_name in enumerate(output_names):
        idx = engine.get_binding_index(output_name)#通过binding_name找到对应的input_id
        dtype = torch_dtype_from_trt(engine.get_binding_dtype(idx))#找到对应的数据类型
        shape = (batch_size,) + tuple(engine.get_binding_shape(idx))#找到对应的形状大小
        device = torch_device_from_trt(engine.get_location(idx))
        output = torch.empty(size=shape, dtype=dtype, device=device)
        #print('&'*10,'device:',device,'idx:',idx,'shape:',shape,'dtype:',dtype,' device:',output.get_device())
        outputs[i] = output
        #print('###line65:',output_name,i,idx,dtype,shape)
        bindings[idx] = output.data_ptr()#绑定输出数据指针
    t2=time.time()
 
    for i, input_name in enumerate(input_names):
        idx =engine.get_binding_index(input_name)
        bindings[idx] = inputs[0].contiguous().data_ptr()#应当为inputs[i]，对应3个输入。但由于我们使用的是单张图片，所以将3个输入全设置为相同的图片。
        #print('#'*10,'input_names:,', input_name,'idx:',idx, inputs[0].dtype,', inputs[0] device:',inputs[0].get_device())
    t3=time.time()
    context.execute_v2(bindings) # 执行推理
    t4=time.time()
        
    
    if len(outputs) == 1:
        outputs = outputs[0] 
    outstr='create Context:%.2f  alloc memory:%.2f  prepare input:%.2f   conext infer:%.2f, total:%.2f'%((t1-t0  )*1000  , (t2-t1)*1000,(t3-t2)*1000,(t4-t3)*1000, (t4-t0)*1000 )
    return outputs[0],outstr

def np_resize_keepRation(img,inp_h, inp_w):
        img_h, img_w = img.shape
        fy=inp_h/img_h
        keep_w = int(img_w* fy )
        Rsize=( keep_w , img_h)
        img = cv2.resize(img, Rsize  )
        #resize后是120,max是160,120-160的地方用边界的值填充
        if keep_w < inp_w:
            img_out = np.zeros((inp_h, inp_w ),dtype=np.uint8)
            img_out[:,:keep_w]=img[:,:]

            img_out[:,keep_w:] = np.tile(img[:,keep_w-1:], inp_w-keep_w)
        else:
            img_out = cv2.resize(img,(inp_w,inp_h))
        return img_out        
        
def recognition_ocr(config, img, model, converter, device,par={}):
    model_mode=par['model_mode'];contextFlag=par['contextFlag']
    if len(img.shape)==3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # github issues: https://github.com/Sierkinhane/CRNN_Chinese_Characters_Rec/issues/211
    h, w = img.shape
    # fisrt step: resize the height and width of image to (32, x)

    img = cv2.resize(img, (0, 0), fx=config.MODEL.IMAGE_SIZE.H / h, fy=config.MODEL.IMAGE_SIZE.H / h, interpolation=cv2.INTER_CUBIC)
    if model_mode=='trt':
        img = np_resize_keepRation(img,par['imgH'], par['imgW'])
    img = np.expand_dims(img,axis=2)


    # normalize
    img = img.astype(np.float32)
    img = (img / 255. - config.DATASET.MEAN) / config.DATASET.STD
    img = img.transpose([2, 0, 1])
    img = torch.from_numpy(img)

    img = img.to(device)
    img = img.view(1, *img.size())
    

    if model_mode=='trt':
        img_input = img.to('cuda:0')
        time2 = time.time()   
        preds,trtstr=OcrTrtForward(model,[img],contextFlag)
    else:
        model.eval()
        preds = model(img)
    _, preds = preds.max(2)
    preds = preds.transpose(1, 0).contiguous().view(-1)

    preds_size = Variable(torch.IntTensor([preds.size(0)]))
    sim_pred = converter.decode(preds.data, preds_size.data, raw=False)
    return sim_pred
class strLabelConverter(object):
    """Convert between str and label.

    NOTE:
        Insert `blank` to the alphabet for CTC.

    Args:
        alphabet (str): set of the possible characters.
        ignore_case (bool, default=True): whether or not to ignore all of the case.
    """

    def __init__(self, alphabet, ignore_case=False):
        self._ignore_case = ignore_case
        if self._ignore_case:
            alphabet = alphabet.lower()
        self.alphabet = alphabet + '-'  # for `-1` index

        self.dict = {}
        for i, char in enumerate(alphabet):
            # NOTE: 0 is reserved for 'blank' required by wrap_ctc
            self.dict[char] = i + 1

    def encode(self, text):
        """Support batch or single str.

        Args:
            text (str or list of str): texts to convert.

        Returns:
            torch.IntTensor [length_0 + length_1 + ... length_{n - 1}]: encoded texts.
            torch.IntTensor [n]: length of each text.
        """

        length = []
        result = []
        decode_flag = True if type(text[0])==bytes else False

        for item in text:

            if decode_flag:
                item = item.decode('utf-8','strict')
            length.append(len(item))
            for char in item:
                index = self.dict[char]
                result.append(index)
        text = result
        return (torch.IntTensor(text), torch.IntTensor(length))

    def decode(self, t, length, raw=False):
        """Decode encoded texts back into strs.

        Args:
            torch.IntTensor [length_0 + length_1 + ... length_{n - 1}]: encoded texts.
            torch.IntTensor [n]: length of each text.

        Raises:
            AssertionError: when the texts and its length does not match.

        Returns:
            text (str or list of str): texts to convert.
        """
        if length.numel() == 1:
            length = length[0]
            assert t.numel() == length, "text with length: {} does not match declared length: {}".format(t.numel(), length)
            if raw:
                return ''.join([self.alphabet[i - 1] for i in t])
            else:
                char_list = []
                for i in range(length):
                    if t[i] != 0 and (not (i > 0 and t[i - 1] == t[i])):
                        char_list.append(self.alphabet[t[i] - 1])
                return ''.join(char_list)
        else:
            # batch mode
            assert t.numel() == length.sum(), "texts with length: {} does not match declared length: {}".format(t.numel(), length.sum())
            texts = []
            index = 0
            for i in range(length.numel()):
                l = length[i]
                texts.append(
                    self.decode(
                        t[index:index + l], torch.IntTensor([l]), raw=raw))
                index += l
            return texts    


def get_alphabets(txtfile  ):
    print(txtfile)
    with open(txtfile,'r') as fp:
        lines=fp.readlines()
    alphas=[x.strip() for x in lines]    
    return "".join(alphas)
def get_cfg(cfg,char_file):
    with open(cfg, 'r') as f:
        #config = yaml.load(f)
        config = yaml.load(f, Loader=yaml.FullLoader)
        config = edict(config)
    config.DATASET.ALPHABETS = get_alphabets(char_file.strip()  )
    config.MODEL.NUM_CLASSES = len(config.DATASET.ALPHABETS)
    return   config   
def custom_mean(x):
    return x.prod()**(2.0/np.sqrt(len(x)))

def contrast_grey(img):
    high = np.percentile(img, 90)
    low  = np.percentile(img, 10)
    return (high-low)/np.maximum(10, high+low), high, low

def adjust_contrast_grey(img, target = 0.4):
    contrast, high, low = contrast_grey(img)
    if contrast < target:
        img = img.astype(int)
        ratio = 200./np.maximum(10, high-low)
        img = (img - low + 25)*ratio
        img = np.maximum(np.full(img.shape, 0) ,np.minimum(np.full(img.shape, 255), img)).astype(np.uint8)
    return img

class NormalizePAD(object):

    def __init__(self, max_size, PAD_type='right'):
        self.toTensor = transforms.ToTensor()
        self.max_size = max_size
        self.max_width_half = math.floor(max_size[2] / 2)
        self.PAD_type = PAD_type

    def __call__(self, img):
        img = self.toTensor(img)
        img.sub_(0.5).div_(0.5)
        c, h, w = img.size()
        Pad_img = torch.FloatTensor(*self.max_size).fill_(0)
        Pad_img[:, :, :w] = img  # right pad
        if self.max_size[2] != w:  # add border Pad
            Pad_img[:, :, w:] = img[:, :, w - 1].unsqueeze(2).expand(c, h, self.max_size[2] - w)

        return Pad_img

class AlignCollate(object):

    def __init__(self, imgH=32, imgW=100, keep_ratio_with_pad=False, adjust_contrast = 0.):
        self.imgH = imgH
        self.imgW = imgW
        self.keep_ratio_with_pad = keep_ratio_with_pad
        self.adjust_contrast = adjust_contrast

    def __call__(self, batch):
        #print('##recongnition.py  line72:  type(batch[0]):',type(batch[0]),batch[0], )
        batch = filter(lambda x: x is not None, batch)
        images = batch

        resized_max_w = self.imgW
        input_channel = 1
        transform = NormalizePAD((input_channel, self.imgH, resized_max_w))

        resized_images = []
        for image in images:
            w, h = image.size
            #### augmentation here - change contrast
            if self.adjust_contrast > 0:
                image = np.array(image.convert("L"))
                image = adjust_contrast_grey(image, target = self.adjust_contrast)
                image = Image.fromarray(image, 'L')

            ratio = w / float(h)
            if math.ceil(self.imgH * ratio) > self.imgW:
                resized_w = self.imgW
            else:
                resized_w = math.ceil(self.imgH * ratio)

            resized_image = image.resize((resized_w, self.imgH), Image.BICUBIC)
            resized_images.append(transform(resized_image))

        image_tensors = torch.cat([t.unsqueeze(0) for t in resized_images], 0)
        return image_tensors