from concurrent.futures import ThreadPoolExecutor
from functools import partial

import cv2
import numpy as np

from mower.models import avatar, portrait
from mower.utils import config
from mower.utils import typealias as tp
from mower.utils.image import cropimg, thres2
from mower.utils.log import logger
from mower.utils.matcher import GOOD_DISTANCE_LIMIT, flann, keypoints
from mower.utils.vector import in_scope, va


def segment_room_select(img: tp.Image) -> list[tp.Scope]:
    "基建房间内选干员"
    line1 = cropimg(img, ((600, 519), (1920, 520)))
    hsv = cv2.cvtColor(line1, cv2.COLOR_RGB2HSV)
    mask = cv2.inRange(hsv, (98, 140, 200), (102, 255, 255))
    line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
    line1[mask > 0] = (255,)
    line1 = thres2(line1, 140)

    last_line = line1[-1]
    prev = last_line[0]
    start = None
    name_x = []
    for i in range(1, line1.shape[1]):
        curr = last_line[i]
        if prev == 0 and curr == 255 and start and i - start > 186:
            name_x.append((i + 415, i + 585))
        elif prev == 255 and curr == 0:
            start = i
        prev = curr

    avatar_y = ((205, 320), (625, 740))
    avatar_p = []
    for x in name_x:
        for y in avatar_y:
            avatar_p.append(tuple(zip(x, y)))

    logger.debug(avatar_p)
    return avatar_p


def segment_team(img: tp.Image) -> list[tp.Scope]:
    "编队界面"
    # TODO: 利用img判断编队缺人的情况
    result = []
    for i in range(6):
        x = 283 + 232 * i
        for y in [204, 637]:
            result.append(((x, y), va((x, y), (180, 110))))
    logger.debug(result)
    return result


def segment_team_select(img: tp.Image) -> list[tp.Scope]:
    "作战编队和训练位选干员"
    line1 = cropimg(img, ((540, 513), (1920, 514)))
    hsv = cv2.cvtColor(line1, cv2.COLOR_RGB2HSV)
    mask = cv2.inRange(hsv, (98, 150, 200), (102, 255, 255))
    line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
    line1[mask > 0] = (255,)
    line1 = thres2(line1, 140)
    last_line = line1[-1]
    prev = last_line[0]
    start = None
    name_x = []
    for i in range(1, line1.shape[1]):
        curr = last_line[i]
        if prev == 255 and curr == 0 and start and i - start > 90:
            name_x.append((i + 450, i + 630))
        elif prev == 0 and curr == 255:
            start = i
        prev = curr

    avatar_y = ((200, 310), (620, 725))
    avatar_p = []
    for x in name_x:
        for y in avatar_y:
            avatar_p.append(tuple(zip(x, y)))

    logger.debug(avatar_p)
    return avatar_p


def match_operator(
    gray: tp.GrayImage, segment: list[tp.Scope], model: list
) -> tuple[tuple[str, tp.Scope]]:
    if len(segment) == 0:
        logger.debug(f"{len(segment)=}")
        return ()

    avatar_p = [[None, 0]] * len(segment)

    mask = np.zeros(gray.shape, dtype=np.uint8)
    for pt1, pt2 in segment:
        cv2.rectangle(mask, pt1, pt2, (255,), -1)

    ori_kp, ori_des = keypoints(gray, mask)

    def match(item):
        name, des = item
        return name, flann.knnMatch(des, ori_des, k=2)

    if config.conf.max_workers > 1:
        with ThreadPoolExecutor(max_workers=config.conf.max_workers) as executor:
            result = executor.map(match, model)
    else:
        result = [match(i) for i in model]

    for name, matches in result:
        good = []
        for pair in matches:
            if len(pair) == 2:
                x, y = pair
                if x.distance < GOOD_DISTANCE_LIMIT * y.distance:
                    good.append(ori_kp[x.trainIdx].pt)
            elif len(pair) == 1:
                x = pair[0]
                good.append(ori_kp[x.trainIdx].pt)
        count = [0] * len(segment)
        for i, scope in enumerate(segment):
            for g in good:
                if in_scope(scope, g):
                    count[i] += 1
        for i, c in enumerate(count):
            if c > avatar_p[i][1]:
                avatar_p[i] = name, c

    op_name = [p[0] for p in avatar_p]
    logger.debug(op_name)

    return tuple(zip(op_name, segment))


match_portrait = partial(match_operator, model=portrait)
match_avatar = partial(match_operator, model=avatar)


def operator_room_select(img: tp.Image) -> tuple[tuple[str, tp.Scope]]:
    "基建房间内选干员"
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    return match_portrait(gray, segment_room_select(img))


def operator_team(img: tp.Image) -> tuple[tuple[str, tp.Scope]]:
    "编队界面"
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    return match_portrait(gray, segment_team(img))


def operator_team_select(img: tp.Image) -> tuple[tuple[str, tp.Scope]]:
    "作战编队和训练位选干员"
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    return match_portrait(gray, segment_team_select(img))