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

# This script is (C) Copyright 2023, P. Lutus: https://www.arachnoid.com
# and is released under the GPL: https://www.gnu.org/licenses/gpl-3.0.en.html

# This is the A1111 version of the morph script

import json
import requests
import io
import base64
import os
import time
from dataclasses import dataclass

class MultiDraw:
  def __init__(self):
    self.server='http://localhost:7860'
    
    self.payload = {
       "seed": 333,
       "width": 1024,
       "height": 1024,
       "sampler_name": 'DPM2 Karras',
       "cfg_scale": 7,
       "steps": 20,
       "batch_size": 1,
       }
      
    override_settings = {
     "sd_model_checkpoint": 'sd_xl_base_1.0.safetensors [31e35c80fc]',
    }
    
    override_payload = {
     "override_settings": override_settings,
     "override_settings_restore_afterwards": False, 
    }
      
    self.payload.update(override_payload)
    
  def display_json(self,title,json):
    print(f'JSON data dump: {title}')
    for key in json:
      print(f'   {key:12s} : {json[key]}')
   
  def draw_save_morph(self,item, draw = True):
    
    filepath = f'{item.savepath}/output_index_{item.index:03d}.png'
    
    # see localhost:7860/docs for full explanation
    
    if draw:
           
      prompt = f'{item.preshared}, ["{item.promptA}":"{item.promptB}":{item.ratio:.03f}], {item.postshared}'
      
      print(f'Prompt: "{prompt}"')
      
      self.payload['prompt'] = prompt
      self.payload['negative_prompt'] = item.negative_prompt
        
      response = requests.post(url=f'{self.server}/sdapi/v1/txt2img', json=self.payload)
            
      r = response.json()
      
      if 'images' in r:
        base64_image = r['images'][0]
        
        # print(f'Successful image generation, saving {filepath}...')
        
        # image info is automatically saved with the file
        with open(filepath, 'wb') as fp:
          fp.write(base64.b64decode(base64_image))
          
      else:
        self.display_json("Error",r['info'])
    return filepath
    

@dataclass
class MorphImageData:
  savepath:str
  preshared:str
  promptA: str
  promptB: str
  postshared: str
  negative_prompt: str
  ratio: float
  nratio: float
  index: int

# interpolation function maps range { xa : xb } -> { ya : yb }

def ntrp(x,xa,xb,ya,yb): return (x-xa) * (yb-ya) / (xb-xa) + ya

def draw_morph(size, savepath, preshared, pa, pb, postshared = "", negative = "", start=0.1, end=0.9):
  md = MultiDraw()
  os.system(f'mkdir {savepath} > /dev/null 2>&1')
  os.system(f'rm {savepath}/*.png > /dev/null 2>&1')
  for index in range(size):
    ratio = 1.0 - ntrp(index,0,size-1,start,end)
    item = MorphImageData(savepath,preshared,pa,pb,postshared,negative,ratio,1.0-ratio,index)
    start_time = time.time()
    # note an available draw argument: True/False
    # this saves time until prompt
    # debugging is complete
    md.draw_save_morph(item)
    duration = float(time.time() - start_time)
    print(f'{savepath}: index {index:03d} ratio:{ratio:7.2f} processing time: {duration:7.03f} sec.')


# cat-dog example:
  
# draw_morph(8,"cat_dog","hyper-realistic","cat","dog", "in forest setting")