# -*- coding: utf-8 -*-
import os
import numpy as np
import cv2


# borrowed from https://github.com/lengstrom/fast-style-transfer/blob/master/src/utils.py
def get_files(img_dir):
    imgs, masks, xmls = list_files(img_dir)
    return imgs, masks, xmls


def list_files(in_path):
    img_files = []
    mask_files = []
    gt_files = []
    for (dirpath, dirnames, filenames) in os.walk(in_path):
        for file in filenames:
            filename, ext = os.path.splitext(file)
            ext = str.lower(ext)
            if ext == '.jpg' or ext == '.jpeg' or ext == '.gif' or ext == '.png' or ext == '.pgm':
                img_files.append(os.path.join(dirpath, file))
            elif ext == '.bmp':
                mask_files.append(os.path.join(dirpath, file))
            elif ext == '.xml' or ext == '.gt' or ext == '.txt':
                gt_files.append(os.path.join(dirpath, file))
            elif ext == '.zip':
                continue
    return img_files, mask_files, gt_files


def save_result(
        img_file: str,
        img: np.ndarray,
        boxes: np.ndarray,
        dirname: str = './result/',
        verticals: np.ndarray = None,
        texts: np.ndarray = None
):
    """
    save text detection result one by one
    Args:
        img_file (str): image file name
        img (array): raw image context
        boxes (array): array of result file
            Shape: [num_detections, 4] for BB output / [num_detections, 4] for QUAD output
        dirname
        verticals
        texts
    Return:
        None
    """
    img = np.array(img)

    # make result file list
    filename, file_ext = os.path.splitext(os.path.basename(img_file))

    # result directory
    res_file = dirname + "res_" + filename + '.txt'
    res_img_file = dirname + "res_" + filename + '.jpg'

    if dirname and not os.path.isdir(dirname):
        os.mkdir(dirname)

    with open(res_file, 'w') as f:
        for i, box in enumerate(boxes):
            poly = np.array(box).astype(np.int32).reshape((-1))
            strResult = ','.join([str(p) for p in poly]) + '\r\n'
            f.write(strResult)

            poly = poly.reshape(-1, 2)
            cv2.polylines(img, [poly.reshape((-1, 1, 2))], True, color=(0, 0, 255), thickness=2)
            pt_color = (0, 255, 255)
            if verticals is not None:
                if verticals[i]:
                    pt_color = (255, 0, 0)

            if texts is not None:
                font = cv2.FONT_HERSHEY_SIMPLEX
                font_scale = 0.5
                cv2.putText(  # noqa
                    img,  # noqa
                    "{}".format(texts[i]),  # noqa
                    (poly[0][0]+1, poly[0][1]+1),  # noqa
                    font,  # noqa
                    font_scale,  # noqa
                    (0, 0, 0),  # noqa
                    thickness=1  # noqa
                )  # noqa
                cv2.putText(
                    img,
                    "{}".format(texts[i]),
                    tuple(poly[0]),
                    font,
                    font_scale,
                    # (0, 255, 255),
                    pt_color,
                    thickness=1
                )

    # Save result image
    cv2.imwrite(res_img_file, img)
