import matplotlib.pyplot as plt
import numpy as np
 
 
def generate_histogram(img, do_print):
  """
  @params: img: can be a grayscale or color image. We calculate the Normalized histogram of this image.
  @params: do_print: if or not print the result histogram
  @return: will return both histogram and the grayscale image 
  """
  if len(img.shape) == 3: # img is colorful, so we convert it to grayscale
    gr_img = np.mean(img, axis=-1)
  else:
    gr_img = img
  '''now we calc grayscale histogram'''
  gr_hist = np.zeros([256])
 
 
  for x_pixel in range(gr_img.shape[0]):
    for y_pixel in range(gr_img.shape[1]):
      pixel_value = int(gr_img[x_pixel, y_pixel])
      gr_hist[pixel_value] += 1
 
 
  '''normalizing the Histogram'''
  gr_hist /= (gr_img.shape[0] * gr_img.shape[1])
  if do_print:
    print_histogram(gr_hist, name="n_h_img", title="Normalized Histogram")
  return gr_hist, gr_img
 
 
 def print_histogram(_histrogram, name, title):
  plt.figure()
  plt.title(title)
  plt.plot(_histrogram, color='#ef476f')
  plt.bar(np.arange(len(_histrogram)), _histrogram, color='#b7b7a4')
  plt.ylabel('Number of Pixels')
  plt.xlabel('Pixel Value')
  plt.savefig("hist_" + name)

def equalize_histogram(img, histo, L):
  eq_histo = np.zeros_like(histo)
  en_img = np.zeros_like(img)
  for i in range(len(histo)):
    eq_histo[i] = int((L - 1) * np.sum(histo[0:i]))
  print_histogram(eq_histo, name="eq_"+str(index), title="Equalized Histogram")
  '''enhance image as well:'''
  for x_pixel in range(img.shape[0]):
    for y_pixel in range(img.shape[1]):
      pixel_val = int(img[x_pixel, y_pixel])
      en_img[x_pixel, y_pixel] = eq_histo[pixel_val]
  '''creating new histogram'''
  hist_img, _ = generate_histogram(en_img, print=False, index=index)
  print_img(img=en_img, histo_new=hist_img, histo_old=histo, index=str(index), L=L)
  return eq_histo

def find_value_target(val, target_arr):
  key = np.where(target_arr == val)[0]
 
 
  if len(key) == 0:
    key = find_value_target(val+1, target_arr)
    if len(key) == 0:
      key = find_value_target(val-1, target_arr)
  vvv = key[0]
  return vvv
 
 
 
 
def match_histogram(inp_img, hist_input, e_hist_input, e_hist_target, _print=True):
  '''map from e_inp_hist to 'target_hist '''
  en_img = np.zeros_like(inp_img)
  tran_hist = np.zeros_like(e_hist_input)
  for i in range(len(e_hist_input)):
    tran_hist[i] = find_value_target(val=e_hist_input[i], target_arr=e_hist_target)
  print_histogram(tran_hist, name="trans_hist_", title="Transferred Histogram")
  '''enhance image as well:'''
  for x_pixel in range(inp_img.shape[0]):
    for y_pixel in range(inp_img.shape[1]):
      pixel_val = int(inp_img[x_pixel, y_pixel])
      en_img[x_pixel, y_pixel] = tran_hist[pixel_val]
  '''creating new histogram'''
  hist_img, _ = generate_histogram(en_img, print=False, index=3)
  print_img(img=en_img, histo_new=hist_img, histo_old=hist_input, index=str(3), L=L)