#-*-coding:utf8-*-
###############################################################################
# Réalise l'insertion et l'extraction, par substitution, d'une image
# dans une autre, dans le domaine DWT. On précise l'ondelette, la matrice DWT.
###############################################################################

import Image as im
import ImageStat
from ImageChops import difference
from BitVector import BitVector
from numpy import *
from os import system

from dwt_marquage import Marquage_DWT
from outilsBase import setBit, getBit, conversion
from suite import Doublement
from chiffrement_image import Chiffrement
from coefficients import Coefficients


class Marquage_DWT_Doublement_IC:

    _nb_bits_coef = 16

    def __init__(self, hote = None, marque = None,
                       matrice_dwt = None, famille = None,
                       position = None, authentification = None,
                       LSB = None, iterations = None, mu = None,
                       Xo = None):
        '''
        Constructeur.
        '''

        # L'hote
        self._hote = hote
        if not self._hote:
            self._hote = raw_input("Quel fichier hôte ? ")
        else:
            print "Hote : ", self._hote

        # La marque
        self._marque = marque
        if not self._marque:
            self._marque = raw_input("Quel marque ? ")
        else:
            print "Marque : ", self._marque

        # Chiffrement
        unChiffrement = Chiffrement(fichier_image = self._marque,
                                    iterations = iterations,
                                    mu = mu,
                                    Xo = Xo)
        unChiffrement.chiffrement()
        self._strategie = unChiffrement.get_strategie()
        self._image = unChiffrement.get_image()


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

        # La famille d'ondelettes
        self._famille = famille
        if not self._famille:
            self._famille = raw_input("Quelle famille d'ondelettes ? ")
        else:
            print "Famille d'ondelettes : ", self._famille

        # La matrice DWT à modifier
        self._matrice_dwt = matrice_dwt
        if not self._matrice_dwt:
            self._matrice_dwt = raw_input("Quelle matrice DWT affecter ? ")
        else:
            print "Matrice DWT : ",self._matrice_dwt
        assert self._matrice_dwt in ['A1','A2','A3','H1','H2','H3',
                                     'D1','D2','D3','V1','V2','V3']

        # Position de l'insertion dans la matrice DWT
        self._position = position
        if not self._position:
            self._position =eval(raw_input("Début de l'insertion ? "))
        else:
            print "Début de l'insertion : ", self._position

        self._authentification = authentification
        if self._authentification == None:
            self._authentification=eval(raw_input("Authentification ? "))
        else:
            print "Authentification : ", str(self._authentification)
        if self._authentification:
            self._MSB = eval(raw_input("Quels MSB ? "))

        # Force de l'insertion
        self._LSB = LSB
        if not self._LSB:
            self._LSB = eval(raw_input("Liste des LSB ? "))
        else:
            print "LSB : ", self._LSB

        print




    def get_strategie(self, Uo, N):
        '''
        Récupère la stratégie des itérations chaotiques.
        '''
        if self._authentification:
            msc = Coefficients().getAllCoefs(self._hote, self._MSB)
            bits = BitVector(bitstring = self._code) ^ msc

        else:
             bits = BitVector(bitstring = self._code)
        doublement = Doublement(Uo, N,
                                generateur = Coefficients()\
                                                 .bit2coef(bits, N))
        suite = doublement.iterateur()
        while True:
            yield suite.next()





    def marquage(self):
        self._unMarquage = Marquage_DWT(hote = self._hote,
                       famille = self._famille)
        matrice = self._unMarquage.get_matrice_DWT(self._matrice_dwt)
        print "Taille de la matrice DWT : ",matrice.shape
        self._matrice_originelle = self._unMarquage.get_matrice_DWT(self._matrice_dwt)

        #self._x0, self._y0 = self._position
        #self._x1, self._y1 = matrice.shape
        #assert self._x0 < matrice.shape[0]
        #assert self._y0 < matrice.shape[1]

        taille = self._nb_bits_coef
        image = []
        for k in range(matrice.shape[0]):
            for l in range(matrice.shape[1]):
                s = 1 if matrice[k,l] >= 0 else -1
                texte = conversion(int(s*matrice[k,l]),2).zfill(taille)
                image += [int(texte[taille-x]) for x in self._LSB]

        assert self._position < len(image)

        compteur = 0

        # L'insertion
        strat = self.get_strategie(self._position, len(image))

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


        # Reconstruction de l'image
        compteur = 0
        for k in range(matrice.shape[0]):
            for l in range(matrice.shape[1]):
                s = 1 if matrice[k,l] >= 0 else -1
                x = [int(d) for d in conversion\
                                 (int(s*matrice[k,l]),2).zfill(taille)]
                for m in self._LSB:
                    x[taille-m] = image[compteur]
                    compteur += 1
                x = ''.join([str(t) for t in x])
                matrice[k,l] = matrice[k,l]%1
                matrice[k,l] += s*int(x,2)

        self._unMarquage.set_matrice_DWT(self._matrice_dwt, matrice)
        self._unMarquage.get_hote_marque().save('lena512marque.png')






    def get_hote_marque(self):
        return self._unMarquage.get_hote_marque()


    def get_matrice_DWT(self):
        return self._unMarquage.get_matrice_DWT(self._matrice_dwt)


    def extraction(self, hote = 'lena_marque.png'):
        print "\n\n==================== Extraction =========================="
        unTatouage = Marquage_DWT_Doublement_IC(hote = hote,
                                          marque = 'invader.png',
                                          famille = 'db1',
                                          position = 2,
                                          authentification = False,
                                          matrice_dwt = 'H2',
                                          LSB = [1],
                                          iterations = 20000,
                                          mu = 4,
                                          Xo = 0.65)
        unTatouage.marquage()
        lena_demarque = unTatouage.get_matrice_DWT()
        M = self._matrice_originelle - lena_demarque

        print "RMS : ",
        print sqrt(sum([k**2 for k in M.flatten().tolist()])/M.size)



if __name__ == '__main__':
    system('clear')
    print "==================== Insertion =========================="
    unTatouage = Marquage_DWT_Doublement_IC(hote = 'lena512.jpg',
                                          marque = 'invader.png',
                                          famille = 'db1',
                                          position = 1,
                                          authentification = False,
                                          matrice_dwt = 'H2',
                                          LSB = [2],
                                          iterations = 20000,
                                          mu = 4,
                                          Xo = 0.65)
    unTatouage.marquage()

    unTatouage.get_hote_marque().save('lena512marque.jpg')
    #unTatouage.get_hote_marque().show()
    #raw_input("Affichage de l'hôte marqué")



    diffs = difference(im.open('lena512.jpg'),unTatouage.get_hote_marque())
    maximum = max(list(diffs.getdata()))
    #diffs.point(lambda x : int(x*255/maximum)).show()
    #raw_input("Différences (facteur de "+str(255./maximum)+")")
    stats = ImageStat.Stat(diffs)
    print "Valeur moyenne des différences : ", stats.mean[0]
    print "Médiane des différences : ", stats.median[0]
    print "Valeurs extrémales : ", stats.extrema[0]
    print "Somme des différences : ", stats.sum[0]
    print "Valeur efficace (RMS) : ", stats.rms[0]
    print "Ecart-type : ", stats.stddev[0]

    # Calcul du PSNR
    from evaluation import Evaluation
    uneEva = Evaluation(im.open('lena512.jpg'),
                        im.open('lena512marque.png'))
    print "PSNR = ",uneEva.PSNR()


    unTatouage.extraction('lena512marque.jpg')
