import os
from google.cloud import vision
import re
import torch
import torchvision
import numpy as np
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import tempfile
import json

def getcredentials():
    secret_key_credential = os.getenv("secret_key")

    with tempfile.NamedTemporaryFile(mode='w+', delete= False, suffix=".json") as temp_file:
        temp_file.write(secret_key_credential)
        tempfile_name = temp_file.name

    return tempfile_name


os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = getcredentials()

##
def info_new_cni(donnees):
    ##
    informations = {}

    # Utilisation d'expressions régulières pour extraire les informations spécifiques
    numero_carte = re.search(r'n° (C\d+)', ' '.join(donnees))
    #prenom_nom = re.search(r'Prénom\(s\)\s+(.*?)\s+Nom\s+(.*?)\s+Signature', ' '.join(donnees))
    nom = re.search(r'Nom\s+(.*?)\s', ' '.join(donnees))
    prenom = re.search(r'Prénom\(s\)\s+(.*?)\s+Nom\s+(.*?)', ' '.join(donnees))
    date_naissance = re.search(r'Date de Naissance\s+(.*?)+(\d{2}/\d{2}/\d{4})', ' '.join(donnees))
    lieu_naissance = re.search(r'Lieu de Naissance\s+(.*?)\s', ' '.join(donnees))
    taille = re.search(r'Sexe Taille\s+(.*?)+(\d+,\d+)', ' '.join(donnees))
    nationalite = re.search(r'Nationalité\s+(.*?)\s+\d+', ' '.join(donnees))
    date_expiration = re.search(r'Date d\'expiration\s+(\d+/\d+/\d+)', ' '.join(donnees))
    sexe = re.search(r'Date de Naissance\s+(.*?)+(\d{2}/\d{2}/\d{4})+(.*)', ' '.join(donnees))

    # Stockage des informations extraites dans un dictionnaire
    if numero_carte:
        informations['Numéro de carte'] = numero_carte.group(1)
    if nom :
        informations['Nom'] = nom.group(1)

    if prenom:
        informations['Prénom'] = prenom.group(1)

    if date_naissance:
        informations['Date de Naissance'] = date_naissance.group(2)
    if lieu_naissance:
        informations['Lieu de Naissance'] = lieu_naissance.group(1)
    if taille:
        informations['Taille'] = taille.group(2)
    if nationalite:
        informations['Nationalité'] = nationalite.group(1)
    if date_expiration:
        informations['Date d\'expiration'] = date_expiration.group(1)
    if sexe :
      informations['sexe'] =  sexe.group(3)[:2]

    return informations

##

def info_ancien_cni(infos):
  """ Extract information in row data of ocr"""

  informations = {}

  immatriculation_patern = r'Immatriculation:\s+(C \d{4} \d{4} \d{2})'
  immatriculation = re.search(immatriculation_patern, ''.join(infos))
  nom = infos[4]
  prenom_pattern = r'Nom\n(.*?)\n'
  prenom = re.search(prenom_pattern, '\n'.join(infos))
  sexe_pattern = r'Prénoms\n(.*?)\n'
  sexe = re.search(sexe_pattern, '\n'.join(infos))
  taille_pattern = r'Sexe\n(.*?)\n'
  taille = re.search(taille_pattern, '\n'.join(infos))
  date_naiss_pattern = r'Taille\s+(.*?)+(\d+/\d+/\d+)' # r'Taille (m)\n(.*?)\n'
  date_naissance = re.search(date_naiss_pattern, ' '.join(infos))
  lieu_pattern = r'Date de Naissance\n(.*?)\n'
  lieu_naissance = re.search(lieu_pattern, '\n'.join(infos))
  valide_pattern = r'Valide jusqu\'au+(.*?)+(\d+/\d+/\d+)'
  validite = re.search(valide_pattern, ' '.join(infos))

  # Stockage des informations extraites dans un dictionnaire
  if immatriculation:
      informations['Immatriculation'] = immatriculation.group(1)
  if nom :
      informations['Nom'] =  infos[4]

  if prenom:
      informations['Prénom'] = prenom.group(1)

  if date_naissance:
      informations['Date de Naissance'] = date_naissance.group(2)
  if lieu_naissance:
      informations['Lieu de Naissance'] = lieu_naissance.group(1)
  if taille:
      informations['Taille'] = taille.group(1)
  
  if validite:
      informations['Date d\'expiration'] = validite.group(2)
  if sexe :
    informations['sexe'] =  sexe.group(1)

  return informations

  ##
def filtrer_elements(liste):
  elements_filtres = []
  for element in liste:
      if element not in ['\r',"RÉPUBLIQUE DE CÔTE D'IVOIRE", "MINISTÈRE DES TRANSPORTS", "PERMIS DE CONDUIRE"]:
          elements_filtres.append(element)
  return elements_filtres

def permis_de_conduite(donnees):
  """ Extraire les information de permis de conduire"""

  informations = {}

  infos = filtrer_elements(donnees)
  
  nom_pattern = r'Nom\n(.*?)\n'
  nom = re.search(nom_pattern, '\n'.join(infos))
  prenom_pattern = r'Prénoms\n(.*?)\n'
  prenom = re.search(prenom_pattern, '\n'.join(infos))
  date_lieu_naissance_patern = r'Date et lieu de naissance\n(.*?)\n'
  date_lieu_naissance = re.search(date_lieu_naissance_patern, '\n'.join(infos))
  date_lieu_delivrance_patern = r'Date et lieu de délivrance\n(.*?)\n'
  date_lieu_delivrance = re.search(date_lieu_delivrance_patern, '\n'.join(infos))
  numero_pattern = r'Numéro du permis de conduire\n(.*?)\n'
  numero = re.search(numero_pattern, '\n'.join(infos))
  restriction_pattern = r'Restriction\(s\)\s+(.*?)+(.*)'
  restriction = re.search(restriction_pattern, ' '.join(infos))
  
  # Stockage des informations extraites dans un dictionnaire
  if nom:
    informations['Nom'] = nom.group(1)

  if prenom : 
     informations['Prenoms'] = prenom.group(1)
  if date_lieu_naissance : 
    informations['Date_et_lieu_de_naissance'] = date_lieu_naissance.group(1)
  if date_lieu_naissance :
    informations['Date_et_lieu_de_délivrance'] = date_lieu_delivrance.group(1)

  informations['Categorie'] = infos[0]
  if numero:
    informations['Numéro_du_permis_de_conduire'] = numero.group(1)

  if restriction:
    informations['Restriction(s)'] = restriction.group(2)

  return informations


# Fonction pour extraire les informations individuelles
def extraire_informations_carte(path, type_de_piece=1):
    """ Detect text in identity card"""

    client = vision.ImageAnnotatorClient()

    with open(path,'rb') as image_file:
      content = image_file.read()
    
    image = vision.Image(content = content)

    # for non dense text
    #response = client.text_detection(image=image)
    #for dense text
    response = client.document_text_detection(image = image)
    texts = response.text_annotations
    ocr_texts = []

    for text in texts:
      ocr_texts.append(f"\r\n{text.description}")

    if response.error.message :
      raise Exception("{}\n For more informations check : https://cloud.google.com/apis/design/errors".format(response.error.message))

    donnees = ocr_texts[0].split('\n')

    if type_de_piece ==1:
      return info_new_cni(donnees)
    elif type_de_piece == 2:
      return info_ancien_cni(donnees)
    elif type_de_piece == 3:
      return permis_de_conduite(donnees)
    else :
      return "Le traitement de ce type de document n'est pas encore pris en charge"

def load_checkpoint(path):
    print('--> Loading checkpoint')
    return torch.load(path,map_location=torch.device('cpu'))

def make_prediction(image_path):
      
    # define the using of GPU or CPU et background training
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    ## load model
    model = load_checkpoint("data/model.pth")
    ## transformation
    test_transforms = A.Compose([
        A.Resize(height=224, width=224, always_apply=True),
        A.Normalize(always_apply=True),
        ToTensorV2(always_apply=True),])

    ## read the image
    image = np.array(Image.open(image_path).convert('RGB'))
    transformed = test_transforms(image= image)
    image_transformed = transformed["image"]
    image_transformed = image_transformed.unsqueeze(0)
    image_transformed = image_transformed.to(device)

    model.eval()
    with torch.set_grad_enabled(False):
      output = model(image_transformed)

    # Post-process predictions
    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    predicted_class = torch.argmax(probabilities).item()
    proba = float(max(probabilities))
    

    return proba, predicted_class