# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""postprocess"""
import argparse
import os

import imageio
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument("--bin_path", type=str, help='bin_path, path to binary files generated by 310 model, default: None')
parser.add_argument("--image_size", type=int, default=256, help='image size, default: image_size.')
parser.add_argument("--output_dir", type=str, default='output_dir',
                    help='output_path, path to store output, default: None')
parser.add_argument("--output_dir_interpolation", type=str, default='output_dir',
                    help='output_path, path to store output, default: None')
args = parser.parse_args()


def convert(img):
    img = np.swapaxes(img, 0, 1)
    img = np.swapaxes(img, 1, 2)
    return img


def post_process(img):
    img = convert(img)
    img = (img + 1.0) / 2.0
    img = np.clip(img, 0, 1.0)
    return img


if __name__ == "__main__":
    bin_path = args.bin_path
    image_size = args.image_size
    content_list = os.listdir(args.bin_path)

    for i in range(0, len(content_list)):
        pic_path = os.path.join(args.bin_path, content_list[i])
        b = np.fromfile(pic_path, dtype=np.float32, count=8 * 3 * image_size * image_size)
        b = np.reshape(b, (8, 3, image_size, image_size))
        # get output images
        output = ()
        style_shifted = convert(b[0])
        std = np.array([0.229, 0.224, 0.225]).reshape([1, 1, 3])
        mean = np.array([0.485, 0.456, 0.406]).reshape([1, 1, 3])
        style = style_shifted * std + mean
        style = np.clip(style, 0, 1.0)
        output = (output + (style,))
        for j in range(1, 7):
            output = ((post_process(b[j]),) + output)
        print(b[7].max())
        print(b[7].min())
        original = post_process(b[7])
        output = ((original,) + output)
        # save images
        content_name = content_list[i].replace(".bin", "")
        image_path_interpolation = os.path.join(args.output_dir_interpolation, content_name) + ".png"
        output_interpolation = np.concatenate(output, axis=1)
        imageio.imsave(image_path_interpolation, output_interpolation)

        image_path = os.path.join(args.output_dir, content_name) + ".png"
        output_single = np.concatenate((output[0], output[7], output[6]), axis=1)
        imageio.imsave(image_path, output_single)
        print("%d / %d , %s \n" % (i + 1, len(content_list), content_name))