from neorl import DE, GWO, SSA, WOA, AEO, MFO, JAYA, HHO, PSO, ES, EPSO
from neorl import EDEV, HCLPSO, BAT, SA, CS, TS, PESA, PESA2
import math, random

#################################
# Define Vessel Function 
#Mixed discrete/continuous/grid
#################################
def Vessel(individual):
    """
    Pressure vesssel design
    x1: thickness (d1)  --> discrete value multiple of 0.0625 in 
    x2: thickness of the heads (d2) ---> categorical value from a pre-defined grid
    x3: inner radius (r)  ---> cont. value between [10, 200]
    x4: length (L)  ---> cont. value between [10, 200]
    """
    
    x=individual.copy()
    x[0] *= 0.0625   #convert d1 to "in" 

    y = 0.6224*x[0]*x[2]*x[3]+1.7781*x[1]*x[2]**2+3.1661*x[0]**2*x[3]+19.84*x[0]**2*x[2];

    g1 = -x[0]+0.0193*x[2];
    g2 = -x[1]+0.00954*x[2];
    g3 = -math.pi*x[2]**2*x[3]-(4/3)*math.pi*x[2]**3 + 1296000;
    g4 = x[3]-240;
    g=[g1,g2,g3,g4]
    
    phi=sum(max(item,0) for item in g)
    eps=1e-5 #tolerance to escape the constraint region
    penality=1e7 #large penality to add if constraints are violated
    
    if phi > eps:  
        fitness=phi+penality
    else:
        fitness=y
    return fitness

def init_sample(bounds):
    #generate an individual from a bounds dictionary
    indv=[]
    for key in bounds:
        if bounds[key][0] == 'int':
            indv.append(random.randint(bounds[key][1], bounds[key][2]))
        elif bounds[key][0] == 'float':
            indv.append(random.uniform(bounds[key][1], bounds[key][2]))
        elif bounds[key][0] == 'grid':
            indv.append(random.sample(bounds[key][1],1)[0])
        else:
            raise Exception ('unknown data type is given, either int, float, or grid are allowed for parameter bounds')   
    return indv
    
ngen=5

for item in ['mixed', 'grid', 'float/int', 'float/grid', 'int/grid', 'float', 'int']:
    bounds = {}
    btype=item  #float, int, grid, float/int, float/grid, int/grid, mixed. 
    
    print(item, 'is running -----')
    if btype=='mixed':
        bounds['x1'] = ['int', 1, 99]
        bounds['x2'] = ['grid', (0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.5625, 0.625)]
        bounds['x3'] = ['float', 10, 200]
        bounds['x4'] = ['float', 10, 200]
        bounds['x5'] = ['grid', ('Hi', 'Bye', 'New')]
        bounds['x6'] = ['int', -5, 5]
    
    elif btype=='int/grid':      
        bounds['x1'] = ['int', 1, 20]
        bounds['x2'] = ['grid', (0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.5625, 0.625)]
        bounds['x3'] = ['int', 10, 200]
        bounds['x4'] = ['int', 10, 200]
        bounds['x5'] = ['grid', ('Hi', 'Bye', 'New')]
        bounds['x6'] = ['int', -5, 5]
    
    elif btype=='float/grid':      
        bounds['x1'] = ['float', 1, 20]
        bounds['x2'] = ['grid', (0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.5625, 0.625)]
        bounds['x3'] = ['float', 10, 200]
        bounds['x4'] = ['float', 10, 200]
        bounds['x5'] = ['grid', ('Hi', 'Bye', 'New')]
        bounds['x6'] = ['float', -5, 5]
        
    elif btype=='float/int':      
        bounds['x1'] = ['int', 1, 20]
        bounds['x2'] = ['float', 1, 20]
        bounds['x3'] = ['int', 10, 200]
        bounds['x4'] = ['float', 10, 200]
        bounds['x5'] = ['float', -5, 5]
        bounds['x6'] = ['int', -5, 5]
    
    elif btype=='float':      
        bounds['x1'] = ['float', 1, 20]
        bounds['x2'] = ['float', 1, 20]
        bounds['x3'] = ['float', 10, 200]
        bounds['x4'] = ['float', 10, 200]
        bounds['x5'] = ['float', -5, 5]
        bounds['x6'] = ['float', -5, 5]
        
    elif btype=='int':      
        bounds['x1'] = ['int', 1, 20]
        bounds['x2'] = ['int', 1, 20]
        bounds['x3'] = ['int', 10, 200]
        bounds['x4'] = ['int', 10, 200]
        bounds['x5'] = ['int', -5, 5]
        bounds['x6'] = ['int', -5, 5]
        
    elif btype=='grid':      
        bounds['x1'] = ['grid', (0.0625, 0.125, 0.375, 0.4375, 0.5625, 0.625)]
        bounds['x2'] = ['grid', (0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.5625, 0.625)]
        bounds['x3'] = ['grid', (1,2,3,4,5)]
        bounds['x4'] = ['grid', (32,64,128)]
        bounds['x5'] = ['grid', ('Hi', 'Bye', 'New')]
        bounds['x6'] = ['grid', ('Cat', 'Dog', 'Bird', 'Fish')]


    npop=70
    x0=[]
    for i in range(npop):
        x0.append(init_sample(bounds))
        
    ########################
    # Setup and evolute GWO
    ########################
    gwo=GWO(mode='min', fit=Vessel, bounds=bounds, nwolves=npop, ncores=1, seed=1)
    x_gwo, y_gwo, gwo_hist=gwo.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_gwo) == y_gwo
    x_gwo, y_gwo, gwo_hist=gwo.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_gwo) == y_gwo
    
    ########################
    # Setup and evolute WOA
    ########################
    woa=WOA(mode='min', bounds=bounds, fit=Vessel, nwhales=npop, a0=1.5, b=1, ncores=1, seed=1)
    x_woa, y_woa, woa_hist=woa.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_woa) == y_woa
    x_woa, y_woa, woa_hist=woa.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_woa) == y_woa

    ########################
    # Setup and evolute SSA
    ########################
    #setup and evolute SSA
    ssa=SSA(mode='min', bounds=bounds, fit=Vessel, nsalps=npop, int_transform='sigmoid', ncores=1, seed=1)
    x_ssa, y_ssa, ssa_hist=ssa.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_ssa) == y_ssa    
    x_ssa, y_ssa, ssa_hist=ssa.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_ssa) == y_ssa
    
    ########################
    # Setup and evolute DE
    ########################
    de=DE(mode='min', bounds=bounds, fit=Vessel, npop=npop, F=0.5, CR=0.7, int_transform='sigmoid', ncores=1, seed=1)
    x_de, y_de, de_hist=de.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_de) == y_de
    x_de, y_de, de_hist=de.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_de) == y_de
    
    ########################
    # Setup and evolute MFO
    ########################
    mfo=MFO(mode='min', bounds=bounds, fit=Vessel, nmoths=npop, int_transform='minmax', ncores=1, seed=1)
    x_mfo, y_mfo, mfo_hist=mfo.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_mfo) == y_mfo
    x_mfo, y_mfo, mfo_hist=mfo.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_mfo) == y_mfo

    ########################
    # Setup and evolute JAYA
    ########################
    jaya=JAYA(mode='min', bounds=bounds, fit=Vessel, npop=npop, int_transform='sigmoid', ncores=1, seed=1)
    x_jaya, y_jaya, jaya_hist=jaya.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_jaya) == y_jaya
    x_jaya, y_jaya, jaya_hist=jaya.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_jaya) == y_jaya
    
    ########################
    # Setup and evolute HHO
    ########################
    hho = HHO(mode='min', bounds=bounds, fit=Vessel, nhawks=npop, 
                      int_transform='minmax', ncores=1, seed=1)
    x_hho, y_hho, hho_hist=hho.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_hho) == y_hho
    
    ########################
    # Setup and evolute PSO
    ########################
    pso=PSO(mode='min', bounds=bounds, fit=Vessel, c1=2.05, c2=2.1, npar=npop,
                    speed_mech='constric', ncores=1, seed=1)
    x_pso, y_pso, pso_hist=pso.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_pso) == y_pso
    x_pso, y_pso, pso_hist=pso.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_pso) == y_pso
    
    ########################
    # Setup and evolute ES 
    ########################
    es = ES(mode='min', fit=Vessel, cxmode='cx2point', bounds=bounds, 
                     lambda_=npop, mu=5, cxpb=0.7, mutpb=0.2, seed=1)
    x_es, y_es, es_hist=es.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_es) == y_es
    x_es, y_es, es_hist=es.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_es) == y_es

    ########################
    # Setup and evolute ES 
    ########################
    #setup and evolute BAT
    bat=BAT(mode='min', bounds=bounds, fit=Vessel, nbats=npop, 
            fmin=0, fmax=1, A=1.0, r0=0.7,
            ncores=1, seed=1)
    x_bat, y_bat, bat_hist=bat.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_bat) == y_bat
    x_bat, y_bat, bat_hist=bat.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_bat) == y_bat

    
    ########################
    # Setup and evolute EPSO
    ########################   
    #setup and evolute EPSO
    epso=EPSO(mode='min', bounds=bounds, g1=int(npop/2), g2=int(npop/2), fit=Vessel, ncores=1, seed=None)
    x_epso, y_epso, epso_hist=epso.evolute(ngen=ngen, LP=1, verbose=0)
    assert Vessel(x_epso) == y_epso
    x_epso, y_epso, epso_hist=epso.evolute(ngen=ngen, x0=x0, LP=1, verbose=0)
    assert Vessel(x_epso) == y_epso

    ########################
    # Setup and evolute EDEV
    ########################  
    #setup and evolute EDEV
    edev=EDEV(mode='min', bounds=bounds, fit=Vessel, npop=npop, ncores=1, seed=1)
    x_edev, y_edev, edev_hist=edev.evolute(ngen=ngen, ng=1, verbose=0)
    assert Vessel(x_edev) == y_edev
    x_edev, y_edev, edev_hist=edev.evolute(ngen=ngen, x0=x0, ng=1, verbose=0)
    assert Vessel(x_edev) == y_edev
    
    ########################
    # Setup and evolute HCLPSO
    ########################  
    #setup and evolute HCLPSO
    hclpso=HCLPSO(mode='min', bounds=bounds, g1=int(npop/2), g2=int(npop/2), fit=Vessel, ncores=1, seed=1)
    x_hclpso, y_hclpso, hclpso_hist=hclpso.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_hclpso) == y_hclpso
    x_hclpso, y_hclpso, hclpso_hist=hclpso.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_hclpso) == y_hclpso

    ########################
    # Setup and evolute SA
    ########################  
    #define a custom moving function
    def my_move(x, **kwargs):
        #-----
        #this function selects two random indices in x and perturb their values
        #-----
        x_new=x.copy()
        #indices=random.sample(range(0,len(x)), 2)
        #for i in indices:
        #    x_new[i] = random.uniform(bounds['x1'][1],bounds['x1'][2])
        x_new=init_sample(sa.bounds)
        
        return x_new

    #setup and evolute parallel SA with `equilibrium` cooling
    sa=SA(mode='min', bounds=bounds, fit=Vessel, chain_size=20, chi=0.2, Tmax=10000, threshold = 1, lmbda=0.05,
          move_func=my_move, reinforce_best='soft', cooling='equilibrium', ncores=4, seed=1)

    x_sa, y_sa, sa_hist=sa.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_sa) == y_sa
    x_sa, y_sa, sa_hist=sa.evolute(ngen=ngen, x0=x0[0:4], verbose=0)
    assert Vessel(x_sa) == y_sa

    ########################
    # Setup and evolute CS
    ########################  
    #setup and evolute CS
    cs = CS(mode = 'min', bounds = bounds, fit = Vessel, ncuckoos = npop, pa = 0.25, seed=None)
    x_cs, y_cs, cs_hist=cs.evolute(ngen = ngen, verbose=0)
    assert Vessel(x_cs) == y_cs
    x_cs, y_cs, cs_hist=cs.evolute(ngen = ngen, x0=x0, verbose=0)
    assert Vessel(x_cs) == y_cs

    ########################
    # Setup and evolute TS
    ########################  
    #setup and evolute TS
    ts=TS(mode = "min", bounds = bounds, fit = Vessel, tabu_tenure=60, 
          penalization_weight = 0.8, swap_mode = "perturb", ncores=1, seed=None)
    x_ts, y_ts, ts_hist=ts.evolute(ngen = ngen, verbose=0)
    assert Vessel(x_ts) == y_ts
    x_ts, y_ts, ts_hist=ts.evolute(ngen = ngen, x0=x0[0], verbose=0)
    assert Vessel(x_ts) == y_ts
    
    ########################
    # Setup and evolute PESA
    ######################## 
    pesa=PESA(mode='min', bounds=bounds, fit=Vessel, npop=npop, mu=40, alpha_init=0.2, 
              alpha_end=1.0, alpha_backdoor=0.1, ncores=1)
    x_pesa, y_pesa, pesa_hist=pesa.evolute(ngen=ngen, verbose=0)
    assert Vessel(x_pesa) == y_pesa
    x_pesa, y_pesa, pesa_hist=pesa.evolute(ngen=ngen, x0=x0, verbose=0)
    assert Vessel(x_pesa) == y_pesa

    ########################
    # Setup and evolute PESA2
    ######################## 
    pesa2=PESA2(mode='min', bounds=bounds, fit=Vessel, npop=npop, nwolves=5, nwhales=5, ncores=1)
    x_pesa2, y_pesa2, pesa2_hist=pesa2.evolute(ngen=ngen, replay_every=2, verbose=0)
    assert Vessel(x_pesa2) == y_pesa2
    x_pesa2, y_pesa2, pesa2_hist=pesa2.evolute(ngen=ngen, x0=x0, replay_every=2, verbose=0)
    assert Vessel(x_pesa2) == y_pesa2

    ########################
    # Setup and evolute AEO
    ########################
    aeo = AEO(mode='min', bounds=bounds, optimizers=[mfo, de, gwo, woa, jaya, hho, pso, es, ssa], gen_per_cycle=3, fit = Vessel)
    x_aeo, y_aeo, aeo_hist = aeo.evolute(3, verbose = 1)
    assert Vessel(x_aeo) == y_aeo
    
    #sys.exit()   #remove to complete the full test