import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from pathlib import Path
import tempfile
import logging

# This mapping helps translate the O*NET 2-digit major group codes
# into human-readable labels for the plot's y-axis.
OCCUPATION_MAJOR_CODES = {
    '11': 'Management',
    '13': 'Business & Financial',
    '15': 'Computer & Mathematical',
    '17': 'Architecture & Engineering',
    '19': 'Life, Physical, & Social Science',
    '21': 'Community & Social Service',
    '23': 'Legal',
    '25': 'Education, Training, & Library',
    '27': 'Arts, Design, & Media',
    '29': 'Healthcare Practitioners',
    '31': 'Healthcare Support',
    '33': 'Protective Service',
    '35': 'Food Preparation & Serving',
    '37': 'Building & Grounds Maintenance',
    '39': 'Personal Care & Service',
    '41': 'Sales & Related',
    '43': 'Office & Admin Support',
    '45': 'Farming, Fishing, & Forestry',
    '47': 'Construction & Extraction',
    '49': 'Installation, Maintenance, & Repair',
    '51': 'Production',
    '53': 'Transportation & Material Moving',
    '55': 'Military Specific',
}


def generate(processed_df: pd.DataFrame):
    """
    Generates a heatmap of the median estimate ratio by occupation and task length quartile.

    This corresponds to 'cell5' from the original analysis notebook. It shows
    how the ratio between upper and lower time estimates varies across
    different occupations and for tasks of different typical lengths (binned
    into quartiles).

    Args:
        processed_df (pd.DataFrame): The preprocessed data. Expected columns:
                                     'lb_estimate_in_minutes',
                                     'ub_estimate_in_minutes', 'onetsoc_major'.

    Returns:
        Path: The path to the generated temporary image file, or None on failure.
    """
    logging.info("Generating heatmap of estimate ratios by occupation and task length...")

    # --- Data Validation and Preparation ---
    required_cols = ['lb_estimate_in_minutes', 'ub_estimate_in_minutes', 'onetsoc_major']
    if not all(col in processed_df.columns for col in required_cols):
        logging.error(f"Missing one or more required columns: {required_cols}. Cannot generate plot.")
        return None

    df = processed_df.copy()

    # Calculate the estimate ratio, handling division by zero and infinity
    df = df[df['lb_estimate_in_minutes'] > 0]
    df['estimate_ratio'] = df['ub_estimate_in_minutes'] / df['lb_estimate_in_minutes']
    df.replace([np.inf, -np.inf], np.nan, inplace=True)
    df.dropna(subset=['estimate_ratio'], inplace=True)

    if df.empty:
        logging.warning("No valid data available for the ratio heatmap.")
        return None

    # 1. Bin lower bounds into quartiles (Q1–Q4)
    # Using duplicates='drop' can help if there are many identical values
    # which can make binning into quantiles fail.
    try:
        df['lb_q'] = pd.qcut(
            df.lb_estimate_in_minutes,
            q=4,
            labels=['Q1 (Shortest)', 'Q2', 'Q3', 'Q4 (Longest)'],
            duplicates='drop'
        )
    except ValueError as e:
        logging.error(f"Could not bin data into quartiles: {e}. There might not be enough unique values.")
        return None


    # 2. Aggregate: median ratio per cell (occupation x task length quartile)
    pivot = df.pivot_table(
        index='onetsoc_major',
        columns='lb_q',
        values='estimate_ratio',
        aggfunc='median'
    )

    # Map the index (onetsoc_major codes) to their corresponding readable labels
    pivot.index = pivot.index.map(OCCUPATION_MAJOR_CODES)
    pivot.dropna(inplace=True) # Drop occupations with no data in some quartiles for a cleaner plot

    if pivot.empty:
        logging.warning("Pivot table is empty after processing. Cannot generate heatmap.")
        return None

    # --- Plotting ---
    try:
        plt.figure(figsize=(12, 10))
        sns.heatmap(
            pivot,
            cmap='RdYlGn_r',  # Red-Yellow-Green (reversed), good for ratios centered around 1
            center=2,         # Center the colormap around a ratio of 2
            annot=True,       # Show the median values in the cells
            fmt='.1f',        # Format annotations to one decimal place
            linewidths=.5,
            cbar_kws={'label': 'Median Upper/Lower Estimate Ratio'}
        )
        plt.xlabel('Task Length (based on lower-bound quartile)', fontsize=12)
        plt.ylabel('Occupation Major Group', fontsize=12)
        plt.title('Typical Estimate Range Width by Occupation and Task Length', fontsize=16)
        plt.tight_layout()

        # --- File Saving ---
        temp_dir = tempfile.gettempdir()
        temp_path = Path(temp_dir) / "ratio_heatmap_by_occupation_and_task_length.png"
        plt.savefig(temp_path, dpi=300)
        logging.info(f"Successfully saved plot to temporary file: {temp_path}")

        return temp_path

    except Exception as e:
        logging.error(f"An error occurred while generating the heatmap: {e}", exc_info=True)
        return None
    finally:
        plt.close()
