"""
## Import Package
"""
import numpy as np
"""
## Linear Equation
"""
def F1(t):
    noise = np.random.normal(loc=0, scale=1)
    a1 = 0.063
    a2 = 5.284
    a3 = 4.887
    a4 = 10.34
    a5 = 105
    return a1*(t**4) - a2*(t**3) + a3*(t**2) + a4*(t**1) + a5 + noise
# number of simulation
n = 1000

# A (matrix) * x (vector with 5 dimensions) = b (vector with 5 dimensions)
A = np.zeros((n, 5))
b = np.zeros((n, 1))

for i in range(n):
    # t wll be a number between 0 and 100
    t = np.random.random()*100
    
    b[i] = F1(t)
    A[i, 0] = t**4
    A[i, 1] = t**3
    A[i, 2] = t**2
    A[i, 3] = t**1
    A[i, 4] = t**0

# Ax = b
x = np.linalg.lstsq(A, b)[0]
x
"""
## Non-Linear Equation
"""
def F2(t, A, B, C, D):
    return A*(t**B) + C*np.cos(D*t) + np.random.normal(0, 1, t.shape[0])
n = 1000
T = np.random.random((n, 1))*100
b2 = F2(T, 0.6, 1.2, 100, 0.4)
"""
### Using genetic algorithm to estimate four parameter we have defined: 0.6, 1.2, 100, 0.4
"""
# genetic algo: first step => encoding
# each row represents a person, gene or solution
population = np.random.randint(0, 2, (10000, 40))
population
# a function to covert gene into parameter
def gene2para(gene):
    A = (np.sum(2**np.arange(10)*gene[0:10]) - 511) / 100
    B = (np.sum(2**np.arange(10)*gene[10:20]) - 511) / 100
    C = (np.sum(2**np.arange(10)*gene[20:30]) - 511)
    D = (np.sum(2**np.arange(10)*gene[30:40]) - 511) / 100
    
    '''
    Because we have known the range of each parameter, and the solution should locate in this range
    A: -5.11 ~ 5.12
    
    '''
    return A, B, C, D
# "error" matrix will store how good the gene is 
error = np.zeros((10000, 1))
for generation in range(10):
    print("#{} Gneartion".format(generation+1))
    for i in range(10000):
        A, B, C, D = gene2para(population[i, :])
        error[i] = np.mean(abs(F2(T, A, B, C, D) - b2))
    
    # sort index of person based on its error
    sort_idx = np.argsort(error[:, 0])
    population = population[sort_idx, :]
    
    # genetic algo: second & third step => Survival of the Fittest
    for i in range(100, 10000):
        father = np.random.randint(0, 100)
        mother = np.random.randint(0, 100)
        while father == mother:
            mother = np.random.randint(0, 100)
        
        mask = np.random.randint(low=0, high=2, size=(1, 40))
        son_gene = np.zeros((40))
        
        mother_gene = population[mother, :]
        father_gene = population[father, :]
        son_gene[mask[0, :] == 1] = father_gene[mask[0, :] == 1]
        son_gene[mask[0, :] == 0] = mother_gene[mask[0, :] == 0]
        population[i, :] = son_gene
        
    
    # genetic algo: forth step => Mutation
    for i in range(1000):
        mutation_person = np.random.randint(0, 10000)
        mutation_gene = np.random.randint(0, 40)
        population[mutation_person, mutation_gene] = 1 - population[mutation_person, mutation_gene]
for i in range(10000):
    A, B, C, D = gene2para(population[i, :])
    error[i] = np.mean(abs(F2(T, A, B, C, D) - b2))
    
# sort index of person based on its error
sort_idx = np.argsort(error[:, 0])
population = population[sort_idx, :]
gene2para(population[0])
# possible improvement
# we should choose better parent (maybe based on its error)
# every person should have probability to live (person who has big error should have small probability instead of 0)
# adaptive mutation rate (high at first, low at final)