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

import argparse
import codecs

from PIL import Image

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

""" sudoku_solve.py: Extract messages from 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.
"""

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


def longToTxt(l):
    hexa = hex(l)[2:].strip('L').encode("utf-8")
    return codecs.decode(hexa, "hex_codec").decode("utf-8")


def base9ToInt(i):
    """ Integer (base9) to Integer (base10) """
    return int(i, 9)


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 tuplesToSize(t):
    """
    Compute size from 9 firsts tuples of list t
    """
    i = 0
    nb = 0
    for l in t[:9][::-1]:
        for elt in l[::-1]:
            nb += elt*(256**i)
            i += 1
    return nb

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

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_out.png"

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_solve.py',
                                 description='Extract data from image with '
                                             'sudoku method.')
parser.add_argument('-src', metavar='imgSrcName', type=str, nargs='?',
                    required=True, help='Input image')
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
s = args.sdk
c1 = args.c1
c2 = args.c2


if __name__ == "__main__":

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

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

    # Size (9 px) Fig 6. General format of embedding data
    sizeTuples = tuplesToSize(imgdata)

    for num, sol in enumerate(SOLUTIONS):
        print("[Solution n° "+str(num+1)+"]")
        expected = subMatrix(sol)  # Fig 2. Sudoku solution after subtracting 1
        refMatrix = sudokuToRefMatrix(expected, w, h)  # Fig 4. Ref. matrix
        data = ""  # tmp flag
        for i in range(sizeTuples):
            px = imgdata[i+9]  # +9 due to size hidding
            data += str(refMatrix[px[c1]][px[c2]])  # extract data

        try:  # If data extracted from refMatrix is printable then print it
            print(longToTxt(base9ToInt(data)))
        except:
            continue