import argparse
from pathlib import Path
import matplotlib.pyplot as plt
import pandas as pd

REPO_ROOT = Path(__file__).resolve().parents[1]
LOG_DIRECTORY = REPO_ROOT / 'results' / 'raw'
OUTPUT_DIRECTORY = REPO_ROOT / 'results' / 'figures'
OUTPUT_DIRECTORY.mkdir(parents=True, exist_ok=True)

ALG = 'proposed'
N_WORKERS = '10'
DATASET = 'fmnist'
PARTITION = 'noniid-labeldir'
MODEL = 'cnn'
SEED = '42'
ALPHAS_TO_PLOT = [0.1, 0.3, 0.5, 0.7, 0.9]
SMOOTHING_WINDOW = 10
FIGSIZE = (8, 5)

colormap = plt.get_cmap('tab10')
colors_for_plot = [colormap(0), colormap(1), colormap(2), colormap(3), colormap(4)]
ls_for_plot = [
    '-',
    '--',
    ':',
    '-.',
    (0, (5, 5)),
]


def parse_args():
    parser = argparse.ArgumentParser(
        description='Plot DuDe-SGD curves on Fashion-MNIST across different alpha values.'
    )
    parser.add_argument('--lr', required=True, help='Learning rate used in the raw CSV filenames, e.g., 0.01 or 0.1.')
    parser.add_argument('--comput-dist', required=True, help='Computation distribution string used in the raw CSV filenames, e.g., normal(1,5).')
    parser.add_argument('--metric', required=True, choices=['test-acc', 'train-loss'], help='Metric to plot.')
    return parser.parse_args()


def smooth_data(data, window=SMOOTHING_WINDOW):
    return data.rolling(window=window, min_periods=1).mean()


def raw_prefix(lr: str, comput_dist: str, alpha_val: float) -> str:
    return (
        f'log_{ALG}-lr{lr}-'
        f'n_workers{N_WORKERS}-'
        f'{comput_dist}-dataset{DATASET}-'
        f'partition{PARTITION}-model{MODEL}-'
        f'alpha{alpha_val}-seed{SEED}'
    )


def output_pdf_name(comput_dist: str, metric: str) -> str:
    # Match the uploaded sample naming style, e.g.
    # proposed-normal(1,5)-fmnist-noniid-labeldir-cnn-test-acc.pdf
    return f'{ALG}-{comput_dist}-{DATASET}-{PARTITION}-{MODEL}-{metric}.pdf'


if __name__ == '__main__':
    args = parse_args()

    if not LOG_DIRECTORY.is_dir():
        raise FileNotFoundError(f"Log directory '{LOG_DIRECTORY}' does not exist.")

    metric_column = 'Train/Loss' if args.metric == 'train-loss' else 'Test/Acc'
    ylabel = 'Training Loss' if args.metric == 'train-loss' else 'Test Accuracy'
    legend_loc = 'upper right' if args.metric == 'train-loss' else 'lower right'

    fig, ax = plt.subplots(figsize=FIGSIZE)

    curve_idx = 0
    for alpha_val in sorted(ALPHAS_TO_PLOT):
        prefix = raw_prefix(args.lr, args.comput_dist, alpha_val)
        csv_path = LOG_DIRECTORY / f'{prefix}-{args.metric}.csv'
        print(f'Trying to load: {csv_path}')

        try:
            df = pd.read_csv(csv_path)
            raw_x = df['Step/Comput'].dropna()
            raw_y = df[metric_column].dropna()
            smoothed_y = smooth_data(raw_y)

            ax.plot(
                raw_x,
                smoothed_y,
                label=f'α={alpha_val}',
                color=colors_for_plot[curve_idx % len(colors_for_plot)],
                linestyle=ls_for_plot[curve_idx % len(ls_for_plot)],
                linewidth=1.5,
            )
            print('File loaded and plotted successfully.')
        except FileNotFoundError:
            print(f'Warning: file not found, skipping: {csv_path}')
        except KeyError:
            print(f"Warning: file '{csv_path}' is missing required columns.")
        except Exception as e:
            print(f"An error occurred while reading '{csv_path}': {e}")

        curve_idx += 1

    ax.set_xlabel('Time Units', fontsize=25)
    ax.set_ylabel(ylabel, fontsize=25)
    ax.legend(fontsize=20, loc=legend_loc)
    ax.grid(True, linestyle=':', alpha=0.7)
    fig.tight_layout()

    output_path = OUTPUT_DIRECTORY / output_pdf_name(args.comput_dist, args.metric)
    fig.savefig(output_path, format='pdf')
    print(f'Saved figure to: {output_path}')
    plt.show()
