import laminate as lam
import constants as c
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap

side_layup = [{'mat' : c.m, 'ori' : 45, 'thi' : 0.25},
              {'mat': c.m, 'ori': -45, 'thi': 0.5},
              {'mat': c.m, 'ori': 45, 'thi': 0.25}]

tb_layup = [{'mat' : c.m, 'ori' : 0, 'thi' : 0.25},
            {'mat': c.m, 'ori': 90, 'thi': 0.5},
            {'mat': c.m, 'ori': 0, 'thi': 0.25}]

def gen_air_width(side_layup):
    h = c.B - 2 * lam.laminateThickness(side_layup)
    return [{'mat' : c.no_mat, 'ori': 0, 'thi': h}]

def gen_air_height(tb_layup):
    h = c.H - 2 * lam.laminateThickness(tb_layup)
    return [{'mat' : c.no_mat, 'ori': 0, 'thi': h}]

def get_sandwitch_Ixx(tb_layup):
    f = lam.laminateThickness(tb_layup)
    return (c.B * f**2 / 12) + 2 * c.B * f * (c.H - f/2)**2

def get_E_hom(layup):
    A = lam.computeA(layup)
    h = lam.laminateThickness(layup)
    return (1/h) * (A[0,0] - ( (A[0,1]**2 / A[1, 1])))

def get_eff_D(layup):
    D = lam.computeD(layup)
    return c.B * (D[0,0] - (D[0,1]**2 / D[1,1]))

def get_eff_G(layup):
    A = lam.computeA(layup)
    return A[2,2]/lam.laminateThickness(layup)

def get_deformation(tb_layup, side_layup):
    Db = get_sandwitch_D(tb_layup, side_layup)
    S = get_eff_G(2 * side_layup) * lam.laminateThickness(2 * side_layup) * (c.H - lam.laminateThickness(tb_layup))

    # From http://www.stmboats.com/articles/sandwich_hb.pdf
    d1 = c.F * c.a * (3 * c.L**2 - 4 * c.a**2) / (48 * Db)
    d2 = c.F * c.a / (2 * S)

    return 2 * (d1 + d2)

def get_sandwitch_D(tb_layup, side_layup):
    f = lam.laminateThickness(tb_layup)
    w = lam.laminateThickness(side_layup)
    h = (c.H - 2 * f)
    b = (c.H - 2 * w)

    Ef = get_E_hom(tb_layup)
    Ec = get_E_hom(side_layup)

    D1 = Ef * c.B * (2 / 3) * ( (c.H / 2)**3 - (h / 2)**3)
    D2 = Ec * (2 / 3) * (c.B - b) * (h / 2)**3
    return D1 + D2
    #return Ef * ((2 / 3) * f**3 + 2 * f * h * (f + h)) + Ec * (2 / 3) * h**3

def hashin(stress):
    #Compute hashin criterion exposure factor, return max(fE_FF, fE_IFF)
    s1, s2, s12 = stress
    XT, YT, XC, YC, SL, ST = c.m['XT'], c.m['YT'], c.m['XC'], c.m['YC'], c.m['S12'], c.m['S23']

    fE_FF = max((s1/XT, -s1/XC))

    if s2 >= 0:
        fE_IFF = (s2/YT)**2 + (s12/ST)**2
    else: # s2 <0:
        fE_IFF = ((s2/(2 * ST))**2 + ( (YC/(2 * ST))**2 - 1) * (s2/YC) + (s12/ST)**2)

    return max((fE_FF, fE_IFF))

def plot_ratio_tb_deflection():
    ratios = np.linspace(0, 1, 50)
    side_layup = [{'mat': c.m, 'ori': 45, 'thi': 0.25},
                  {'mat': c.m, 'ori': -45, 'thi': 0.5},
                  {'mat': c.m, 'ori': 45, 'thi': 0.25}]
    deflections = np.zeros(len(ratios))
    for i, r in enumerate(ratios):
        tb_layup = [{'mat': c.m, 'ori': 0, 'thi': r/2},
                    {'mat': c.m, 'ori': 90, 'thi': (1 - r)},
                    {'mat': c.m, 'ori': 0, 'thi': r/2}]
        deflections[i] = get_deformation(tb_layup, side_layup)

    plt.plot(ratios, deflections)
    plt.xlabel('Top/bot ratios (0/90)')
    plt.ylabel(r'$\delta$')
    plt.show()

def plot_ori_side_deflection(tb_ori=10):
    tb_layup = [{'mat' : c.m, 'ori' :tb_ori, 'thi' : 1},
                {'mat': c.m, 'ori': -tb_ori, 'thi': 2},
                {'mat': c.m, 'ori':tb_ori, 'thi': 1}]

    t_list = np.linspace(0, 90, 100)
    deflections = np.zeros(len(t_list))
    for i, t in enumerate(t_list):
        side_layup = [{'mat': c.m, 'ori': t, 'thi': 1},
                    {'mat': c.m, 'ori': -t, 'thi': 2},
                    {'mat': c.m, 'ori': t, 'thi': 1}]
        deflections[i] = get_deformation(tb_layup, side_layup)

    plt.plot(t_list, deflections, label=tb_ori)
    plt.xlabel('side orientation (+/-)')
    plt.ylabel(r'$\delta$')

def plot_ori_tb_deflection():
    side_layup = [{'mat': c.m, 'ori': 45, 'thi': 0.25},
                  {'mat': c.m, 'ori': -45, 'thi': 0.5},
                  {'mat': c.m, 'ori': 45, 'thi': 0.25}]

    t_list = np.linspace(0, 90, 50)
    h = 1
    deflections = np.zeros(len(t_list))
    for i, t in enumerate(t_list):
        tb_layup = [{'mat': c.m, 'ori': t, 'thi': 1},
                      {'mat': c.m, 'ori': -t, 'thi': 2},
                      {'mat': c.m, 'ori': t, 'thi': 1}]
        deflections[i] = get_deformation(tb_layup, side_layup)

    plt.plot(t_list, deflections)
    plt.xlabel('Top/bot orientation (+/-)')
    plt.ylabel(r'$\delta$')
    plt.show()

def plot_ratio_side_deflection():
    ratios = np.linspace(0, 1, 50)
    tb_layup = [{'mat': c.m, 'ori': 0, 'thi': 0.25},
                  {'mat': c.m, 'ori': 90, 'thi': 0.5},
                  {'mat': c.m, 'ori': 0, 'thi': 0.25}]
    deflections = np.zeros(len(ratios))
    for i, r in enumerate(ratios):
        side_layup = [{'mat': c.m, 'ori': 9.5, 'thi': r / 4},
                      {'mat': c.m, 'ori': -9.5, 'thi': r / 4},
                      {'mat': c.m, 'ori': -45, 'thi': (1 - r)/6},
                      {'mat': c.m, 'ori': 45, 'thi': (1 - r) * 2 / 3},
                      {'mat': c.m, 'ori': -45, 'thi': (1 - r)/6},
                      {'mat': c.m, 'ori': -9.5, 'thi': r / 4},
                      {'mat': c.m, 'ori': 9.5, 'thi': r / 4}]
        deflections[i] = get_deformation(tb_layup, side_layup)

    plt.plot(ratios, deflections)
    plt.xlabel('Side ratios ')
    plt.ylabel(r'$\delta$')
    plt.show()

def get_stress():
    I = get_sandwitch_Ixx(tb_layup)
    max_layer_stress = 0
    for layer in tb_layup:
        t = layer['ori']
        sigma_max = (c.H * c.F * c.a) / (2 * I)
        stress = np.dot(lam.T2Ds(t), [sigma_max, 0, 0])
        if stress > max_layer_stress:
            max_layer_stress = stress

    return max_layer_stress

def get_shear(tb_layup):
    I = get_sandwitch_Ixx(tb_layup)
    Q = (c.B / 2) * ((c.H / 2)**2 - ((c.H / 2) - lam.laminateThickness(tb_layup))**2)
    return c.F * Q / (I * c.B)

def get_max_exposure(h_side, h_tb, side_ori=9.5):
    side_layup = [{'mat': c.m, 'ori': side_ori, 'thi': h_side/4},
                  {'mat': c.m, 'ori': -side_ori, 'thi': h_side/2},
                  {'mat': c.m, 'ori': side_ori, 'thi': h_side/4}]

    tb_layup = [{'mat': c.m, 'ori': 0, 'thi':  0.455 * h_tb},
                {'mat': c.m, 'ori': 90, 'thi': 0.05 * h_tb},
                {'mat': c.m, 'ori': 0, 'thi':  0.455 * h_tb}]

    f = lam.laminateThickness(tb_layup)
    w = lam.laminateThickness(side_layup)
    h = c.H - 2 * f
    b = c.B - 2 * w
    D = get_sandwitch_D(tb_layup, side_layup)

    Ef = get_E_hom(tb_layup)
    Ec = get_E_hom(side_layup)

    stress_f_max = - c.F * c.a * c.H * Ef / (2 * D)
    max_stress_f = [stress_f_max, 0, 0]

    stress_c_interface = - c.F * (c.H - h) * c.a * Ec / (2 * D) * (c.B / (c.B - b))
    shear_c_interface = - ( c.F / (2 * D) ) * Ef * f * (f + h) * (c.B / (c.B - b))
    max_stress_c = [stress_c_interface, 0, shear_c_interface]

    shear_center = - ( c.F / (2 * D) ) * (Ec * (h / 2)**2 + Ef * f * (f + h)) * (c.B / (c.B - b))
    max_shear = [0, 0, shear_center]

    max_face_fE = 0
    for layer in tb_layup:
        t = layer['ori']
        layer_stress = np.dot(lam.T2Ds(t), max_stress_f)
        this_face_fE = hashin(layer_stress)
        if this_face_fE > max_face_fE:
            max_face_fE = this_face_fE

    max_interface_fE = 0
    max_center_fE = 0
    for layer in side_layup:
        t = layer['ori']
        interface_stress = np.dot(lam.T2Ds(t), max_stress_c)
        center_stress = np.dot(lam.T2Ds(t), max_shear)

        center_fE = hashin(center_stress)
        if center_fE > max_center_fE:
            max_center_fE = center_fE

        interface_fE = hashin(interface_stress)
        if interface_fE > max_interface_fE:
            max_interface_fE = interface_fE

    return max((max_face_fE, max_center_fE, max_interface_fE))

def plot_envelope(side_ori):
    h_tb_list = np.linspace(0.1, 20, 50)
    h_side_list = np.zeros(len(h_tb_list))

    for i, h_tb in enumerate(h_tb_list):
        print(h_tb)
        h_side = 0.5
        expo = get_max_exposure(h_side, h_tb, side_ori=side_ori)
        while expo > 1/c.safety_factor:
            if h_side > 250:
                break
            h_side += 0.01
            expo = get_max_exposure(h_side, h_tb, side_ori=side_ori)

        h_side_list[i] = h_side

    area = 2 * h_tb_list * c.B + 2 * h_side_list * (c.H - 2 * h_tb_list)
    print(area)
    mass = area * c.L * c.m['rho']/1000

    ax.plot(h_tb_list, h_side_list, label='± '+str(side_ori), color=cmap(side_ori/90))
    twn.plot(h_tb_list, mass, linestyle='--', color=cmap(side_ori/90))
    #twn.set_ylabel(r'total area [mm$^2$]')
    print('±', side_ori)


def plot_expo():
    side_layup = [{'mat': c.m, 'ori': 45, 'thi': 1e-10},
                  {'mat': c.m, 'ori': -45, 'thi': 1e-10},
                  {'mat': c.m, 'ori': 45, 'thi': 1e-10}]
    tb_expo = []
    side_shear_expo = []
    side_stress_expo = []
    fE_max = 0
    h_list = np.linspace(0.75, 5, 50)
    for h in h_list:
        tb_layup = [{'mat': c.m, 'ori': 0, 'thi': 0.35 * h},
                    {'mat': c.m, 'ori': 15, 'thi': 0.1 * h},
                    {'mat': c.m, 'ori': -15, 'thi': 0.1 * h},
                    {'mat': c.m, 'ori': 15, 'thi': 0.1 * h},
                    {'mat': c.m, 'ori': 0, 'thi': 0.35 * h}]

        D = get_sandwitch_D(tb_layup, side_layup)
        global_strain = [- c.F * c.H / (2 * D), 0, 0]
        tau_max = get_shear(tb_layup)

        tb_fE = 0
        for layer in tb_layup:
            t = layer['ori']
            layer_strain = np.dot(lam.T2De(t), global_strain)
            layer_stress = np.dot(lam.Q2D(layer['mat']), layer_strain)
            this_tb_fE = hashin(layer_stress)
            if this_tb_fE > tb_fE:
                tb_fE = this_tb_fE

        side_shear_fE = 0
        side_stress_fE = 0
        for layer in side_layup:
            t = layer['ori']
            layer_strain = np.dot(lam.T2De(t), global_strain)
            layer_stress = np.dot(lam.Q2D(layer['mat']), layer_strain)
            center_stress = np.dot(lam.T2Ds(t), [0, 0, tau_max])

            layer_shear_fE = hashin(center_stress)
            if layer_shear_fE > side_shear_fE:
                side_shear_fE = layer_shear_fE

            layer_stress_fE = hashin(layer_stress)
            if layer_stress_fE > side_stress_fE:
                side_stress_fE = layer_stress_fE

        if max(side_shear_fE, side_stress_fE, tb_fE) > fE_max:
            fE_max = max(side_shear_fE, tb_fE)

        side_shear_expo.append(side_shear_fE)
        side_stress_expo.append(side_stress_fE)
        tb_expo.append(tb_fE)

    plt.plot(h_list, tb_expo, label='tb')
    plt.plot(h_list, side_shear_expo, label='side shear', linestyle='--')
    plt.plot(h_list, side_stress_expo, label='side stress', linestyle='--')
    plt.legend()
    plt.show()

cmap = get_cmap('viridis')
fig, ax = plt.subplots(1,1)
twn = ax.twinx()

plot_envelope(9.5)
plot_envelope(30)
plot_envelope(45)
plot_envelope(65)
plot_envelope(80)

plt.xlabel('h top/bot [mm]')
ax.set_ylabel('wall thickness [mm]')
plt.title(r'$f_E =$ '+str(round(1/c.safety_factor, 2)))
twn.set_yscale('log')
ax.legend()
plt.show()

plot_ori_side_deflection(10)
plot_ori_side_deflection(40)
plot_ori_side_deflection(60)
plot_ori_side_deflection(0)
plt.legend()
plt.show()
#plot_ori_tb_deflection()
#plot_ratio_tb_deflection()
#plot_ratio_side_deflection()