2024-10-15 10:42:23,275 [INFO] Loading data from local_battle_data.json
2024-10-15 10:43:11,962 [INFO] Data loaded successfully with 1799991 records
2024-10-15 10:43:11,965 [INFO] Creating advanced scatter plot
2024-10-15 10:43:14,296 [ERROR] An error occurred while creating the scatter plot: 'rank'
2024-10-15 10:43:14,296 [ERROR] Process terminated due to an error: 'rank
import os
import sys
import math
import json
import logging
from typing import Any

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from tqdm import tqdm
import requests
from json import JSONDecodeError
# Configure pandas display options
pd.options.display.float_format = "{:.2f}".format

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)

# Constants
DATA_URL: str = "https://storage.googleapis.com/arena_external_data/public/clean_battle_20240814_public.json"
LOCAL_FILE_NAME: str = "local_battle_data.json"
OUTPUT_FOLDER: str = "output_plots"


def download_data(url: str, local_path: str) -> None:
    """
    Downloads JSON data from the specified URL and saves it locally.

    Args:
        url (str): The URL to download data from.
        local_path (str): The local file path to save the downloaded data.

    Raises:
        requests.HTTPError: If the HTTP request to download data fails.
        Exception: If any other exception occurs during the download.
    """
    try:
        logging.info(f"Starting download from {url}")
        response = requests.get(url, stream=True)
        response.raise_for_status()
        total_size = int(response.headers.get('content-length', 0))
        with open(local_path, 'wb') as file, tqdm(
            desc=local_path,
            total=total_size,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for data in response.iter_content(chunk_size=1024):
                size = file.write(data)
                bar.update(size)
        logging.info(f"Data downloaded successfully and saved to {local_path}")
    except requests.HTTPError as http_err:
        logging.error(f"HTTP error occurred while downloading data: {http_err}")
        raise
    except Exception as err:
        logging.error(f"An error occurred while downloading data: {err}")
        raise



def load_data(local_path: str) -> pd.DataFrame:
    """
    Loads JSON data from a local file into a pandas DataFrame and sorts it by timestamp.

    Args:
        local_path (str): The local file path to load data from.

    Returns:
        pd.DataFrame: The loaded and sorted battles DataFrame.

    Raises:
        FileNotFoundError: If the local file does not exist.
        JSONDecodeError: If the file content is not valid JSON.
        Exception: If any other exception occurs during data loading.
    """
    try:
        logging.info(f"Loading data from {local_path}")
        with open(local_path, 'r',encoding="utf-8") as file:
            battles = pd.read_json(file).sort_values(by="tstamp", ascending=True)
        logging.info(f"Data loaded successfully with {len(battles)} records")
        return battles
    except FileNotFoundError as fnf_err:
        logging.error(f"File not found: {fnf_err}")
        raise
    except JSONDecodeError as json_err:
        logging.error(f"JSON decoding error: {json_err}")
        raise
    except Exception as err:
        logging.error(f"An error occurred while loading data: {err}")
        raise

def create_advanced_scatter_plot(battles: pd.DataFrame, output_folder: str) -> None:
    """
    Creates an advanced scatter plot with interactive controls and saves it as an HTML file.

    The scatter plot visualizes battles with:
    - X-axis: Model A
    - Y-axis: Model B
    - Marker size varies according to the battle rank
    - Interactive play and pause buttons to animate over time

    Args:
        battles (pd.DataFrame): The battles DataFrame containing battle records.
        output_folder (str): The directory where the plot will be saved.

    Raises:
        Exception: If any error occurs during plot creation or saving.
    """
    try:
        logging.info("Creating advanced scatter plot")

        # Ensure output directory exists
        os.makedirs(output_folder, exist_ok=True)

        # Example: Suppose we have a 'rank' column to determine scatter sizes
        # If not present, we'll create a dummy 'rank' based on 'turn' as an example
        if 'rank' not in battles.columns:
            battles['rank'] = battles['turn'].rank(method='dense').astype(int)

        # Aggregate data as needed, for example by model pairs and timestamp
        aggregated = (
            battles.groupby(['model_a', 'model_b', 'tstamp'])
            .size()
            .reset_index(name='battle_count')
        )

        # Normalize rank for marker sizes
        aggregated['size'] = aggregated['rank'].apply(lambda x: 10 + (x * 5))

        # Create scatter plot
        fig = px.scatter(
            aggregated,
            x='model_a',
            y='model_b',
            animation_frame='tstamp',
            size='size',
            hover_name='battle_count',
            title='Battle Outcomes Over Time',
            labels={
                'model_a': 'Model A',
                'model_b': 'Model B',
                'tstamp': 'Timestamp',
                'size': 'Rank'
            },
            size_max=60,
            color='rank',
            color_continuous_scale='Viridis'
        )

        # Update layout for better readability
        fig.update_layout(
            title={
                'text': "Advanced Battle Scatter Plot",
                'x': 0.5,
                'xanchor': 'center'
            },
            xaxis_title="Model A",
            yaxis_title="Model B",
            legend_title="Rank",
            height=700
        )

        # Add play and pause buttons
        fig.update_layout(
            updatemenus=[
                {
                    "buttons": [
                        {
                            "args": [None, {"frame": {"duration": 500, "redraw": True},
                                            "fromcurrent": True, "transition": {"duration": 300}}],
                            "label": "Play",
                            "method": "animate"
                        },
                        {
                            "args": [[None], {"frame": {"duration": 0, "redraw": False},
                                              "mode": "immediate",
                                              "transition": {"duration": 0}}],
                            "label": "Pause",
                            "method": "animate"
                        }
                    ],
                    "direction": "left",
                    "pad": {"r": 10, "t": 87},
                    "showactive": False,
                    "type": "buttons",
                    "x": 0.1,
                    "xanchor": "right",
                    "y": 0,
                    "yanchor": "top"
                }
            ]
        )

        # Save the plot to the output folder
        output_path = os.path.join(output_folder, "advanced_scatter_plot.html")
        fig.write_html(output_path)
        logging.info(f"Advanced scatter plot saved to {output_path}")

    except Exception as err:
        logging.error(f"An error occurred while creating the scatter plot: {err}")
        raise


def main() -> None:
    """
    Main function to orchestrate data download, loading, visualization, and saving.

    Raises:
        Exception: Propagates exceptions from called functions.
    """
    try:
        # Step 1: Download the latest data
        # download_data(DATA_URL, LOCAL_FILE_NAME)

        # Step 2: Load the data into a DataFrame
        battles_df = load_data(LOCAL_FILE_NAME)

        # Step 3: Create advanced scatter plot and save it
        create_advanced_scatter_plot(battles_df, OUTPUT_FOLDER)

        logging.info("All tasks completed successfully.")

    except Exception as err:
        logging.error(f"Process terminated due to an error: {err}")
        sys.exit(1)


if __name__ == "__main__":
    main()