"""Module for computing energy release rates of tilted fracture tests."""

# Third-party imports
import numpy as np

from uncertainties import ufloat, wrap

import weac

def GIc(layers, phi, a, p, t, E, G, L=1000, C1=4.4):
    """
    Calculate mode I fracture toughness.

    Parameters
    ----------
    layers : ndarray
        Slab layering as list of densities (kg/m^3) and thicknesses (mm).
    E : float
        Weak-layer young modulus (MPa).
    G : float
        Weak-layer shear modulus (MPa).
    phi : float
        Slope angle (degrees).
    t : float
        Weak-layer thickness (mm).
    a : float
        Cut length (mm).
    p : float
        Surface line load (N/mm).
    L : int, optional
        PST lengths (mm). Default is 1000 mm.
    C1 : float, optional
        Slab young modulus density parametrization exponen. Default is 4.4.

    Returns
    -------
    float
        Mode I fracture toughness (J/m^2).
    """

    # Compute Poisson's ratio
    nu = E/2/G-1

    # Initialize system
    pst = weac.Layered(system='-pst')
    # Set weak-layer properties
    pst.set_foundation_properties(t=t, E=E, nu=nu, update=True)
    # Recalculate layering with specified exponent
    pst.set_beam_properties(layers=layers, C1=C1, update=True)
    # Set gravitational surface line load (N/mm)
    pst.set_surface_load(p=p)

    # Compute segementation with crack as unsupported segement
    segments = pst.calc_segments(L=L, a=a)['crack']
    # Assemble system of linear equations and solve the boundary-value
    # problem for free constants
    C = pst.assemble_and_solve(phi=phi, **segments)

    # Energy release rate in J/m^2
    Gdif = 1e3*pst.gdif(C=C, phi=phi, **segments)

    # Return only mode I
    return Gdif[1]

def GIIc(layers, phi, a, p, t, E, G, L=1000, C1=4.4):
    """
    Calculate mode I fracture toughness.

    Parameters
    ----------
    layers : ndarray
        Slab layering as list of densities (kg/m^3) and thicknesses (mm).
    E : float
        Weak-layer young modulus (MPa).
    G : float
        Weak-layer shear modulus (MPa).
    phi : float
        Slope angle (degrees).
    t : float
        Weak-layer thickness (mm).
    a : float
        Cut length (mm).
    p : float
        Surface line load (N/mm).
    L : int, optional
        PST lengths (mm). Default is 1000 mm.
    C1 : float, optional
        Slab young modulus density parametrization exponen. Default is 4.4.

    Returns
    -------
    float
        Mode II fracture toughness (J/m^2).
    """

    # Compute Poisson's ratio
    nu = E/2/G-1
 
    # Initialize system
    pst = weac.Layered(system='-pst')
    # Set weak-layer properties
    pst.set_foundation_properties(t=t, E=E, nu=nu, update=True)
    # Recalculate layering with specified exponent
    pst.set_beam_properties(layers=layers, C1=C1, update=True)
    # Set gravitational surface line load (N/mm)
    pst.set_surface_load(p=p)

    # Compute segementation with crack as unsupported segement
    segments = pst.calc_segments(L=L, a=a)['crack']
    # Assemble system of linear equations and solve the boundary-value
    # problem for free constants
    C = pst.assemble_and_solve(phi=phi, **segments)
    # Energy release rate in J/m^2
    Gdif = 1e3*pst.gdif(C=C, phi=phi, **segments)

    # Return only mode II
    return Gdif[2]

def complement_df(
        df, Ewl, da=10, dt=1, dphi=2, dp=0.025, nu=0.25, C1=4.4):
    """
    Calculate mode I and mode II fracture toughness for all rows in dataframe.

    Parameters
    ----------
    df : pd.DataFrame
        Data set with slab layering, weak-layer properties, and cut lengths.
    Ewl : float
        Young's modulus of the weak layer.
    da : float, optional
        Cut length standard deviation (mm). Default is 10 mm.
    dt : float, optional
        Weak-layer thickness standard deviation (mm). Default is 1 mm.
    dphi : float, optional
        Slope angle standard deviation (degree). Default is 2 degrees.
    dp : float, optional
        Surface line load (N/mm) standard deviation. Default is 0.025 (2.5%).
    nu : float, optional
        Poisson's ratio of the weak layer. Default is 0.25.
    C1 : float, optional
        Exponent of Young-modulus parametrization. Default is 4.4.
    """
    
    # Wrap energy-release-rate functions for use with uncertainties
    wGIc = wrap(GIc)
    wGIIc = wrap(GIIc)
    
    # Make sure we don't modify the original dataframe
    df = df.copy(deep=True)
    
    # Assign uncertainties
    a = df.rc.apply(lambda a: ufloat(a, da, 'a'))
    t = df.wl_thickness.apply(lambda t: ufloat(t, dt, 't'))
    p = df.surface_lineload.apply(lambda p: ufloat(p, p*dp, 'p'))
    phi = df.slope_incl.apply(lambda phi: ufloat(phi, dphi, 'phi'))
    
    # Initialize output arrays
    G1c, G2c = [], []
    
    # Loop thorugh all rows in dataframe
    for i, exp in df.iterrows():
        
        # Compute weak-layer shear modulus
        Gwl = Ewl/(2*(1+nu))
        
        # Unpack layering
        layers = exp.layers_mean
        
        # Compute fracture toughnesses
        G1c.append(wGIc(layers, phi[i], a[i], p[i], t[i], Ewl, Gwl, exp.L, C1))
        G2c.append(wGIIc(layers, phi[i], a[i], p[i], t[i], Ewl, Gwl, exp.L, C1))
        
    # Write results to dataframe
    df['GIc'] = np.array(G1c)
    df['GIIc'] = np.array(G2c)
    df['Gc'] = df['GIc'] + df['GIIc']
    df['Gii/G'] = df['GIIc']/df['Gc']
    
    # Update columns with uncertainties
    df['slope_incl'] = phi
    df['surface_lineload'] = p
    df['rc'] = a
    
    # Return updated dateframe
    return df