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

DEFAULT_RESULTS_DIR = Path('./results/raw')
DEFAULT_SEEDS = [3407, 114514]
DEFAULT_FIGSIZE = (16, 3)
DEFAULT_RECT = [
    [0.05, 0.18, 0.2, 0.70],
    [0.2875, 0.18, 0.2, 0.70],
    [0.525, 0.18, 0.2, 0.70],
    [0.7625, 0.18, 0.2, 0.70],
]

PLOT_ORDER = [
    'Synchronous SGD',
    'MIFA',
    'FedBuff',
    'Delayed SGD',
    'Uniform ASGD',
    'Shuffled ASGD',
    r'DuDe-SGD',
]

colormap = plt.get_cmap('tab10')
colors = [colormap(i) for i in range(7)]
colors = [colors[2], colors[6], colors[1], colors[0], colors[3], colors[5], colors[4]]
line_styles = ['-.', '-.', '--', '-', ':', '--', '-']


def build_alg_prefix(sync_buffer: int):
    return {
        'Delayed SGD': lambda tau_std, alpha: f'async-lr0.005-normal(1,{tau_std})-alpha{alpha}',
        'FedBuff': lambda tau_std, alpha: f'buffer(5)-lr0.005-normal(1,{tau_std})-alpha{alpha}',
        'Synchronous SGD': lambda tau_std, alpha: f'buffer({sync_buffer})-lr0.01-normal(1,{tau_std})-alpha{alpha}',
        'Uniform ASGD': lambda tau_std, alpha: f'concurrency(5)-lr0.01-normal(1,{tau_std})-alpha{alpha}',
        'Shuffled ASGD': lambda tau_std, alpha: f'shuffle-lr0.01-normal(1,{tau_std})-alpha{alpha}',
        'MIFA': lambda tau_std, alpha: f'mifa-lr0.01-normal(1,{tau_std})-alpha{alpha}',
        r'DuDe-SGD': lambda tau_std, alpha: f'proposed-lr0.01-normal(1,{tau_std})-alpha{alpha}',
    }


def default_settings(n_workers: int):
    if n_workers == 10:
        return ['1-0.1', '5-0.1', '1-0.5', '5-0.5']
    if n_workers == 30:
        return ['1-0.05', '5-0.05', '1-0.1', '5-0.1']
    return ['1-0.1', '5-0.1', '1-0.5', '5-0.5']


def parse_args():
    parser = argparse.ArgumentParser(
        description='Display CIFAR comparison figures from seed-level CSV files in results/raw/.')
    parser.add_argument('--results-dir', type=Path, default=DEFAULT_RESULTS_DIR,
                        help='Directory containing {dataset}-{metric}-{n_workers}-{seed}.csv files, typically results/raw/.'))
    parser.add_argument('--dataset', type=str, default='cifar',
                        help='Dataset tag used in the merged CSV filenames, e.g., cifar.')
    parser.add_argument('--n-workers', type=int, default=10, help='Number of workers used in the experiments.')
    parser.add_argument('--seeds', type=int, nargs='+', default=DEFAULT_SEEDS,
                        help='Random seeds to include when averaging the curves.')
    parser.add_argument('--metric', choices=['acc', 'loss', 'both'], default='both',
                        help='Which metric to display.')
    parser.add_argument('--settings', type=str, nargs='+', default=None,
                        help="Optional list of setting keys in the form 'tauStd-alpha', e.g., 1-0.1 5-0.1 1-0.5 5-0.5.")
    return parser.parse_args()


def smooth_data(data, window=20):
    mean = data.rolling(window).mean()
    std = data.rolling(window).std()
    return mean, std


def comparison_path(results_dir: Path, dataset: str, metric: str, n_workers: int, seed: int) -> Path:
    return results_dir / f'{dataset}-{metric}-{n_workers}-{seed}.csv'


def load_frames(results_dir: Path, dataset: str, metric: str, n_workers: int, seeds):
    frames = {}
    for seed in seeds:
        path = comparison_path(results_dir, dataset, metric, n_workers, seed)
        if path.exists():
            frames[seed] = pd.read_csv(path, encoding='ISO-8859-1')
        else:
            print(f'Warning: comparison file not found and will be skipped: {path}')
    if not frames:
        raise FileNotFoundError(
            f'No comparison CSV files were found in {results_dir} for metric={metric}.')
    return frames


def extend_initial(raw: pd.DataFrame, column: str, initial_value: float) -> pd.DataFrame:
    if raw.empty:
        return raw
    first_step = raw['Step/Comput'].min()
    if pd.isna(first_step):
        return raw
    first_step = int(first_step)
    if first_step <= 0:
        return raw
    prefix = pd.DataFrame({'Step/Comput': range(0, first_step), column: [initial_value] * first_step})
    return pd.concat([prefix, raw], ignore_index=True)


def collect_metric_data(frames, tau_std, alpha, metric_name, alg_prefix):
    x_data_dict = {alg: {} for alg in PLOT_ORDER}
    step_col_candidates = ['Step/Comput/Cul', 'Step/Comput']
    initial_value = 0.0976 if metric_name == 'Test/Acc' else 2.320514949
    need_initial_extension = {'Synchronous SGD', 'MIFA', 'FedBuff', 'Shuffled ASGD'}

    for seed, df in frames.items():
        step_col = next((col for col in step_col_candidates if col in df.columns), None)
        if step_col is None:
            raise KeyError('Comparison CSV must contain Step/Comput or Step/Comput/Cul.')

        for alg in PLOT_ORDER:
            column = f"{alg_prefix[alg](tau_std, alpha)} - {metric_name}"
            if column not in df.columns:
                print(f"Warning: column '{column}' is missing in seed {seed}; skipping it.")
                continue

            raw = df[[step_col, column]].dropna(axis=0, how='any').copy()
            raw = raw.rename(columns={step_col: 'Step/Comput'})
            raw = raw.sort_values('Step/Comput')
            raw = raw.drop_duplicates(subset=['Step/Comput'], keep='last')

            if alg in need_initial_extension:
                raw = extend_initial(raw, column, initial_value)

            for x, y in zip(raw['Step/Comput'], raw[column]):
                x_int = int(x)
                x_data_dict[alg].setdefault(x_int, {})[seed] = float(y)

    return x_data_dict


def plot_metric(metric_name, ylabel, frames, settings, alg_prefix):
    fig = plt.figure(figsize=DEFAULT_FIGSIZE)
    handles_for_legend = {}

    for i, key in enumerate(settings):
        ax = fig.add_axes(DEFAULT_RECT[i])
        tau_std, alpha = key.split('-', 1)
        x_data_dict = collect_metric_data(frames, tau_std, alpha, metric_name, alg_prefix)

        for idx, alg in enumerate(PLOT_ORDER):
            zorder = idx if alg != r'DuDe-SGD' else 10
            line_width = 2.0 if alg == r'DuDe-SGD' else 1.5

            x_values = np.array(sorted(x_data_dict[alg].keys()))
            x_values = x_values[x_values < 8000]
            if len(x_values) == 0:
                continue

            y_all = {x: np.array(list(x_data_dict[alg][x].values())) for x in x_values}
            y_mean = [np.mean(v) for v in y_all.values()]
            x_smooth, _ = smooth_data(pd.Series(x_values), window=10)
            y_mean_smooth, y_std_smooth = smooth_data(pd.Series(y_mean), window=10)
            min_len = min(len(x_smooth), len(y_mean_smooth))
            x_smooth = x_smooth[:min_len]
            y_mean_smooth = y_mean_smooth[:min_len]
            y_std_smooth = y_std_smooth[:min_len]

            interval = 10
            (line,) = ax.plot(
                x_smooth[::interval],
                y_mean_smooth[::interval],
                linestyle=line_styles[idx],
                color=colors[idx],
                linewidth=line_width,
                label=alg,
                zorder=zorder,
            )
            handles_for_legend.setdefault(alg, line)
            ax.fill_between(
                x_smooth[::interval],
                y_mean_smooth[::interval] - y_std_smooth[::interval],
                y_mean_smooth[::interval] + y_std_smooth[::interval],
                color=colors[idx],
                alpha=0.15,
            )

        ax.set_xlim(0, 6000)
        if metric_name == 'Test/Acc':
            ax.set_ylim(0.1, 0.6)
        ax.grid(True)
        ax.set_xlabel('Time Units', fontsize=12)
        ax.set_title(rf'$\sigma={tau_std}$, $\alpha={alpha}$', fontsize=12)
        if i == 0:
            ax.set_ylabel(ylabel, fontsize=12)

    ordered_handles = [handles_for_legend[alg] for alg in PLOT_ORDER if alg in handles_for_legend]
    ordered_labels = [alg for alg in PLOT_ORDER if alg in handles_for_legend]
    if ordered_handles:
        fig.legend(ordered_handles, ordered_labels, ncol=min(4, len(ordered_labels)),
                   loc='upper center', bbox_to_anchor=(0.5, 1.04), fontsize=11)
    fig.tight_layout(rect=[0, 0, 1, 0.90])
    plt.show()


def main():
    args = parse_args()
    settings = args.settings if args.settings is not None else default_settings(args.n_workers)
    alg_prefix = build_alg_prefix(args.n_workers)

    if args.metric in ('loss', 'both'):
        loss_frames = load_frames(args.results_dir, args.dataset, 'loss', args.n_workers, args.seeds)
        plot_metric('Train/Loss', 'Training Loss', loss_frames, settings, alg_prefix)

    if args.metric in ('acc', 'both'):
        acc_frames = load_frames(args.results_dir, args.dataset, 'acc', args.n_workers, args.seeds)
        plot_metric('Test/Acc', 'Test Accuracy', acc_frames, settings, alg_prefix)


if __name__ == '__main__':
    main()
