"""
Processing of Elastic Modulus, Shear Modulus, Poisson's Ratio
from tensile frame testing.

Refer to Standards:
    ASTM D3039 and D3518

Running script:
    1. Script has been verified to work with the following versions:

        python==3.9.25
        matplotlib==3.9.4
        npTDMS==1.10.0
        numpy==2.0.2
        pandas==2.3.3
        PyYAML==6.0.3
        scipy==1.13.1

        Without versions, libraries can be installed as:
        `python3 -m pip install numpy nptdms matplotlib pandas scipy pyyaml`

    2. Download and extract the tar file of data to be in the same folder as
    the present script. `tar -xzvf static_coupon_data.tar.gz`

    3. Material specifications can be looked up via the panel number included
    with the outputs. For tests conducted at 0 and 90 degrees, the elastic
    modulus and Poisson's ratio are of interest. Tests at 45 degrees are
    generally for the shear modulus.

    4. *.csv files for different cases are saved in an `outputs` folder at the
    same location as this script.

Data Preparation Notes:
    1. Text file notes are included in some coupon folders about repeat tests
    and/or potential damage to coupons before tensile testing.
"""

import os
import numpy as np
from nptdms import TdmsFile

import glob

import matplotlib.pyplot as plt

import pandas as pd

from scipy.signal import butter, filtfilt


###############################################################################
######## Inputs                                                        ########
###############################################################################

show_plots = False

default_start_threshold = 1000 # start strain (ue)
default_end_threshold = 3000 # End strain (ue)

default_max_samples = int(1e30) # maximum number of samples to include

lbf_to_N = 4.4482216 # N/lbf

filter_strain = True
high_freq = 50 # for filter
fs = 1000 # Hz
transverse_sign = -1 # measurement is positive in compression, so flip sign

filedir = os.path.dirname(os.path.realpath(__file__)) # get path to this file
data_folder = os.path.join(filedir, 'static_coupon_data')


case_dict = {
    'Epoxy' :
    {
        'UD0' : {},
        'UD45' : {},
        'UD90' : {'start_threshold' : 500, 'end_threshold' : 1000},
        'BX0_90' : {},
        'BX45' : {},
        'TX0' : {'max_samples' : int(1e5)},
        'TX45' : {'max_samples' : int(1e5)},
        'TX90' : {'max_samples' : int(1e5), 'end_threshold' : 2000},
    },
    'Elium' :
    {
        'UD0' : {},
        'UD45' : {},
        'UD90' : {},
        'BX0' : {},
        'BX45' : {},
    },
    'PECAN' :
    {
        'UD0' :  {},
        'UD45' : {},
        'UD90' : {},
        'BX0' : {},
        'BX45' : {},
    },
    'CFRP' :
    {
        '3mm_Epoxy' : {},
    },
}

###############################################################################
###### Verify Data is Located Correctly                                  ######
###############################################################################

if not os.path.exists(data_folder):
    raise Exception('Data folder not found. Please download and extract '
                    'static_coupon_data.tar.gz to the same folder as this '
                    'script.')

###############################################################################
######## Coupon Sizing                                                 ########
###############################################################################

uni_dict = {
    '0deg' : {
        'number' : [25, 26, 27, 28, 29],
        'width' : [14.9, 14.855, 14.835, 14.87, 14.8425],
        'thickness' : [2.2725, 2.385, 2.2775, 2.325, 2.355],
        },
    'p45deg' : {
        'number' : [30, 31, 32, 33, 34],
        'width' : [14.8775, 14.8675, 14.8975, 14.8625, 14.845],
        'thickness' : [2.3225, 2.3475, 2.3075, 2.3475, 2.34],
        },
    '90deg' : {
        'number' : [35, 36, 37, 38, 39],
        'width' : [14.865, 14.8475, 14.93, 14.86, 14.8625],
        'thickness' : [2.3775, 2.3625, 2.5025, 2.33, 2.4175],
        },
    }

bx_dict = {
    '0deg' : {
        'number' : [25, 26, 27, 28, 29],
        'width' : [14.7525, 14.835, 14.8275, 14.805, 14.8475],
        'thickness' : [1.775, 1.7725, 1.8325, 1.815, 1.8475],
        },
    'p45deg' : {
        'number' : [30, 31, 32, 33, 34],
        'width' : [14.8325, 14.83, 14.8425, 14.8275, 14.8375],
        'thickness' : [1.84, 1.8475, 1.8425, 1.84, 1.8575],
        },
    }

tx_dict = {
    '0deg' : {
        'number' : [1, 2, 3, 4, 5],
        'width' : [14.8025, 14.8125, 14.805, 14.795, 14.77],
        'thickness' : [1.705, 1.7425, 1.8175, 1.725, 1.7475],
        },
    'm45deg' : {
        'number' : [16, 17, 18, 19, 20],
        'width' : [14.8075, 14.7875, 14.7775, 14.82, 14.7975],
        'thickness' : [1.8425, 1.765, 1.8, 1.8125, 1.815],
        },
    '90deg' : {
        'number' : [6, 7, 8, 9, 10],
        'width' : [14.77, 14.7825, 14.805, 14.825, 14.7675],
        'thickness' : [1.8125, 1.78, 1.8075, 1.7825, 1.805],
        },
    }


carbon_dict = {
    '0deg' : {
        'number' : [7, 8, 9, 10, 11, 12],
        'width' : [14.84, 14.8925, 14.855, 14.8575, 14.83, 14.8775],
        'thickness' : [2.98, 3.0175, 2.9825, 2.98, 2.9775, 2.98],
        },
    'load_frame' : '250 kN'
    }


ud_pecan0_dict = {
    '0deg' : {
        'number' : [6, 7, 8, 9, 10],
        'width' : [14.8625, 14.8175, 14.875, 15.0125, 14.785],
        'thickness' : [2.2775, 2.2875, 2.2825, 2.1975, 2.27],
        },
    'load_frame' : '12 kN'
    }


ud_pecan90_dict = {
    '90deg' : {
        'number' : [6, 7, 8, 9, 10],
        'width' : [14.87, 14.8875, 14.865, 14.8825, 14.9325],
        'thickness' : [2.365, 2.355, 2.34, 2.4325, 2.34],
        },
    'load_frame' : '12 kN'
    }

ud_pecan45_dict = {
    'p45deg' : {
        'number' : [6, 7, 8, 9, 10],
        'width' : [14.8275, 14.815, 14.83, 14.8025, 14.82],
        'thickness' : [2.2525, 2.25, 2.225, 2.2375, 2.25],
        'adjust_max' : {
            '7' : 56000 #56000 is good here
            }
        },
    'load_frame' : '12 kN'
    }


bx_pecan0_dict = {
    '0deg' : {
        'number' : [6, 7, 8, 9, 10],
        'width' : [14.8325, 14.77, 14.79, 14.805, 14.7875],
        'thickness' : [1.8, 1.8325, 1.8575, 1.835, 1.855],
        },
    'load_frame' : '12 kN'
    }

bx_pecan45_dict = {
    '45deg' : {
        'number' : [7, 8, 9, 10, 11, 12],
        'width' : [14.8375, 14.775, 14.8225, 14.7875, 14.81, 14.7925],
        'thickness' : [1.8575, 1.89, 1.9, 1.8525, 1.8525, 1.8725],
        },
    'load_frame' : '12 kN'
    }

epoxy_ud_dict = {
    '0deg' : {
        'number' : [13, 14, 15, 16, 17, 18],
        'width' : [14.7925, 14.825, 14.7775, 14.8, 14.785, 14.755],
        'thickness' : [2.2575, 2.25, 2.2275, 2.325, 2.24, 2.225],
        },
    '90deg' : {
        'number' : [11, 12, 13, 14, 15, 16],
        'width' : [14.8825, 14.815, 14.8875, 14.8175, 14.86, 14.875],
        'thickness' : [2.245, 2.2525, 2.245, 2.26, 2.245, 2.2825],
        },

    'load_frame' : '12 kN'
    }


epoxy_ud_pm45_dict = {
    'pm45deg' : {
        'number' : [1,2,3,4,5,6],
        'width' : [14.82, 14.8225, 14.8525, 14.8325, 14.855, 14.8075],
        'thickness' : [2.5475, 2.5225, 2.535, 2.5125, 2.5075, 2.495],
        },
    'load_frame' : '12 kN'
    }

epoxy_bx_dict = {
    '0deg' : {
        'number' : [4, 6, 8],
        'width' : [15.5775, 15.5375, 15.5375],
        'thickness' : [1.8275, 1.81, 1.84],
        },
    '90deg' : {
        'number' : [4, 6, 8],
        'width' : [15.46, 15.54, 15.4925],
        'thickness' : [1.8225, 1.8375, 1.8125],
        },
    'p45deg' : {
        'number' : [4, 6, 8],
        'width' : [15.2575, 15.56, 15.5425],
        'thickness' : [1.745, 1.8425, 1.8575],
        },
    'm45deg' : {
        'number' : [4, 6, 8],
        'width' : [15.51, 15.475, 15.4875],
        'thickness' : [1.82, 1.7875, 1.79],
        },
    'load_frame' : '12 kN'
    }

coupon_sizes = {
    '2025003' : epoxy_bx_dict,
    '2025004' : epoxy_ud_dict,
    '2025008' : uni_dict,
    '2025009' : bx_dict,
    '2025011' : tx_dict,
    'PCF001' : carbon_dict,
    '2025018' : ud_pecan90_dict,
    '2025019' : ud_pecan0_dict,
    '2025020' : ud_pecan45_dict,
    '2025021' : bx_pecan0_dict,
    '2025023' : bx_pecan45_dict,
    '2025036' : epoxy_ud_pm45_dict,
    }


def process_data_set(folder_name, start_threshold, end_threshold, max_samples):

    ###########################################################################
    ######## Parse Files                                               ########
    ###########################################################################
    filenames = sorted(glob.glob(f"{folder_name}/*.tdms"))

    parsed_files = [None] * len(filenames)

    # Iterate over each filename
    for ind,filepath in enumerate(filenames):

        curr_file = os.path.basename(filepath)

        dash_components = curr_file.split('-')
        under_components = dash_components[-1].split('_')

        parsed_files[ind] = {
            'filename': curr_file,
            'panel': dash_components[1], # panel number
            'angle': dash_components[2], # coupon cut angle
            'number': int(under_components[0]), # coupon number
            'date' : '-'.join(under_components[1:4]),
            'time' : ':'.join(under_components[4:7]),
        }

        parsed_freq = float(under_components[-1][:-7])

        assert parsed_freq == fs, 'Sample frequency from filename' \
            + ' is not what is assumed for filtering.'

    ###########################################################################
    ######## Load Data                                                 ########
    ###########################################################################

    recorded_data = [None] * len(filenames)

    processed_results = [None] * len(filenames)

    for ind,filepath in enumerate(filenames):

        tdms_file = TdmsFile.read(filepath)

        # Identify parameters of the tests
        panel = parsed_files[ind]['panel']
        angle = parsed_files[ind]['angle']
        coupon_num = parsed_files[ind]['number']

        size_ind = np.where(np.array(coupon_sizes[panel][angle]['number'])
                            == coupon_num)[0][0]

        width = coupon_sizes[panel][angle]['width'][size_ind] * 1e-3
        thickness = coupon_sizes[panel][angle]['thickness'][size_ind] * 1e-3


        curr_max_samples = max_samples
        if 'adjust_max' in coupon_sizes[panel][angle].keys():

            if '{}'.format(coupon_num) in coupon_sizes[panel][angle]\
                ['adjust_max']:

                curr_max_samples = coupon_sizes[panel][angle]\
                    ['adjust_max']['{}'.format(coupon_num)]

        # list of channels
        # tdms_file.groups()[0].channels()

        force_sign = 1.0

        if 'load_frame' in coupon_sizes[panel].keys():
            if coupon_sizes[panel]['load_frame'].upper() == '250 kN'.upper():

                strain_key = 'Axial Strain 250 kN'
                load_key = 'Load 250kN'
                force_sign = -1.0

            elif coupon_sizes[panel]['load_frame'].upper() == '12 kN'.upper():

                strain_key = 'Axial Strain 12 kN'
                load_key = 'Load 12kN'

            else:
                raise
        else:
            strain_key = 'Axial Strain'
            load_key = 'Load 12kN'

        recorded_data[ind] = {
            'axial_data' : tdms_file.groups()[0][strain_key]\
                [:curr_max_samples],
            'axial_units' :
                tdms_file.groups()[0][strain_key].properties['Units'],
            'transverse_data' : transverse_sign * \
                tdms_file.groups()[0]['Transverse Strain'][:curr_max_samples],
            'transverse_units' :
                tdms_file.groups()[0]['Transverse Strain'].properties['Units'],
            'force_data' : force_sign*tdms_file.groups()[0][load_key]\
                [:curr_max_samples],
            'force_units' :
                tdms_file.groups()[0][load_key].properties['Units'],
            }


        ######################
        # Filter data to eliminate noise

        if filter_strain:

            high = high_freq / (0.5 * fs)
            filter_order = 3

            b, a = butter(filter_order, high, btype='lowpass')

            recorded_data[ind]['axial_data'] \
                = filtfilt(b, a, recorded_data[ind]['axial_data'])

            recorded_data[ind]['transverse_data'] \
                = filtfilt(b, a, recorded_data[ind]['transverse_data'])

            recorded_data[ind]['force_data'] \
                = filtfilt(b, a, recorded_data[ind]['force_data'])

        ######################
        # Find the Index with Max Strain

        max_ind = np.argmax(recorded_data[ind]['axial_data'])

        start_ind = np.where(recorded_data[ind]['axial_data'] \
                             > start_threshold)[0][0]

        # find the last index (before the peak strain) where the ending
        # threshold is met
        end_ind = np.where(recorded_data[ind]['axial_data'][:max_ind] \
                             < end_threshold)[0][-1]

        ######################
        # Linear regression results

        axial_curr = recorded_data[ind]['axial_data'][start_ind:end_ind]
        force_curr = recorded_data[ind]['force_data'][start_ind:end_ind]
        trans_curr = recorded_data[ind]['transverse_data'][start_ind:end_ind]

        force_fit,ss_res,_,_ = np.linalg.lstsq(
                    np.vstack((np.ones_like(axial_curr),
                              axial_curr)).T,
                    force_curr,
                    rcond=None)

        ss_tot = ((force_curr
                  - force_curr.mean())**2).sum()

        rsq_force = 1 - ss_res / ss_tot

        trans_fit,ss_res,_,_ = np.linalg.lstsq(
                    np.vstack((np.ones_like(axial_curr),
                              axial_curr)).T,
                    trans_curr,
                    rcond=None)

        ss_tot = ((trans_curr
                  - trans_curr.mean())**2).sum()

        rsq_trans = 1 - ss_res / ss_tot


        ######################
        # Shear modulus fit

        shear_stress = force_curr * lbf_to_N / 2 / width / thickness
        shear_strain = axial_curr/1e6 - (trans_curr / (width*1e3))

        shear_fit,ss_res,_,_ = np.linalg.lstsq(
                    np.vstack((np.ones_like(shear_strain),
                              shear_strain)).T,
                    shear_stress,
                    rcond=None)

        ss_tot = ((shear_stress - shear_stress.mean())**2).sum()

        rsq_shear = 1 - ss_res / ss_tot

        ######################
        # Format results

        processed_results[ind] = {
            'panel' : panel,
            'angle' : angle,
            'coupon_number' : coupon_num,
            # 'axial_lbf/ue' : force_fit[1],
            # 'trans_mm/ue' : trans_fit[1],
            'elastic_mod_Pa' : force_fit[1] * lbf_to_N * 1e6 / width \
                / thickness,
            'rsq_elastic_mod' : rsq_force[0],
            'poissons' : -1*(trans_fit[1] / (width*1e3)) * 1e6,
            'rsq_poissons' : rsq_trans[0],
            'shear_mod_at_45deg_Pa' : shear_fit[1],
            'rsq_shear' : rsq_shear[0]
            }

        ######################
        if show_plots:
            plt.plot(recorded_data[ind]['axial_data'],
                     recorded_data[ind]['force_data'])

            plt.xlabel(recorded_data[ind]['axial_units'])
            plt.ylabel(recorded_data[ind]['force_units'])

            plt.title(os.path.basename(filepath))
            plt.show()


            plt.plot(recorded_data[ind]['axial_data'],
                     recorded_data[ind]['transverse_data'])

            plt.xlabel(recorded_data[ind]['axial_units'])
            plt.ylabel(recorded_data[ind]['transverse_units'])

            plt.title(os.path.basename(filepath))
            plt.show()


            plt.plot(recorded_data[ind]['axial_data'])

            bounds = [recorded_data[ind]['axial_data'].min(),
                      recorded_data[ind]['axial_data'].max()]

            plt.plot([start_ind, start_ind], bounds)
            plt.plot([end_ind, end_ind], bounds)

            plt.plot([0, len(recorded_data[ind]['axial_data'])],
                     [start_threshold, start_threshold])

            plt.plot([0, len(recorded_data[ind]['axial_data'])],
                     [end_threshold, end_threshold])

            plt.ylabel(recorded_data[ind]['axial_units'])

            plt.title(os.path.basename(filepath))
            plt.show()




            plt.plot(axial_curr, force_curr)
            plt.xlabel(recorded_data[ind]['axial_units'])
            plt.ylabel('Force ' + recorded_data[ind]['force_units'])
            plt.title('Fitted Data ' + os.path.basename(filepath))
            plt.show()


            plt.plot(axial_curr, trans_curr)
            plt.xlabel(recorded_data[ind]['axial_units'])
            plt.ylabel('Transverse ' + recorded_data[ind]['transverse_units'])
            plt.title('Fitted Data ' + os.path.basename(filepath))
            plt.show()

            plt.plot(shear_strain*1e6, shear_stress/1e6)
            plt.xlabel('Shear strain [ue]')
            plt.ylabel('Shear Stress [MPa]')
            plt.show()


        # axial_curr = recorded_data[ind]['axial_data'][start_ind:end_ind]
        # force_curr = recorded_data[ind]['force_data'][start_ind:end_ind]
        # trans_curr = recorded_data[ind]['transverse_data'][start_ind:end_ind]



    ###########################################################################
    ######## Output CSV of Data                                        ########
    ###########################################################################

    output_cols = ['panel', 'angle', 'coupon_number', 'elastic_mod_Pa',
                   'rsq_elastic_mod', 'poissons', 'rsq_poissons',
                   'shear_mod_at_45deg_Pa', 'rsq_shear']


    # output_csv = './pandas_csv_out.csv'
    df = pd.DataFrame(processed_results, columns=output_cols)
    df.to_csv(output_fname, index=False)


    print('{} Coupons Reported'.format(df['elastic_mod_Pa'].shape[0]))

    print('Avg. Elastic Modulus: {:.4f} GPa'.format(
        df['elastic_mod_Pa'].mean()/1e9))

    print('Avg. Poisson Ratio: {:.4f}'.format(df['poissons'].mean()))

    print('Shear Modulus for 45 degree orientation: {:.4f} GPa\n'.format(
        df['shear_mod_at_45deg_Pa'].mean()/1e9))




###############################################################################
######## Loop through and actually process all data files              ########
###############################################################################

# Make the output folder
os.makedirs(os.path.join(filedir, 'outputs'), exist_ok=True)

# Process data to output folder
for top_folder in case_dict.keys():
    for second_folder in case_dict[top_folder].keys():

        folder_name = os.path.join(data_folder, top_folder, second_folder)

        output_fname = os.path.join(filedir, 'outputs',
                            './{}_{}.csv'.format(top_folder, second_folder))

        curr_dict = case_dict[top_folder][second_folder]

        start_threshold = curr_dict.get('start_threshold',
                                        default_start_threshold)

        end_threshold = curr_dict.get('end_threshold',
                                        default_end_threshold)

        max_samples = curr_dict.get('max_samples',
                                        default_max_samples)

        print('Processing Data for {} {}'.format(top_folder, second_folder))

        process_data_set(folder_name, start_threshold,
                         end_threshold, max_samples)
