import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import sys, os, pathlib

assert len(sys.argv) == 3, f'wrong args, syntax is: python {sys.argv[0]} <result_input_path> <dynamic (0|1)>'

file = sys.argv[1]  #'./results/RCV1.C4.csv'
loop = bool(int(sys.argv[2]))

plotname = pathlib.Path(file).name.replace(".csv", ".png")

# plot the data
fig, axs = plt.subplots(3)

try:
    while True:
        df = pd.read_csv(file, sep='\t')

        xs = df['it']

        y_r = df['R']
        y_rhat = df['Rhat']
        y_rhatCC = df['RhatCC']
        label='R'
        axs[0].plot(xs, y_rhat, label='$\hat{'+label+'}$')
        axs[0].plot(xs, y_rhatCC, label='$\hat{'+label+'}_{CC}$')
        axs[0].plot(xs, y_r, label=label)
        axs[0].legend()
        axs[0].grid()

        y_r = df['te-prev']
        y_rhat = df['te-estim']
        y_rhatCC = df['te-estimCC']
        label='P'
        axs[1].plot(xs, y_rhat, label='$\hat{'+label+'}$')
        axs[1].plot(xs, y_rhatCC, label='$\hat{'+label+'}_{CC}$')
        axs[1].plot(xs, y_r, label=label)
        axs[1].legend()
        axs[1].grid()

        y_ae = df['AE']
        y_ae_cc = df['AE_CC']
        axs[2].plot(xs, y_ae, label='AE')
        axs[2].plot(xs, y_ae_cc, label='AE-CC')
        axs[2].legend()
        axs[2].grid()

        os.makedirs('./plots', exist_ok=True)
        plt.savefig(f'./plots/{plotname}')

        if not loop:
            break
        else:
            plt.pause(.5)
            axs[0].cla()
            axs[1].cla()
            axs[2].cla()

except KeyboardInterrupt:
    print("\n[exit]")