from functools import partial

import cv2
import numpy as np
from skimage.metrics import structural_similarity

from mower.static import avatar, portrait
from mower.utils import typealias as tp
from mower.utils.image import cmatch, cropimg, loadres, thres2
from mower.utils.log import logger
from mower.utils.matcher import GOOD_DISTANCE_LIMIT, flann, search_params
from mower.utils.vector import in_scope

FAST = cv2.FastFeatureDetector_create()
ORB = cv2.ORB_create(nfeatures=1000000, nlevels=1)

FLANN_INDEX_LSH = 6
index_params = dict(
    algorithm=FLANN_INDEX_LSH,
    table_number=1,
    key_size=18,
    multi_probe_level=1,
)
flann1 = cv2.FlannBasedMatcher(index_params, search_params)


def fast_keypoints(img: tp.GrayImage, mask: tp.GrayImage | None = None):
    return ORB.compute(img, FAST.detect(img, mask))


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, (95, 230, 200), (105, 255, 230))
    line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
    line1[mask > 0] = (255,)
    mask = cv2.inRange(hsv, (10, 0, 0), (20, 255, 255))
    line1[mask > 0] = (0,)
    line1 = thres2(line1, 100)
    line = line1[-1]
    prev = line[0]
    start = None
    name_x = []
    for i in range(1, line1.shape[1]):
        curr = line[i]
        if prev == 0 and curr == 255 and start and i - start > 186:
            name_x.append((i + 434, i + 574))
        elif prev == 255 and curr == 0:
            start = i
        prev = curr

    avatar_y = ((242, 317), (662, 737))
    avatar_p = []
    for x in name_x:
        for y in avatar_y:
            avatar_p.append(tuple(zip(x, y)))

    empty = loadres("choose_agent/empty")
    if avatar_p and cmatch(cropimg(img, avatar_p[-1]), empty, 20):
        avatar_p.pop()

    logger.debug(avatar_p)
    return avatar_p


def segment_team(img: tp.Image) -> list[tp.Scope]:
    "编队界面"
    result = [
        ((190 + 232 * x, 235 + 433 * y), (330 + 232 * x, 310 + 433 * y))
        for x in range(6)
        for y in range(2)
    ] + [((1582, 668), (1722, 743))]
    res = loadres("choose_agent/no_oper")
    for i in range(13):
        if cmatch(cropimg(img, result[i]), res):
            gray = cropimg(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), result[i])
            res_gray = cv2.cvtColor(res, cv2.COLOR_RGB2GRAY)
            ssim = structural_similarity(gray, res_gray)
            logger.debug(f"cmatch+SSIM:{i=} {ssim=}")
            if ssim >= 0.9:
                break
    return result[:i]


def segment_team_select(img: tp.Image) -> list[tp.Scope]:
    "作战编队选干员"
    line1 = cropimg(img, ((545, 510), (1920, 511)))
    # 选中时的蓝色边框变白，信赖排序时的橙色变黑
    hsv = cv2.cvtColor(line1, cv2.COLOR_RGB2HSV)
    mask = cv2.inRange(hsv, (95, 230, 200), (105, 255, 230))
    line1 = cv2.cvtColor(line1, cv2.COLOR_RGB2GRAY)
    line1[mask > 0] = (255,)
    mask = cv2.inRange(hsv, (10, 0, 0), (20, 255, 255))
    line1[mask > 0] = (0,)
    line1 = thres2(line1, 100)
    line = line1[-1]
    prev = line[0]
    start = None
    name_x = []
    for i in range(1, line1.shape[1]):
        curr = line[i]
        if prev == 0 and curr == 255 and start and i - start > 90:
            name_x.append((i + 381, i + 521))
        elif prev == 255 and curr == 0:
            start = i
        prev = curr

    avatar_y = ((229, 304), (645, 720))
    avatar_p = []
    for x in name_x:
        for y in avatar_y:
            avatar_p.append(tuple(zip(x, y)))

    empty = loadres("choose_agent/empty")
    if avatar_p and cmatch(cropimg(img, avatar_p[-1]), empty, 20):
        avatar_p.pop()

    logger.debug(avatar_p)
    return avatar_p


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

    w = model["w"]
    h = model["h"]
    col = model["col"]
    name_list = model["name"]
    pt2 = model["pt"]
    des2 = model["des"]

    mask = np.zeros((1080, 1920), dtype=np.uint8)
    for seg in segment:
        cv2.rectangle(mask, seg[0], seg[1], (255,), -1)
    kp1, des1 = fast_keypoints(gray, mask)

    result = [{} for _ in range(len(segment))]
    for match in flann.knnMatch(des1, des2, k=2):
        if len(match) != 2:
            continue
        m, n = match
        if m.distance > GOOD_DISTANCE_LIMIT * n.distance:
            continue
        x, y = pt2[m.trainIdx]
        x, y = int(x / w), int(y / h)
        name = name_list[y * col + x]
        for i, seg in enumerate(segment):
            if in_scope(seg, kp1[m.queryIdx].pt):
                if name in result[i]:
                    result[i][name] += 1
                else:
                    result[i][name] = 1
                break
    result = [list(r.items()) for r in result]
    result = [max(r, key=lambda p: p[1]) if r else None for r in result]
    logger.debug(result)

    bad_index = []
    for i, r in enumerate(result):
        if r is None or r[1] < 10:
            bad_index.append(i)
    if bad_index:
        mask = np.zeros((1080, 1920), dtype=np.uint8)
        for i in bad_index:
            cv2.rectangle(mask, segment[i][0], segment[i][1], (255,), -1)
        kp1, des1 = fast_keypoints(gray, mask)

        result1 = [{} for _ in range(len(bad_index))]
        for match in flann1.knnMatch(des1, des2, k=2):
            if len(match) != 2:
                continue
            m, n = match
            if m.distance > GOOD_DISTANCE_LIMIT * n.distance:
                continue
            x, y = pt2[m.trainIdx]
            x, y = int(x / w), int(y / h)
            name = name_list[y * col + x]
            for i, j in enumerate(bad_index):
                if in_scope(segment[j], kp1[m.queryIdx].pt):
                    if name in result1[i]:
                        result1[i][name] += 1
                    else:
                        result1[i][name] = 1
                    break
        result1 = [list(r.items()) for r in result1]
        result1 = [max(r, key=lambda p: p[1]) if r else None for r in result1]
        logger.debug(result1)

        for i, r in enumerate(result1):
            result[bad_index[i]] = r

    return list(zip([r[0] if r else None for r in result], segment))


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


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


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


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