from collections import defaultdict
from pathlib import Path

import ast
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend for better PDF output
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import random
import seaborn as sns
import tensorflow
import textwrap


# Configure matplotlib for high-quality PDF output with ICLR 2025 specifications
plt.rcParams['pdf.fonttype'] = 42  # Use TrueType fonts for PDF
plt.rcParams['ps.fonttype'] = 42   # Use TrueType fonts for PostScript
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Times', 'Times New Roman', 'DejaVu Serif', 'Computer Modern Roman']
plt.rcParams['mathtext.fontset'] = 'cm'
plt.rcParams['axes.linewidth'] = 1.2
plt.rcParams['lines.linewidth'] = 1.5
# ICLR 2025 uses 11pt base font, but images are scaled to 0.48*textwidth
# So we need larger fonts to compensate for the scaling
plt.rcParams['axes.labelsize'] = 22
plt.rcParams['xtick.labelsize'] = 20
plt.rcParams['ytick.labelsize'] = 20
plt.rcParams['legend.fontsize'] = 20
plt.rcParams['figure.titlesize'] = 24


# Supported dataset choices (alphabetical)
DATASETS = [
    'breast_cancer',
    'custom',
    'diabetes',
    'digits',
    'glass',
    'iris',
    'leaf',
    'mnist',
    'parkinsons',
    'wine',
]

# Paul Tol's muted color scheme
ID_TO_COLOR = {
    "Diabetes": "#332288",   # Dark blue
    "Glass": "#117733",      # Dark green
    "Leaf": "#CC6677",       # Rose
    "Parkinsons": "#AA4499", # Purple
    "Wine": "#999933",       # Greenish yellow

    "SHAP": "#CC6677",       # Rose
    "LIME": "#AA4499",       # Purple

    "Group": "#332288",  # Dark blue
    "Aggregated": "#117733", # Dark green
    "Individual": "#CC6677", # Rose
    "Inflated": "#CC6677"    # Rose
}

def set_random_seed(seed):
    np.random.seed(seed)
    random.seed(seed)
    tensorflow.random.set_seed(seed)


def count_json_files(directory):
    count = 0

    for file in Path(directory).iterdir():
        if not file.is_file() or file.suffix != '.json':
            continue

        count += 1

    return count


def pairplot(stats, file_name):
    # Convert to DataFrame
    stats = sorted(stats, key=lambda x: x["Dataset"])
    data = pd.DataFrame(stats)

    # Create a categorical version of δ for coloring
    delta_categories = ['Very Low', 'Low', 'Medium', 'High', 'Very High']
    data['Delta_Category'] = pd.cut(data['Delta'], bins=5, labels=delta_categories)

    # Select numeric columns for pairplot
    numeric_cols = data.select_dtypes(include='number').columns.tolist()

    # Create PairGrid
    grid = sns.PairGrid(data, vars=numeric_cols, hue='Delta_Category', palette='winter')

    # Plot scatter points
    grid.map_diag(sns.kdeplot)
    grid.map_offdiag(sns.scatterplot)

    # Draw lines connecting points for each δ across all pairs
    for delta_val, group in data.groupby('Delta'):
        for i, x_var in enumerate(numeric_cols):
            for j, y_var in enumerate(numeric_cols):
                if i != j:
                    ax = grid.axes[j, i]  # Row j, Col i
                    ax.plot(group[x_var], group[y_var], alpha=0.5)

    plt.suptitle("Pair Plot with Lines by Delta", y=1.02, fontsize=20)
    plt.tight_layout()
    plt.savefig(file_name, dpi=600, bbox_inches='tight', format='pdf')
    plt.close()


def lineplot(stats, x_col_name, y_col_name, file_name, index_col_name="ID", title=None, log_scale=False):
    # Convert to DataFrame
    stats = sorted(stats, key=lambda x: x["Dataset"])

    # Capitalize dataset names for better presentation
    for stat in stats:
        stat["Dataset"] = stat["Dataset"].title()

    data = pd.DataFrame(stats)

    # Get unique id values and sort them
    unique_ids = sorted(data[index_col_name].unique())

    # Create the line plot
    plt.figure(figsize=(10, 6), dpi=600)

    # Plot a line for each id value
    delta_data = []

    for id in unique_ids:
        id_data = data[data[index_col_name] == id]

        if id_data.empty:
            continue

        plot_data = id_data.groupby(x_col_name)[y_col_name].mean().reset_index()
        plot_data = plot_data.sort_values(x_col_name)

        if id.startswith('δ'):
            delta_data.append(plot_data)
        else:
            plt.plot(plot_data[x_col_name],
                    plot_data[y_col_name],
                    marker='o',
                    label=f'{id}',
                    color=get_color(id, unique_ids),
                    alpha=0.7,
                    linewidth=2,
                    markersize=8)

    if delta_data:
        delta_data = pd.concat(delta_data, ignore_index=True)
        plot_data = delta_data.groupby(x_col_name)[y_col_name].mean().reset_index()
        plot_data = plot_data.sort_values(x_col_name)

        plt.plot(plot_data[x_col_name],
                plot_data[y_col_name],
                marker='o',
                label='Group',
                color=get_color('Group', unique_ids),
                alpha=0.7,
                linewidth=2,
            markersize=8)

    plt.xlabel(x_col_name, fontsize=22)
    plt.ylabel(y_col_name, fontsize=22)

    if log_scale:
        plt.yscale('log')

    if title is None:
        title = f'{x_col_name} vs {y_col_name}'

    # Break the title into multiple lines when it is too long
    title = "\n".join(textwrap.wrap(title, width=50))

    plt.title(title, fontsize=24)
    plt.legend(fontsize=20)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(file_name, dpi=600, bbox_inches='tight', format='pdf')
    plt.close()


def print_grouped_means(stats, group_by_cols, value_cols):
    """
    Print the mean of specified columns, grouped by a tuple of columns as a 2D table.
    Args:
        stats (list of dict): The statistics data.
        group_by_cols (tuple): A tuple of two column names (x_axis_col, y_axis_col).
        value_cols (list of str): The columns to compute means for.
    """
    x_col, y_col = group_by_cols

    # Group by both columns
    grouped = defaultdict(list)
    for stat in stats:
        key = (stat[x_col], stat[y_col])
        grouped[key].append(stat)

    # Calculate means for each group
    group_means = {}
    for (x_val, y_val), group_stats in grouped.items():
        means = {}
        for col in value_cols:
            values = [s[col] for s in group_stats]
            means[col] = np.mean(values)
        group_means[(x_val, y_val)] = means

    # Get unique values for each axis
    x_values = sorted(set(key[0] for key in group_means.keys()))
    y_values = sorted(set(key[1] for key in group_means.keys()))

    # Print results for each value column
    for col in value_cols:
        print(f"\nMean '{col}' grouped by '{x_col}' (columns) and '{y_col}' (rows):")

        # Create the 2D table
        rows = []

        # Header row
        header = [f"{y_col}"] + [str(x_val) for x_val in x_values]
        rows.append(header)

        # Data rows
        for y_val in y_values:
            row = [str(y_val)]
            for x_val in x_values:
                if (x_val, y_val) in group_means:
                    mean_val = group_means[(x_val, y_val)][col]
                    row.append(f"{mean_val:.4f}")
                else:
                    row.append("N/A")  # No data for this combination
            rows.append(row)

        # Calculate column widths for alignment
        col_widths = []
        for i in range(len(header)):
            max_width = max(len(str(row[i])) for row in rows)
            col_widths.append(max_width)

        # Print the table
        for row in rows:
            formatted_row = []
            for i, cell in enumerate(row):
                formatted_row.append(str(cell).ljust(col_widths[i]))
            print("  ".join(formatted_row))
        print()  # Extra line between tables


def get_color(id, unique_ids):
    if not id.startswith('δ'):
        try:
            return ID_TO_COLOR[id]
        except:
            return ID_TO_COLOR[-1]

    deltas_ids = [id for id in unique_ids if id.startswith('δ')]
    deltas_ids.sort()
    index = deltas_ids.index(id)

    try:
        return ID_TO_COLOR['δ'][index]
    except:
        return ID_TO_COLOR['δ'][-1]


def load_csv_files(file_or_directory):
    """Load CSV files from a file path or directory and return as DataFrame."""
    path = Path(file_or_directory)

    if path.is_file() and path.suffix == '.csv':
        # Single CSV file
        try:
            df = pd.read_csv(path)
            return parse_csv_dataframe(df)
        except Exception as e:
            print(f"Error loading {path}: {e}")
            return pd.DataFrame()
    elif path.is_dir():
        # Directory containing CSV files
        csv_files = list(path.glob("*.csv"))

        if not csv_files:
            raise FileNotFoundError(f"No CSV files found in {path}")

        dfs = []
        for csv_file in csv_files:
            try:
                df = pd.read_csv(csv_file)
                dfs.append(parse_csv_dataframe(df))
            except Exception as e:
                print(f"Error loading {csv_file}: {e}")

        if dfs:
            return pd.concat(dfs, ignore_index=True)
        else:
            return pd.DataFrame()
    else:
        raise FileNotFoundError(f"Path {path} is not a valid CSV file or directory")


def count_csv_files(file_or_directory):
    """Count the number of rows in CSV file(s)."""
    path = Path(file_or_directory)

    if path.is_file() and path.suffix == '.csv':
        # Single CSV file
        try:
            df = pd.read_csv(path)
            return len(df)
        except Exception as e:
            print(f"Error counting rows in {path}: {e}")
            return 0
    elif path.is_dir():
        # Directory containing CSV files
        csv_files = list(path.glob("*.csv"))
        total_rows = 0
        for csv_file in csv_files:
            try:
                df = pd.read_csv(csv_file)
                total_rows += len(df)
            except Exception as e:
                print(f"Error counting rows in {csv_file}: {e}")
        return total_rows
    else:
        return 0


def parse_csv_dataframe(df):
    """Parse a DataFrame and convert string representations to proper data types."""
    df_parsed = df.copy()

    # Handle different data types
    for column in df_parsed.columns:
        if column in ['samples', 'sample']:
            # Parse list of lists or list
            df_parsed[column] = df_parsed[column].apply(lambda x:
                ast.literal_eval(x) if isinstance(x, str) else x)

        elif column == 'explanations':
            # Parse list of strings
            df_parsed[column] = df_parsed[column].apply(lambda x:
                ast.literal_eval(x) if isinstance(x, str) else x)

        elif column == 'inflated_explanations':
            # Parse list of lists of [lower_bound, upper_bound] pairs
            df_parsed[column] = df_parsed[column].apply(lambda x:
                parse_inflated_explanations(x) if isinstance(x, str) else x)

        elif column in ['custom_dim', 'label', 'sample_index', 'seed']:
            # Convert to int
            df_parsed[column] = pd.to_numeric(df_parsed[column], errors='coerce').fillna(0).astype(int)

        elif column in ['epsilon', 'delta', 'computation_time']:
            # Convert to float
            df_parsed[column] = pd.to_numeric(df_parsed[column], errors='coerce').fillna(0.0)

    return df_parsed


def parse_inflated_explanations(value):
    """Parse inflated_explanations field specifically."""
    try:
        if isinstance(value, str):
            parsed = ast.literal_eval(value)
            # Handle the case where it's [[[...]]] - extract the inner list
            if isinstance(parsed, list) and len(parsed) == 1 and isinstance(parsed[0], list):
                return parsed[0]
            else:
                return parsed
        else:
            return value
    except:
        return value


# Backward compatibility - replace JSON functions with CSV equivalents
def load_json_files(directory):
    """Load CSV files instead of JSON files for backward compatibility."""
    return load_csv_files(directory)


def count_json_files(directory):
    """Count CSV files instead of JSON files for backward compatibility."""
    return count_csv_files(directory)
