# -* coding:UTF-8 -*
# !/usr/bin/env python
import numpy as np
import gradio as gr

import os
from PIL import Image, ImageDraw, ImageOps

from codeinterpreterapi import CodeInterpreterSession


def imGreyAlpha(im):
    if isinstance(im,np.ndarray):
        im = Image.fromarray(im)

    # im = ImageOps.invert(im)
    
    grey = im.convert('L') # 转成灰度
    px = grey.load() # 获取灰度数组

    w,h = im.size # 获取图片大小
    im2 = Image.new(mode="RGBA", size=(w, h), color=(255,255,255,0)) # 新建图片
    draw = ImageDraw.Draw(im2) # 获取绘制句柄
    

    # 遍历像素点
    for i in range(w):
        for j in range(h):
            alpha = px[i, j] # 获取灰度作为alpha值
            if alpha == 255:
                draw.point((i,j), fill=(0,0,0,0))
            # elif alpha == 255:
            #      draw.point((i,j), fill=(0,0,0,255))
            elif alpha <= 50:
                draw.point((i,j), fill=(255,255,255,255))
            else:
                pt = im.getpixel((i,j))
                draw.point((i,j), fill=(pt[0] ,pt[1],pt[2],255-alpha))
                # draw.point((i, j), fill=(255, 255, 255, 255 -alpha)) # 填充像素点

    
    return im2 # 返回图像


def codeinterpreter(openai_key,prompt,filterBG, files):

    fileList = []
    if files != None:
        for idx, file in enumerate(files):
            fileList.append(file)

    if openai_key == "123321":
        openai_key = os.getenv("openaikey")
    else:
        return ['openai key must set!',None]

    with CodeInterpreterSession(model="gpt-3.5-turbo",openai_api_key=openai_key) as session:
    #async with CodeInterpreterSession(model="gpt-3.5-turbo",openai_api_key="") as session:
        response = session.generate_response_sync(prompt, fileList,True)

        images = []

        for _file in response.files:
            if filterBG:
                img = _file.get_image()
                img2 = imGreyAlpha(img)
                images.append(img2)
            else:
                images.append(_file.get_image())

        return [response.content,images]
        

with gr.Blocks() as app:
    with gr.Tab(label="codeinterpreter"):
        with gr.Row():
            with gr.Column():
                inp1=gr.Textbox(label="openai_key")
                inp2=gr.Textbox(label="prompt",info="input the prompt")
                inp3=gr.Checkbox(label="do_filter_bg?", info="do background filter?")
                inp4=gr.Files()
                btn = gr.Button(value="Submit")
            with gr.Column():
                out1=gr.Textbox(label="result")
                out2=gr.Gallery()
        gr.Examples([["Plot the nvidea stock vs microsoft stock over the last 6 months."],
               ["Plot a sin wave and show it to me."],
               ["贵州茅台最近半年走势"],
               ["Plot the bitcoin chart of 2023 YTD"]],
                [inp2])
    with gr.Tab(label="image background filter"):
        with gr.Row():
            with gr.Column():
                inp10=gr.Image(label="original image")
                btn10 = gr.Button(value="Submit")
            with gr.Column():
                out10=gr.Image(label="result")
    btn.click(codeinterpreter, inputs=[inp1,inp2,inp3,inp4], outputs=[out1,out2],api_name="getresult")
    btn10.click(imGreyAlpha,inputs=[inp10],outputs=[out10],api_name="imageFilter")

app.launch()