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

import argparse
import codecs
import random
import sys

from PIL import Image

# http://infohost.nmt.edu/tcc/help/lang/python/examples/sudoku/
# (Ported to python3)
from sudosolver import *

""" sudoky.py: Generate steganographic images using sudoku method based on
Chang et al.’s steganographic scheme and "Steganography using Sudoku Puzzle"
paper from Roshan Shetty B R, Rohith J, Mukund V, Rohan Honwade. This script
use a reference matrix as large as the input image to minify distortion. It
doesn't affect the extracting process.
"""

__author__ = "Alex GARRIDO (Zeecka)"
__credits__ = ["Chang et al.", "Roshan Shetty B R", "Rohith J",
               "Mukund V", "Rohan Honwade"]
__license__ = "WTFPL"


def txtToLong(t):
    hexa = codecs.encode(t.encode("utf-8"), "hex_codec")
    return int(hexa.decode("utf-8"), 16)


def intToBase9(i):
    """ Integer (base10) to Integer (base9) """
    if i == 0:
        return '0'
    nums = []
    while i:
        i, r = divmod(i, 9)
        nums.append(str(r))
    return ''.join(reversed(nums))


def addToSolutions(x):
    """
    SudokuSolver solutionObserver handler
    Add each matrice solution to SOLUTIONS
    """
    global SOLUTIONS
    M = []
    for rowx in range(9):
        l = []
        for colx in range(9):
            cell = x.get(rowx, colx)
            l.append(cell)
        M.append(l)
    SOLUTIONS.append(M)


def subMatrix(M, n=1):
    """
    Substract n to each elt in an 2D array
    (substract 1 by default)
    """
    return [[x-n for x in l] for l in M]


def sudokuToRefMatrix(s, w, h):
    """
    Convert sudoku "s" to a reference matrix with
    size of "size"
    """
    M = []
    for i in range(h):
        M.append((s[i % len(s)] * ((w//9)+1))[:w])
    return M


def sizeToTuples(t):
    """
    Compute size of tuple and return it with 9 tuples of 3 integer
    Used to hide size in 9 px
    ie. size = 1337,
    return [(0,0,0),(0,0,0),...,(0,0,0),(0,5,57)
    because 1337 = 62*(256^0) + 5*(256^1)
    """
    N = len(t)
    Nbin = bin(N)[2:].zfill(8*3*9)  # Convert as binary
    l = [int(Nbin[x:x+8], 2) for x in range(0, len(Nbin), 8)]
    return [(l[i], l[i+1], l[i+2]) for i in range(0, 9*3, 3)]


def manhattanDist(p1, p2):
    """
    USED TO HIDE DATA ONLY
    """
    return abs(p1[0]-p2[0])+abs(p1[1]-p2[1])


def getIndexHide(x, y, r, n):
    """
    USED TO HIDE DATA ONLY
    Return nearest index from x and y value equal to n (in r as refMatrix)
    This function is used to compute new value of embbeding image
    Fig 5. Illustration of CE H (solid line), CE V (dotted line),
    CE B (dashed line)
    """

    # Horizontal line
    Xh, Yh = 0, 0
    if x > 3 and x < 252:
        CEh = [(x+i, y) for i in range(-4, 5) if r[y][x+i] == n]
    elif x <= 3:
        CEh = [(x+i, y) for i in range(10) if r[y][i] == n]  # Overflow l.
    else:  # x >= 252
        CEh = [(x+i, y) for i in range(247, 256) if r[y][i] == n]  # Ov. r.

    # Vertical line
    if y > 3 and y < 252:
        CEv = [(x, y+i) for i in range(-4, 5) if r[y+i][x] == n]
    elif x <= 3:
        CEv = [(x, y+i) for i in range(10) if r[i][x] == n]  # Overflow t.
    else:  # x >= 252
        CEv = [(x, y+i) for i in range(247, 256) if r[i][x] == n]  # Ov. b.

    # Square (Get top left of the square first)
    xmod = x - (x % 3)
    ymod = y - (y % 3)
    CEb = [(xmod+i, ymod+j) for i in range(3)
           for j in range(3) if r[ymod+j][xmod+i] == n]

    # Get nearest between CEh, CEv, Ceb
    newX, newY = CEb[0]
    if manhattanDist(CEh[0], (x, y)) < manhattanDist(CEb[0], (x, y)):
        newX, newY = CEh[0]
    if manhattanDist(CEv[0], (x, y)) < manhattanDist(CEh[0], (x, y)):
        newX, newY = CEv[0]
    if manhattanDist(CEv[0], (x, y)) < manhattanDist(CEb[0], (x, y)):
        newX, newY = CEv[0]

    return newX, newY


###############################################################################

# Sudoku with 5 solutions
s = "..4...3.." \
    ".8...2..9" \
    "7..9...6." \
    "..879...." \
    "2..4.6..3" \
    "....319.." \
    ".3..69.18" \
    "1..8...3." \
    "..6...2.."

imgSrcName = "lena.png"
imgDstName = "lena_out.png"

hidden = "Bravo ! Vous pouvez valider le challenge avec le flag suivant: " \
       "APRK{*5UD0KU_M4ST3R*}.\n" \
       "Congratz ! You can validate the challenge with the following flag: " \
       "APRK{*5UD0KU_M4ST3R*}.\n"

c1 = 0  # Channel 1 is Red ==> 0
c2 = 1  # Channel 2 is Green ==> 1


SOLUTIONS = []  # List of solutions for the sudoku "s"

parser = argparse.ArgumentParser(prog='sudoku.py',
                                 description='Hide data in image with '
                                             'sudoku method.')
parser.add_argument('-src', metavar='imgSrcName', type=str, nargs='?',
                    required=True, help='Embbed image')
parser.add_argument('-out', metavar='imgDstName', type=str, nargs='?',
                    required=True, help='Output image')
parser.add_argument('-txt', metavar='hidden', type=str, nargs='?',
                    default=hidden, help='Hidden text')
parser.add_argument('--sdk', metavar='s', type=str, nargs='?', default=s,
                    help='Sudoku used for reference matrix. Use a 81 chars '
                         'long string with . as unknow value')
parser.add_argument('--c1', metavar='c1', type=int, default=0, nargs='?',
                    help='Indice of hidden channel 1 (default is 0 for red)')
parser.add_argument('--c2', metavar='c2', type=int, default=1, nargs='?',
                    help='Indice of hidden channel 2 (default is 1 for green)')

args = parser.parse_args()
imgSrcName = args.src
imgDstName = args.out
hidden = args.txt
s = args.sdk
c1 = args.c1
c2 = args.c2

if __name__ == "__main__":

    if imgDstName.lower().endswith("jpg") or \
       imgDstName.lower().endswith("jpeg"):
        print("Output image can't be in jpg format due to jpeg degradation")
        sys.exit()

    slv = SudokuSolver(s, solutionObserver=addToSolutions)
    slv.solve()

    imgSrc = Image.open(imgSrcName)
    w, h = imgSrc.size
    imgdata = list(imgSrc.getdata())  # List of (r,g,b)

    imgDst = Image.new(imgSrc.mode, imgSrc.size)
    imgDstData = []

    expectedNb = random.randint(0, len(SOLUTIONS)-1)  # Random solution
    expected = SOLUTIONS[expectedNb]  # Random expected solution

    expected = subMatrix(expected)  # Fig 2. Sudoku solution, subtracting 1
    refMatrix = sudokuToRefMatrix(expected, w, h)  # Fig 4. Reference matrix

    secret = intToBase9(txtToLong(hidden))  # Convert flag to base9

    # START EXTRACT
    # Size (9 px) Fig 6. General format of embedding data
    sizeTuples = sizeToTuples(secret)

    for i in range(9):
        imgDstData.append(sizeTuples[i])

    for i in range(len(imgdata)-9):  # III - A. Data embedding
        px = imgdata[i+9]  # +9 due to size hidding

        # M(R,G) = M(gi,gi+1) = M(Y,X)
        # Note, "chosen as X-axis and Y-axis" in paper should be reverse
        # https://en.wikipedia.org/wiki/Matrix_(mathematics)
        gi = px[c1]  # gi = y = c1 (default Red)
        gip1 = px[c2]  # gi+1 = x = c2 (default Green)

        # Xnew, Ynew for secret[i] value with smallest distortion from p1, p2
        # Note: Red = Y and Green = X because "gi = R and gi+1 = G."
        t = list(imgdata[i+9])
        if i < len(secret):  # Still data to embbed
            Xnew, Ynew = getIndexHide(gip1, gi, refMatrix, int(secret[i]))
            t[c1] = Ynew
            t[c2] = Xnew
        imgDstData.append(tuple(t))  # Append tuple to list

    imgDst.putdata(imgDstData)  # Append list to image data
    imgDst.save(imgDstName, subsampling=0, quality=100)