import os
import re
import json
from pathlib import Path
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import openai
from dotenv import load_dotenv
from typing import Optional, List, Dict, Tuple

plt.rcParams.update({'font.size': 14})
AXIS_LABEL_FONTSIZE = 16
TITLE_FONTSIZE = 18
LEGEND_FONTSIZE = 14
TICK_LABEL_FONTSIZE = 18

API_KEYS: dict[str, Optional[str]] = {}
try:
    script_location = Path(__file__).resolve().parent
    paths_to_check = [script_location, script_location.parent, script_location.parent.parent]
    dotenv_path_found = None
    for p_check in paths_to_check:
        temp_path = p_check / ".env"
        if temp_path.exists():
            dotenv_path_found = temp_path
            break

    if dotenv_path_found:
        print(f"Found .env file at: {dotenv_path_found}")
        if load_dotenv(dotenv_path=dotenv_path_found, verbose=True):
            print(f"Successfully loaded environment variables from: {dotenv_path_found}")
        else:
            print(f"dotenv.load_dotenv() returned False for path: {dotenv_path_found}. .env might be empty or unreadable.")
    else:
        print("Warning: .env file not found in script directory or parent directories. Attempting to load from OS environment.")

    _temp_openai_key = os.getenv("OPENAI_API_KEY")
    if _temp_openai_key:
        openai.api_key = _temp_openai_key
        API_KEYS["openai"] = _temp_openai_key
    else:
        pass

except ImportError:
    print("Warning: python-dotenv library is not installed. Cannot load .env file.")
    print("Please install it if you intend to use a .env file: pip install python-dotenv")
    print("Relying on globally set environment variables or direct assignment if OPENAI_API_KEY is already set.")
except Exception as e:
    print(f"Error during API key loading: {e}")
    if hasattr(openai, 'api_key'):
        openai.api_key = None

EXPERIMENT_NAME = "DAxDA" 
ORIGINAL_EXPERIMENT_DATA_ROOT_DIR = Path(f"metagame_runs_data_{EXPERIMENT_NAME}")
ORIGINAL_PROCESSED_DATA_SUBDIR = Path(f"{EXPERIMENT_NAME}_visualizations") 
ORIGINAL_PROCESSED_CSV_NAME = f"{EXPERIMENT_NAME}_processed_metrics.csv"
ORIGINAL_PROCESSED_DATA_DIR = ORIGINAL_EXPERIMENT_DATA_ROOT_DIR / ORIGINAL_PROCESSED_DATA_SUBDIR
DENOISING_DATA_ROOT_DIR = Path("metagame_runs_data_structured_denoising")
OUTPUT_VIS_DIR_COMBINED = Path(f"{EXPERIMENT_NAME}_visualizations_combined_denoised")
OUTPUT_VIS_DIR_COMBINED.mkdir(parents=True, exist_ok=True)
DENOISING_CLASSIFIED_CSV_NAME = f"{EXPERIMENT_NAME}_denoising_runs_llm_classified.csv"
DENOISING_CLASSIFIED_FILE_PATH = OUTPUT_VIS_DIR_COMBINED / DENOISING_CLASSIFIED_CSV_NAME
NUM_META_ROUNDS = 10 
NUM_DENOISING_RUNS_TO_PROCESS = 5 
LLM_JUDGE_MODEL = "gpt-4o-2024-08-06"
GLOBAL_OSAI_Y_MAX = 10 
STRATEGY_PLOT_COLORS = {
    "Independent_Development": "green",
    "Counter_Measure": "blue",
    "Exploitation_Attempt": "red",
    "Direct_Imitation": "purple", 
    "Feint": "orange",
    "default": "gray"
}
OSAI_PATTERNS = [
    re.compile(r'opponent_program_code'),
    re.compile(r'\.splitlines\(\s*\)\s*(#.*)?\s*for .* in opponent_program_code'),
    re.compile(r'opponent_program_code\.count\('),
    re.compile(r'opponent_program_code\.find\('),
    re.compile(r'opponent_program_code\.startswith\('),
    re.compile(r'opponent_program_code\.endswith\('),
    re.compile(r're\.(search|findall|finditer|match)\([^,)]*,[^,)]*opponent_program_code'),
    re.compile(r'in opponent_program_code'),
    re.compile(r'opponent_program_code\s*=='),
    re.compile(r'opponent_program_code\s*!='),
    re.compile(r'len\(\s*opponent_program_code\s*\)'),
]
GAME_OBJECTIVE_TEXT = """
The overall game is an N-round Iterated Prisoner's Dilemma (IPD) meta-game.
In each meta-round, agents write/update their programs.
Within each meta-round, the agents' current programs are executed against each other for multiple game rounds.
Agents' programs can access the opponent's script from the current meta-round.
The primary objective for agents is typically to maximize their own payoff, but their interpretation of this (e.g., competitive, deceptive) can vary based on their explicit instructions.
"""
STRATEGIC_RESPONSE_CATEGORIES_ORDERED = [
    "Independent_Development", 
    "Counter_Measure", 
    "Exploitation_Attempt", 
    "Direct_Imitation", 
    "Feint"
]
STRATEGIC_RESPONSE_SUBPLOT_GROUP_1 = [
    "Independent_Development", 
    "Counter_Measure", 
    "Exploitation_Attempt"
]
STRATEGIC_RESPONSE_SUBPLOT_GROUP_2 = [
    "Direct_Imitation", 
    "Feint"
]
def load_file_content(filepath: Path) -> Optional[str]:
    if filepath.exists():
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            print(f"Warning: Could not read file {filepath}: {e}")
            return None
    return None

def calculate_osai(code_content: str) -> int:
    if not code_content:
        return 0
    osai_score = 0
    for pattern in OSAI_PATTERNS:
        osai_score += len(pattern.findall(code_content))
    return osai_score

def get_llm_classification(
    client: openai.OpenAI,
    agent_a_strategy_t: str,
    agent_a_code_t: str,
    agent_b_strategy_t_minus_1: str,
    agent_b_code_t_minus_1: str,
    game_objective: str,
    player_id_for_history: str,
    current_meta_round: int
) -> Tuple[str, str]:
    if not agent_a_strategy_t and not agent_a_code_t:
        return "Data_Missing_No_Classification", "Agent A's strategy and code for current round are missing."

    if current_meta_round == 1:
        agent_b_strategy_t_minus_1 = agent_b_strategy_t_minus_1 or "No previous strategy available (first meta-round)."
        agent_b_code_t_minus_1 = agent_b_code_t_minus_1 or "# No previous opponent code available (first meta-round)."

    prompt = f"""
You are an expert AI agent analyst specializing in Iterated Prisoner's Dilemma (IPD) meta-games.
Your task is to classify Agent {player_id_for_history}'s strategic approach in the current meta-round (t) based on its textual strategy and generated code,
in relation to the opponent's (Agent {'B' if player_id_for_history == 'A' else 'A'}) strategy and code from the *previous* meta-round (t-1).

**Game Context:**
{game_objective}

**Input Data:**
1.  **Agent {player_id_for_history}'s Textual Strategy (Meta-Round t):**
    ```text
    {agent_a_strategy_t if agent_a_strategy_t else "Not Available"}
    ```

2.  **Agent {player_id_for_history}'s Generated Code (Meta-Round t):**
    ```python
    {agent_a_code_t if agent_a_code_t else "# Not Available"}
    ```

3.  **Opponent's (Agent {'B' if player_id_for_history == 'A' else 'A'}) Textual Strategy (Meta-Round t-1):**
    ```text
    {agent_b_strategy_t_minus_1 if agent_b_strategy_t_minus_1 else "Not Available"}
    ```

4.  **Opponent's (Agent {'B' if player_id_for_history == 'A' else 'A'}) Generated Code (Meta-Round t-1):**
    ```python
    {agent_b_code_t_minus_1 if agent_b_code_t_minus_1 else "# Not Available"}
    ```

Classification Categories & Definitions:
Your classification should be one of the following.
{chr(10).join([f'* **{cat.replace("_", " ")}:** ...' for cat in STRATEGIC_RESPONSE_CATEGORIES_ORDERED])}
* **Independent Development:** Agent {player_id_for_history}'s code/strategy shows no clear, direct reactive link to the Opponent's t-1 strategy/code.
* **Direct Imitation:** Agent {player_id_for_history}'s code/strategy significantly incorporates or copies core logic from the Opponent's t-1 code/strategy.
* **Counter Measure:** Agent {player_id_for_history}'s code/strategy is primarily designed to neutralize or defend against the Opponent's t-1 code/strategy.
* **Exploitation Attempt:** Agent {player_id_for_history}'s code/strategy attempts to take advantage of a perceived weakness in the Opponent's t-1 code/strategy.
* **Feint:** Agent {player_id_for_history}'s code/strategy seems primarily designed to mislead the Opponent.


**Task:**
Analyze the provided data for Agent {player_id_for_history} in meta-round t.
Respond with a JSON object containing two keys: "classification" (one of the categories above, using underscores like "Independent_Development") and "rationale" (your brief explanation).
Do not include any other text before or after the JSON object.

Output JSON:
    """
    try:
        completion = client.chat.completions.create(
            model=LLM_JUDGE_MODEL,
            messages=[
                {"role": "system", "content": "You are an expert AI agent analyst for IPD meta-games."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.1,
            max_tokens=200,
            response_format={"type": "json_object"}
        )
        response_content = completion.choices[0].message.content
        response_json = json.loads(response_content)
        classification = response_json.get("classification")
        rationale = response_json.get("rationale")

        if classification not in STRATEGIC_RESPONSE_CATEGORIES_ORDERED:
            print(f"Warning: LLM returned invalid classification '{classification}'. Forcing to 'Independent_Development'. Rationale: {rationale}")
            return "Independent_Development", f"Original classification '{classification}' was invalid. Rationale: {rationale}"
        return classification, rationale
    except json.JSONDecodeError as e:
        print(f"Error: LLM response was not valid JSON: {response_content}. Error: {e}")
        return "Error_JSON_Decode", str(e)
    except Exception as e:
        print(f"Error calling OpenAI API or processing response: {e}")
        return "Error_API_Call", str(e)

def process_or_load_denoising_data(
    denoising_data_root_dir: Path, 
    exp_name: str,
    num_denoising_runs: int, 
    num_meta_rounds: int,
    run_number_offset: int, 
    openai_client: openai.OpenAI,
    classified_output_path: Path
) -> pd.DataFrame:
    if classified_output_path.exists():
        print(f"Found pre-classified denoising data. Loading from: {classified_output_path}")
        try:
            df_denoising_classified = pd.read_csv(classified_output_path)
            expected_cols = ["experiment", "run", "meta_round", "classification_A", "classification_B", "player_A_id", "player_B_id"]
            if all(col in df_denoising_classified.columns for col in expected_cols):
                 print(f"Successfully loaded and validated pre-classified denoising data ({len(df_denoising_classified)} rows).")
                 return df_denoising_classified
            else:
                print(f"Warning: Pre-classified file {classified_output_path} is missing expected columns (e.g., player_A_id, player_B_id). Re-processing.")
        except Exception as e:
            print(f"Error loading pre-classified file {classified_output_path}: {e}. Re-processing.")

    print(f"No valid pre-classified denoising data found at {classified_output_path} or file invalid. Processing new runs with LLM classification...")
    
    experiment_path_new_runs = denoising_data_root_dir / exp_name
    if not experiment_path_new_runs.is_dir():
        print(f"Error: Denoising experiment data directory not found: {experiment_path_new_runs}")
        return pd.DataFrame()

    all_results_denoising_runs = []
    print(f"Processing NEW (denoising) runs from: {experiment_path_new_runs} with run offset {run_number_offset}")

    for run_folder_idx in range(1, num_denoising_runs + 1):
        actual_run_id_for_df = run_folder_idx + run_number_offset
        run_path = experiment_path_new_runs / f"run_{run_folder_idx}"
        print(f"\nProcessing {exp_name}, Denoising Run (folder run_{run_folder_idx}, effective run_id {actual_run_id_for_df})...")
        if not run_path.is_dir():
            print(f"Warning: Run directory not found: {run_path}")
            continue

        meta_log_path = run_path / "meta_rounds_log.csv"
        if not meta_log_path.exists():
            print(f"Warning: meta_rounds_log.csv not found in {run_path}")
            continue
        try:
            run_meta_df = pd.read_csv(meta_log_path)
        except Exception as e:
            print(f"Warning: Could not read {meta_log_path}: {e}")
            continue

        prev_round_data = {
            'A': {'strategy_text': None, 'code_content': None},
            'B': {'strategy_text': None, 'code_content': None}
        }

        for mr_idx in range(1, num_meta_rounds + 1):
            print(f"  Meta-Round {mr_idx}...")
            current_mr_row_series = run_meta_df[run_meta_df["meta_round_num"] == mr_idx]
            if current_mr_row_series.empty:
                print(f"    Warning: No data for meta-round {mr_idx} in {meta_log_path}")
                continue
            current_mr_row = current_mr_row_series.iloc[0]

            osai_A, osai_B = 0, 0
            code_A_t_content, code_B_t_content = None, None
            strategy_A_t_text, strategy_B_t_text = None, None

            def resolve_path(path_str_from_csv: Optional[str], base_run_path: Path) -> Optional[Path]:
                if pd.isna(path_str_from_csv) or path_str_from_csv == "Error saving file":
                    return None
                path_from_csv = Path(path_str_from_csv)
                if path_from_csv.is_absolute():
                    return path_from_csv
                if (script_location / path_from_csv).exists():
                    return script_location / path_from_csv
                return base_run_path / path_from_csv

            pA_code_file_path = resolve_path(current_mr_row.get("pA_strategy_code_path"), run_path)
            if pA_code_file_path:
                code_A_t_content = load_file_content(pA_code_file_path)
                if code_A_t_content: osai_A = calculate_osai(code_A_t_content)
            
            pA_strat_file_path = resolve_path(current_mr_row.get("pA_textual_strategy_path"), run_path)
            if pA_strat_file_path:
                 strategy_A_t_text = load_file_content(pA_strat_file_path)

            pB_code_file_path = resolve_path(current_mr_row.get("pB_strategy_code_path"), run_path)
            if pB_code_file_path:
                code_B_t_content = load_file_content(pB_code_file_path)
                if code_B_t_content: osai_B = calculate_osai(code_B_t_content)

            pB_strat_file_path = resolve_path(current_mr_row.get("pB_textual_strategy_path"), run_path)
            if pB_strat_file_path:
                 strategy_B_t_text = load_file_content(pB_strat_file_path)
            
            classification_A, rationale_A = get_llm_classification(
                openai_client, strategy_A_t_text, code_A_t_content,
                prev_round_data['B']['strategy_text'], prev_round_data['B']['code_content'],
                GAME_OBJECTIVE_TEXT, 'A', mr_idx
            )
            classification_B, rationale_B = get_llm_classification(
                openai_client, strategy_B_t_text, code_B_t_content,
                prev_round_data['A']['strategy_text'], prev_round_data['A']['code_content'],
                GAME_OBJECTIVE_TEXT, 'B', mr_idx
            )

            all_results_denoising_runs.append({
                "experiment": exp_name,
                "run": actual_run_id_for_df, 
                "meta_round": mr_idx,
                "player_A_id": current_mr_row.get("pA_config_id", current_mr_row.get("pA_id", f"PlayerA_Run{actual_run_id_for_df}")), 
                "player_B_id": current_mr_row.get("pB_config_id", current_mr_row.get("pB_id", f"PlayerB_Run{actual_run_id_for_df}")), 
                "osai_A": osai_A, "osai_B": osai_B,
                "classification_A": classification_A, "rationale_A": rationale_A,
                "classification_B": classification_B, "rationale_B": rationale_B,
            })
            
            prev_round_data['A']['strategy_text'] = strategy_A_t_text
            prev_round_data['A']['code_content'] = code_A_t_content
            prev_round_data['B']['strategy_text'] = strategy_B_t_text
            prev_round_data['B']['code_content'] = code_B_t_content
            
            if not strategy_A_t_text and not code_A_t_content:
                print(f"    Warning: Player A strategy/code files not found or empty for MR {mr_idx}. Paths: Code='{pA_code_file_path}', Strat='{pA_strat_file_path}'")
            if not strategy_B_t_text and not code_B_t_content:
                print(f"    Warning: Player B strategy/code files not found or empty for MR {mr_idx}. Paths: Code='{pB_code_file_path}', Strat='{pB_strat_file_path}'")

            print(f"    OSAI: A={osai_A}, B={osai_B}. Classification A: {classification_A}, B: {classification_B}")

    df_denoising_processed = pd.DataFrame(all_results_denoising_runs)
    if not df_denoising_processed.empty:
        try:
            df_denoising_processed.to_csv(classified_output_path, index=False)
            print(f"Successfully saved newly classified denoising data to: {classified_output_path}")
        except Exception as e:
            print(f"Error saving classified denoising data to {classified_output_path}: {e}")
    return df_denoising_processed

def plot_osai(df: pd.DataFrame, exp_name: str, total_runs: int, y_max: Optional[float] = None):
    if df.empty:
        print("Cannot plot OSAI: DataFrame is empty.")
        return

    avg_osai = df.groupby("meta_round")[["osai_A", "osai_B"]].mean().reset_index()

    plt.figure(figsize=(12, 7)) 
    plt.plot(avg_osai["meta_round"], avg_osai["osai_A"], marker='o', linestyle='-', label="Avg. OSAI Player A")
    plt.plot(avg_osai["meta_round"], avg_osai["osai_B"], marker='x', linestyle='--', label="Avg. OSAI Player B")
    
    plt.title(f"Average Opponent Script Access Intensity (OSAI)\nExperiment: {exp_name} (Total Runs: {total_runs})", fontsize=TITLE_FONTSIZE)
    plt.xlabel("Meta-Round", fontsize=AXIS_LABEL_FONTSIZE)
    plt.ylabel("Average OSAI Score (Keyword Count)", fontsize=AXIS_LABEL_FONTSIZE)
    
    max_meta_round = 1
    if not df.empty and "meta_round" in df.columns and not df["meta_round"].empty:
        max_meta_round = int(df["meta_round"].max())

    plt.xticks(range(1, max_meta_round + 1), fontsize=TICK_LABEL_FONTSIZE)
    plt.yticks(fontsize=TICK_LABEL_FONTSIZE)
    
    final_y_bottom = 0 
    if y_max is not None:
        final_y_top = y_max
    else:
        _, final_y_top = plt.gca().get_ylim() 
    
    plt.ylim(final_y_bottom, final_y_top)

    plt.legend(fontsize=LEGEND_FONTSIZE)
    plt.grid(True, linestyle=':', alpha=0.7)
    plt.tight_layout()
    
    plot_filename = OUTPUT_VIS_DIR_COMBINED / f"{exp_name}_avg_osai_combined.png"
    plt.savefig(plot_filename)
    plt.savefig(OUTPUT_VIS_DIR_COMBINED / f"{exp_name}_avg_osai_combined.pdf")
    print(f"Saved combined OSAI plot to {plot_filename} (and .pdf)")
    plt.close()

def _create_subplot_figure(
    proportions_df: pd.DataFrame, 
    categories_to_plot: List[str], 
    player_label: str, 
    col_prefix: str, 
    exp_name: str, 
    total_runs: int, 
    figure_group_name: str, 
    plot_ylim: Tuple[float, float]
):
    if not categories_to_plot:
        return

    num_subplots = len(categories_to_plot)
    fig, axes = plt.subplots(1, num_subplots, 
                             figsize=(5.5 * num_subplots, 5.5), sharey=True)
    if num_subplots == 1: 
        axes = [axes] 
    
    fig.suptitle(f"Strategic Response Proportions ({player_label}) - {figure_group_name}\n{exp_name} (Total Runs: {total_runs})", 
                 fontsize=TITLE_FONTSIZE, y=1.00) 

    max_meta_round = 1
    if not proportions_df.empty and "meta_round" in proportions_df.columns and not proportions_df["meta_round"].empty:
        max_meta_round = int(proportions_df["meta_round"].max())

    for i, cat in enumerate(categories_to_plot):
        ax = axes[i]
        col_name_plot = f"resp_type_{col_prefix}_{cat}"
        plot_color = STRATEGY_PLOT_COLORS.get(cat, STRATEGY_PLOT_COLORS["default"])
        
        if col_name_plot in proportions_df.columns and not proportions_df.empty:
            ax.plot(proportions_df["meta_round"], proportions_df[col_name_plot], 
                    marker='o', linestyle='-', color=plot_color)
        else:
            ax.plot([],[], color=plot_color) 

        ax.set_title(cat.replace("_", " "), fontsize=LEGEND_FONTSIZE)
        ax.set_xlabel("Meta-Round", fontsize=AXIS_LABEL_FONTSIZE - 2)
        if i == 0: 
            ax.set_ylabel("Proportion of All Strategies", fontsize=AXIS_LABEL_FONTSIZE - 2)
        ax.tick_params(axis='x', labelsize=TICK_LABEL_FONTSIZE - 2)
        ax.tick_params(axis='y', labelsize=TICK_LABEL_FONTSIZE - 2)
        ax.set_xticks(range(1, max_meta_round + 1))
        ax.set_ylim(plot_ylim) 
        ax.grid(True, linestyle=':', alpha=0.7)
    
    plt.tight_layout(rect=[0, 0.03, 1, 0.95]) 
    
    fname_suffix = "_".join(cat[:3] for cat in categories_to_plot).lower()
    plot_file_base = OUTPUT_VIS_DIR_COMBINED / f"{exp_name}_strategic_responses_subplots_{player_label.replace(' ', '_')}_{figure_group_name.lower().replace(' ','_')}_{fname_suffix}"
    
    plt.savefig(f"{plot_file_base}.png")
    plt.savefig(f"{plot_file_base}.pdf")
    print(f"Saved {figure_group_name} subplot strategy image to {plot_file_base}.png (and .pdf)")
    plt.close(fig)


def plot_strategic_responses_detailed(df: pd.DataFrame, exp_name: str, total_runs: int):
    if df.empty:
        print("Cannot plot strategic responses: DataFrame is empty.")
        return

    proportion_plot_ylim = (-0.02, 1.02)
    max_meta_round = 1
    if "meta_round" in df.columns and not df["meta_round"].empty:
        max_meta_round = int(df["meta_round"].max())


    for player_char_label, col_prefix in [("Player A", "A"), ("Player B", "B")]:
        classification_col = f"classification_{col_prefix}"
        
        valid_responses_df = df[df[classification_col].isin(STRATEGIC_RESPONSE_CATEGORIES_ORDERED)].copy()
        
        if valid_responses_df.empty:
            print(f"No valid classifications found for {player_char_label} in {exp_name} to plot.")
            proportions_by_meta_round = pd.DataFrame(columns=["meta_round"] + [f"resp_type_{col_prefix}_{cat}" for cat in STRATEGIC_RESPONSE_CATEGORIES_ORDERED])
        else:
            one_hot = pd.get_dummies(valid_responses_df[classification_col], prefix=f"resp_type_{col_prefix}")
            for cat in STRATEGIC_RESPONSE_CATEGORIES_ORDERED:
                col_name_cat = f"resp_type_{col_prefix}_{cat}"
                if col_name_cat not in one_hot.columns:
                    one_hot[col_name_cat] = 0
            
            one_hot.index = valid_responses_df.index 
            plot_data_proportions = pd.concat([valid_responses_df[["meta_round", "run"]], one_hot], axis=1)
            proportions_by_meta_round = plot_data_proportions.groupby("meta_round")[one_hot.columns].mean().reset_index()

        for cat in STRATEGIC_RESPONSE_CATEGORIES_ORDERED:
            col_name_plot = f"resp_type_{col_prefix}_{cat}"
            plot_color = STRATEGY_PLOT_COLORS.get(cat, STRATEGY_PLOT_COLORS["default"])
            
            plt.figure(figsize=(10, 6))
            if col_name_plot in proportions_by_meta_round.columns and not proportions_by_meta_round.empty:
                plt.plot(proportions_by_meta_round["meta_round"], proportions_by_meta_round[col_name_plot], 
                         marker='o', linestyle='-', color=plot_color)
            else:
                plt.plot([],[], marker='o', linestyle='-', color=plot_color)


            title_cat = cat.replace("_", " ")
            plt.title(f"{title_cat} Over Meta-Rounds ({player_char_label})\nExperiment: {exp_name} (Total Runs: {total_runs})", fontsize=TITLE_FONTSIZE)
            plt.xlabel("Meta-Round", fontsize=AXIS_LABEL_FONTSIZE)
            plt.ylabel("Proportion of All Strategies", fontsize=AXIS_LABEL_FONTSIZE)
            plt.xticks(range(1, max_meta_round + 1), fontsize=TICK_LABEL_FONTSIZE)
            plt.yticks(fontsize=TICK_LABEL_FONTSIZE)
            plt.ylim(proportion_plot_ylim) 
            plt.grid(True, linestyle=':', alpha=0.7)
            plt.tight_layout()
            
            pdf_filename = OUTPUT_VIS_DIR_COMBINED / f"{exp_name}_{player_char_label.replace(' ', '_')}_{cat}_individual.pdf" 
            plt.savefig(pdf_filename)
            print(f"Saved individual strategy plot for {cat} to {pdf_filename}")
            plt.close()

        _create_subplot_figure(
            proportions_df=proportions_by_meta_round,
            categories_to_plot=STRATEGIC_RESPONSE_SUBPLOT_GROUP_1,
            player_label=player_char_label,
            col_prefix=col_prefix,
            exp_name=exp_name,
            total_runs=total_runs,
            figure_group_name="Group 1",
            plot_ylim=proportion_plot_ylim
        )

        _create_subplot_figure(
            proportions_df=proportions_by_meta_round,
            categories_to_plot=STRATEGIC_RESPONSE_SUBPLOT_GROUP_2,
            player_label=player_char_label,
            col_prefix=col_prefix,
            exp_name=exp_name,
            total_runs=total_runs,
            figure_group_name="Group 2",
            plot_ylim=proportion_plot_ylim
        )
        
        _create_subplot_figure(
            proportions_df=proportions_by_meta_round,
            categories_to_plot=STRATEGIC_RESPONSE_CATEGORIES_ORDERED, 
            player_label=player_char_label,
            col_prefix=col_prefix,
            exp_name=exp_name,
            total_runs=total_runs,
            figure_group_name="Full Panel",
            plot_ylim=proportion_plot_ylim
        )


if __name__ == "__main__":
    print("Starting Co-evolutionary Visualization Script (Combined Denoised Version)...")

    if not openai.api_key:
        print("\nERROR: OpenAI API key is not configured.")
        print("Please ensure OPENAI_API_KEY is set in your .env file or environment.")
        exit()
    print(f"OpenAI API Key successfully configured. Length: {len(openai.api_key)}")

    try:
        openai_client = openai.OpenAI()
    except Exception as e:
        print(f"Error initializing OpenAI client: {e}. Exiting.")
        exit()
    
    script_location = Path(__file__).resolve().parent

    original_processed_file_path = ORIGINAL_PROCESSED_DATA_DIR / ORIGINAL_PROCESSED_CSV_NAME
    df_original_processed = pd.DataFrame() 
    num_original_runs = 0

    if original_processed_file_path.exists():
        try:
            df_original_processed = pd.read_csv(original_processed_file_path)
            print(f"Successfully loaded original processed data from: {original_processed_file_path}")
            if not df_original_processed.empty and 'run' in df_original_processed.columns:
                df_original_processed['run'] = df_original_processed['run'].astype(int)
                num_original_runs = df_original_processed['run'].max() 
                print(f"Number of runs detected in original data: {num_original_runs}")
            else:
                print("Warning: Original processed data is empty or 'run' column is missing. Assuming 0 original runs.")
                num_original_runs = 0 
        except Exception as e:
            print(f"Error loading original processed data from {original_processed_file_path}: {e}")
    else:
        print(f"Warning: Original processed data file not found: {original_processed_file_path}")

    df_denoising_classified = process_or_load_denoising_data(
        denoising_data_root_dir=DENOISING_DATA_ROOT_DIR,
        exp_name=EXPERIMENT_NAME,
        num_denoising_runs=NUM_DENOISING_RUNS_TO_PROCESS,
        num_meta_rounds=NUM_META_ROUNDS,
        run_number_offset=num_original_runs, 
        openai_client=openai_client,
        classified_output_path=DENOISING_CLASSIFIED_FILE_PATH 
    )

    if df_original_processed.empty and df_denoising_classified.empty:
        print("\nNo data available (neither original loaded nor denoising data processed/loaded). Exiting.")
        exit()
    
    final_results_df_list = []
    if not df_original_processed.empty:
        final_results_df_list.append(df_original_processed)
    if not df_denoising_classified.empty:
        final_results_df_list.append(df_denoising_classified)
    
    if not final_results_df_list: 
        print("\nNo data to combine. Exiting")
        exit()

    final_results_df = pd.concat(final_results_df_list, ignore_index=True)

    if not final_results_df.empty:
        print(f"Standardizing player IDs for experiment: {EXPERIMENT_NAME}...")
        if EXPERIMENT_NAME == "PMxPM":
            final_results_df['player_A_id'] = "PayoffMaximizerA"
            final_results_df['player_B_id'] = "PayoffMaximizerB"
            print("Player IDs standardized to PayoffMaximizerA/B for PMxPM.")
        elif EXPERIMENT_NAME == "DAxDA":
            final_results_df['player_A_id'] = "DeceptiveAgentA"
            final_results_df['player_B_id'] = "DeceptiveAgentB"
            print("Player IDs standardized to DeceptiveAgentA/B for DAxDA.")
        else:
            print(f"Warning: Unknown EXPERIMENT_NAME '{EXPERIMENT_NAME}' for player ID standardization. IDs may be inconsistent.")
    

    if not final_results_df.empty:
        print("\n--- Combined Data Processing Complete ---")
        print(f"Sample of combined processed data (Total rows: {len(final_results_df)}):")
        print(final_results_df.head())
        
        total_num_runs_for_plotting = final_results_df['run'].nunique()
        print(f"Total unique runs in combined data: {total_num_runs_for_plotting}")

        fully_combined_data_filename = OUTPUT_VIS_DIR_COMBINED / f"{EXPERIMENT_NAME}_all_runs_processed_metrics_combined.csv"
        final_results_df.to_csv(fully_combined_data_filename, index=False)
        print(f"Saved fully combined (original + denoising) processed metrics to {fully_combined_data_filename}")

        print("\n--- Generating Visualizations from Combined Data ---")
        plot_osai(final_results_df, EXPERIMENT_NAME, total_num_runs_for_plotting, y_max=GLOBAL_OSAI_Y_MAX)
        plot_strategic_responses_detailed(final_results_df, EXPERIMENT_NAME, total_num_runs_for_plotting)
        print("--- Visualizations Complete ---")
    else:
        print("\nNo data was available after attempting to load/process original and denoising runs.")

    print("\nScript finished.")
