import os
import time
import logging
import numpy as np
import keras.optimizers
import cv2
import gradio as gr
import Network
import utls

class ImageProcessor:

    def __init__(self, model_name):
        self.model_name = model_name
        self.mbllen = Network.build_mbllen((None, None, 3))
        self.mbllen.load_weights(f'{model_name}.h5')
        self.opt = keras.optimizers.Adam(lr=2 * 1e-04, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
        self.mbllen.compile(loss='mse', optimizer=self.opt)

    def process_image(self, image, highpercent, lowpercent, gamma, maxrange):
        img_A = np.array(image) / 255.0
        img_A = img_A[np.newaxis, :]

        maxrange /= 10.
        hsvgamma = gamma / 10.

        out_pred = self.mbllen.predict(img_A)

        fake_B = out_pred[0, :, :, :3]
        fake_B = self.adjust_image(fake_B, maxrange, highpercent, lowpercent, hsvgamma)

        return fake_B

    def adjust_image(self, fake_B, maxrange, highpercent, lowpercent, hsvgamma):
        gray_fake_B = fake_B[:, :, 0] * 0.299 + fake_B[:, :, 1] * 0.587 + fake_B[:, :, 1] * 0.114
        percent_max = sum(sum(gray_fake_B >= maxrange)) / sum(sum(gray_fake_B <= 1.0))
        max_value = np.percentile(gray_fake_B[:], highpercent)
        if percent_max < (100-highpercent)/100.:
            scale = maxrange / max_value
            fake_B = fake_B * scale
            fake_B = np.minimum(fake_B, 1.0)

        gray_fake_B = fake_B[:,:,0]*0.299 + fake_B[:,:,1]*0.587 + fake_B[:,:,1]*0.114
        sub_value = np.percentile(gray_fake_B[:], lowpercent)
        fake_B = (fake_B - sub_value) * (1./(1-sub_value))

        imgHSV = cv2.cvtColor(fake_B, cv2.COLOR_RGB2HSV)
        H, S, V = cv2.split(imgHSV)
        S = np.power(S, hsvgamma)
        imgHSV = cv2.merge([H, S, V])
        fake_B = cv2.cvtColor(imgHSV, cv2.COLOR_HSV2RGB)
        fake_B = np.minimum(fake_B, 1.0)

        return fake_B

if __name__ == "__main__":
    logging.basicConfig(filename='image_processor.log', level=logging.INFO)

    image_processor = ImageProcessor(model_name='Syn_img_lowlight_withnoise')
    iface = gr.Interface(
    image_processor.process_image,
    [
        gr.inputs.Image(shape=(None, None)),  # Corrected line
        gr.inputs.Slider(85, 100, default=95, label="High Percent"),
        gr.inputs.Slider(0, 100, default=5, label="Low Percent"),
        gr.inputs.Slider(6, 10, default=8, label="Gamma"),
        gr.inputs.Slider(0, 20, default=8, label="Max Range"),
    ],
    gr.outputs.Image(type="numpy"),
)



    iface.launch()