import numpy as np
import scipy.optimize as opt
import params as p
import computations as comp
import os
import matplotlib.pyplot as plt
import laminate as lam

WETLAND_CRADLE = False
BALANCED = True
SYMMETRIC = True

def saverun(x0, result, layup):
    if result.success and result.fun < 1e6:
        ppm = result.fun
        x = [r for r in result.x]
        layers = len(layup)
        h = lam.laminateThickness(layup)
        length = get_max_length(layup)

        if WETLAND_CRADLE is False:
            filename = 'opt_results_bedrock/'+str(layers)+'layers_'+str(round(ppm, 0))[:-2]
        else:
            filename = 'opt_results_wetland/' + str(layers) + 'layers_' + str(round(ppm, 0))[:-2]

        if os.path.isfile(filename+'.txt'):
            filename += '_2'

        file_str = 'wetland cradle : ' + str(WETLAND_CRADLE) + '\n' +\
                    'kr/meter : ' + str(ppm) +'\n' +\
                    'max length [m] : ' + str(round(length,0)/1000)[:-2] +'\n' +\
                    'Solution : ' + str(x) +'\n' +'Initial guess :' + str(x0) +'\n' +\
                    'Layers : ' + str(layers) +'\n' +\
                    'Total thickness [mm] : ' + str(h) +'\n\n'
        file_str += layup_to_str(layup)

        i = 3
        while os.path.isfile(filename+'.txt'):
            if WETLAND_CRADLE is False:
                filename = 'opt_results_bedrock/' + str(layers) + 'layers_' + str(round(ppm, 0))[:-2]+'_'+str(i)
            else:
                filename = 'opt_results_wetland/' + str(layers) + 'layers_' + str(round(ppm, 0))[:-2]+'_'+str(i)
            i+=1
        filename += '.txt'
        with open(filename, 'w') as file:
            file.write(file_str)

        print('Saved')
        print(file_str)
        print('In', filename)

    else:
        ppm = result.fun
        x = [r for r in result.x]
        layers = len(layup)
        h = lam.laminateThickness(layup)
        length = get_max_length(layup)

        if WETLAND_CRADLE is False:
            filename = 'opt_results_bedrock/'+str(layers)+'layers_'+str(round(ppm, 0))[:-2]
        else:
            filename = 'opt_results_wetland/' + str(layers) + 'layers_' + str(round(ppm, 0))[:-2]

        file_str = 'wetland cradle : ' + str(WETLAND_CRADLE) + '\n' + \
                   'kr/meter : -' + '\n' + \
                   'max length [m] : -' + '\n' + \
                   'Solution : DID NOT CONVERGE' + '\n' + \
                   'Initial guess :' + str(x0) + '\n' + \
                   'Layers : ' + str(layers) + '\n' + \
                   'Total thickness : -' + '\n\n'

        i = 3
        while os.path.isfile(filename + '.txt'):
            if WETLAND_CRADLE is False:
                filename = 'opt_results_bedrock/' + str(layers) + 'layers_' + str(round(ppm, 0))[:-2]+'_'+str(i)
            else:
                filename = 'opt_results_wetland/' + str(layers) + 'layers_' + str(round(ppm, 0))[:-2]+'_'+str(i)
            i += 1
        filename += '.txt'
        with open(filename, 'w') as file:
            file.write(file_str)

        print('DID NOT CONVERGE')


def layup_to_str(layup):
    layup_str = 'Balanced : '+str(BALANCED)+'\n' +\
                'Symmetric : '+str(SYMMETRIC)+'\n'+\
                'Layer      ori     thi\n'
    for i, l in enumerate(layup):
        layup_str += str(i+1)+'          '+str(l['ori'])+'       '+str(l['thi'])+'\n'

    return layup_str

def get_max_length(layup): #Hardcoded maximum length
    max_L = 30e3

    L = 5000
    tol = 1e-10
    dL = 1000

    while abs(dL) > tol:

        while not comp.check_failure(layup, L):

            if L >= max_L:
                return max_L

            L += dL
            dL *= 2
        dL *= 0.1

        while comp.check_failure(layup, L):
            if L <= 1000:
                return 0
            while L <= dL:
                dL *= 0.1
            L -= dL
        dL *= 0.1
    return L

def get_price_per_meter(layup, L):
    if L == 0:
        return 1e10
    h = lam.laminateThickness(layup)
    V = (np.pi * (p.R_inner+ h)**2 - np.pi*(p.R_inner)**2) * L
    m = V * p.m['rho'] * 1e3

    if WETLAND_CRADLE is True:
        return (m * p.prod_cost + p.wetland_cradle_cost) / (L / 1000)
    elif WETLAND_CRADLE is False:
        return (m * p.prod_cost + p.rock_cradle_cost) / (L / 1000)
    elif WETLAND_CRADLE is None:
        return (m * p.prod_cost + p.avg_cradle_cost) / (L / 1000)

def gen_sym_layup(args, ntimes):
    t_list, d_list = args[:len(args)//2], args[len(args)//2:]
    layup = [dict() for i in range(len(t_list))]
    for l in range(len(t_list)):
        layup[l] = {'mat':p.m, 'ori': t_list[l], 'thi': d_list[l]}

    layup += layup[::-1] #Enforcing symmetry
    layup *= ntimes
    return layup

def gen_sym_bal_layup(args, ntimes):
    t_list, d_list = args[:len(args)//2], args[len(args)//2:]
    layup = [dict() for i in range(len(t_list)*2)]
    for l in range(len(t_list)):
        layup[2 * l] = {'mat':p.m, 'ori': t_list[l], 'thi':d_list[l]}
        layup[2*l + 1] = {'mat':p.m, 'ori': -t_list[l], 'thi':d_list[l]} #Enforcing balance

    layup += layup[::-1] #Enforcing symmetry
    layup *= ntimes
    return layup

def gen_layup(args, ntimes):
    if BALANCED and SYMMETRIC:
        layup = gen_sym_bal_layup(args, ntimes)
    elif SYMMETRIC and not BALANCED:
        layup = gen_sym_layup(args, ntimes)
    else:
        raise NotImplementedError('The case balanced = '+str(BALANCED)+', symmetric = '+str(SYMMETRIC)+' is not implemented!')

    return layup

def optimize_target(args, ntimes=1):
    layup = gen_layup(args, ntimes)

    L = get_max_length(layup)
    ppm = get_price_per_meter(layup, L)
    print(round(L, 2), round(ppm, 2), [round(x,2) for x in args])
    return ppm

def optimize_thickness_target(d_list, theta_list, ntimes=1):
    args = np.concatenate((theta_list, d_list))
    layup = gen_layup(args, ntimes)
    L = get_max_length(layup)
    ppm = get_price_per_meter(layup, L)
    print(round(L, 2), round(ppm, 2), [round(x, 2) for x in args])
    return ppm

def optimize_angles_target(theta_list, d_list, ntimes=1):
    args = np.concatenate((theta_list, d_list))

    layup = gen_layup(args, ntimes)

    L = get_max_length(layup)
    ppm = get_price_per_meter(layup, L)
    print(round(L, 2), round(ppm, 2), [round(x, 2) for x in args])
    return ppm

def optimize_layup(ntimes):
    theta_bound = [(15, 85)] #limit of angles in each layer
    thi_bound = (0.1, 5)

    x0 = [30, 85, 15, 1.5, 1.5, 1.5]
    bounds = theta_bound * (len(x0)//2) + [thi_bound] * (len(x0)//2)
    print(bounds, x0)
    init_layup = gen_layup(x0, ntimes)
    print('Running optimize with initial guess (ntimes = ', ntimes,')', sep='')
    print(layup_to_str(init_layup))

    sol = opt.minimize(optimize_target, x0=x0, bounds=bounds, args=ntimes)
    result_layup = gen_layup(sol.x, ntimes)
    saverun(x0, sol, result_layup)

def optimize_thickness(ntimes):
    thi_bound = (0.1, 5)

    x0 = [1.5]
    theta_list = [85]
    bounds = [thi_bound] * len(x0)

    init_layup = gen_layup(theta_list+x0, ntimes)
    print('Running optimize with initial guess (ntimes = ', ntimes, ')', sep='')
    print(layup_to_str(init_layup))

    sol = opt.minimize(optimize_thickness_target, x0=x0, bounds=bounds, args=(theta_list, ntimes))
    result_layup = gen_layup(np.concatenate((np.array(theta_list), sol.x)), ntimes)

    saverun(x0, sol, result_layup)

def optimize_orientation(ntimes):
    theta_bounds = [(15, 85), (15, 85)]

    x0 = [57, 15]
    d_list = [1.4, 0.7]

    init_layup = gen_layup(x0+d_list, ntimes)
    print('Running optimize with initial guess (ntimes = ', ntimes, ')', sep='')
    print(layup_to_str(init_layup))

    sol = opt.minimize(optimize_angles_target, x0=x0, bounds=theta_bounds, args=(d_list, ntimes))
    result_layup = gen_layup(np.concatenate((sol.x, np.array(d_list))), ntimes)

    saverun(x0, sol, result_layup)

# for i in (1, 2, 3, 4):
#     optimize_layup(i)
#
# WETLAND_CRADLE = True
#
# for i in (1, 2, 3, 4):
#     optimize_layup(i)


# bal, sym: 10113.75 2958.08 [85.0, 84.53, 0.32, 0.2], ntimes = 2
#