import argparse
import glob
import multiprocessing as multi
import pathlib

import matplotlib.pyplot as plt
import numpy as np
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:
        pool.map(
            plot_single_data, data_files, chunksize=chunksize)
    return


def plot_single_data(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')

    for i in range(6):
        plt.figure(figsize=(14., 7.))
        plt.plot(
            t_diff * ltec[:, i], strain[:, i], '.', ms=1,
            label=f"{data_directory}_{i}")
        plt.xlabel('LTEC * T_diff')
        plt.ylabel('Strain')
        plt.xlim(-6., 6.)
        plt.ylim(-6., 6.)
        plt.title(data_directory)
        plt.legend()
        save_file_name = data_directory / f"ltec_strain_{i}.pdf"
        plt.savefig(save_file_name)
        print(f"Plot saved in: {save_file_name}")
        plt.close()

    return


if __name__ == '__main__':
    main()
