import cv2
import torch
import onnx
import onnxruntime
import numpy as np

import time

# codeformer converted to onnx
# using https://github.com/redthing1/CodeFormer


class CodeFormerEnhancer:
    def __init__(self, model_path="CodeFormer/weights/CodeFormer/codeformer.pth", device='cpu'):
        model = onnx.load(model_path)
        session_options = onnxruntime.SessionOptions()
        session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
        providers = ["CPUExecutionProvider"]
        if device == 'cuda':
            providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),"CPUExecutionProvider"]
        self.session = onnxruntime.InferenceSession("codeformer.onnx", sess_options=session_options, providers=providers)

    def enhance(self, img, w=0.9):
        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
        img = img.astype(np.float32)[:,:,::-1] / 255.0
        img = img.transpose((2, 0, 1))
        nrm_mean = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
        nrm_std = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
        img = (img - nrm_mean) / nrm_std

        img = np.expand_dims(img, axis=0)

        out = self.session.run(None, {'x':img.astype(np.float32), 'w':np.array([w], dtype=np.double)})[0]
        out = (out[0].transpose(1,2,0).clip(-1,1) + 1) * 0.5
        out = (out * 255)[:,:,::-1]

        return out.astype('uint8')