from abc import abstractmethod, abstractproperty
class GeneticAlgorithm():
    """Genetic Algorithm Interface."""
    @abstractmethod
    def create_population(self):
        pass
    @abstractmethod
    def selection(self):
        pass
    @abstractmethod
    def mutation(self):
        pass
    @abstractmethod
    def cross(self):
        pass
    @abstractmethod
    def call(self):
        pass

import numpy as np
import random
class Individ:
    """Generate a new individ."""
    def __init__(self, y1, y2, length):
        self.length = length
        self.A=np.random.randint(y1, y2, self.length)    # Declare the attribute "A" of our individ (this will be the line 010..101).
        self.fit=0                                       # Declare the attribute "fit" of the individ (this will be the sum of the elements of our line 010..101) (for now, we will assign it the value 0)

class SimpleSegmentationGA(GeneticAlgorithm):
    """Genetic Algorithm Implementation."""
    def __init__(self, pop_num, lenght, delta_x, gray):
        self.pop_num = pop_num
        self.lenght = lenght    # Length of individ.
        self.delta_x =  delta_x # Value of h.
        self.gray = 256 - gray # Inversion.

def create_population(self, y1, y2):
  """Creates a population, Dad, Mom and 4 sons"""
  self.y1 = y1
  self.y2 = y2
  self.population = []                                               # Declare an array "population" in which we will store a population of 5 individuals.
  for i in range(self.pop_num):
    c=Individ(self.y1, self.y2, self.lenght)                       # Create a new individual.
    self.population.append(c)                                      # Add the i-th unique individual to the array (fill in our population)
    self.mother = Individ(self.y1, self.y2, self.lenght)               # Initialize the variables with which we will work: mom, dad, 4 sons ..
    self.father = Individ(self.y1, self.y2, self.lenght)
    self.son1 = Individ(self.y1, self.y2, self.lenght)
    self.son2 = Individ(self.y1, self.y2, self.lenght)
    self.son3 = Individ(self.y1, self.y2, self.lenght)
    self.son4 = Individ(self.y1, self.y2, self.lenght)
    self.par_and_sons = []                                             #.. and an array of individs "Parents and Children" in which we will store
    for j in range(self.pop_num*3):                                    # Initialize our array of "Parents and Sons" with random individs.
      self.par_and_sons.append(Individ(self.y1, self.y2, self.lenght))

def cross(self):
  """Crosses the best individs with each other."""
  for i in range(self.pop_num):                                      # Put in the first pop_num elements of the "Parents and Sons" array our entire input population.
    self.par_and_sons[i].A=self.population[i].A.copy()
    random.shuffle(self.population)                                    # Shuffle population.
    tt=0                                                               # The counter that is needed to implement a non-trivial crossing.
    for s in range(0,self.pop_num,2):                                  # From 0 to pop_num with step 2. That is. here we take pop_num / 2 pairs of parents.
      self.mother.A=self.population[tt+int(self.pop_num/2)].A        # Let the last pop_num / 2 individuals of our population be our mothers.
      self.father.A=self.population[tt].A                            # And let first pop_num / 2 individuals of our population be dads.
      tt=tt+1
      ran=random.random()
      for n in range(self.lenght):                                   # Crossover.
        if random.random()>0.5:
          self.son1.A[n] = self.father.A[n]
          self.son2.A[self.lenght-1-n] = self.father.A[n]
          self.son3.A[n] = self.mother.A[n]
          self.son4.A[self.lenght-1-n] = self.mother.A[n]
        else:
          self.son1.A[n] = self.mother.A[n]
          self.son2.A[self.lenght-1-n] = self.mother.A[n]
          self.son3.A[n] = self.father.A[n]
          self.son4.A[self.lenght-1-n] = self.father.A[n]
        self.par_and_sons[self.pop_num+2*s].A = self.son1.A.copy()
        self.par_and_sons[self.pop_num+2*s+1].A = self.son2.A.copy()
        self.par_and_sons[self.pop_num+2*s+2].A = self.son3.A.copy()
        self.par_and_sons[self.pop_num+2*s+3].A = self.son4.A.copy()

def mutation(self):
  """Mutates individuals with a 20% probability."""
  for r in range(self.pop_num*3, 5):                                 # Mutation.
    for w in range(0,self.lenght):
      if random.random()<0.2:
        self.par_and_sons[r].A[w] = self.par_and_sons[r].A[w] + np.random.randint(-20, 20)  # Offset + -20 pixels.

def selection(self):
  """Sorts by fit and selects the best pop_num individs."""
  for i in range(self.pop_num*3):                                    # It is important. Next, we will rank the array of parents and children in ascending order of survivability (sum (fit)).
    self.par_and_sons[i].fit = SimpleSegmentationGA.fitness_function(self.gray, self.delta_x, self.lenght, self.par_and_sons[i].A)
    #  Sort.
    self.par_and_sons = sorted(self.par_and_sons, key=lambda individ: individ.fit)
    self.population=self.par_and_sons[:self.pop_num].copy()

def fitness(image, delta_x, length, individ):
    """Responsible for calculating the "fit" (counts the amount)."""
    summa = 0
    sum_vrt = 0
    for i in range(length):
        sum_ = np.sum(image[individ[i], i*delta_x:i*delta_x+delta_x])
        if i>0:
            if individ[i]>individ[i-1]:
                sum_vrt = np.sum(image[individ[i-1]:individ[i], i*delta_x])
            else:
                sum_vrt = np.sum(image[individ[i]:individ[i-1], i*delta_x])
        summa=summa + sum_ + sum_vrt
    return summa

def call(self):
  """Calls other functions and returns the selected population."""
  self.cross()
  self.mutation()
  self.selection()
  return  self.population[0]

import numpy as np
import random, time
import cv2
from scipy.ndimage import gaussian_filter
from scipy.signal import find_peaks
from numba import njit
import pymp
@njit
def fitness(image, delta_x, length, individ):
    """Responsible for calculating the "fit" (counts the amount)."""
    summa = 0
    sum_vrt = 0
    for i in range(length):
        sum_ = np.sum(image[individ[i], i*delta_x:i*delta_x+delta_x])
        if i>0:
            if individ[i]>individ[i-1]:
                sum_vrt = np.sum(image[individ[i-1]:individ[i], i*delta_x])
            else:
                sum_vrt = np.sum(image[individ[i]:individ[i-1], i*delta_x])
        summa=summa + sum_ + sum_vrt
    return summa
def find_peaks_(image):
    """Calculates ranges of random numbers for our individs."""
    height, width = image.shape[:2]
    img_matrix = [sum(i)/len(i) for i in image]
    x=[i for i in range(height)]
    y = [255-i for i in img_matrix]
    y = gaussian_filter(y, sigma=20)
    maxs, _ = find_peaks(y)
    maxs = maxs.tolist()
    return maxs
def run(gaObj, peaks, epoch, parallel=True, number_cpu=4):
    """Can choose how to run."""
    gen_lines = pymp.shared.list()
    if parallel:
        with pymp.Parallel(number_cpu) as p:
            range_peaks = p.range(len(peaks)-1)
            gen_lines = run_genetic_algorithm(gaObj, peaks, epoch, range_peaks, gen_lines)
    else:
        range_peaks = range(len(peaks)-1)
        gen_lines = run_genetic_algorithm(gaObj, peaks, epoch, range_peaks, gen_lines)
    return np.moveaxis(np.array(gen_lines), 0, 1)    # Swap line and epoch axes.
def draw_lines(image, lines, delta_x, imagename, epoch='last', color=(0,0,255), thickness = 3, IMG_FOLDER='output'):
    """Draws the result of a genetic algorithm."""
    image_copy = image.copy()
    for j in lines:
        for i in range(len(j)):
            x1 = i*delta_x
            x2 = i*delta_x+delta_x
            y1 = j[i]
            y2 = j[i]
            start_point = (x1, y1)
            if i>0:
                cv2.line(image_copy, preview_point, start_point, color, thickness)
            end_point = (x2, y2)
            cv2.line(image_copy, start_point, end_point, color, thickness)
            preview_point = end_point
    cv2.imwrite(f"{IMG_FOLDER}/{imagename[:-4]}_gen_line_{epoch}.jpg", image_copy)
def run_genetic_algorithm(gaObj, peaks, epoch, range_peaks, gen_lines):
    """Runs a genetic alghoritm."""
    for line in range_peaks:
        gaObj.create_population(peaks[line], peaks[line+1])
        epoch_line = pymp.shared.list()
        for p in range(epoch):
            st_time = time.time()
            gen_line = gaObj.call()
            epoch_line.append(gen_line.A)            # For draw results each epoch, add results of each epoch.
            print(f'Line = {line}, Epoch = {p}, fit = {gen_line.fit}, Time = {time.time()-st_time}')
        gen_lines.append(epoch_line)
    return gen_lines
def crop_lines(image, lines, delta_x, imagename, IMG_FOLDER='crops'):
    """Crop the lines."""
    for i in range(-1,len(lines)):
        height, width = image.shape[:2]
        if i == (-1):
            first_line = [[i*delta_x, 0] for i, el in enumerate(lines[i])]
        else:
            first_line = [[i*delta_x, el] for i, el in enumerate(lines[i])]
        first_line = first_line[::-1] # Reverse first line.
        if i < len(lines)-1:
            second_line = [[i*delta_x, el] for i, el in enumerate(lines[i+1])]
        else:
            second_line = [[i*delta_x, height] for i, el in enumerate(lines[i])]
        points = np.array([first_line + second_line])
        mask = np.zeros((height, width), dtype=np.uint8)
        cv2.fillPoly(mask, points, (255))
        res = cv2.bitwise_and(image,image,mask = mask)
        rect = cv2.boundingRect(points) # Returns (x,y,w,h) of the rect
        im2 = np.full((res.shape[0], res.shape[1], 3), (0, 255, 0), dtype=np.uint8 ) # You can also use other colors or simply load another image of the same size
        maskInv = cv2.bitwise_not(mask)
        colorCrop = cv2.bitwise_or(im2,im2,mask = maskInv)
        finalIm = res + colorCrop
        cropped = finalIm[rect[1]: rect[1] + rect[3], rect[0]: rect[0] + rect[2]]
        cv2.imwrite(f"{IMG_FOLDER}/croped_{imagename[:-4]}_{i}.png", cropped)

import cv2
import glob, os
from utils import *
from SimpleSegmentationGA import SimpleSegmentationGA
if __name__ == '__main__':
    # Hyper parameters for genetic alghoritms.
    POP_NUM=80
    DELTAX=50 # Value of h.
    EPOCH=100
    IMAGES_PATH = 'input'
    # IMAGE_PATH = '7.jpg'
    images = glob.glob(f'{IMAGES_PATH}/*.jpg')
    for imagename in images:
        fstart_time = time.time()
        # Image operations.
        image = cv2.imread(imagename)
        height, width = image.shape[:2]
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # Calculate individ's lenght.
        length_individ = int(width/DELTAX)
        # Calculate the range of random numbers for an individ.
        peaks = find_peaks_(gray)
        # Create object of class.
        ssga = SimpleSegmentationGA(POP_NUM, length_individ, DELTAX, gray)
        start_time = time.time()
        # Run genetic algorithm.
        lines = run(ssga, peaks, EPOCH, parallel=True)
        print('Time run =', time.time()-start_time)
        last_epoch_lines = sorted(lines[-1], key = lambda x:x[0]) # Sort for crop. Because when using parallel, the lines lose consistency.
        start_time = time.time()
        # Draw the result of the genetic alghoritm. lines[-1] -> result of last epoch.
        draw_lines(image, last_epoch_lines, DELTAX, os.path.basename(imagename))
        print('Time draw_lines =', time.time()-start_time)
        start_time = time.time()
        # Crop the result of genetic alghoritm.
        crop_lines(image, last_epoch_lines, DELTAX, os.path.basename(imagename))
        print('Time crop_lines =', time.time()-start_time)
        print('Time for one image =', time.time()-fstart_time)