"""
Processes LDV data based on yaml file settings for each coupon.


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 files and data to the same folder:
        a. `process_ldv.py` - this script that processes the dynamic data
        b. `ldv_pars.yaml` - definition of tests and parameters.
        c. `dynamic_utils.py` - functions used in processing
        d. `dynamic_coupon_data.tar.gz` - raw data, extract with
           `tar -xzvf dynamic_coupon_data.tar.gz`

    3. `filter_data = False` and `mode_list = None` will process all of the
    data. Additional options to reduce to a desired subset are described below.

    4. Run this script. A csv with all of the data will be produced with the
    name defined by `output_csv` below.

    5. At the end of the script, a plot of all damping measurements together
    and some text output is provided for reference. Formal analysis should
    utilize the output csv file to appropriately restrict data to pass quality
    metrics (e.g., area_error).

"""

import os
import shutil
import numpy as np

import warnings

import matplotlib as mpl
import matplotlib.pyplot as plt

import pandas as pd

from nptdms import TdmsFile

import yaml

import dynamic_utils as utils

mpl.style.use('seaborn-v0_8-colorblind')


###############################################################################
###### Inputs and Settings                                               ######
###############################################################################

filedir = os.path.dirname(os.path.realpath(__file__)) # get path to this file

base_folder = os.path.join(filedir, 'dynamic_coupon_data') # data folder

# settings file
yaml_file = os.path.join(filedir, 'ldv_pars.yaml')

# output details for all tests
output_csv = os.path.join(filedir, 'full_dynamic_data.csv')

# flag to copy all of the data files to a new folder location
# (used to prepare data upload)
copy_data = False
copy_dest = os.path.join(filedir, 'data_copy/')

# Generate about 4 plots per file processed.
# Showing plots generally slows down processing significantly.
show_plots = False
show_final_plots = False

# this is filter the list of coupons, always filter time series.
filter_data = False # if False, the will run all data measurements

# if filter_data == True, these filters can be used to look at only a subset of
# the data
angle_filter = None # options: [0, 90, -45, 45]
resin_filter = None # options: ['BP20', 'elium', 'epoxy', 'aluminum', 'balsa', 'pultruded', 'PECAN']
fiber_filter = None # options: ['triax', 'biax', 'uniax', 'carbon']
index_filter = None # range(8) # index full list before filtering
post_filter_index = None # index list after filtering

# Mode list is applied seperately from filtering, must be strings
mode_list = None # ['1', '2', 'T1'] # None sets the default to be all modes.


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

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

###############################################################################
###### Load Yaml File                                                    ######
###############################################################################

with open(yaml_file, 'r') as file:
    yaml_settings = yaml.safe_load(file)

###############################################################################
###### Filter to desired coupons                                         ######
###############################################################################

coupon_list = yaml_settings['coupons']

# code will be faster if this is a reasonable upper bound on number of tests.
estimated_tests = 3000

if mode_list is None:
    mode_list = ['1', '2', 'T1']

if filter_data:

    if index_filter is not None:
        coupon_list = [coupon_list[ind] for ind in index_filter]

    # Filter on angles
    if angle_filter is not None:
        coupon_list = [coupon for coupon in coupon_list
                       if coupon.get('angle') in angle_filter]

    # Get the material specs
    mat_list = [yaml_settings['general']['panels'][coupon['panel']]['name']
                for coupon in coupon_list]

    mat_mask = [None] * len(mat_list)

    for ind,mat in enumerate(mat_list):

        if resin_filter is not None:
            resin_include = any(word.upper() in mat.upper()
                                for word in resin_filter)
        else:
            resin_include = True

        if fiber_filter is not None:
            fiber_include = any(word.upper() in mat.upper()
                                for word in fiber_filter)
        else:
            fiber_include = True

        mat_mask[ind] = resin_include and fiber_include

    coupon_list = [coupon for ind,coupon in enumerate(coupon_list)
                   if mat_mask[ind]]

    if post_filter_index is not None:
        coupon_list = [coupon_list[ind] for ind in post_filter_index]

###############################################################################
###### Preprocess Coupon Number to be a list in all cases                ######
###############################################################################

for coupon in coupon_list:
    if not isinstance(coupon['coupon'], list):
        coupon['coupon'] = [coupon['coupon']]

###############################################################################
###### Processes Data                                                    ######
###############################################################################

test_ind = 0

output_cols = ['panel', 'angle', 'resin', 'fiber', 'date', 'test_location',
               'bolt_torque',
               'coupon', 'mode', 'index', 'filename', 'infused_date',
               'freq_hz', 'zeta_frac_crit_lsq', 'eta_loss_factor_lsq',
               'lsq_rsq',
               'area_error', 'zeta_frac_crit_log_dec']

# Expected file naming format, can be overwritten by yaml
default_format = '{date:}/ringdown_{panel:}-{angle:}deg-{coupon:}' \
    + '_mode{mode:}_{index}.{ext}'

# df = pd.DataFrame(columns=columns)
rows_list = estimated_tests * [None]

for coupon_dict in coupon_list:

    for coupon_num in coupon_dict['coupon']:

        for mode_name in mode_list:

            mode_key = 'mode' + mode_name

            if mode_key in coupon_dict.keys():

                for index in coupon_dict[mode_key]['file_indices']:

                    if 'fileformat' in coupon_dict.keys():
                        curr_format = coupon_dict['fileformat']
                    else:
                        curr_format = default_format

                    angle_str = coupon_dict.get('angle')
                    if angle_str == 45:
                        angle_str = 'p45'
                    elif angle_str == -45:
                        angle_str = 'm45'

                    fname = curr_format.format(date=coupon_dict.get('date'),
                                                panel=coupon_dict.get('panel'),
                                                angle=angle_str,
                                                coupon=coupon_num,
                                                mode=mode_name,
                                                index=index,
                                                ext=coupon_dict.get('ext'))

                    if angle_str is None:
                        fname = fname.replace('-Nonedeg', '')

                    path_fname = os.path.join(base_folder, fname)

                    # Processes Data
                    if fname.split('.')[-1].upper() == 'UNV':

                        # unv data was saved in units of m/s
                        data_dict = utils.load_unv(path_fname)


                        time = data_dict[3]['time']
                        signal = data_dict[3]['data']

                        assert data_dict[3]['type'].upper() == 'VELOCITY', \
                            'Need to handle output block ind to find velocity.'

                        max_amp = coupon_dict[mode_key].get('max_amp', 0.035)
                        min_amp = coupon_dict[mode_key].get('min_amp', 0.001)
                        units = 'm/s'

                    elif fname.split('.')[-1].upper() == 'TDMS':

                        # tdms files were saved in mm/s
                        tdms_file = TdmsFile.read(path_fname)

                        in_channel = coupon_dict.get('channel',
                                                     'cDAQ1Mod1/ai2')

                        signal = tdms_file.groups()[0][in_channel][:]

                        dt = tdms_file.groups()[0]\
                            [in_channel].properties['wf_increment']

                        unit_string = tdms_file.groups()[0][in_channel]\
                            .properties['unit_string']

                        assert unit_string == 'mm/s', \
                            'Not getting expected units of mm/s.'

                        if 'Applied Sensitivity mm/s/V' \
                            in tdms_file._properties.keys():
                                assert tdms_file._properties\
                                    ['Applied Sensitivity mm/s/V'] == 12.5, \
                                    'Not at expected sensitivity.'
                        else:
                            if '{}'.format(coupon_dict.get('date')) \
                                > '2025-08-14':

                                # Know that sensitivity was not correctly
                                # named in the first date of NREL testing.
                                # For future tests, raise a warning.
                                warnings.warn('Sensitivity not in expected '\
                                              + 'file format.')

                        if 'Data Clipped' in tdms_file._properties.keys():
                            assert not tdms_file._properties['Data Clipped'], \
                                'Data file clipped: {}.'.format(fname)

                        time = np.arange(signal.shape[0])*dt

                        max_amp = coupon_dict[mode_key].get('max_amp', 35.0)
                        min_amp = coupon_dict[mode_key].get('min_amp', 1.0)
                        units = 'mm/s'

                    else:
                        assert False, 'Unrecognized file extension.' \
                            + ' Expect unv or tdms.'

                    freq, zeta, zeta_lsq, lsq_rsq, norm_area_error \
                        = utils.filter_log_dec(time, signal,
                              low_freq=coupon_dict[mode_key]['filter_low'],
                              high_freq=coupon_dict[mode_key]['filter_high'],
                              show_plots=show_plots,
                              title=fname,
                              filter_order=coupon_dict[mode_key].get(
                                  'filter_order', 3),
                              start_offset=coupon_dict[mode_key].get(
                                  'start_offset', 0.0),
                              max_amp=max_amp,
                              min_amp=min_amp,
                              units=units)

                    panel_dict = yaml_settings['general']['panels'] \
                        [coupon_dict['panel']]

                    if coupon_dict['date'] \
                        in yaml_settings['general']['location_dates']['cu']:

                        test_location = 'CU'
                    elif coupon_dict['date'] \
                        in yaml_settings['general']['location_dates']['nlr1']:

                        test_location = 'nlr1'

                    elif coupon_dict['date'] \
                        in yaml_settings['general']['location_dates']['nlr2']:

                        test_location = 'nlr2'
                    else:
                        assert False, 'Date not in test location list.'

                    # get bolt torque:
                    coupon_torque = coupon_dict.get('bolt_torque', 'default')

                    # The coupon_torque could be a value with units, or it
                    # could be high/low/default and need to look up the value.
                    output_torque = yaml_settings['general']['bolt_torque']\
                        .get(coupon_torque, coupon_torque)

                    # Format outputs
                    curr_data = {
                        'panel' : coupon_dict['panel'],
                        'angle' : coupon_dict.get('angle'),
                        'resin' : panel_dict.get('resin'),
                        'fiber' : panel_dict.get('fiber'),
                        'date' : '{}'.format(coupon_dict.get('date')),
                        'test_location': test_location,
                        'bolt_torque' : output_torque,
                        'coupon' : coupon_num,
                        'infused_date' : panel_dict.get('infused'),
                        'mode' : mode_name,
                        'index' : index,
                        'filename' : fname,
                        'freq_hz' : freq,
                        'zeta_frac_crit_lsq' : zeta_lsq,
                        'eta_loss_factor_lsq' : 2*zeta_lsq,
                        'lsq_rsq' : lsq_rsq,
                        'area_error' : norm_area_error,
                        'zeta_frac_crit_log_dec' : zeta,
                        }

                    if test_ind >= len(rows_list):
                        # extend the row list dynamically when necessary
                        rows_list += len(rows_list) * [None]

                    rows_list[test_ind] = curr_data
                    test_ind += 1

                    if np.mod(test_ind, 10) == 0:
                        print('Finished processing {: 4d} measurements.'.format(
                            test_ind))

                    if copy_data:
                        new_path_fname = os.path.join(copy_dest, fname)

                        # create folder to the new path
                        os.makedirs(os.path.dirname(new_path_fname),
                                    exist_ok=True)

                        # copy file
                        shutil.copy(path_fname, new_path_fname)


                    # df = pd.concat([df, pd.DataFrame([row])], ignore_index=True)

df = pd.DataFrame(rows_list[:test_ind], columns=output_cols)

print('Total of {} damping measurements.'.format(test_ind))

###############################################################################
###### Plotting                                                          ######
###############################################################################

df.to_csv(output_csv, index=False)


all_freq = df['freq_hz'].to_numpy()
all_damp = df['zeta_frac_crit_log_dec'].to_numpy()
all_damp_lsq = df['zeta_frac_crit_lsq'].to_numpy() # preferred value
all_lsq_rsq = df['lsq_rsq'].to_numpy()
mode_number = df['mode'].to_numpy()
area_error = df['area_error'].to_numpy()


if show_final_plots:

    plt.plot(all_freq, np.array(all_damp)*2, 'x',
             label='Log Dec')

    plt.plot(all_freq, np.array(all_damp_lsq)*2, '+',
             label='Least Squares (Prefer)')

    plt.xlabel('Frequency [Hz]')
    plt.ylabel('Loss Factor [-]')

    plt.xlim((0, np.maximum(105, all_freq.max()+10)))
    plt.ylim((0, 2.2*np.maximum(all_damp.max(), all_damp_lsq.max())))

    ax = plt.gca()
    ax.tick_params(which='major', left=True, right=True, top=True, bottom=True,
                   direction='in')

    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))

    plt.title('All Data')

    plt.show()

    #########
    # Mask on R^2 value of the linear fit of the exponential

    mask = np.logical_and(np.array(all_lsq_rsq) > 0.90,
                          area_error < 0.02)

    plt.plot(all_freq[mask], all_damp[mask]*2, 'x',
             label='Log Dec')

    plt.plot(all_freq[mask], all_damp_lsq[mask]*2, '+',
             label='Least Squares (Prefer)')

    plt.xlabel('Frequency [Hz]')
    plt.ylabel('Loss Factor [-]')

    plt.xlim((0, np.maximum(105, all_freq.max()+10)))
    plt.ylim((0, 2.2*np.maximum(all_damp.max(), all_damp_lsq.max())))

    ax = plt.gca()
    ax.tick_params(which='major', left=True, right=True, top=True, bottom=True,
                   direction='in')

    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))

    plt.title('Masked on R^2 and Area Error')
    plt.show()

print('Following summary information includes all data points even those that'
      + ' fail quality metrics.')

if (mode_number == '1').sum() > 0:

    print('Eta Mode 1 range: {:.4f}-{:.4f} [%]'.format(
        2*100*(all_damp_lsq[mode_number == '1'].min()),
        2*100*(all_damp_lsq[mode_number == '1'].max())))


    print('Eta Mode 1 (mean, min, max): {:.3f} & {:.3f} & {:.3f} [%]'.format(
        2*100*(all_damp_lsq[mode_number == '1'].mean()),
        2*100*(all_damp_lsq[mode_number == '1'].min()),
        2*100*(all_damp_lsq[mode_number == '1'].max())))

    print('Mode 1 - number of reported data points: {:d} tests'.format(
        (mode_number == 1).sum() + (mode_number == '1').sum()))


    print('Frequency Mode 1 range: {:.2f}-{:.2f} [Hz]'.format(
        all_freq[mode_number == '1'].min(),
        all_freq[mode_number == '1'].max()))

if (mode_number == '2').sum() > 0:

    print('\nEta Mode 2 range: {:.4f}-{:.4f} [%]'.format(
        2*100*(all_damp_lsq[mode_number == '2'].min()),
        2*100*(all_damp_lsq[mode_number == '2'].max())))

    print('Frequency Mode 2 range: {:.2f}-{:.2f} [Hz]'.format(
        all_freq[mode_number == '2'].min(),
        all_freq[mode_number == '2'].max()))

if (mode_number == 'T1').sum() > 0:

    print('\nEta Mode T1 range: {:.4f}-{:.4f} [%]'.format(
        2*100*(all_damp_lsq[mode_number == 'T1'].min()),
        2*100*(all_damp_lsq[mode_number == 'T1'].max())))

    print('Frequency Mode T1 range: {:.2f}-{:.2f} [Hz]'.format(
        all_freq[mode_number == 'T1'].min(),
        all_freq[mode_number == 'T1'].max()))