#!/usr/bin/env python3

## Copyright 2018 Intel Corporation
## SPDX-License-Identifier: Apache-2.0

import os
import time
import numpy as np
import torch

from config import *
from util import *
from dataset import *
from model import *
from color import *
from result import *

# Inference function object
class Infer(object):
  def __init__(self, cfg, device, result=None, is_aux=False):
    # Load the result config
    result_dir = get_result_dir(cfg, result)
    if not os.path.isdir(result_dir):
      error('result does not exist')
    self.result_cfg = load_config(result_dir)
    self.device = device
    self.features = self.result_cfg.features
    self.main_feature = get_main_feature(self.features)
    self.aux_features = get_aux_features(self.features)
    self.all_channels = get_dataset_channels(self.features)
    self.num_main_channels = len(get_dataset_channels(self.main_feature))
    self.is_aux = is_aux

    # Initialize the model
    self.model = get_model(self.result_cfg)
    self.model.to(device)

    # Load the checkpoint
    checkpoint = load_checkpoint(result_dir, device, cfg.num_epochs, self.model)
    self.epoch = checkpoint['epoch']

    # Infer in FP16 if the model was trained with mixed precision
    if self.result_cfg.precision == 'mixed':
      self.model.half()
    if device.type == 'cpu':
      self.model.float() # CPU does not support FP16, so convert it back to FP32

    # Initialize the transfer function
    if not self.is_aux or self.main_feature != 'z':
      self.transfer = get_transfer_function(self.result_cfg)
    else:
      self.transfer = None # already applied to auxiliary depth

    # Set the model to evaluation mode
    self.model.eval()

    # Initialize auxiliary feature inference
    self.aux_infers = {}
    if self.aux_features:
      for aux_result in set(cfg.aux_results):
        aux_infer = Infer(cfg, device, aux_result, is_aux=True)
        if (aux_infer.main_feature not in self.aux_features) or aux_infer.aux_features:
          error(f'result {aux_result} does not correspond to an auxiliary feature')
        self.aux_infers[aux_infer.main_feature] = aux_infer

  # Inference function
  def __call__(self, input, exposure=1.):
    image = input.clone()

    # Apply the transfer function to the main feature
    color = image[:, 0:self.num_main_channels, ...]
    if self.main_feature == 'hdr':
      color *= exposure
    if self.transfer:
      color = self.transfer.forward(color)
    image[:, 0:self.num_main_channels, ...] = color

    # Pad the output
    shape = image.shape
    image = F.pad(image, (0, round_up(shape[3], self.model.alignment) - shape[3],
                          0, round_up(shape[2], self.model.alignment) - shape[2]))

    # Prefilter the auxiliary features
    for aux_feature, aux_infer in self.aux_infers.items():
      aux_channels = get_dataset_channels(aux_feature)
      aux_channel_indices = get_channel_indices(aux_channels, self.all_channels)
      aux = image[:, aux_channel_indices, ...]
      aux = aux_infer(aux)
      image[:, aux_channel_indices, ...] = aux

    # Filter the main feature
    if self.result_cfg.precision == 'mixed' and self.device.type != 'cpu':
      image = image.half()
    if self.main_feature == 'sh1':
      # Iterate over x, y, z
      image = torch.cat([self.model(torch.cat((image[:, i:i+3, ...], image[:, 9:, ...]), 1)) for i in [0, 3, 6]], 1)
    else:
      image = self.model(image)
    image = image.float()

    # Unpad the output
    image = image[:, :, :shape[2], :shape[3]]

    # Sanitize the output
    image = torch.clamp(image, min=0.)

    # Apply the inverse transfer function
    if self.transfer:
      image = self.transfer.inverse(image)

    if self.main_feature == 'hdr':
      image /= exposure
    else:
      image = torch.clamp(image, max=1.)

    return image

def main():
  # Parse the command line arguments
  cfg = parse_args(description='Performs inference on a dataset using the specified training result.')

  # Initialize the PyTorch device
  device = init_device(cfg)

  # Initialize the inference function
  infer = Infer(cfg, device)
  print('Result:', cfg.result)
  print('Epoch:', infer.epoch)

  # Initialize the dataset
  data_dir = get_data_dir(cfg, cfg.input_data)
  image_sample_groups = get_image_sample_groups(data_dir, infer.features)

  # Iterate over the images
  print()
  output_dir = os.path.join(cfg.output_dir, cfg.input_data)
  metric_sum = {metric : 0. for metric in cfg.metric}
  metric_count = 0

  # Saves an image in different formats
  def save_images(path, image, image_srgb, num_channels, feature_ext=infer.main_feature):
    if feature_ext == 'sh1':
      # Iterate over x, y, z
      for i, axis in [(0, 'x'), (3, 'y'), (6, 'z')]:
        save_images(path, image[:, i:i+3, ...], image_srgb[:, i:i+3, ...], num_channels, 'sh1' + axis)
      return

    image      = tensor_to_image(image)
    image_srgb = tensor_to_image(image_srgb)
    filename_prefix = path + '.' + feature_ext + '.'
    for format in cfg.format:
      if format in {'exr', 'pfm', 'phm', 'hdr'}:
        # Transform to original range
        if infer.main_feature in {'sh1', 'nrm'}:
          image = image * 2. - 1. # [0..1] -> [-1..1]
        save_image(filename_prefix + format, image, num_channels=num_channels[feature_ext])
      else:
        save_image(filename_prefix + format, image_srgb, num_channels=num_channels[feature_ext])

  with torch.inference_mode():
    for group, input_names, target_name in image_sample_groups:
      # Create the output directory if it does not exist
      output_group_dir = os.path.join(output_dir, os.path.dirname(group))
      if not os.path.isdir(output_group_dir):
        os.makedirs(output_group_dir)

      # Load metadata for the images if it exists
      tonemap_exposure = 1.
      metadata = load_image_metadata(os.path.join(data_dir, group))
      if metadata:
        tonemap_exposure = metadata['exposure']
        save_image_metadata(os.path.join(output_dir, group), metadata)

      # Load the target image if it exists
      if target_name:
        target, target_num_channels = load_image_features(os.path.join(data_dir, target_name), infer.main_feature)
        target = image_to_tensor(target, batch=True).to(device)
        target_srgb = transform_feature(target, infer.main_feature, 'srgb', tonemap_exposure)

      # Iterate over the input images
      for input_name in input_names:
        print(input_name, '...', end='', flush=True)

        # Load the input image
        input, input_num_channels = load_image_features(os.path.join(data_dir, input_name), infer.features)

        # Compute the autoexposure value
        exposure = autoexposure(input) if infer.main_feature == 'hdr' else 1.

        # Infer
        input = image_to_tensor(input, batch=True).to(device)
        output = infer(input, exposure)

        input = input[:, 0:infer.num_main_channels, ...] # keep only the main feature
        input_srgb  = transform_feature(input,  infer.main_feature, 'srgb', tonemap_exposure)
        output_srgb = transform_feature(output, infer.main_feature, 'srgb', tonemap_exposure)

        # Compute metrics
        metric_str = ''
        if target_name and cfg.metric:
          for metric in cfg.metric:
            value = compare_images(output_srgb, target_srgb, metric)
            metric_sum[metric] += value
            if metric_str:
              metric_str += ', '
            metric_str += f'{metric}={value:.4f}'
          metric_count += 1

        # Save the input and output images
        output_suffix = cfg.result if cfg.output_suffix is None else cfg.output_suffix
        output_name = input_name + '.' + output_suffix
        if cfg.num_epochs:
          output_name += f'_{epoch}'
        if cfg.save_all:
          save_images(os.path.join(output_dir, input_name), input, input_srgb, input_num_channels)
        save_images(os.path.join(output_dir, output_name), output, output_srgb, input_num_channels)

        # Print metrics
        if metric_str:
          metric_str = ' ' + metric_str
        print(metric_str)

      # Save the target image if it exists
      if cfg.save_all and target_name:
        save_images(os.path.join(output_dir, target_name), target, target_srgb, target_num_channels)

  # Print summary
  if metric_count > 0:
    metric_str = ''
    for metric in cfg.metric:
      value = metric_sum[metric] / metric_count
      if metric_str:
        metric_str += ', '
      metric_str += f'{metric}_avg={value:.4f}'
    print()
    print(f'{cfg.result}: {metric_str} ({metric_count} images)')

if __name__ == '__main__':
  main()