#!/usr/bin/env python
# coding: utf-8

# In[1]:


#Tous les codes sont basés sur l'environnement suivant
#python 3.7
#opencv 3.1.0
#pytorch 1.4.0

import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import cv2
import matplotlib.pyplot as plt
import numpy as np
import random
import math
import pickle
import random
from PIL import Image
import sys


# In[3]:


#Les fonctions dans ce bloc ne sont pas utilisées par le réseau, mais certaines fonctions d'outils


def tensor_imshow(im_tensor,cannel):
    b,c,h,w=im_tensor.shape
    if c==1:
        plt.imshow(im_tensor.squeeze().detach().numpy())
    else:
        plt.imshow(im_tensor.squeeze().detach().numpy()[cannel,:])

#    Obtenez des données d'entraînement
#    frag,vt=get_training_fragment(frag_size,image)
#    frag est un patch carrée de taille (frag_size*frag_size) a partir du image(Son emplacement est aléatoire)
#    vt est la vérité terrain de la forme Dirac.
def get_training_fragment(frag_size,im):
    h,w,c=im.shape
    n=random.randint(0,int(h/frag_size)-1)
    m=random.randint(0,int(w/frag_size)-1) 
    shape=frag_size/4
    vt_h=math.ceil((h+1)/shape)
    vt_w=math.ceil((w+1)/shape)
    vt=np.zeros([vt_h,vt_w])
    vt_h_po=round((vt_h-1)*(n*frag_size/(h-1)+(n+1)*frag_size/(h-1))/2)
    vt_w_po=round((vt_w-1)*(m*frag_size/(w-1)+(m+1)*frag_size/(w-1))/2)
    vt[vt_h_po,vt_w_po]=1
    vt = np.float32(vt)
    vt=torch.from_numpy(vt.reshape(1,1,vt_h,vt_w))
    
    return im[n*frag_size:(n+1)*frag_size,m*frag_size:(m+1)*frag_size,:],vt

#    Cette fonction convertit l'image en variable de type Tensor.
#    Toutes les données de calcul du réseau sont de type Tensor
#    Img.shape=[Height,Width,Channel]
#    Tensor.shape=[Batch,Channel,Height,Width]
def img2tensor(im):
    im=np.array(im,dtype="float32")
    tensor_cv = torch.from_numpy(np.transpose(im, (2, 0, 1)))
    im_tensor=tensor_cv.unsqueeze(0)
    return im_tensor

#   Trouvez les coordonnées de la valeur maximale dans une carte de corrélation
#   x,y=show_coordonnee(carte de corrélation)
def show_coordonnee(position_pred):
    map_corre=position_pred.squeeze().detach().numpy()
    h,w=map_corre.shape
    max_value=map_corre.max()
    coordonnee=np.where(map_corre==max_value)
    return coordonnee[0].mean()/h,coordonnee[1].mean()/w

#   Filtrer les patchs en fonction du nombre de pixels noirs dans le patch
#   Si seuls les pixels non noirs sont plus grands qu'une certaine proportion(seuillage), revenez à True, sinon False
def test_fragment32_32(frag,seuillage):
    a=frag[:,:,0]+frag[:,:,1]+frag[:,:,2]
    mask = (a == 0)
    arr_new = a[mask]
    if arr_new.size/a.size<=(1-seuillage):
        return True
    else:
        return False
    
#    Ces deux fonctions permettent de sauvegarder le réseau dans un fichier
#    ou de load le réseau stocké à partir d'un fichier
def save_net(file_path,net):
    pkl_file = open(file_path, 'wb')
    pickle.dump(net,pkl_file)
    pkl_file.close()
def load_net(file_path):   
    pkl_file = open(file_path, 'rb')
    net= pickle.load(pkl_file)
    pkl_file.close()
    return net


# In[4]:


#    Les fonctions de ce bloc sont utilisées pour construire le réseau

#   Créer un poids de type DeepMatch comme valeur initiale de Conv1 (non obligatoire)
def ini():
    kernel=torch.zeros([8,3,3,3])
    array_0=np.array([[1,2,1],[0,0,0],[-1,-2,-1]],dtype='float32')
    array_1=np.array([[2,1,0],[1,0,-1],[0,-1,-2]],dtype='float32')
    array_2=np.array([[1,0,-1],[2,0,-2],[1,0,-1]],dtype='float32')
    array_3=np.array([[0,-1,-2],[1,0,-1],[2,1,0]],dtype='float32')
    array_4=np.array([[-1,-2,-1],[0,0,0],[1,2,1]],dtype='float32')
    array_5=np.array([[-2,-1,0],[-1,0,1],[0,1,2]],dtype='float32')
    array_6=np.array([[-1,0,1],[-2,0,2],[-1,0,1]],dtype='float32')
    array_7=np.array([[0,1,2],[-1,0,1],[-2,-1,0]],dtype='float32')
    for i in range(3):
        kernel[0,i,:]=torch.from_numpy(array_0)
        kernel[1,i,:]=torch.from_numpy(array_1)
        kernel[2,i,:]=torch.from_numpy(array_2)
        kernel[3,i,:]=torch.from_numpy(array_3)
        kernel[4,i,:]=torch.from_numpy(array_4)
        kernel[5,i,:]=torch.from_numpy(array_5)
        kernel[6,i,:]=torch.from_numpy(array_6)
        kernel[7,i,:]=torch.from_numpy(array_7)
    return torch.nn.Parameter(kernel,requires_grad=True) 

#   Calculer le poids initial de la couche convolutive add
#   n, m signifie qu'il y a n * m sous-patches dans le patch d'entrée
#   Par exemple, le patch d'entrée est 16 * 16, pour les patchs 4 * 4 de la première couche, n = 4, m = 4
#   pour les patchs 8 * 8 de la deuxième couche, n = 2, m = 2
def kernel_add_ini(n,m):
    input_canal=int(n*m)
    output_canal=int(n/2)*int(m/2)
    for i in range(int(n/2)):
        for j in range(int(m/2)):
            kernel_add=np.zeros([1,input_canal],dtype='float32')
            kernel_add[0,i*2*m+j*2]=1
            kernel_add[0,i*2*m+j*2+1]=1
            kernel_add[0,(i*2+1)*m+j*2]=1
            kernel_add[0,(i*2+1)*m+j*2+1]=1
            if i==0 and j==0:
                add=torch.from_numpy(kernel_add.reshape(1,input_canal,1,1))
            else:
                add_=torch.from_numpy(kernel_add.reshape(1,input_canal,1,1))
                add=torch.cat((add,add_),0)
    return torch.nn.Parameter(add,requires_grad=False) 

#   Calculer le poids initial de la couche convolutive shift
#   shift+add Peut réaliser l'étape de l'agrégation
#   Voir ci-dessus pour les paramètres n et m. 
#   Pour des étapes plus détaillées, veuillez consulter mon rapport de stage
def kernel_shift_ini(n,m):
    input_canal=int(n*m)
    output_canal=int(n*m)
    
    kernel_shift=torch.zeros([output_canal,input_canal,3,3])
    
    array_0=np.array([[1,0,0],[0,0,0],[0,0,0]],dtype='float32')
    array_1=np.array([[0,0,1],[0,0,0],[0,0,0]],dtype='float32')
    array_2=np.array([[0,0,0],[0,0,0],[1,0,0]],dtype='float32')
    array_3=np.array([[0,0,0],[0,0,0],[0,0,1]],dtype='float32')
    
    kernel_shift_0=torch.from_numpy(array_0)
    kernel_shift_1=torch.from_numpy(array_1)
    kernel_shift_2=torch.from_numpy(array_2)
    kernel_shift_3=torch.from_numpy(array_3)
    
    
    for i in range(n):
        for j in range(m):
            if i==0 and j==0:
                kernel_shift[0,0,:]=kernel_shift_0
            else:
                if i%2==0 and j%2==0:
                    kernel_shift[i*m+j,i*m+j,:]=kernel_shift_0
                if i%2==0 and j%2==1:
                    kernel_shift[i*m+j,i*m+j,:]=kernel_shift_1
                if i%2==1 and j%2==0:
                    kernel_shift[i*m+j,i*m+j,:]=kernel_shift_2
                if i%2==1 and j%2==1:
                    kernel_shift[i*m+j,i*m+j,:]=kernel_shift_3
                    
    return torch.nn.Parameter(kernel_shift,requires_grad=False) 

#   Trouvez le petit patch(4 * 4) dans la n ème ligne et la m ème colonne du patch d'entrée
#   Ceci est utilisé pour calculer la convolution et obtenir la carte de corrélation
def get_patch(fragment,psize,n,m):
    return fragment[:,:,n*psize:(n+1)*psize,m*psize:(m+1)*psize]
###################################################################################################################
class Net(nn.Module):
    def __init__(self,frag_size,psize):
        super(Net, self).__init__()
       
        h_fr=frag_size
        w_fr=frag_size
        
        n=int(h_fr/psize) #    n*m patches dans le patch d'entrée
        m=int(w_fr/psize)
        
        self.conv1 = nn.Conv2d(3,8,kernel_size=3,stride=1,padding=1)
        #    Si vous souhaitez initialiser Conv1 avec les poids de DeepMatch, exécutez la ligne suivante
        #    self.conv1.weight=ini()
        self.Relu = nn.ReLU(inplace=True)
        self.maxpooling=nn.MaxPool2d(3,stride=2, padding=1)
        
        self.shift1=nn.Conv2d(n*m,n*m,kernel_size=3,stride=1,padding=1)
        self.shift1.weight=kernel_shift_ini(n,m)
        self.add1 = nn.Conv2d(n*m,int(n/2)*int(m/2),kernel_size=1,stride=1,padding=0)
        self.add1.weight=kernel_add_ini(n,m)
        
        n=int(n/2)
        m=int(m/2)
        if n>=2 and m>=2:#   Si n=m=1，Notre réseau n'a plus besoin de plus de couches pour agréger les cartes de corrélation
            self.shift2=nn.Conv2d(n*m,n*m,kernel_size=3,stride=1,padding=1)
            self.shift2.weight=kernel_shift_ini(n,m)
            self.add2 = nn.Conv2d(n*m,int(n/2)*int(m/2),kernel_size=1,stride=1,padding=0)
            self.add2.weight=kernel_add_ini(n,m)
        
        n=int(n/2)
        m=int(m/2)
        if n>=2 and m>=2:
            self.shift3=nn.Conv2d(n*m,n*m,kernel_size=3,stride=1,padding=1)
            self.shift3.weight=kernel_shift_ini(n,m)
            self.add3 = nn.Conv2d(n*m,int(n/2)*int(m/2),kernel_size=1,stride=1,padding=0)
            self.add3.weight=kernel_add_ini(n,m)
              
    def get_descripteur(self,img,using_cuda):
        #   Utilisez Conv1 pour calculer le descripteur,
        descripteur_img=self.Relu(self.conv1(img))
        b,c,h,w=descripteur_img.shape
        couche_constante=0.5*torch.ones([1,1,h,w])
        if using_cuda:
            couche_constante=couche_constante.cuda()
        #    Ajouter une couche constante pour éviter la division par 0 lors de la normalisation
        descripteur_img=torch.cat((descripteur_img,couche_constante),1)
        #    la normalisation
        descripteur_img_norm=descripteur_img/torch.norm(descripteur_img,dim=1)
        return descripteur_img_norm
    
    def forward(self,img,frag,using_cuda):
        psize=4
        #   Utilisez Conv1 pour calculer le descripteur,
        descripteur_input1=self.get_descripteur(img,using_cuda)
        descripteur_input2=self.get_descripteur(frag,using_cuda)
        
        b,c,h,w=frag.shape
        n=int(h/psize)
        m=int(w/psize)
        
        #######################################
        #    Calculer la carte de corrélation par convolution pour les n*m patchs plus petit.
        for i in range(n):
            for j in range(m):
                if i==0 and j==0:
                    map_corre=F.conv2d(descripteur_input1,get_patch(descripteur_input2,psize,i,j),padding=2)
                else:
                    a=F.conv2d(descripteur_input1,get_patch(descripteur_input2,psize,i,j),padding=2)
                    map_corre=torch.cat((map_corre,a),1)
        ########################################
        #    Étape de polymérisation
        map_corre=self.maxpooling(map_corre)
        map_corre=self.shift1(map_corre)
        map_corre=self.add1(map_corre)
        
        #########################################
        #    Répétez l'étape d'agrégation jusqu'à obtenir le graphique de corrélation du patch d'entrée
        n=int(n/2)
        m=int(m/2)
        if n>=2 and m>=2:
            map_corre=self.maxpooling(map_corre)
            map_corre=self.shift2(map_corre)
            map_corre=self.add2(map_corre)
        
        
        n=int(n/2)
        m=int(m/2)
        if n>=2 and m>=2:
            map_corre=self.maxpooling(map_corre)
            map_corre=self.shift3(map_corre)
            map_corre=self.add3(map_corre)
        
        
        b,c,h,w=map_corre.shape
        #   Normalisation de la division par maximum
        map_corre=map_corre/(map_corre.max())
        #   Normalisation SoftMax
        #map_corre=(F.softmax(map_corre.reshape(1,1,h*w,1),dim=2)).reshape(b,c,h,w)
        return map_corre


# In[5]:


def run_net(net,img,frag,frag_size,using_cuda):
    h,w,c=frag.shape
    n=int(h/frag_size)
    m=int(w/frag_size)
    frag_list=[]
    #####################################
    #    Obtenez des patchs carrés des fragments et mettez-les dans la frag_list
    for i in range(n):
        for j in range(m):
            frag_32=frag[i*frag_size:(i+1)*frag_size,j*frag_size:(j+1)*frag_size]
            if test_fragment32_32(frag_32,0.6):
                frag_list.append(frag_32)
    img_tensor=img2tensor(img)
    ######################################
    if using_cuda:
        img_tensor=img_tensor.cuda()
        
    coordonnee_list=[]
    #######################################
    #    Utilisez le réseau pour calculer les positions de tous les patchs dans frag_list[]
    #    Mettez le résultat du calcul dans coordonnee_list[]
    for i in range(len(frag_list)):
        frag_tensor=img2tensor(frag_list[i])
        if using_cuda:
            frag_tensor=frag_tensor.cuda()
        res=net.forward(img_tensor,frag_tensor,using_cuda)
        if using_cuda:
            res=res.cpu()
        po_h,po_w=show_coordonnee(res)
        coordonnee_list.append([po_h,po_w])
    h_img,w_img,c=img.shape
    position=[]
    for i in range(len(coordonnee_list)):
        x=int(round(h_img*coordonnee_list[i][0]))
        y=int(round(w_img*coordonnee_list[i][1]))
        position.append([x,y])
    return position


# In[10]:


if __name__=='__main__':
    
    #     La taille du patch d'entrée est de 16*16
    frag_size=16
    #     La taille du plus petit patch dans réseau est de 4 *4 fixée
    psize=4
    using_cuda=True
    
    
    net=Net(frag_size,psize)
    
    #     Pour chaque fresque, le nombre d'itérations est de 1000
    itera=1000
    
    if using_cuda:
        net=net.cuda()
    
    #    Choisissez l'optimiseur et la fonction de coût
    optimizer = torch.optim.Adam(net.parameters())
    loss_func = torch.nn.MSELoss()
    
    #    Dans le processus d'apprentissage du réseau,le changement d'erreur est placé dans loss_value=[] 
    #    et le changement de Conv1 poids est placé dans para_value[]
    loss_value=[]
    para_value=[]
    ####################################################training_net
    
    #Les données d'entraînement sont 6 fresques
    for n in range(6):
        im_path="./fresque"+str(n)+".ppm"
        img_training=cv2.imread(im_path)
        h,w,c=img_training.shape
        
        #    Si la peinture murale est trop grande, sous-échantillonnez-la et rétrécissez-la
        while h*w>(1240*900):
            img_training=cv2.resize(img_training,(int(h/2),int(w/2)),interpolation=cv2.INTER_CUBIC)
            h,w,c=img_training.shape
        im_tensor=img2tensor(img_training)
        
        if using_cuda:
            im_tensor=im_tensor.cuda()
        for i in range(itera):
            #     Tous les 100 cycles, enregistrez le changement de poids
            if i%100==0:
                para=net.conv1.weight
                para=para.detach().cpu()
                para_value.append(para)
            frag,vt=get_training_fragment(frag_size,img_training)
            frag_tensor=img2tensor(frag)
            if using_cuda:
                vt=vt.cuda()
                frag_tensor=frag_tensor.cuda()
            #     Utilisez des patchs et des fresques de données d'entraînement pour faire fonctionner le réseau
            frag_pred=net.forward(im_tensor,frag_tensor,using_cuda)
            b,c,h,w=vt.shape
            #    Utilisez la fonction de coût pour calculer l'erreur
            err_=loss_func(vt,frag_pred)
            #    Utilisez l'optimiseur pour ajuster le poids de Conv1
            optimizer.zero_grad()
            err_.backward(retain_graph=True)
            optimizer.step()
            
            loss_value.append(err_.tolist())
    
            del frag_tensor,frag_pred,err_,vt
            torch.cuda.empty_cache()


# In[7]:


len(loss_value)


# In[11]:


plt.plot(loss_value)


# In[12]:


file_path="./net_trainned6000"
save_net(file_path,net)