#!/usr/bin/env python
"""
sd api txt2img benchmark
"""
import asyncio
import base64
import io
import json
import time
from PIL import Image
import sdapi
from util import Map, log


options = Map({
    'restore_faces': False,
    'prompt': 'photo of two dice on a table',
    'negative_prompt': 'foggy, blurry',
    'steps': 20,
    'batch_size': 1,
    'n_iter': 1,
    'seed': -1,
    'sampler_name': 'Euler a',
    'cfg_scale': 0,
    'width': 512,
    'height': 512
})


# batch = [1, 1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128]
batch = [1, 1, 2, 4, 8, 12, 16]
oom = 0


async def txt2img():
    t0 = time.perf_counter()
    data = {}
    try:
        data = await sdapi.post('/sdapi/v1/txt2img', options)
    except Exception:
        return -1
    if 'error' in data:
        return -1
    if 'info' in data:
        info = Map(json.loads(data['info']))
    else:
        return 0
    log.debug({ 'info': info })
    for i in range(len(data['images'])):
        data['images'][i] = Image.open(io.BytesIO(base64.b64decode(data['images'][i].split(',',1)[0])))
        log.debug({ 'image': data['images'][i].size })
    t1 = time.perf_counter()
    return t1 - t0


def memstats():
    mem = sdapi.getsync('/sdapi/v1/memory')
    cpu = mem.get('ram', 'unavailable')
    gpu = mem.get('cuda', 'unavailable')
    if 'active' in gpu:
        gpu['session'] = gpu.pop('active')
    if 'reserved' in gpu:
        gpu.pop('allocated')
        gpu.pop('reserved')
        gpu.pop('inactive')
    if 'events' in gpu:
        global oom # pylint: disable=global-statement
        oom = gpu['events']['oom']
        gpu.pop('events')
    return cpu, gpu


def gb(val: float):
    return round(val / 1024 / 1024 / 1024, 2)


async def main():
    log.info({ 'benchmark': { 'batch-sizes': batch } })
    sdapi.quiet = True
    await sdapi.session()
    await sdapi.interrupt()
    opts = await sdapi.get('/sdapi/v1/options')
    opts = Map(opts)
    log.info({ 'options': {
        'resolution': [options.width, options.height],
        'model': opts.sd_model_checkpoint,
        'vae': opts.sd_vae,
        'hypernetwork': opts.sd_hypernetwork,
        'sampler': options.sampler_name,
        'preview': opts.show_progress_every_n_steps
    } })
    cpu, gpu = memstats()
    log.info({ 'system': { 'cpu': cpu, 'gpu': gpu }})
    for i in range(len(batch)):
        if oom > 0:
            continue
        options['batch_size'] = batch[i]
        ts = await txt2img()
        if ts > 0:
            await asyncio.sleep(0)
            cpu, gpu = memstats()
            if i == 0:
                log.info({ 'warmup': round(ts, 2) })
            else:
                peak = gpu['session']['peak'] if 'session' in gpu else 0
                log.info({ 'batch': batch[i], 'its': round(options.steps / (ts / batch[i]), 2), 'img': round(ts / batch[i], 2), 'wall': round(ts, 2), 'peak': gb(peak), 'oom': oom > 0 })
        else:
            await asyncio.sleep(10)
            cpu, gpu = memstats()
            log.info({ 'batch': batch[i], 'result': 'error', 'gpu': gpu, 'oom': oom > 0 })
            break
    if oom > 0:
        log.info({ 'benchmark': 'ended with oom so you should probably restart your automatic server now' })
    await sdapi.close()


if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        log.warning({ 'interrupted': 'keyboard request' })
        sdapi.interruptsync()