import base64
import json
import requests

"""
Explanation: Querying Clipper with an Image
This notebook shows how to setup Clipper to accept querying with image. We will deploy a function that compute the size of the image using PIL image library and query it with image data.
End of explanation
"""


raw_bytes = open('imgs/clipper-logo.png', "rb").read()
raw_bytes[:10]

"""
Explanation: Using input_type = "bytes"
An image can be read a string of raw bytes. Clipper can take any base64 encoded bytes as input and decode it to send to your model. 

Below, we define the query function and the model that calculate image size.
Query Function
Our input JSON will contains a base64 encoded string. This can be done in the following three step procedure:

Read in the raw bytes.
End of explanation
"""


encoded_bytes = base64.b64encode(raw_bytes)
print(type(encoded_bytes))
encoded_bytes[:10]

"""
Explanation: Encode the raw bytes to base64 encoding, and note that the result is still bytes
End of explanation
"""


encoded_string = encoded_bytes.decode()
print(type(encoded_string))
encoded_string[:10]

"""
Explanation: Turn the bytes into string so json can serialize it
End of explanation
"""


def query(addr, filename):
    url = "http://%s/image-example/predict" % addr
    req_json = json.dumps({
        "input":
        base64.b64encode(open(filename, "rb").read()).decode() # bytes to unicode
    })
    headers = {'Content-type': 'application/json'}
    r = requests.post(url, headers=headers, data=req_json)
    print(r.json())

"""
Explanation: Finally our query function will look like this:
End of explanation
"""


def image_size(imgs):
    """
    Input: 
    - imgs : (np.ndarray) of shape (n, d). n is the number of data in this batch
             d is the length of the bytes as numpy int8 array.  
    Output:
    - sizes : List[Tuple(int, int),...]
    """
    import base64
    import io
    import os
    import PIL.Image
    import tempfile
  
    num_imgs = len(imgs)
    sizes = []
    for i in range(num_imgs):
        # Create a temp file to write to
        tmp = tempfile.NamedTemporaryFile('wb', delete=False, suffix='.png')
        tmp.write(io.BytesIO(imgs[i]).getvalue())
        tmp.close()
        
        # Use PIL to read in the file and compute size
        size = PIL.Image.open(tmp.name, 'r').size
        
        # Remove the temp file
        os.unlink(tmp.name) 

        sizes.append(size)
    return sizes

"""
Explanation: Model Function
The model will takes in an array of byte arrays. The bytes are already decoded for you.
End of explanation
"""


from clipper_admin import ClipperConnection, DockerContainerManager
from clipper_admin.deployers import python as python_deployer

clipper_conn = ClipperConnection(DockerContainerManager())
clipper_conn.start_clipper()


python_deployer.create_endpoint(
    clipper_conn=clipper_conn, 
    name="image-example", 
    input_type="bytes", 
    func=image_size, 
    pkgs_to_install=['pillow']
)

"""
Explanation: Start Clipper and Deploy Function
We can now start Clipper with Docker and deploy our function.
Be sure to add pillow as external package by passing it in to pkgs_to_install parameter in deployer function!
End of explanation
"""


query(clipper_conn.get_query_addr(), 'imgs/clipper-logo.png')

"""
Explanation: Try it out!
End of explanation
"""


def query_json(addr, filename, image_format):
    url = "http://%s/image-example-string/predict" % addr
    req_json = json.dumps({
        "input":
        json.dumps({
            'data': base64.b64encode(open(filename, "rb").read()).decode(),
            'format': image_format
        })
    })
    headers = {'Content-type': 'application/json'}
    r = requests.post(url, headers=headers, data=req_json)
    print(r.json())

def image_size_json(imgs):
    """
    Input: 
    - imgs : an array of strings 
    Output:
    - sizes : List[Tuple(int, int),...]
    """
    import base64
    import io
    import os
    import PIL.Image
    import tempfile
    import json
  
    num_imgs = len(imgs)
    sizes = []
    for i in range(num_imgs):
        # Deserialize the query
        data = json.loads(imgs[i])
        image_format = data['format']
        image_bytes = data['data'].encode()
        
        # Create a temp file to write to
        tmp = tempfile.NamedTemporaryFile('wb', delete=False, suffix='.{}'.format(image_format))
        tmp.write(io.BytesIO(base64.b64decode(image_bytes)).getvalue())
        tmp.close()
        
        # Use PIL to read in the file and compute size
        size = PIL.Image.open(tmp.name, 'r').size
        
        # Remove the temp file
        os.unlink(tmp.name) 

        sizes.append(size)
    return sizes

python_deployer.create_endpoint(
    clipper_conn=clipper_conn, 
    name="image-example-string", 
    input_type="strings", 
    func=image_size_json, 
    pkgs_to_install=['pillow']
)

query_json(clipper_conn.get_query_addr(), 'imgs/clipper-logo.jpg', 'jpg')

query_json(clipper_conn.get_query_addr(), 'imgs/clipper-logo.png', 'png')

"""
Explanation: Using input_type = "strings"
Sometime you need to pass in metadata about the image like the file format. In the example below, we will serialize our image data as base64 encoded string and make it part of our input query. We will input a json serialized string as input and decoded in the model function.
End of explanation
"""


clipper_conn.stop_all()

"""
Explanation: Shutdown Clipper
Finally, we will stop all clipper containers
End of explanation
"""


