import numpy as np
import onnxruntime as ort
from PIL import Image
import os
import math

def out(output):
    image = output.squeeze()
    image = (image * 0.5 + 0.5) * 255
    image = image.astype(np.uint8)
    image = np.transpose(image, (1, 2, 0))
    return Image.fromarray(image, "RGB").resize((512, 512), Image.LANCZOS)

def generate_gan():
    return Image.new("RGB", (512, 512), color="gray")

def generate_progan():
    model_path = os.path.join("model", "batik_progan.onnx")
    session = ort.InferenceSession(model_path)

    noise = np.random.randn(1, 512, 1, 1).astype(np.float32)
    alpha = np.array([1.0], dtype=np.float32)

    output = session.run(None, {
        'z': noise,
        'alpha': alpha,
    })

    return out(output[0])

def generate_dcgan():
    model_path = os.path.join("model", "batik_dcgan.onnx")
    session = ort.InferenceSession(model_path)

    noise = np.random.randn(1, 512, 1, 1).astype(np.float32)
    input_name = session.get_inputs()[0].name
    
    output = session.run(None, {input_name: noise})
    return out(output[0])

def generate_stylegan():
    model_path = os.path.join("model", "batik_stylegan.onnx")
    session = ort.InferenceSession(model_path)

    LATENT_FEATURES = 512
    RESOLUTION = 256
    LAST_INDEX = math.log2(RESOLUTION) - 2

    z = np.random.randn(1, LATENT_FEATURES).astype(np.float32)
    alpha = np.array([1.0], dtype=np.float32)
    steps = np.array([LAST_INDEX], dtype=np.int64)

    output = session.run(None, {
        'z': z,
        'alpha': alpha,
        'steps': steps
    })
    return out(output[0])