import requests
import json
import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from adjustText import adjust_text
from load_results import download_results, load_results_into_models, BenchmarkResults

def ensure_image_dir():
    """Ensure the images directory exists relative to this script."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    image_dir = os.path.join(script_dir, 'images')
    os.makedirs(image_dir, exist_ok=True)
    return image_dir

def get_theoretical_ratio(dataset_name: str) -> float | None:
    """Get theoretical compression ratio for a dataset if applicable."""
    if dataset_name.startswith('bernoulli-'):
        try:
            p = float(dataset_name.split('-')[1])
            probs = np.array([p, 1-p])
            entropy = -np.sum(probs * np.log2(probs))  # Shannon entropy
            return 8 / entropy  # 8 bits per byte divided by entropy
        except (IndexError, ValueError):
            return None
    elif dataset_name.startswith('gaussian-q'):
        try:
            sigma = float(dataset_name.split('-q')[1])
            # Calculate probabilities for integers in a reasonable range around mean (0)
            # For 99.7% coverage, use ±4*sigma
            range_max = int(4 * sigma)
            x = np.arange(-range_max, range_max + 1)
            # Calculate probabilities using normal distribution PDF
            probs = (1/(sigma * np.sqrt(2*np.pi))) * np.exp(-(x**2)/(2*sigma**2))
            # Normalize probabilities to sum to 1
            probs = probs / np.sum(probs)
            # Calculate entropy
            nonzero_probs = probs[probs > 0]  # Avoid log(0)
            entropy = -np.sum(nonzero_probs * np.log2(nonzero_probs))
            return 16 / entropy  # 8 bits per byte divided by entropy
        except (IndexError, ValueError):
            return None
    return None

def plot_single_compression_ratio_comparison(results: BenchmarkResults, dataset_name: str, ax: Axes):
    """Create a horizontal bar plot for a specific dataset."""
    # Filter for the specified dataset
    dataset_results = [r for r in results.results if r.dataset == dataset_name]

    # Get unique algorithms and their best compression ratios
    algorithm_ratios = {}
    for r in dataset_results:
        current_ratio = r.compression_ratio
        if r.algorithm not in algorithm_ratios or current_ratio < algorithm_ratios[r.algorithm]:
            algorithm_ratios[r.algorithm] = current_ratio

    # Sort by compression ratio
    sorted_items = sorted(algorithm_ratios.items(), key=lambda x: x[1])
    algorithms = [item[0] for item in sorted_items]
    ratios = [item[1] for item in sorted_items]

    # Create horizontal bar plot
    ax.barh(algorithms, ratios)
    ax.set_xlabel('Compression Ratio')
    ax.set_ylabel('Algorithm')
    ax.set_title(f'Compression Ratios for {dataset_name}')

    # Add theoretical limit line if applicable
    theoretical_ratio = get_theoretical_ratio(dataset_name)
    if theoretical_ratio is not None:
        ax.axvline(x=theoretical_ratio, color='red', linestyle='--')
        ax.legend()

def figure_tradeoff_comparison(results: BenchmarkResults, datasets: list[str]):
    """Create scatter plots showing compression ratio vs speed tradeoffs."""
    # Create two figures - one for encode, one for decode
    for speed_type in ['encode', 'decode']:
        n_plots = len(datasets)
        fig, axes = plt.subplots(1, n_plots, figsize=(10*n_plots, 8))
        fig.suptitle(f'Compression Ratio vs {speed_type.capitalize()} Speed Tradeoff')

        # Handle single dataset case
        if n_plots == 1:
            axes = [axes]

        for dataset_name, ax in zip(datasets, axes):
            # Filter for specified dataset
            dataset_results = [r for r in results.results if r.dataset == dataset_name]

            # Get best results for each algorithm
            algorithm_metrics = {}
            for r in dataset_results:
                speed = r.encode_mb_per_sec if speed_type == 'encode' else r.decode_mb_per_sec
                if r.algorithm not in algorithm_metrics:
                    algorithm_metrics[r.algorithm] = {
                        'ratio': r.compression_ratio,
                        'speed': speed
                    }
                else:
                    # Update if this instance has better trade-off score
                    current_score = speed / r.compression_ratio
                    existing_score = (algorithm_metrics[r.algorithm]['speed'] /
                                   algorithm_metrics[r.algorithm]['ratio'])
                    if current_score > existing_score:
                        algorithm_metrics[r.algorithm] = {
                            'ratio': r.compression_ratio,
                            'speed': speed
                        }

            # Extract data for plotting
            ratios = [metrics['ratio'] for metrics in algorithm_metrics.values()]
            speeds = [metrics['speed'] for metrics in algorithm_metrics.values()]
            algorithms = list(algorithm_metrics.keys())

            # Create scatter plot
            ax.scatter(ratios, speeds, alpha=0.6)

            texts = [ax.text(ratios[i], speeds[i], alg, fontsize=9) for i, alg in enumerate(algorithms)]
            adjust_text(texts,
                       arrowprops=dict(arrowstyle='->', color='gray', linewidth=0.5),
                       ensure_inside_axes=False,
                       ax=ax)

            # Customize plot
            ax.set_xlabel('Compression Ratio (higher is better)')
            ax.set_ylabel(f'{speed_type.capitalize()} Speed MB/sec (higher is better)')
            ax.set_title(f'{dataset_name}')
            ax.grid(True, alpha=0.3)

            # Add theoretical limit line if applicable
            theoretical_ratio = get_theoretical_ratio(dataset_name)
            if theoretical_ratio is not None:
                ax.axvline(x=theoretical_ratio, color='red', linestyle='--')
                ax.legend()

        plt.tight_layout()
        # Save figure
        plt.savefig(os.path.join(ensure_image_dir(), f'iid_tradeoff_{speed_type}_{"-".join(datasets)}.png'))
        plt.close()

def figure_speeds_comparison(results: BenchmarkResults, datasets: list[str]):
    """Create horizontal bar plots comparing encode/decode speeds for the given datasets."""
    # Create figure with enough subplots for all datasets
    n_plots = len(datasets)
    fig, axes = plt.subplots(1, n_plots, figsize=(10*n_plots, 6))

    # Handle single dataset case
    if n_plots == 1:
        axes = [axes]

    for dataset_name, ax in zip(datasets, axes):
        # Filter for specified dataset
        dataset_results = [r for r in results.results if r.dataset == dataset_name]

        # Get unique algorithms and their speeds
        algorithm_speeds = {}
        for r in dataset_results:
            if r.algorithm not in algorithm_speeds:
                algorithm_speeds[r.algorithm] = {
                    'encode': r.encode_mb_per_sec,
                    'decode': r.decode_mb_per_sec
                }
            else:
                # Update if this instance has better average speed
                current_avg = (r.encode_mb_per_sec + r.decode_mb_per_sec) / 2
                existing_avg = (algorithm_speeds[r.algorithm]['encode'] +
                              algorithm_speeds[r.algorithm]['decode']) / 2
                if current_avg > existing_avg:
                    algorithm_speeds[r.algorithm] = {
                        'encode': r.encode_mb_per_sec,
                        'decode': r.decode_mb_per_sec
                    }

        # Sort alphabetically by algorithm name
        sorted_items = sorted(algorithm_speeds.items(), key=lambda x: x[0])
        algorithms = [item[0] for item in sorted_items]
        encode_speeds = [item[1]['encode'] for item in sorted_items]
        decode_speeds = [item[1]['decode'] for item in sorted_items]

        # Set up positions for the bars
        y_pos = np.arange(len(algorithms))
        bar_height = 0.35

        # Create bars
        ax.barh(y_pos - bar_height/2, encode_speeds, bar_height,
                label='Encode Speed', color='royalblue', alpha=0.8)
        ax.barh(y_pos + bar_height/2, decode_speeds, bar_height,
                label='Decode Speed', color='seagreen', alpha=0.8)

        # Customize plot
        ax.set_yticks(y_pos)
        ax.set_yticklabels(algorithms)
        ax.set_xlabel('Speed (MB/sec)')
        ax.set_title(f'Encode/Decode Speeds for {dataset_name}')
        ax.legend()

        # Add grid for readability
        ax.grid(True, axis='x', alpha=0.3)

    plt.tight_layout()
    # Save figure
    plt.savefig(os.path.join(ensure_image_dir(), f'iid_speeds_comparison_{"-".join(datasets)}.png'))
    plt.close()

def figure_compression_ratios_comparison(results: BenchmarkResults, datasets: list[str]):
    """Create horizontal bar plots comparing compression ratios for given datasets."""
    # Create figure with enough subplots for all datasets
    n_plots = len(datasets)
    fig, axes = plt.subplots(1, n_plots, figsize=(10*n_plots, 6))

    # Handle single dataset case
    if n_plots == 1:
        axes = [axes]

    # Plot comparisons for each dataset
    for dataset, ax in zip(datasets, axes):
        plot_single_compression_ratio_comparison(results, dataset, ax)

    plt.tight_layout()
    # Save figure
    plt.savefig(os.path.join(ensure_image_dir(), f'iid_compression_ratios_{"-".join(datasets)}.png'))
    plt.close()

def main():
    try:
        # Download the results
        raw_data = download_results()

        # Load into pydantic models
        results = load_results_into_models(raw_data)

        print(f"Successfully loaded {len(results.results)} benchmark results")

        datasets_list = [
            ['bernoulli-0.1', 'bernoulli-0.5'],
            ['gaussian-q2', 'gaussian-q5']
        ]

        for datasets in datasets_list:
            print("\nGenerating compression ratios plot...")
            figure_compression_ratios_comparison(results, datasets)

            print("\nGenerating speed comparison plot...")
            figure_speeds_comparison(results, datasets)

            print("\nGenerating tradeoff plots...")
            figure_tradeoff_comparison(results, datasets)

        return results

    except requests.RequestException as e:
        print(f"Error downloading results: {e}")
        return None
    except json.JSONDecodeError as e:
        print(f"Error parsing JSON: {e}")
        return None
    except Exception as e:
        print(f"Error loading data into models: {e}")
        return None

if __name__ == '__main__':
    main()