import numpy as np
from funcy import identity

from raven_utils.constant import PROPERTY, TARGET, INPUTS
from raven_utils.decode import decode_target, target_mask
from raven_utils.image import EXIST, COR, draw_images, add_text
from raven_utils.render.rendering import render_panels
from raven_utils.render_ import TYPES, SIZES
from raven_utils.tools import none, is_model, filter_keys
from raven_utils.uitls import get_val_index



def draw_board(images, target=None, predict=None, image=None, desc=None, layout=None, break_=20):
    if image != "target" and predict is not None:
        image = images[predict:predict + 1]
    elif images is None and target is not None:
        image = images[target:target + 1]
    # image = False to not draw anything
    border = [{COR: target - 8, EXIST: list(range(4)) if predict is None else (1, 3)}] + [{COR: p, EXIST: (0, 2)} for p
                                                                                          in none(predict)]

    boards = []
    boards.append(draw_images(
        np.concatenate([images[:8], image[None] if len(image.shape) == 3 else image]) if image is not None else images[
                                                                                                                :8]))
    if layout == 1:
        i = draw_images(images[8:], column=4, border=border)
        if break_:
            i = np.concatenate([np.zeros([break_, i.shape[1], 1]), i], axis=0)
        boards.append(i)

    else:
        boards.append(
            draw_images(np.concatenate([images[8:], predict]) if predict is not None else images[8:], column=4,
                        border=target - 8))
    full_board = draw_images(boards, grid=False)
    if desc:
        full_board = add_text(full_board, desc)
    return full_board


def draw_boards(images, target=None, predict=None,image=None, desc=None,  layout=None):
    boards = []
    for i, im in enumerate(images):
        boards.append(draw_board(im, target[i][0] if target is not None else None,
                                 predict[i] if predict is not None else None,
                                 image[i] if image is not None else None,
                                 desc[i] if desc is not None else None, layout=layout))
    return boards


def draw_from_generator(generator, predict=None, no=1, indexes=None, layout=1):
    data,_ = val_sample(generator, no, indexes)
    return draw_raven(data, predict=predict, pre_fn=generator.data.data["inputs"].fn, layout=layout)


def val_sample(generator, no=1, indexes=None):
    if indexes is None:
        indexes = get_val_index(base=no)
    data = generator.data[indexes]
    return data, indexes

def render_from_model(data,predict,pre_fn=identity):
    data = filter_keys(data, PROPERTY, reverse=True)
    if is_model(predict) or str(type(predict)) == "<class 'tensorflow.python.saved_model.load.Loader._recreate_base_user_object.<locals>._UserObject'>":
        predict = predict(data)
    pro = np.array(target_mask(predict['predict_mask'].numpy()) * predict["predict"].numpy(), dtype=np.int8)
    return pre_fn(render_panels(pro, target=False)[None])[0]

def draw_raven(data, predict=None, pre_fn=identity, layout=1):
    if is_model(predict):
        d = filter_keys(data, PROPERTY, reverse=True)
        # tmp change
        res = predict(d)
        pro = np.array(target_mask(res['mask'].numpy()) * res["predict"].numpy(),dtype=np.int8)
        predict = pre_fn(render_panels(pro, target=False)[None])[0]
    target = data[TARGET]
    target_index = data["index"]
    images = data[INPUTS]
    # np.equal(res['predict'], pro[:,:102]).sum()

    if hasattr(predict, "shape"):
        if len(predict.shape) > 3:
            # iamges
            image = predict
            # todo create index and output based on image
            predict = None
            predict_index = None
        elif len(predict.shape) == 3:
            image = render_panels(predict, target=False)
            # Create index based on predict.
            predict_index = None
        else:
            image = images[predict]
            predict_index = predict
            predict = target
    else:
        image = K.gather(images, target_index[:, 0])
        predict_index = None
        predict = None

    # elif not(hasattr(target,"shape") and len(target.shape) > 3):
    #     if hasattr(target,"shape") and target.shape[-1] == OUTPUT_SIZE:
    #         pro = target
    #         predict = render_panels(pro)
    #     elif hasattr(target,"shape") and target.shape[-1] == FEATURE_NO:
    #         # pro = target
    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
    #     else:
    #         pro = np.zeros([no, OUTPUT_SIZE], dtype="int")
    #         # predict = [None] * no
    #         predict = render_panels(data[TARGET])

    image = draw_boards(images, target=target_index, predict=predict_index,image=image,  desc=None,
                        layout=layout)

    all_rules = extract_rules(data[PROPERTY])
    target_desc = get_desc(target)
    if predict is not None:
        predict_desc = get_desc(predict)
    else:
        predict_desc = [None] * len(target_desc)
    for a, po, to in zip(all_rules, predict_desc, target_desc):
        # fl(predict_desc[-1])
        if po is None:
            po = [None] * len(to)
        for p, t in zip(po, to):
            a.extend(
                [" ".join([str(i) for i in t])] + (
                    [" ".join([str(i) for i in p]), ""] if p is not None else []
                )
            )
        # a.extend([""] + [] + [""] + [" ".join(fl(p))])

    # image = draw_boards(data[INPUTS],target=data["index"], predict=predict[:no], desc=all_rules, no=no,layer=layer)
    return lu([(i, j) for i, j in zip(image, all_rules)])


def extract_rules(data):
    all_rules = []
    for d in data:
        rules = []
        for j, rule_group in enumerate(d.findAll("Rule_Group")):
            # rules_all.append(rule_group['id'])
            for j, rule in enumerate(rule_group.findAll("Rule")):
                rules.append(f"{rule['attr']} - {rule['name']}")
            rules.append("")
        all_rules.append(rules)
    return all_rules


def get_desc(target, exist=None, types=TYPES, sizes=SIZES):
    decoded = decode_target(target)
    exist = decoded[1] if exist is None else exist
    taken = np.stack(take(decoded[2], np.array(exist, dtype=bool))).T

    figures_no = np.sum(exist, axis=-1)
    desc = np.split(taken, np.cumsum(figures_no))[:-1]
    # figures_no = np.sum(exist, axis=-1)
    # div = np.split(desc, np.cumsum(figures_no))[:-1]
    result = []
    for pd in desc:
        r = []
        for p in pd:
            r.append([p[0], sizes[p[1]], types[p[2]]])
        result.append(r)

    return result