import os
from os.path import join, exists
import argparse
import pathlib
from enum import Enum

import click
import numpy as np
import pandas as pd

import download_data
import dataframe
import plotter
from model_types import ModelTypes, model_types_map


class DeepAugmentModelTypes(Enum):
    STANDARD = ModelTypes.STANDARD.value
    LP_ADV = ModelTypes.LP_ADV.value
    ROBUST_INTV = ModelTypes.ROBUST_INTV.value
    MORE_DATA = ModelTypes.MORE_DATA.value
    DEEPAUGMENT = ('_nolegend_', 'tab:brown', 150, 0.9, 's')


def get_model_type(df_row):
    if 'deepaugment' in df_row.name.lower():
        return DeepAugmentModelTypes.DEEPAUGMENT
    elif model_types_map[df_row.name] == ModelTypes.MORE_DATA:
        return DeepAugmentModelTypes.MORE_DATA
    elif model_types_map[df_row.name] == ModelTypes.LP_ADV:
        return DeepAugmentModelTypes.LP_ADV
    elif model_types_map[df_row.name] == ModelTypes.ROBUST_INTV:
        return DeepAugmentModelTypes.ROBUST_INTV
    elif model_types_map[df_row.name] == ModelTypes.STANDARD:
        return DeepAugmentModelTypes.STANDARD


def show_in_plot(df_row):
    model_name, model_type = df_row.name.lower(), df_row.model_type
    return 'subsample' not in model_name and df_row['val-on-imagenet-r-classes'] >= 75


def use_for_line_fit(df_row):
    model_name, model_type, in_plot = df_row.name.lower(), df_row.model_type, df_row.show_in_plot
    return 'aws' not in model_name and 'batch64' not in model_name and model_type is DeepAugmentModelTypes.STANDARD and in_plot


@click.command()
@click.option('--x_axis', type=str, default='val-on-imagenet-r-classes')
@click.option('--y_axis', type=str, default='imagenet-r')
@click.option('--transform', type=str, default='logit')
@click.option('--num_bootstrap_samples', type=int, default=1000) 
@click.option('--output_dir', type=str, default=str((pathlib.Path(__file__).parent / '../outputs').resolve()))
@click.option('--output_file_dir', type=str, default=str((pathlib.Path(__file__).parent / '../paper/appendix').resolve()))
@click.option('--skip_download', is_flag=True, type=bool)
def generate_xy_plot(x_axis, y_axis, transform, num_bootstrap_samples, output_dir, output_file_dir, skip_download):

    if skip_download:
        filename = join(output_dir, 'grid_df.pkl')
        if not exists(filename):
            raise Exception(f'Downloaded data not found at {filename}. Please run python src/plotting/download_data.py first')
        df = pd.read_pickle(filename)
    else:
        df = download_data.download_plotting_data(output_dir, store_data=True, verbose=True)

    df, df_metadata = dataframe.extract_metadata(df)
    df, df_metadata = dataframe.replace_10percent_with_metadata(df, df_metadata)
    df, df_metadata = dataframe.aggregate_corruptions_with_metadata(df, df_metadata)

    df = prepare_df_for_plotting(df, df_metadata, [x_axis, y_axis])
    df = plotter.add_plotting_data(df, [x_axis, y_axis])

    # auto set xlim and ylim based on visible points
    df_visible = df[df.show_in_plot == True]
    xlim = [df_visible[x_axis].min() - 1, df_visible[x_axis].max() + 0.5]
    ylim = [df_visible[y_axis].min() - 1, df_visible[y_axis].values.max() + 1]

    fig, _ = plotter.model_scatter_plot(df, x_axis, y_axis, xlim, ylim, DeepAugmentModelTypes, 
                                         transform=transform, tick_multiplier=5, num_bootstrap_samples=num_bootstrap_samples,
                                         title='Distribution Shift to ImageNet-R', x_label='ImageNet (class-subsampled)', y_label='ImageNet-R', 
                                         figsize=(12, 8), include_legend=False, return_separate_legend=False)
    
    os.makedirs(output_file_dir, exist_ok=True)

    fig.savefig(join(output_file_dir, f'imagenet_r.pdf'), dpi='figure', bbox_inches='tight', pad_inches=0.1)
    print(f"Plot saved to {join(output_file_dir, f'imagenet_r.pdf')}")


def prepare_df_for_plotting(df, df_metadata, columns):
    assert set(columns).issubset(set(df.columns))

    df = df[columns]
    df_metadata = df_metadata[[x+'_dataset_size' for x in columns]]
    df = df.merge(df_metadata, right_index=True, left_index=True)
    df = df.dropna()

    df['model_type'] = df.apply(get_model_type, axis=1)
    df['show_in_plot'] = df.apply(show_in_plot, axis=1)
    df['use_for_line_fit'] = df.apply(use_for_line_fit, axis=1)

    return df


if __name__ == '__main__':
    generate_xy_plot()
