visualization.py

import matplotlib
matplotlib.use('Agg') # Non-interactive backend
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import os

COLOR_MAP = {
    1: '#2ecc71',      # Green (P1)
    2: '#f1c40f',      # Yellow (P2)
    3: '#e67e22',      # Orange (P3)
    "3 (VIP)": '#d35400', 
    "FAIL": '#c0392b', # Red
    "RANDOM": '#8e44ad', 
    "Mimo": '#95a5a6'  
}

def clean_priority(val):
    try:
        if isinstance(val, (int, float)):
            val = int(val)
        return str(val).replace('.0', '')
    except:
        return str(val)

def try_int(x):
    try: return int(x)
    except: return x

def generate_report(df, output_folder, config_manager=None):
    """
    Generates graphs and saves them to the output folder.
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Load capacities from config
    initial_capacities = {}
    if config_manager:
        initial_capacities = config_manager.get_project_capacities()

    # Prepare data
    df = df.copy()
    df['Priorita_Clean'] = df['Priorita'].apply(clean_priority)
    
    # Setup styles
    plt.style.use('dark_background')
    sns.set_style("darkgrid", {"axes.facecolor": "#2d2d2d", "grid.color": "#444444"})
    plt.rcParams.update({
        'font.size': 12,
        'figure.facecolor': '#1E1E1E',
        'axes.facecolor': '#2d2d2d',
        'savefig.facecolor': '#1E1E1E',
        'text.color': 'white',
        'axes.labelcolor': 'white',
        'xtick.color': 'white',
        'ytick.color': 'white'
    })

    # --- GRAPH 1: Pie Chart (Satisfaction) ---
    try:
        plt.figure(figsize=(10, 8))
        counts = df['Priorita_Clean'].value_counts()
        colors = [COLOR_MAP.get(try_int(x), '#95a5a6') for x in counts.index]
        
        plt.pie(counts, labels=counts.index, autopct='%1.1f%%', startangle=140, colors=colors, pctdistance=0.85, explode=[0.05]*len(counts))
        plt.title('Celková spokojenost studentů', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_folder, '1_spokojenost.png'), dpi=150)
        plt.close()
    except Exception as e:
        print(f"Error generating graph 1: {e}")

    # --- GRAPH 2: Bar Chart (Classes vs Priorities) ---
    try:
        plt.figure(figsize=(14, 8))
        cross_tab = pd.crosstab(df['Trida'], df['Priorita_Clean'])
        # Sort columns safely
        cols = sorted(cross_tab.columns, key=lambda x: float(x) if str(x).replace('.','').isdigit() else 99)
        cross_tab = cross_tab[cols]
        bar_colors = [COLOR_MAP.get(try_int(c), '#333333') for c in cross_tab.columns]

        ax = cross_tab.plot(kind='bar', stacked=True, color=bar_colors, figsize=(14, 8), edgecolor='white')
        plt.title('Rozdělení priorit v jednotlivých třídách', fontsize=16, fontweight='bold')
        plt.xlabel('Třída')
        plt.ylabel('Počet studentů')
        plt.legend(title='Priorita', bbox_to_anchor=(1.05, 1), loc='upper left')

        for c in ax.containers:
            ax.bar_label(c, label_type='center', color='white', fontweight='bold')

        plt.tight_layout()
        plt.savefig(os.path.join(output_folder, '2_tridy_priority.png'), dpi=150)
        plt.close()
    except Exception as e:
        print(f"Error generating graph 2: {e}")

    # --- GRAPH 3: Heatmap (Projects vs Classes) ---
    try:
        plt.figure(figsize=(16, 10))
        heatmap_data = pd.crosstab(df['Vybrany_Projekt'], df['Trida'])
        sns.heatmap(heatmap_data, annot=True, fmt='d', cmap='YlGnBu', linewidths=.5)
        plt.title('Heatmapa: Rozložení tříd do projektů', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_folder, '3_heatmapa.png'), dpi=150)
        plt.close()
    except Exception as e:
        print(f"Error generating graph 3: {e}")

    # --- GRAPH 4: Capacity (Optional but good) ---
    if initial_capacities:
        try:
            plt.figure(figsize=(14, 10))
            proj_counts = df['Vybrany_Projekt'].value_counts()
            proj_df = pd.DataFrame({'Projekt': proj_counts.index, 'Pocet': proj_counts.values})
            proj_df['Kapacita'] = proj_df['Projekt'].apply(lambda x: initial_capacities.get(x, 0))
            proj_df['Naplněnost %'] = (proj_df['Pocet'] / proj_df['Kapacita'] * 100).round(1)
            proj_df = proj_df.sort_values('Naplněnost %', ascending=False)

            ax = sns.barplot(x='Pocet', y='Projekt', data=proj_df, color='#3498db')

            for index, p in enumerate(ax.patches):
                proj_name = ax.get_yticklabels()[index].get_text()
                cap = initial_capacities.get(proj_name, 0)
                width = p.get_width()
                plt.axvline(x=cap, ymin=index/len(proj_df), ymax=(index+1)/len(proj_df), color='red', linestyle='--', linewidth=2)
                
                # Safe text placement
                try:
                    percent = proj_df[proj_df["Projekt"]==proj_name]["Naplněnost %"].values[0]
                except:
                    percent = 0
                
                ax.text(width + 0.5, p.get_y() + p.get_height()/2, 
                        f'{int(width)}/{cap} ({percent}%)', 
                        va='center', fontweight='bold')

            plt.title('Obsazenost projektů vs. Kapacita', fontsize=16, fontweight='bold')
            plt.xlabel('Počet studentů')
            plt.tight_layout()
            plt.savefig(os.path.join(output_folder, '4_obsazenost.png'), dpi=150)
            plt.close()
        except Exception as e:
            print(f"Error generating graph 4: {e}")

    return True