#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: viz.py
# Credit: zxytim

import numpy as np
import os, sys
import io
import cv2
from .fs import mkdir_p
from .argtools import shape2d

try:
    import matplotlib.pyplot as plt
except ImportError:
    pass

__all__ = ['pyplot2img', 'build_patch_list', 'pyplot_viz',
        'dump_dataflow_images', 'interactive_imshow']

def pyplot2img(plt):
    buf = io.BytesIO()
    plt.axis('off')
    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
    buf.seek(0)
    rawbuf = np.frombuffer(buf.getvalue(), dtype='uint8')
    im = cv2.imdecode(rawbuf, cv2.IMREAD_COLOR)
    buf.close()
    return im

def pyplot_viz(img, shape=None):
    """ use pyplot to visualize the image
        Note: this is quite slow. and the returned image will have a border
    """
    plt.clf()
    plt.axes([0,0,1,1])
    plt.imshow(img)
    ret = pyplot2img(plt)
    if shape is not None:
        ret = cv2.resize(ret, shape)
    return ret

def minnone(x, y):
    if x is None: x = y
    elif y is None: y = x
    return min(x, y)

def interactive_imshow(img, lclick_cb=None, rclick_cb=None, **kwargs):
    """
    :param lclick_cb: a callback(img, x, y) for left click
    :param kwargs: can be {key_cb_a ... key_cb_z: callback(img)}
    """
    name = 'random_window_name'
    cv2.imshow(name, img)

    def mouse_cb(event, x, y, *args):
        if event == cv2.EVENT_LBUTTONUP and lclick_cb is not None:
            lclick_cb(img, x, y)
        elif event == cv2.EVENT_RBUTTONUP and rclick_cb is not None:
            rclick_cb(img, x, y)
    cv2.setMouseCallback(name, mouse_cb)
    key = chr(cv2.waitKey(-1) & 0xff)
    cb_name = 'key_cb_' + key
    if cb_name in kwargs:
        kwargs[cb_name](img)
    elif key == 'q':
        cv2.destroyWindow(name)
    elif key == 'x':
        sys.exit()
    elif key == 's':
        cv2.imwrite('out.png', img)

def build_patch_list(patch_list,
        nr_row=None, nr_col=None, border=None,
        max_width=1000, max_height=1000,
        shuffle=False, bgcolor=255,
        viz=False, lclick_cb=None):
    """
    Generate patches.
    :param patch_list: bhw or bhwc images in [0,255]
    :param border: defaults to 0.1 * min(image_width, image_height)
    :param nr_row, nr_col: rows and cols of the grid
    :parma max_width, max_height: if nr_row/col are not given, use this to infer the rows and cols
    :param shuffle: shuffle the images
    :param bgcolor: background color
    :param viz: use interactive imshow to visualize the results
    :param lclick_cb: only useful when viz=True. a callback(patch, idx)
    """
    # setup parameters
    patch_list = np.asarray(patch_list)
    if patch_list.ndim == 3:
        patch_list = patch_list[:,:,:,np.newaxis]
    assert patch_list.ndim == 4 and patch_list.shape[3] in [1, 3], patch_list.shape
    if shuffle:
        np.random.shuffle(patch_list)
    if lclick_cb is not None:
        viz = True
    ph, pw = patch_list.shape[1:3]
    if border is None:
        border = int(0.1 * min(ph, pw))
    mh, mw = max(max_height, ph + border), max(max_width, pw + border)
    if nr_row is None:
        nr_row = minnone(nr_row, max_height / (ph + border))
    if nr_col is None:
        nr_col = minnone(nr_col, max_width / (pw + border))

    canvas = np.zeros((nr_row * (ph + border) - border,
             nr_col * (pw + border) - border,
             patch_list.shape[3]), dtype='uint8')

    def draw_patch(plist):
        cur_row, cur_col = 0, 0
        canvas.fill(bgcolor)
        for patch in plist:
            r0 = cur_row * (ph + border)
            c0 = cur_col * (pw + border)
            canvas[r0:r0+ph, c0:c0+pw] = patch
            cur_col += 1
            if cur_col == nr_col:
                cur_col = 0
                cur_row += 1

    nr_patch = nr_row * nr_col
    start = 0

    def lclick_callback(img, x, y):
        if lclick_cb is None:
            return
        x = x // (pw + border)
        y = y // (pw + border)
        idx = start + y * nr_col + x
        if idx < end:
            lclick_cb(patch_list[idx], idx)

    while True:
        end = start + nr_patch
        cur_list = patch_list[start:end]
        if not len(cur_list):
            return
        draw_patch(cur_list)
        if viz:
            interactive_imshow(canvas, lclick_cb=lclick_callback)
        yield canvas
        start = end

def dump_dataflow_images(df, index=0, batched=True,
        number=1000, output_dir=None,
        scale=1, resize=None, viz=None,
        flipRGB=False, exit_after=True):
    """
    :param df: a DataFlow
    :param index: the index of the image component
    :param batched: whether the component contains batched images or not
    :param number: how many datapoint to take from the DataFlow
    :param output_dir: output directory to save images, default to not save.
    :param scale: scale the value, usually either 1 or 255
    :param resize: (h, w) or Nne, resize the images
    :param viz: (h, w) or None, visualize the images in grid with imshow
    :param flipRGB: apply a RGB<->BGR conversion or not
    :param exit_after: exit the process after this function
    """
    if output_dir:
        mkdir_p(output_dir)
    if viz is not None:
        viz = shape2d(viz)
        vizsize = viz[0] * viz[1]
    if resize is not None:
        resize = tuple(shape2d(resize))
    vizlist = []

    df.reset_state()
    cnt = 0
    while True:
        for dp in df.get_data():
            if not batched:
                imgbatch = [dp[index]]
            else:
                imgbatch = dp[index]
            for img in imgbatch:
                cnt += 1
                if cnt == number:
                    if exit_after:
                        sys.exit()
                    else:
                        return
                if scale != 1:
                    img = img * scale
                if resize is not None:
                    img = cv2.resize(img, resize)
                if flipRGB:
                    img = img[:,:,::-1]
                if output_dir:
                    fname = os.path.join(output_dir, '{:03d}.jpg'.format(cnt))
                    cv2.imwrite(fname, img)
                if viz is not None:
                    vizlist.append(img)
            if viz is not None and len(vizlist) >= vizsize:
                patch = next(build_patch_list(
                    vizlist[:vizsize],
                    nr_row=viz[0], nr_col=viz[1], viz=True))
                vizlist = vizlist[vizsize:]

if __name__ == '__main__':
    import cv2
    imglist = []
    for i in range(100):
        fname = "{:03d}.png".format(i)
        imglist.append(cv2.imread(fname))
    for idx, patch in enumerate(build_patch_list(
            imglist, max_width=500, max_height=200)):
        of = "patch{:02d}.png".format(idx)
        cv2.imwrite(of, patch)
