#-*- coding:utf-8 -*-
######################################################################
# Ce script permet d'insérer une image NB dans une image, et de le récupérer.
# L'insertion est de type Doublement. Les étapes sont les suivantes :
#       - On demande à l'utilisateur l'image à insérer.
#       - On la chiffre
#       - On demande la fenêtre d'insertion.
#       - On effectue l'insertion en substituant, dans la fenêtre, le bit du
#          message encodé à celui préexistant.
#
# On propose alors diverses attaques.
#
# Puis, on passe à l'extraction...
#       - On récupère les bits du code, par soustraction avec l'image
#         d'origine.
# Cette méthode ne nécessite donc pas l'image d'origine pour récupérer la marque.
######################################################################

import binascii
import struct
import Image as im
import ImageStat

from ImageChops import difference
from math import sqrt
from numpy import array
from os import system, listdir
from sys import argv
from BitVector import BitVector
from copy import deepcopy
from ConfigParser import ConfigParser

from outilsBase import getBit, setBit, conversion
from attaque import Attaque
from suite import Doublement
from chiffrement_image import Chiffrement
from coefficients import Coefficients
from evaluation import Evaluation


OK = 100
PAS_OK = 99


def f(L):
    return [int(not k) for k in L]



class Spatial_Doublement:

    _bits_par_mode = {'1' : 1, 'L' : 8}


    def __init__(self, hote = '', watermark = '',
                       config = None,
                       mu = None, Xo = None,
                       Resultat = 'resultat.txt',
                       authentification = None,
                       x0 = None, y0 = None,
                       LSB = None,
                       iterations_chiffrement = None):
        '''
        Constructeur.

        Récupère le code à insérer, et les paramètres d'insertion.
        '''

        # Initialisation
        self._watermark = watermark
        self._liste = []
        self._hote = hote
        self._resultat = open(Resultat,'w')
        self._authentification = authentification
        self._x0, self._y0 = x0,y0
        self._iterations = iterations_chiffrement
        self._config = config

        if LSB != None:
            self._bits_faibles = LSB

        # Chiffrement
        unChiffrement = Chiffrement(fichier_image = self._watermark,
                                    mu = mu, Xo = Xo,
                                    authentification = self._authentification,
                                    iterations = self._iterations)

        unChiffrement.chiffrement()
        [self._mu, self._Xo] = unChiffrement.get_parametres()
        self._iterations = unChiffrement.get_iterations()

        self._strategie = unChiffrement.get_strategie()
        self._watermark = unChiffrement.get_image()


        # Récupération des bits de l'image chiffrée
        self._code = [k/255 for k in list(self._watermark.getdata())]

        # Récupération des paramètres d'insertion
        self._get_parametres_insertion()







    def _get_parametres_insertion(self):
        '''
        Récupère les paramètres d'insertion, à savoir le pixel de début de
        doublement.
        '''

        '''
        if self._authentification == None:
            self._authentification = eval(raw_input("Authentification (True/False) ? "))
            if self._authentification:
                self._bits_forts = eval(raw_input("Liste des MSB : "))
        '''

        if self._x0 == None:
            # Rectangle d'insertion
            print "Début du marquage ? "
            self._x0, self._y0 = eval(raw_input("Réponse : "))

        if '_bits_faibles' not in dir(self):
            self._bits_faibles = eval(raw_input("Liste des LSB : "))

            #taille = im.open('lena.png').size[0]
            print "\nCapacité d'insertion :",len(self._bits_faibles)*taille*taille
            self._x1, self._y1 = taille, taille









    def get_strategie(self):
        '''
        Récupère la stratégie des itérations chaotiques.
        '''
        hote = im.open(self._hote)
        N = hote.size[0]
        N *= hote.size[1]
        N *= len(self._bits_faibles)

        '''
        if self._authentification:
            msc = Coefficients().getAllCoefs(self._hote, self._bits_forts)
            bits = BitVector(bitstring = self._code) ^ msc

        else:
             bits = BitVector(bitstring = self._code)
        doublement = Doublement(1, N,
                                generateur = Coefficients().bit2coef(bits, N))'''

        doublement = Doublement(1, N, generateur = self._strategie)

        suite = doublement.iterateur()
        while True:
            yield suite.next()





    def insertion(self, nom = ''):
        '''
        L'insertion, à proprement parlé.

        On commence par regarder si l'image hote est suffisament grande. Puis,
        on réalise l'insertion, on montre l'image marquée, et on la sauvegarde
        dans lena_marque.png.
        '''
        # Initialisation
        hote = im.open(self._hote)
        mode = self._bits_par_mode[hote.mode]
        image = []
        for k in range(hote.size[0]):
            for l in range(hote.size[1]):
                texte = conversion(hote.getpixel((k,l)),2).zfill(mode)
                image += [int(texte[mode-x]) for x in self._bits_faibles]


        # L'insertion
        strat = self.get_strategie()
        X = set()
        if len(image)%4 == 0:
            image.append(0)

        for x in range(len(self._code)):
            suivant = strat.next()
            #y = f(image)
            #image[suivant] = y[suivant]
            image[suivant] = int(not image[suivant])

        # Reconstruction de l'image
        image_nouvelle = im.new(hote.mode, hote.size)
        compteur = 0
        for k in range(hote.size[0]):
            for l in range(hote.size[1]):
                x = [int(d) for d in conversion(hote.getpixel((k,l)),2).zfill(mode)]
                for m in self._bits_faibles:
                    x[mode-m] = image[compteur]
                    compteur += 1
                x = ''.join([str(t) for t in x])
                image_nouvelle.putpixel((k,l), int(x,2))

        # Affichage et sauvegarde de l'hôte marqué
        #print "On montre l'image, que l'on sauvegarde ensuite"
        image_nouvelle.save(nom)
        #image_nouvelle.show()
        #raw_input("  (Appuyer sur une touche pour continuer)")

        uneEva = Evaluation(im.open(self._hote), im.open(nom))
        #print "Différences Lena/Lena marque"
        #uneEva.visualisationDifferences()
        #raw_input("Differences")
        self._resultat.write("PSNR = "+str(uneEva.PSNR()))




    def extraction(self, hote = '', original = None, 
                         watermark = None, LSB = None, mu = None,
                         Xo = None):
        '''
        Réalise l'extraction de l'image.

        On récupère le code extrait, qui est la chaîne de bits du message encodé
        par Reed-Solomon.
        '''
        #print "\n================== Extraction =================="
        #print "Hote, dans laquelle réaliser l'insertion : ",hote

        uneAction = Spatial_Doublement(hote = hote,
                           watermark = watermark,
                           mu = mu,
                           Xo = Xo,
                           LSB = LSB,
                           authentification = self._authentification,
                           x0 = self._x0, y0 = self._y0,
                           iterations_chiffrement = self._iterations)

        #print "Sauvegarde dans lena_demarque.png"
        tempo_="tempo_"+argv[3]+".png"
        uneAction.insertion(nom = tempo_)

        image1 = im.open(original)
        image2 = im.open(tempo_)

        if self._config.get('EVALUATION','type') == "differences":
            compteur = 0
            diff = difference(image1, image2)
            for k in range(image1.size[0]):
                for l in range(image1.size[0]):
                    if diff.getpixel((k,l)) >0 :
                        compteur += 1
            dd = "\nNombre de différences (entre lena et lena démarque) :"+str(compteur)+" sur "
            dd += str(image1.size[0]*image1.size[1])+'.\n'
            self._resultat.write(dd)
            if compteur < self._config.getfloat("DIFFERENCES","seuil"):
                print "OK"
                exit(OK)
            else:
                print "PAS OK"
                exit(PAS_OK)







if __name__ == '__main__':
    config = ConfigParser()
    config.read([argv[1]])

    repertoire_images = config.get('GENERAL','repertoire_images')+'/'
    repertoire_temporaire = config.get('GENERAL','repertoire_temporaire')+'/'
    
    if repertoire_temporaire not in listdir('.'):
        system("mkdir "+repertoire_temporaire)
    
    nom = ''.join(argv[2].split('.')[:-1])
    extension = argv[2].lstrip(nom)

    image_hote = repertoire_images+argv[2]
    image_marquee = repertoire_temporaire+nom+"_marquee"+extension
    image_attaquee = repertoire_temporaire+nom+"_attaquee"+extension

    authentification = config.getboolean('GENERAL','authentification')
    WATERMARK = config.get('GENERAL','watermark')
    
    Resultat = config.get('GENERAL','repertoire_resultat')+'/'+config.get('GENERAL','resultat')

    uneAction = Spatial_Doublement(hote = image_hote,
                    watermark = WATERMARK,
                    config = config,
                    LSB = eval(config.get('INSERTION','LSB')),
                    mu = config.getfloat("CHIFFREMENT_INSERTION","mu"),
                    Xo = config.getfloat("CHIFFREMENT_INSERTION","Xo"),
                    authentification = authentification,
                    x0 = 1, y0 = 1,
                    Resultat = Resultat,
                    iterations_chiffrement = 20000)
    uneAction.insertion(nom = image_marquee)

    # Attaque éventuelle 
    if config.getboolean('ATTAQUE','attaque'):
        uneAttaque = Attaque(nom = image_marquee,
                             configuration = argv[1])
        #uneAttaque.show()
        uneAttaque.save(image_attaquee)
        hote = image_attaquee
    else:
        hote = image_marquee

    # Extraction et evaluation
    uneAction.extraction(hote = hote,
                    original = image_hote,
                    mu = config.getfloat("CHIFFREMENT_EXTRACTION","mu"),
                    Xo = config.getfloat("CHIFFREMENT_EXTRACTION","Xo"),
                    LSB = eval(config.get('EXTRACTION','LSB')),
                    watermark = WATERMARK)

    # Nettoyage
    system("rm tempo_"+argv[3]+".png")
    system("rm "+repertoire_temporaire+"*")
