import argparse
import os
import platform
import struct
import subprocess
import time
from collections import namedtuple
from typing import List

import cv2
import numpy as np
import torch.multiprocessing as mp
from numba import njit

import blender.histogram_blend as histogram_blend
from blender.guide import (BaseGuide, ColorGuide, EdgeGuide, PositionalGuide,
                           TemporalGuide)
from blender.poisson_fusion import poisson_fusion
from blender.video_sequence import VideoSequence
from flow.flow_utils import flow_calc
from src.video_util import frame_to_video

OPEN_EBSYNTH_LOG = False
MAX_PROCESS = 8

os_str = platform.system()

if os_str == 'Windows':
    ebsynth_bin = '.\\deps\\ebsynth\\bin\\ebsynth.exe'
elif os_str == 'Linux':
    ebsynth_bin = './deps/ebsynth/bin/ebsynth'
elif os_str == 'Darwin':
    ebsynth_bin = './deps/ebsynth/bin/ebsynth.app'
else:
    print('Cannot recognize OS. Run Ebsynth failed.')
    exit(0)


@njit
def g_error_mask_loop(H, W, dist1, dist2, output, weight1, weight2):
    for i in range(H):
        for j in range(W):
            if weight1 * dist1[i, j] < weight2 * dist2[i, j]:
                output[i, j] = 0
            else:
                output[i, j] = 1
            if weight1 == 0:
                output[i, j] = 0
            elif weight2 == 0:
                output[i, j] = 1


def g_error_mask(dist1, dist2, weight1=1, weight2=1):
    H, W = dist1.shape
    output = np.empty_like(dist1, dtype=np.byte)
    g_error_mask_loop(H, W, dist1, dist2, output, weight1, weight2)
    return output


def create_sequence(base_dir, beg, end, interval, key_dir):
    sequence = VideoSequence(base_dir, beg, end, interval, 'video', key_dir,
                             'tmp', '%04d.png', '%04d.png')
    return sequence


def process_one_sequence(i, video_sequence: VideoSequence):
    interval = video_sequence.interval
    for is_forward in [True, False]:
        input_seq = video_sequence.get_input_sequence(i, is_forward)
        output_seq = video_sequence.get_output_sequence(i, is_forward)
        flow_seq = video_sequence.get_flow_sequence(i, is_forward)
        key_img_id = i if is_forward else i + 1
        key_img = video_sequence.get_key_img(key_img_id)
        for j in range(interval - 1):
            i1 = cv2.imread(input_seq[j])
            i2 = cv2.imread(input_seq[j + 1])
            flow_calc.get_flow(i1, i2, flow_seq[j])

        guides: List[BaseGuide] = [
            ColorGuide(input_seq),
            EdgeGuide(input_seq,
                      video_sequence.get_edge_sequence(i, is_forward)),
            TemporalGuide(key_img, output_seq, flow_seq,
                          video_sequence.get_temporal_sequence(i, is_forward)),
            PositionalGuide(flow_seq,
                            video_sequence.get_pos_sequence(i, is_forward))
        ]
        weights = [6, 0.5, 0.5, 2]
        for j in range(interval):
            # key frame
            if j == 0:
                img = cv2.imread(key_img)
                cv2.imwrite(output_seq[0], img)
            else:
                cmd = f'{ebsynth_bin} -style {os.path.abspath(key_img)}'
                for g, w in zip(guides, weights):
                    cmd += ' ' + g.get_cmd(j, w)

                cmd += (f' -output {os.path.abspath(output_seq[j])}'
                        ' -searchvoteiters 12 -patchmatchiters 6')
                if OPEN_EBSYNTH_LOG:
                    print(cmd)
                subprocess.run(cmd,
                               shell=True,
                               capture_output=not OPEN_EBSYNTH_LOG)


def process_sequences(i_arr, video_sequence: VideoSequence):
    for i in i_arr:
        process_one_sequence(i, video_sequence)


def run_ebsynth(video_sequence: VideoSequence):

    beg = time.time()

    processes = []
    mp.set_start_method('spawn')

    n_process = min(MAX_PROCESS, video_sequence.n_seq)
    cnt = video_sequence.n_seq // n_process
    remainder = video_sequence.n_seq % n_process

    prev_idx = 0

    for i in range(n_process):
        task_cnt = cnt + 1 if i < remainder else cnt
        i_arr = list(range(prev_idx, prev_idx + task_cnt))
        prev_idx += task_cnt
        p = mp.Process(target=process_sequences, args=(i_arr, video_sequence))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

    end = time.time()

    print(f'ebsynth: {end-beg}')


@njit
def assemble_min_error_img_loop(H, W, a, b, error_mask, out):
    for i in range(H):
        for j in range(W):
            if error_mask[i, j] == 0:
                out[i, j] = a[i, j]
            else:
                out[i, j] = b[i, j]


def assemble_min_error_img(a, b, error_mask):
    H, W = a.shape[0:2]
    out = np.empty_like(a)
    assemble_min_error_img_loop(H, W, a, b, error_mask, out)
    return out


def load_error(bin_path, img_shape):
    img_size = img_shape[0] * img_shape[1]
    with open(bin_path, 'rb') as fp:
        bytes = fp.read()

    read_size = struct.unpack('q', bytes[:8])
    assert read_size[0] == img_size
    float_res = struct.unpack('f' * img_size, bytes[8:])
    res = np.array(float_res,
                   dtype=np.float32).reshape(img_shape[0], img_shape[1])
    return res


def process_seq(video_sequence: VideoSequence,
                i,
                blend_histogram=True,
                blend_gradient=True):

    key1_img = cv2.imread(video_sequence.get_key_img(i))
    img_shape = key1_img.shape
    interval = video_sequence.interval
    beg_id = video_sequence.get_sequence_beg_id(i)

    oas = video_sequence.get_output_sequence(i)
    obs = video_sequence.get_output_sequence(i, False)

    binas = [x.replace('jpg', 'bin') for x in oas]
    binbs = [x.replace('jpg', 'bin') for x in obs]

    obs = [obs[0]] + list(reversed(obs[1:]))
    inputs = video_sequence.get_input_sequence(i)
    oas = [cv2.imread(x) for x in oas]
    obs = [cv2.imread(x) for x in obs]
    inputs = [cv2.imread(x) for x in inputs]
    flow_seq = video_sequence.get_flow_sequence(i)

    dist1s = []
    dist2s = []
    for i in range(interval - 1):
        bin_a = binas[i + 1]
        bin_b = binbs[i + 1]
        dist1s.append(load_error(bin_a, img_shape))
        dist2s.append(load_error(bin_b, img_shape))

    lb = 0
    ub = 1
    beg = time.time()
    p_mask = None

    # write key img
    blend_out_path = video_sequence.get_blending_img(beg_id)
    cv2.imwrite(blend_out_path, key1_img)

    for i in range(interval - 1):
        c_id = beg_id + i + 1
        blend_out_path = video_sequence.get_blending_img(c_id)

        dist1 = dist1s[i]
        dist2 = dist2s[i]
        oa = oas[i + 1]
        ob = obs[i + 1]
        weight1 = i / (interval - 1) * (ub - lb) + lb
        weight2 = 1 - weight1
        mask = g_error_mask(dist1, dist2, weight1, weight2)
        if p_mask is not None:
            flow_path = flow_seq[i]
            flow = flow_calc.get_flow(inputs[i], inputs[i + 1], flow_path)
            p_mask = flow_calc.warp(p_mask, flow, 'nearest')
            mask = p_mask | mask
        p_mask = mask

        # Save tmp mask
        # out_mask = np.expand_dims(mask, 2)
        # cv2.imwrite(f'mask/mask_{c_id:04d}.jpg', out_mask * 255)

        min_error_img = assemble_min_error_img(oa, ob, mask)
        if blend_histogram:
            hb_res = histogram_blend.blend(oa, ob, min_error_img,
                                           (1 - weight1), (1 - weight2))

        else:
            # hb_res = min_error_img
            tmpa = oa.astype(np.float32)
            tmpb = ob.astype(np.float32)
            hb_res = (1 - weight1) * tmpa + (1 - weight2) * tmpb

        # cv2.imwrite(blend_out_path, hb_res)

        # gradient blend
        if blend_gradient:
            res = poisson_fusion(hb_res, oa, ob, mask)
        else:
            res = hb_res

        cv2.imwrite(blend_out_path, res)
    end = time.time()
    print('others:', end - beg)


def main(args):
    global MAX_PROCESS
    MAX_PROCESS = args.n_proc

    video_sequence = create_sequence(f'{args.name}', args.beg, args.end,
                                     args.itv, args.key)
    if not args.ne:
        run_ebsynth(video_sequence)
    blend_histogram = True
    blend_gradient = args.ps
    for i in range(video_sequence.n_seq):
        process_seq(video_sequence, i, blend_histogram, blend_gradient)
    if args.output:
        frame_to_video(args.output, video_sequence.blending_dir, args.fps,
                       False)
    if not args.tmp:
        video_sequence.remove_out_and_tmp()


def process_video(name, output=None, fps=30, beg=1, end=101, itv=10, key='keys0', n_proc=8, ps=False, ne=False, tmp=False):
    """
    Process a video with the specified parameters.

    :param name: Path to input video.
    :param output: Path to output video.
    :param fps: The FPS of output video.
    :param beg: The index of the first frame to be stylized.
    :param end: The index of the last frame to be stylized.
    :param itv: The interval of key frame.
    :param key: The subfolder name of stylized key frames.
    :param n_proc: The max process count.
    :param ps: Use poisson gradient blending if True.
    :param ne: Do not run ebsynth if True (use previous ebsynth output).
    :param tmp: Keep temporary output if True.
    """
    Args = namedtuple('Args', ['name', 'output', 'fps', 'beg', 'end', 'itv', 'key', 'n_proc', 'ps', 'ne', 'tmp'])
    args = Args(name, output, fps, beg, end, itv, key, n_proc, ps, ne, tmp)

    main(args)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('name', type=str, help='Path to input video')
    parser.add_argument('--output',
                        type=str,
                        default=None,
                        help='Path to output video')
    parser.add_argument('--fps',
                        type=float,
                        default=30,
                        help='The FPS of output video')
    parser.add_argument('--beg',
                        type=int,
                        default=1,
                        help='The index of the first frame to be stylized')
    parser.add_argument('--end',
                        type=int,
                        default=101,
                        help='The index of the last frame to be stylized')
    parser.add_argument('--itv',
                        type=int,
                        default=10,
                        help='The interval of key frame')
    parser.add_argument('--key',
                        type=str,
                        default='keys0',
                        help='The subfolder name of stylized key frames')
    parser.add_argument('--n_proc',
                        type=int,
                        default=8,
                        help='The max process count')
    parser.add_argument('-ps',
                        action='store_true',
                        help='Use poisson gradient blending')
    parser.add_argument(
        '-ne',
        action='store_true',
        help='Do not run ebsynth (use previous ebsynth output)')
    parser.add_argument('-tmp',
                        action='store_true',
                        help='Keep temporary output')

    args = parser.parse_args()
    main(args)