import FebOpt, subprocess
import numpy as np
from scipy.interpolate import interp1d
from scipy.optimize import fmin_l_bfgs_b
from scoop import futures
import matplotlib.pyplot as plt
import random, string, pickle
from deap import base
from deap import creator
from deap import tools

scale = [(1.0,0.1),(0.15,0.1),(5e-3,1e-4)]
IND_SIZE = 3    #number of degrees of freedom 
creator.create("FitnessMin",base.Fitness, weights=(-1.0,))
creator.create("Individual",list,fitness=creator.FitnessMin)

def evaluate(x):
    subprocess.call('(export OMP_NUM_THREADS=1; /home/scott/FEBio/febio.lnx64 -i '+x[0]+' -silent) > /dev/null',shell=True)
    logname = string.replace(x[0],'.feb','.log')
    results = r.parseLog(name=logname)
    if "ERROR" in results:
        return [1e64]
    if 'opt' not in x[0]:
        subprocess.call('rm '+x[0]+' '+logname,shell=True)
    xi = np.arange(0.0,results[-1,0],0.1)
    y = interp1d(x[2][0],x[2][1])
    yi = y(xi)
    obj_fit = interp1d(results[:,0],results[:,1])
    obji = -360*obj_fit(xi)
    plt.plot(xi,yi,'b')
    plt.plot(xi,obji,'r')
    plt.savefig('img/opt'+str("%04d" % x[1])+'.png',bbox_inches=0)
    plt.clf()
    diff = yi-obji
    yrange = abs(np.max(yi)-np.min(yi))
    f = np.max(abs(diff))/yrange
    g  = np.linalg.norm(diff)/np.sqrt(yi.size)/yrange
    return [(f+g)/2]

def bfgsObj(x,target):
    print("-- Current Material Properties --")
    print(' E: '), x[0]*scale[0][0]+scale[0][1]
    print(' v'), x[1]*scale[1][0]+scale[1][1]
    print(' perm1'),x[2]*scale[2][0]+scale[2][1]
    gx = []
    dx = []
    for i in xrange(len(x)):
        dmy = np.array(x)
        dx.append(0.0001*(1+dmy[i]))
        dmy[i] = dmy[i]+dx[i]
        gx.append(dmy)
    for i in xrange(len(x)):
        dmy = np.array(x)
        dmy[i] = dmy[i]+dx[i]
        gx.append(dmy)

    models = []
    for i in xrange(2*len(x)+1):
        name = str(i)
        if i < 1:
            makeModels(x,name)
        else:
            makeModels(gx[i-1],name)
        models.append((name,target))
    f = list(futures.map(solveFeb,models))

    df = np.array(f[1:])
    df = (-3*f[0]+4*df[0:len(x)]-df[len(x):])/(2*np.array(dx))
    print 'Current objective: ',f[0]
    print 'Current gradient: ',df
    return f[0], df

def solveFeb(x):
    subprocess.call('(export OMP_NUM_THREADS=1; /home/scott/FEBio/febio.lnx64 -i '+x[0]+'.feb -silent) > /dev/null',shell=True)
    logname = x[0]+'.log'
    results = r.parseLog(name=logname)
    if "ERROR" in results:
        print("FE model crashed during BFGS - Aborting")
        raise SystemExit
    subprocess.call('rm '+x[0]+'.feb '+logname,shell=True)
    xi = np.arange(0.0,results[-1,0],0.1)
    y = interp1d(x[1][0],x[1][1])
    yi = y(xi)
    obj_fit = interp1d(results[:,0],results[:,1])
    obji = -360*obj_fit(xi)
    diff = yi-obji
    yrange = abs(np.max(yi)-np.min(yi))
    f = np.max(abs(diff))/yrange
    g  = np.linalg.norm(diff)/np.sqrt(yi.size)/yrange
    return (f+g)/2

toolbox = base.Toolbox()
toolbox.register("map",futures.map)
toolbox.register("attribute",random.random)
toolbox.register("individual",tools.initRepeat, creator.Individual,toolbox.attribute, n=IND_SIZE)
toolbox.register("population",tools.initRepeat, list, toolbox.individual)
toolbox.register("evaluate",evaluate)
toolbox.register("mate",tools.cxUniform,indpb=0.2)
toolbox.register("mutate",tools.mutShuffleIndexes,indpb=0.2)
toolbox.register("select",tools.selTournament, tournsize=3)

def materialUpdate(dmy,r):
    M = 2.0
    ksi = 10.0
    w = [[1.0,0.4,0.4],
            [1.0]*3,
            [0.4,0.4,1.0]]
    kw = [[1.6,0.6,0.6],[1.0,1.0,1.0],[0.6,0.6,1.2]]
    r.updateMaterials({
        'ecs':{'E':dmy[0],
            'v':dmy[1],
            'ksi':np.array([w[0][0]*ksi,w[0][1]*ksi,w[0][2]*ksi]),
            'perm1':np.array([kw[0][0]*dmy[2],kw[0][1]*dmy[2],kw[0][2]*dmy[2]]),
            'M':np.array([M]*3)},
        'ecsp':{'E':dmy[0],
            'v':dmy[1],
            'ksi':np.array([w[0][0]*ksi,w[0][1]*ksi,w[0][2]*ksi]),
            'perm1':np.array([kw[0][0]*dmy[2],kw[0][1]*dmy[2],kw[0][2]*dmy[2]]),
            'M':np.array([M]*3)},
        'ecm':{'E':dmy[0],
            'v':dmy[1],
            'ksi':np.array([w[1][0]*ksi,w[1][1]*ksi,w[1][2]*ksi]),
            'perm1':np.array([kw[1][0]*dmy[2],kw[1][1]*dmy[2],kw[1][2]*dmy[2]]),
            'M':np.array([M]*3)},
        'ecmp':{'E':dmy[0],
            'v':dmy[1],
            'ksi':np.array([w[1][0]*ksi,w[1][1]*ksi,w[1][2]*ksi]),
            'perm1':np.array([kw[1][0]*dmy[2],kw[1][1]*dmy[2],kw[1][2]*dmy[2]]),
            'M':np.array([M]*3)},
        'ecd':{'E':dmy[0],
            'v':dmy[1],
            'ksi':np.array([w[2][0]*ksi,w[2][1]*ksi,w[2][2]*ksi]),
            'perm1':np.array([kw[2][0]*dmy[2],kw[2][1]*dmy[2],kw[2][2]*dmy[2]]),
            'M':np.array([M]*3)},
        'ecdp':{'E':dmy[0],
            'v':dmy[1],
            'ksi':np.array([w[2][0]*ksi,w[2][1]*ksi,w[2][2]*ksi]),
            'perm1':np.array([kw[2][0]*dmy[2],kw[2][1]*dmy[2],kw[2][2]*dmy[2]]),
            'M':np.array([M]*3)}})

def makeModels(x,name):
    #transforms x as x_new = scale[i][0]*x[i] + scale[i][1]
    dmy = []
    for i in xrange(len(x)):
        dmy.append(x[i]*scale[i][0]+scale[i][1])
    dmy = np.array(dmy)
    materialUpdate(dmy,r)
    r.changeLog(name+'.log')
    r.writeModel(name=name+'.feb')

r = FebOpt.FebOpt('axi3.feb',obj_var='Fz',obj_var_type='rigid_body_data',obj_var_range='9')

def main(N=35,NGEN=3,CXPB=0.8,MUTPB=0.05):
    hof = tools.HallOfFame(1)
    fid = open('exp.pkl','rb')
    exp_dat = pickle.load(fid)
    fid.close()
    t = exp_dat['Time'][u'segment_001']['LoadCell']
    t[0] = 0.0
    peak_loc = [77,9154,18266] #indices of location where displacement first converges
    disp_start = [0,9050,18130]
    d = exp_dat['Load'][u'segment_001']['LoadCell']
    d[0] = 0.0
    d = d - 1.53
    d[0:77] = np.linspace(0.0,d[77],77)
    for i in xrange(len(peak_loc)):
        m = d[peak_loc[i]]-d[disp_start[i]]
        m /= (t[peak_loc[i]] - t[disp_start[i]])
        for j in xrange(peak_loc[i]-disp_start[i]):
            d[disp_start[i]+j] = d[disp_start[i]]+m*j*.1
    target = (t,d)
    pop = toolbox.population(n=N)
    model_names = []
    evalcnt = 0
    for i in xrange(len(pop)):
        makeModels(pop[i],str(i))
        model_names.append((str(i)+'.feb',evalcnt,target))
        evalcnt += 1
    fitnesses = list(toolbox.map(toolbox.evaluate,model_names))
    hof.update(pop)
    for ind, fit in zip(pop,fitnesses):
        ind.fitness.values = fit
    for g in range(NGEN):
        print("-- Generation %i --" %g)
        #select the next generation
        offspring = toolbox.select(pop,len(pop))
        # Clone the selected individuals
        offspring = map(toolbox.clone,offspring)

        # Apply crossover and mutation on the offspring
        for child1, child2 in zip(offspring[::2], offspring[1::2]):
            if random.random() < CXPB:
                toolbox.mate(child1,child2)
                del child1.fitness.values
                del child2.fitness.values

        for mutant in offspring:
            if random.random() < MUTPB:
                toolbox.mutate(mutant)
                del mutant.fitness.values

        invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
        model_names = []
        for i in xrange(len(invalid_ind)):
            makeModels(invalid_ind[i],str(i))
            model_names.append((str(i)+'.feb',evalcnt,target))
            evalcnt += 1

        fitnesses = list(toolbox.map(toolbox.evaluate,model_names))
        for ind, fit in zip(invalid_ind, fitnesses):
            ind.fitness.values = fit
        pop[:] = offspring
        hof.update(pop)
        #adaptive evolution probabilities
        '''
        fmax = np.max(pop)
        fmean = np.mean(pop)
        fprime = 0
        if fprime > fmean:
            CXPB = k1*(fmax-fprime)/(fmax-fmean)
            MUTP = k2*(fmax-f)/(fmax-fmean)
        else:
            CXPB = k3
            MUTP = k4
        '''
        fits = [ind.fitness.values[0] for ind in pop]
        print(" Current fitness Statistics: ")
        print(" Min %s" % min(fits))
        print(" Max %s" % max(fits))
        print(" Mean %s" % np.mean(fits))
        print(" Standard Deviation %s" % np.std(fits))
        if np.std(fits) < 1e-5:
            break

    fits = [ind.fitness.values[0] for ind in pop]
    #best = fits.index(min(fits))
    best_ind = hof[0]
    '''
    print("-- Final Fitness Statistics --")
    print(" Size of final population %s" % len(fits))
    print(" Min %s" % fits[best])
    print(" Max %s" % max(fits))
    print(" Mean %s" % np.mean(fits))
    print(" Standard Deviation %s" % np.std(fits))
    '''
    print("-- Material Constants of Best Individual --")
    print(' E: '), best_ind[0]*scale[0][0]+scale[0][1]
    print(' v'), best_ind[1]*scale[1][0]+scale[1][1]
    print(' perm1'),best_ind[2]*scale[2][0]+scale[2][1]

    print('\n-- Starting BFGS optimization from GA optimum --')
    bounds = [(0,1),(0,1),(0,1)]
    x0 = best_ind
    x,f,d = fmin_l_bfgs_b(bfgsObj,x0,fprime=None,args=(target,),
            bounds=bounds,iprint=0,factr=1e12,m=100)
    # write the optimized model
    makeModels(x,"axi1_opt")
    print("-- Material properties determined by BFGS --")
    print(' E: '), x[0]*scale[0][0]+scale[0][1]
    print(' v'), x[1]*scale[1][0]+scale[1][1]
    print(' perm1'),x[2]*scale[2][0]+scale[2][1]
    # solve the optimized model and plot it
    print '-- Solving optimized model --'
    toolbox.evaluate(["axi3_opt.feb",evalcnt,target])
if __name__ == '__main__':
    main()