import argparse
import glob
import multiprocessing as multi
import pathlib

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import siml


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'data_directory',
        type=pathlib.Path,
        help='Data directory path')
    args = parser.parse_args()

    data_files = glob.glob(
        str(args.data_directory / '**/global_ltec_array.npy'),
        recursive=True)
    max_process = siml.util.determine_max_process()
    chunksize = max(
        len(data_files) // max_process // 16, 1)

    with multi.Pool(max_process) as pool:
        errors = pool.map(
            calculate_error, data_files, chunksize=chunksize)
    df = pd.DataFrame(errors)
    df.to_csv(
        args.data_directory / 'ltec_strain_error.csv', index=False,
        header=['data_directory', 'abs_error', 'rela_error', 'max_abs_strain'])

    return


def calculate_error(data_file):
    data_directory = pathlib.Path(data_file).parent

    ltec = np.load(data_file)
    t_diff = np.load(data_directory / 'nodal_t_diff.npy')
    strain = np.load(data_directory / 'nodal_strain_array.npy')
    estimate = ltec * t_diff
    rmse = np.mean((estimate - strain)**2)**.5
    scale = np.mean(strain**2)**.5
    print(f"Finish process: {data_directory}")
    return data_directory, rmse, rmse / scale, np.max(np.abs(strain))


if __name__ == '__main__':
    main()
