# Пример работы с IANPR2 из Python
import os
import cv2
import numpy as np
import json
import ianpr2class as ia2
import logging

logging.basicConfig(level=logging.DEBUG)

colorsD = [
    (0, 0, 255),
    (0, 255, 255),
    (0, 255, 0),
    (255, 0, 0),
    (255, 255, 0),
    (255, 0, 255)
]

objectNames = [ "number","person","car","moto","bus","truck" ]

def drawImage(img,outFilename,jsonResult,jsonAddResult,writefile = True):
    font = cv2.FONT_HERSHEY_SIMPLEX 
    numbers = []
    for j in jsonResult:        
        w = img.shape[1]
        if j.get("x") == None:
            continue
        x1 = j["x"]
        x2 = j["x"] + j["width"]
        y1 = j["y"]
        y2 = j["y"] + j["height"]
        numbers.append([int(x1),int(y1)])
        img = cv2.rectangle(img, (int(x1),int(y1)), (int(x2),int(y2)), (0,0,255),2)
        s = '{:.2f}'.format(j["score"])
        s += '|{}:{:.1f}{}'.format(j["strings"][0],j["confidence"][0],j["templates"][0])
        yt = int(y1)-10
        if yt < 30:
            yt = int(y2)+30
        textSize = cv2.getTextSize(s, fontFace=font, fontScale=1, thickness=2)
        if x1 + textSize[0][0] > w:
            x1 = w - textSize[0][0]
        img = cv2.putText(img,s,(int(x1),yt),font,1,(0,0,255),2)

    for j in jsonAddResult:
        if j.get("x") == None:
            continue
        x1 = j["x"]
        x2 = j["x"] + j["width"]
        y1 = j["y"]
        y2 = j["y"] + j["height"]
        img = cv2.rectangle(img, (int(x1),int(y1)), (int(x2),int(y2)), colorsD[j["classid"]%6],2)
        if j["imagenum"] > -1:            
            img = cv2.line(img,(int(x1),int(y1)),(numbers[j["imagenum"]]), colorsD[j["classid"]%6])

        s = objectNames[j["classid"]%6] + ":" + '{:.2f}'.format(j["score"])
        yt = int(y1)-10
        if yt < 30:
            yt = int(y2)+30
        textSize = cv2.getTextSize(s, fontFace=font, fontScale=1, thickness=2)
        if x1 + textSize[0][0] > img.shape[1]:
            x1 = img.shape[1] - textSize[0][0]
        img = cv2.putText(img,s,(int(x1),yt),font,1,colorsD[j["classid"]%6],2)

    if writefile:
        cv2.imwrite(outFilename,img)

def extractObjectsByClassid(image_array, json_data, target_classid=0):
    """
    Вырезает объекты с заданным classid из изображения
    """

    # Парсим JSON если строка
    if isinstance(json_data, str):
        data = json.loads(json_data)
    else:
        data = json_data
    
    extracted_images = []
    
    # Извлекаем список объектов
    objects = data.get("images", [[]])[0]
    
    for obj in objects:
        if obj.get("classid") == target_classid:
            x = obj["x"]
            y = obj["y"]
            width = obj["width"]
            height = obj["height"]
            score = obj.get("score", 0)
            
            # Вырезаем область (убеждаемся, что координаты в пределах изображения)
            h, w = image_array.shape[:2]
            x1 = max(0, x)
            y1 = max(0, y)
            x2 = min(w, x + width)
            y2 = min(h, y + height)
            
            # Вырезаем изображение
            cropped = image_array[y1:y2, x1:x2]
            
            if cropped.size > 0:  # Проверяем, что вырезанная область не пуста
                extracted_images.append(cropped)
    
    return extracted_images

# 1 изображение в пакете и возвращение дополнительного результата
def iANPR2Image(config, imgPath, code):  
    if ianpr2.initiANPR2Object(config,lic):
        if ianpr2.setSettings(settings):
            # Объект готов принимать изображения на распознавание                            
            img = cv2.imread(imgPath)
            if ianpr2.putImage(img):
                if ianpr2.inference():
                    s = ianpr2.getResult(code)
                    print(s) # JSON в строке
                    sAdd = ianpr2.getAddResult()
                    print(sAdd)
                    if (len(s) > 1):
                        d = json.loads(s)
                        add = [{}]
                        if (len(sAdd) > 1):
                            dAdd = json.loads(sAdd)
                            if dAdd.get("images"):
                                add = dAdd["images"][0]
                        if d.get("images"):
                            drawImage(img,"out.png",d["images"][0],add)

# 1 изображение First и Second шаги отдельно
def iANPR2ImageFS(config, imgPath, code):  
    if ianpr2.initiANPR2Object(config,lic):
        if ianpr2.setSettings(settings):
            # Объект готов принимать изображения на распознавание                            
            img = cv2.imread(imgPath)
            if ianpr2.putImage(img):
                if ianpr2.inferenceFirst():
                    sAdd = ianpr2.getAddResult()
                    print(sAdd)
                    objects = extractObjectsByClassid(img,sAdd)
                    for im in objects:
                        if ianpr2.putImage(im):
                            if ianpr2.inferenceSecond():
                                s = ianpr2.getResult(code)
                                print(s) # JSON в строке

# 4 изображения в пакете (только GPU)
def iANPR2Images4(config, imgPaths):  
    if ianpr2.initiANPR2Object(config,lic):
        if ianpr2.setSettings(settings):
            # Объект готов принимать изображения на распознавание                            
            imgs = []
            for p in imgPaths:
                img = cv2.imread(p)
                imgs.append(img)
                ianpr2.putImage(img)                    
            if ianpr2.inference():
                s = ianpr2.getResult()
                print(s) # JSON в строке
                if (len(s) > 1):
                    d = json.loads(s)
                    if d.get("images"):
                        for i in range(len(param1)):
                            drawImage(imgs[i],"out"+str(i)+".png",d["images"][i],[{}])

# видео
def iANPR2Video(config, videoPath,code):  
    if ianpr2.initiANPR2Object(config,lic):
        if ianpr2.setSettings(settings):
            # Видео
            cap = cv2.VideoCapture(videoPath)
            out = None
            while cap.isOpened():
                ret, frame = cap.read()         
                if ret:
                    if out is None:
                        out = cv2.VideoWriter('outvideo.avi',cv2.VideoWriter_fourcc('D','I','V','X'), 30, (frame.shape[1],frame.shape[0]))                     
                    image = frame.copy()        
                    if ianpr2.putImage(image):
                        if ianpr2.inference():
                            s = ianpr2.getResult(code)                            
                            if len(s) > 1:
                                print(s)
                                d = json.loads(s)
                                if d.get("images"):
                                    drawImage(image,"out.png",d["images"][0],[{}],False)
                            s = ianpr2.getMemResult(code)                            
                            if len(s) > 1:
                                print(s)
                                d = json.loads(s)
                                if d.get("numbers"):
                                    l = len(d["numbers"][0])
                                    font = cv2.FONT_HERSHEY_SIMPLEX 
                                    st = []
                                    textSize = []
                                    w = 0
                                    h = 0
                                    for i in d["numbers"][0]:
                                        st.append(i["string"]+":"+i["template"])
                                        textSize.append(cv2.getTextSize(st[-1], fontFace=font, fontScale=1, thickness=2))                                        
                                        if textSize[-1][0][0] > w:
                                            w = textSize[-1][0][0]
                                        h += textSize[-1][0][1]
                                    image = cv2.rectangle(image, (0,0), (w,h), (200,200,200),cv2.FILLED)
                                    for i in range(l):
                                        image = cv2.putText(image,st[i],(0,(i+1)*h//l),font,1,(0,0,0),2)
                     
                    cv2.imshow("iANPR2Video",image)
                    out.write(image)
                    c = cv2.waitKey(1)
                    if c & 0xFF == ord('q'):
                        break       
                else:
                    break


# видео траектория
def iANPR2VideoTrajectory(config, videoPath,code):  
    settings["mem"] = 0
    settings["templateCountries"]=["ru"]
    points = []
    if ianpr2.initiANPR2Object(config,lic):
        if ianpr2.setSettings(settings):
            # Видео
            cap = cv2.VideoCapture(videoPath)
            out = None
            while cap.isOpened():
                ret, frame = cap.read()         
                if ret:
                    if out is None:
                        out = cv2.VideoWriter('outvideo.avi',cv2.VideoWriter_fourcc('D','I','V','X'), 30, (frame.shape[1],frame.shape[0]))      
                        ianpr2.initTrajectory(frame.shape[1],frame.shape[0],10000, 100, 1)
                        points.append((0,int(frame.shape[0] * 0.4)))
                        points.append((frame.shape[1]-1,int(frame.shape[0] * 0.4)))
                        points.append((0,int(frame.shape[0] * 0.45)))
                        points.append((frame.shape[1]-1,int(frame.shape[0] * 0.45)))                                        
                        ianpr2.createLineIntersection(points[0][0],points[0][1],points[1][0],points[1][1],
                                                      points[2][0],points[2][1],points[3][0],points[3][1])
                    image = frame.copy()        
                    if ianpr2.putImage(image):
                        if ianpr2.inference():
                            sF = ianpr2.getResult(code) 
                            print(sF)
                            sAdd = ianpr2.getAddResult()
                            print("ADD=",sAdd)
                            if len(sF) > 1 or len(sAdd) > 1:                                
                                d = None
                                dAdd = None
                                if len(sF) > 1:
                                    d = json.loads(sF)   
                                if len(sAdd) > 1:
                                    dAdd = json.loads(sAdd)
                                add = [{}]
                                num = [{}]
                                if not dAdd is None and dAdd.get("images"): 
                                    add = dAdd["images"][0]
                                if not d is None and d.get("images"):
                                    num = d["images"][0]                                    
                                drawImage(image,"out.png",num,add,False)                            

                            ianpr2.trajectoryProcess(sF,sAdd,code)
                            trajRes = ianpr2.getTrajectoryResult(code)
                            print(trajRes)
                            

                            # Нарисовать траекторию
                            if len(trajRes) > 1:
                                d = json.loads(trajRes)
                                if d.get("trajectories"):   
                                    font = cv2.FONT_HERSHEY_SIMPLEX 
                                    st = []
                                    textSize = []
                                    w = 0
                                    h = 0
                                    l = len(d["trajectories"])
                                    for i,t in enumerate(d["trajectories"]):
                                        st.append(t["strNumber"] + ":" + t["templateNumber"] + ":" + t["zone"] + ":" + \
                                            str(t["weights"]) + "|ID:" + str(t["ID"]) + "|Class:" + str(t["classId"]))
                                        textSize.append(cv2.getTextSize(st[-1], fontFace=font, fontScale=1, thickness=2))                                        
                                        if textSize[-1][0][0] > w:
                                            w = textSize[-1][0][0]
                                        h += textSize[-1][0][1]
                                        # рисуем точки
                                        for j,p in enumerate(t["points"]):
                                            image = cv2.circle(image,(p[0],p[1]),3,(0,0,255),2)
                                            if j > 1:
                                                image = cv2.line(image,(p[0],p[1]),(t["points"][j-1][0],t["points"][j-1][1]),(0,0,255))
                                    image = cv2.rectangle(image, (0,0), (w,h), (200,200,200),cv2.FILLED)
                                    sInt = ""
                                    for i in range(l):
                                        image = cv2.putText(image,st[i],(0,(i+1)*h//l),font,1,(0,0,0),2)
                                        # Пересечение траекторий
                                        res = ianpr2.getLineIntersection(i)
                                        if res == 1:
                                            sInt += st[i] + ":up->down;";
                                        if res == 2:
                                            sInt += st[i] + ":down->up;";
                                    if sInt != "":
                                        cv2.putText(image,sInt,(0, int(frame.shape[0] * 0.9)),font,1,(0,0,255),2)

                            # Нарисовать линии                                 
                            image = cv2.line(image,points[0],points[1],(255,0,0)) 
                            image = cv2.line(image,points[2],points[3],(0,255,0)) 
                            

                                        
                     
                    cv2.imshow("iANPR2Video",image)
                    out.write(image)
                    c = cv2.waitKey(1)
                    if c & 0xFF == ord('q'):
                        break       
                else:
                    break

# Начало программы
if __name__ == "__main__":
    ianpr2 = ia2.iANPR2( './libianpr2.7.so' ) 


    # Настройки
    settings = {
        "detectConfThresh":0.25,
	    "minPlateHeight": 12,
	    "minPlateWidth": 60,	
	    "templateCountries":[
            "ru","kg","am","kz","by","az","tj","uz"
            ],
	    "turnNumber":1, # 0 or 1
	    "memoryNumberFrames":8,
	    "memoryNumberRepeat":4,
	    "minSymbolsPlate": 6,
	    "maxSymbolsPlate": 15,
	    "mem":1, # 0 or 1
        "expandedResult":1,# 0 or 1
	    "typeResultNumberStr":0, # 0-3
        "bigSymbols":1,
        "detectConfPerson":0.5,
        "detectConfVehicle":0.5
    }

    # Читаем лицензию
    lic = ""
    with open("lic2.key", "r") as f:
        # итерация по строкам
        for line in f:
            lic = line.strip()
            break

    # Распознавание 1 изображения - выдача расширенного результата с дополнительными объектами
    iANPR2Image("ianpr2dataconfig.json","image2.jpg","utf-8")
    #iANPR2Image("ianpr2dataconfig.json","084COAA.jpg","utf-8")    
    #iANPR2Image("ianpr2dataconfig.json","19039991.jpg","utf-8")
    #iANPR2Image("ianpr2dataconfig.json","19039991.jpg","cp1251")    
    #iANPR2ImageFS("ianpr2dataconfig.json","image2.jpg","utf-8")
    #iANPR2ImageFS("ianpr2dataconfig.json","image1.jpg","utf-8")
    #iANPR2ImageFS("ianpr2dataconfig.json","image3.jpg","utf-8")

    # Распознавание 4х изображений
    #iANPR2Images4("ianpr2dataconfigb4.json",["gr6.jpg","empty.jpg","gr6.jpg","image2.jpg"])

    # Распознавание видео-файла
    #iANPR2Video("ianpr2dataconfig.json","D:\\data\\auto\\video.mp4","cp1251")
    #iANPR2Video("ianpr2dataconfig.json","bypolicetest.avi","cp1251")
    #iANPR2Video("ianpr2dataconfig.json","736_12fps.mp4","cp1251")

    # Распознавание видео-файла вместе с траекторией
    #iANPR2VideoTrajectory("ianpr2dataconfig.json","d:/2.MP4","cp1251")
    #iANPR2VideoTrajectory("ianpr2dataconfig.json","D:/data/video/v1.avi","cp1251")
    #iANPR2VideoTrajectory("ianpr2dataconfig.json","bypolicetest.avi","utf-8")

