#-*-coding:utf8-*-
import Image as im
import pywt
from numpy import *


class Marquage_DWT:

    _matrices_DWT = {'A1' : None,
                     'H1' : None, 'D1' : None, 'V1' : None,
                     'A2' : None,
                     'H2' : None, 'D2' : None, 'V2' : None,
                     'A3' : None,
                     'H4' : None, 'D3' : None, 'V3' : None
                     }

    def __init__(self, hote = 'lena.png', famille = 'db1'):
        '''
        Constructeur
        '''
        self._hote = im.open(hote)

        self._hote_marque = im.new(self._hote.mode,
                                   self._hote.size)
        self._hote_marque.putdata(self._hote.getdata())

        self._matrice_hote = array(self._hote.getdata()).\
                                            reshape(self._hote.size)

        self._famille = famille

        self._get_matrices()




    def _get_matrices(self):

        A1,(H1,V1,D1) = pywt.wavedec2\
                        (self._matrice_hote,
                        self._famille,
                        level = 1)

        self._matrices_DWT['A1'] = A1
        self._matrices_DWT['H1'] = H1
        self._matrices_DWT['V1'] = V1
        self._matrices_DWT['D1'] = D1

        A2,(H2,V2,D2),(H1,V1,D1) = pywt.wavedec2\
                        (self._matrice_hote,
                        self._famille,
                        level = 2)

        self._matrices_DWT['A2'] = A2
        self._matrices_DWT['H2'] = H2
        self._matrices_DWT['V2'] = V2
        self._matrices_DWT['D2'] = D2

        A3,(H3,V3,D3), (H2,V2,D2), (H1,V1,D1) = pywt.wavedec2\
                        (self._matrice_hote,
                        self._famille,
                        level = 3)

        self._matrices_DWT['A3'] = A3
        self._matrices_DWT['H3'] = H3
        self._matrices_DWT['V3'] = V3
        self._matrices_DWT['D3'] = D3



    def get_matrice_DWT(self, matrice):
        assert matrice in self._matrices_DWT.keys()
        return self._matrices_DWT[matrice]


    def set_matrice_DWT(self, matrice, valeur):
        assert matrice in self._matrices_DWT.keys()
        self._matrices_DWT[matrice] = valeur
        self._reconstruit_image(matrice)


    def _reconstruit_image(self, matrice_modifiee):
        if '1' in matrice_modifiee:
            self._hote_marque.putdata(pywt.waverec2((\
                        self._matrices_DWT['A1'],
                        (self._matrices_DWT['H1'],
                         self._matrices_DWT['V1'],
                         self._matrices_DWT['D1'])),
                        self._famille).flatten())
        elif '2' in matrice_modifiee:
            self._hote_marque.putdata(pywt.waverec2((\
                        self._matrices_DWT['A2'],
                        (self._matrices_DWT['H2'],
                         self._matrices_DWT['V2'],
                         self._matrices_DWT['D2']),
                        (self._matrices_DWT['H1'],
                         self._matrices_DWT['V1'],
                         self._matrices_DWT['D1'])),
                        self._famille).flatten())
        elif '3' in matrice_modifiee:
            self._hote_marque.putdata(pywt.waverec2((\
                        self._matrices_DWT['A3'],
                        (self._matrices_DWT['H3'],
                         self._matrices_DWT['V3'],
                         self._matrices_DWT['D3']),
                        (self._matrices_DWT['H2'],
                         self._matrices_DWT['V2'],
                         self._matrices_DWT['D2']),
                        (self._matrices_DWT['H1'],
                         self._matrices_DWT['V1'],
                         self._matrices_DWT['D1'])),
                        self._famille).flatten())


    def get_hote_marque(self):
        return self._hote_marque


if __name__ == '__main__':
    print "Mise à zéro des matrices d'approximation"
    for k in ['A1','A2','A3']:
        print "  Matrice "+k
        unMarquage = Marquage_DWT(hote = 'lena.png',
                       famille = 'db1')
        matrice_nulle = zeros(unMarquage.get_matrice_DWT(k).shape)
        unMarquage.set_matrice_DWT(k,matrice_nulle)
        unMarquage.get_hote_marque().show()
        raw_input('   (Appuyez sur une touche)')

    print
    print "Mise à zéro des matrices de détail D"
    for k in ['D1','D2','D3']:
        print "  Matrice "+k
        unMarquage = Marquage_DWT(hote = 'lena.png',
                       famille = 'db1')
        matrice_nulle = zeros(unMarquage.get_matrice_DWT(k).shape)
        unMarquage.set_matrice_DWT(k,matrice_nulle)
        unMarquage.get_hote_marque().show()
        raw_input('   (Appuyez sur une touche)')
