#!/usr/bin/env python
# coding: utf-8

# This notebook explores how to expires assets in an experiment after a certain date.
# 
# First, we will make sure we have all of the necessary Python modules:

# In[1]:


get_ipython().run_line_magic('pip', 'install -U "comet_ml>=3.44.0" aitk.utils Pillow')


# Next, we import all of the modules we will need:

# In[2]:


# Import Python modules:
import random
import os
import glob
import urllib
import zipfile
import datetime
import random
import json

# Additional modules:
from PIL import Image
from aitk.utils import gallery
import comet_ml


# For this demo, we'll create some fun images. First we need to download them, so we write a general download and unzip function:

# In[3]:


def download(url, filename=None):
    filename = filename if filename is not None else os.path.basename(url)
    basename, ext = os.path.splitext(filename)
    g = urllib.request.urlopen(url, timeout=5)
    with open(filename, "wb") as f:
        f.write(g.read())
    if ext == ".zip":
        with zipfile.ZipFile(filename, "r") as zip_ref:
            zip_ref.extractall(basename)


# And we download a dataset of image parts. These are based on the images and code from:
# 
# https://github.com/pixegami/pixel-punk-avatars

# In[23]:


download("https://github.com/dsblank/pixel-punk-avatars/raw/main/images/images.zip")


# We should now have a directory called "images" that contains faces, eyes, hair, and accessories.
# 
# Now, we make sure that we have our Comet API key set:

# In[5]:


comet_ml.login()


# We write a simple image generation function:

# In[24]:


def generate_image():
    bg_color = random.choice([(120, 150, 180), (255, 225, 150)])
    image = Image.new("RGBA", (24, 24), bg_color)
    for layer, probability in [
        ("0_face", 1.0),
        ("1_eye", 1.0),
        ("2_hair", 0.8),
        ("3_accessory", 0.15),
    ]:
        if random.random() > probability:
            continue
        items = glob.glob(f"images/{layer}/*.png")
        layer_image = Image.open(random.choice(items))
        image = Image.alpha_composite(image, layer_image)
    return image


# Every time we call it, we get a different image:

# In[25]:


image = generate_image()
image.resize((500, 500), resample=Image.Resampling.NEAREST)


# Now, let's call it 625 times to generate a nice dataset:

# In[26]:


images = [generate_image() for i in range(625)]


# A nice function from `aitk.utils` is the `gallery()` function that generates a gallery of images from a list of images:

# In[27]:


gallery(images)


# Now, we create an experiment to log all of the images.

# In[10]:


experiment = comet_ml.start(project_name="pixel-punk-avatars")


# We'll need some support functions to help us:

# In[11]:


def random_day(n=100):
    """
    Return a random day n days in the future.
    """
    today = datetime.date.today()
    return add_days(today, random.randint(1, n))


def add_days(day, n):
    """
    Method to do some date math.
    """
    return day + datetime.timedelta(days=n)


def day_to_timestamp(day):
    """
    Turn a day into an integer timestamp
    """
    date = datetime.datetime(day.year, day.month, day.day)
    return date.timestamp()


def timestamp_to_day(timestamp):
    """
    Turn a timestamp into a day
    """
    return datetime.datetime.fromtimestamp(timestamp)


# The reason for the timestamp is that we can't log an experation day directly, but need to encode it as a number.
# 
# Now we log each image, storing the date to expire in the metadata of the image. Note that for this demo we create random dates between 1 and 100 days in the future.

# In[12]:


for image in images:
    experiment.log_image(image, metadata={"expires_on": day_to_timestamp(random_day())})


# and call `experiment.end()` because we are in a Jupyter Notebook.

# In[13]:


experiment.end()


# We can check to see that the images were loaded by examining the Graphics tab in the Come UI:

# In[30]:


experiment.display(tab="graphics")


# Now we are ready to expire those images past a certain date. We'll use the Comet API to get the assets.

# In[14]:


api = comet_ml.API()


# We write a function to delete the image assets given an experiment key (id) and date:

# In[31]:


def delete_experiment_assets(experiment_key, expires_on):
    api_experiment = api.get_experiment_by_key(experiment_key)
    count = 0
    for asset_data in api_experiment.get_asset_list():
        if asset_data["type"] == "image":
            asset_metadata = json.loads(asset_data["metadata"])
            asset_id = asset_data["assetId"]
            timestamp = asset_metadata["expires_on"]
            if timestamp > expires_on:
                print(f"Deleting {asset_id}")
                api_experiment.delete_asset(asset_id)
                count += 1
    print(f"{count} assets expired")


# For this demo, let's pick a day 90 days in the future:

# In[21]:


# A date 90 days in the future:
expires_on = day_to_timestamp(add_days(datetime.date.today(), 90))


# And now, we simply call the function given the experiment.id from above, and the date 90 days from now:

# In[22]:


delete_experiment_assets(experiment.id, expires_on)


# That should have deleted about 10% of the 625 images, or about 62 images give or take.

# If you run the `delete_experiment_assets()` again with the same arguments, it shouldn't delete any assets.

# In[ ]:


delete_experiment_assets(experiment.id, expires_on)


# Now, every once in a while you can call `delete_experiment_assets()` with a paticular date, and you'll delete all of the assets that should be "expired".