import stc
import cv2
import random
import numpy as np
import gradio as gr
from PIL import Image
from scipy import signal


title = "Steganography"
description =  '''Explore hiding messages in images using content adaptive steganography and STCs (Syndrome Trellis Codes).
               http://dde.binghamton.edu/download/syndrome/ .
               We use HILL https://ieeexplore.ieee.org/document/7025854 .
               Python implementation adapted from Daniel Lerch's https://github.com/daniellerch/pySTC .
               To encode:
               Drag and drop a PNG file, write a message and enter a key (remember the key), the generated image has the secret message encoded.
               To decode:
               Drag and drop the stego file that you just generated, enter the key.
               Note that this software is supplied "as is," without any services or security guaranties.
               '''

def HILL(input_image, operation, message, key):
    tmp_name = str(random.randint(100,500))
    try:
        buffer = input_image
        I = cv2.imdecode(np.frombuffer(buffer, np.uint8), 1)
        I = cv2.cvtColor(I,cv2.COLOR_BGR2GRAY)
        cv2.imwrite('tmp/'+tmp_name+'.png',I)
    except Exception as e:
        print(e)
        raise ValueError('Unable to read image')

    if operation == 'decode':
        try:
            stc.extract('tmp/'+tmp_name+'.png', key, 'tmp/'+tmp_name+'.txt')
            return 'tmp/'+tmp_name+'.txt'
        except:
            raise ValueError('Unable to decode')

    else:
        H = np.array(
           [[-1,  2, -1],
            [ 2, -4,  2],
            [-1,  2, -1]])
        L1 = np.ones((3, 3)).astype('float32')/(3**2)
        L2 = np.ones((15, 15)).astype('float32')/(15**2)
        costs = signal.convolve2d(I, H, mode='same', boundary='symm')
        costs = abs(costs)
        costs = signal.convolve2d(costs, L1, mode='same', boundary='symm')
        costs = 1/costs
        costs = signal.convolve2d(costs, L2, mode='same', boundary='symm')
        costs[costs == np.inf] = 1
        stc.embed('tmp/'+tmp_name+'.png', costs, message, key, 'tmp/'+tmp_name+'.png')
        return 'tmp/'+tmp_name+'.png'

iface = gr.Interface(fn=HILL,
                    inputs=[gr.components.File(type='binary'), gr.components.Radio(['encode', 'decode']), gr.components.Textbox(), gr.components.Textbox()],
                    outputs=gr.components.File(),
                    examples=[['tmp/8825.png', 'encode', 'This is a secret message', 'secret-key'],
                             ['tmp/9390.png', 'encode', 'This is another secret message', 'secret-key-2']],
                    title=title,
                    description=description)
iface.launch()