#!/usr/bin/env python3
"""
Plot KEGG annotation results
"""
import sys
import logging
from pathlib import Path
from typing import Dict, List
from collections import OrderedDict

from ..config import config

LOG = logging.getLogger(__name__)


# ============================================================================
# PARSING FUNCTIONS
# ============================================================================

def parse_ko_hierarchy_file(keg_file: Path) -> Dict[str, Dict[str, List[str]]]:
    """
    Parse KEGG KO hierarchy file (like ko00001.keg).
    
    Returns: {category: {pathway: [ko_ids]}}
    """
    if not keg_file.exists():
        LOG.error(f"KEGG file not found: {keg_file}")
        return {}
    
    result = OrderedDict()
    current_category = ""
    current_pathway = ""
    line_count = 0
    d_line_count = 0  # Count D lines with actual KOs
    
    LOG.info(f"Parsing KEGG KO hierarchy file: {keg_file}")
    
    with open(keg_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.rstrip('\n')
            if not line.strip():
                continue
            
            line_count += 1
            
            # Skip comment/special lines
            if line.startswith('#') or line.startswith('!') or line.startswith('+'):
                continue
            
            # Get tag (first character)
            tag = line[0] if line else ''
            
            # Category (A level) - format: A09100 Metabolism
            if tag == 'A':
                # Extract category name (after the number)
                # Format: A09100 Metabolism
                parts = line.split(maxsplit=1)
                if len(parts) >= 2:
                    current_category = parts[1].strip()
                    if current_category not in result:
                        result[current_category] = OrderedDict()
                        LOG.debug(f"Line {line_num}: Found category '{current_category}'")
                continue
            
            # Pathway (B level) - format: B   09101 Carbohydrate metabolism
            if tag == 'B':
                # Check if this is a real pathway or just an empty B line
                if len(line.strip()) > 1:  # Not just "B"
                    # Extract pathway name
                    parts = line.split(maxsplit=1)
                    if len(parts) >= 2:
                        current_pathway = parts[1].strip()
                        
                        if current_category and current_pathway:
                            if current_category not in result:
                                result[current_category] = OrderedDict()
                            
                            if current_pathway not in result[current_category]:
                                result[current_category][current_pathway] = []
                                LOG.debug(f"Line {line_num}: Found pathway '{current_pathway}' under '{current_category}'")
                continue
            
            # Sub-pathway (C level) - format: C    00010 Glycolysis / Gluconeogenesis [PATH:ko00010]
            if tag == 'C':
                # Extract sub-pathway name (remove [PATH:...])
                # Format: C    00010 Glycolysis / Gluconeogenesis [PATH:ko00010]
                if '[' in line:
                    line = line.split('[')[0]  # Remove everything after [
                
                parts = line.split(maxsplit=1)
                if len(parts) >= 2:
                    pass  # We don't use C level for plotting
                continue
            
            # KO entry (D level) - format: D      K00844  HK; hexokinase [EC:2.7.1.1]
            if tag == 'D' and current_category and current_pathway:
                # Ensure dicts exist
                if current_category not in result:
                    result[current_category] = OrderedDict()
                if current_pathway not in result[current_category]:
                    result[current_category][current_pathway] = []
                
                # Extract KO ID (first word after D and spaces)
                # Format: D      K00844  HK; hexokinase [EC:2.7.1.1]
                parts = line.split()
                if len(parts) >= 2:
                    # Second part is usually the KO ID (like K00844)
                    ko_id = parts[1]
                    
                    # Validate it looks like a KO ID (starts with K followed by numbers)
                    if (ko_id.startswith('K') and 
                        (ko_id[1:].replace('_', '').replace('.', '').isdigit() or
                         any(c.isdigit() for c in ko_id[1:]))):
                        result[current_category][current_pathway].append(ko_id)
                        d_line_count += 1
                        
                        if LOG.level <= logging.DEBUG and d_line_count <= 3:
                            LOG.debug(f"Line {line_num}: Added KO '{ko_id}' to '{current_pathway}'")
                continue
    
    # Calculate statistics
    total_categories = len(result)
    total_pathways = sum(len(pathways) for pathways in result.values())
    total_kos = sum(
        len(kos)
        for pathways in result.values()
        for kos in pathways.values()
    )
    
    LOG.info(f"Parsed {line_count} non-empty lines")
    LOG.info(f"Found {total_categories} categories, {total_pathways} pathways, {d_line_count} D-lines with KOs, {total_kos} unique KOs")
    
    # Show summary
    if result:
        LOG.info("Top categories:")
        for i, (category, pathways) in enumerate(list(result.items())[:5], 1):
            pathway_count = len(pathways)
            ko_count = sum(len(k) for k in pathways.values())
            LOG.info(f"  {i}. {category}: {pathway_count} pathways, {ko_count} KOs")
    else:
        LOG.warning("No data parsed from KEGG file")
    
    # Check if we actually found any KOs
    if total_kos == 0:
        LOG.warning("No KOs found in the file. This might be a structure-only file.")
        # Still return the structure if we found categories/pathways
        if total_categories > 0 or total_pathways > 0:
            LOG.info("Returning hierarchical structure without KO counts")
    
    return result


def parse_gene_annotation_file(keg_file: Path) -> Dict[str, Dict[str, List[str]]]:
    """
    Parse gene annotation KEGG file (from make-keg command).
    
    Returns: {category: {pathway: [gene_ids]}}
    """
    if not keg_file.exists():
        LOG.error(f"KEGG file not found: {keg_file}")
        return {}
    
    result = OrderedDict()
    current_category = ""
    current_pathway = ""
    line_count = 0
    d_line_count = 0
    
    LOG.info(f"Parsing gene annotation KEGG file: {keg_file}")
    
    with open(keg_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.rstrip('\n')
            if not line.strip():
                continue
            
            line_count += 1
            tag = line[0] if line else ''
            
            # Skip comment/special lines
            if line.startswith('#') or line.startswith('!') or line.startswith('+'):
                continue
            
            # Category (A level)
            if tag == 'A':
                # Remove HTML tags if present
                content = line[1:].lstrip()
                if '<b>' in content and '</b>' in content:
                    content = content.replace('<b>', '').replace('</b>', '')
                current_category = content.strip()
                
                if current_category and current_category not in result:
                    result[current_category] = OrderedDict()
                    LOG.debug(f"Line {line_num}: Found category '{current_category}'")
                continue
            
            # Pathway (B level)
            if tag == 'B':
                # Remove HTML tags if present
                content = line[1:].lstrip()
                if '<b>' in content and '</b>' in content:
                    content = content.replace('<b>', '').replace('</b>', '')
                current_pathway = content.strip()
                
                if current_category and current_pathway:
                    if current_category not in result:
                        result[current_category] = OrderedDict()
                    
                    if current_pathway not in result[current_category]:
                        result[current_category][current_pathway] = []
                        LOG.debug(f"Line {line_num}: Found pathway '{current_pathway}' under '{current_category}'")
                continue
            
            # Protein (D level) - only if we have a category and pathway
            if tag == 'D' and current_category and current_pathway:
                # Ensure dicts exist
                if current_category not in result:
                    result[current_category] = OrderedDict()
                if current_pathway not in result[current_category]:
                    result[current_category][current_pathway] = []
                
                # Try to extract gene ID
                gene_id = None
                d_line_count += 1
                
                # Method 1: Tab-separated
                if '\t' in line:
                    parts = line.split('\t')
                    if parts and len(parts) > 0:
                        first_part = parts[0].strip()
                        subparts = first_part.split()
                        if len(subparts) >= 2:
                            gene_id = subparts[1]
                
                # Method 2: Space-separated
                if not gene_id and ' ' in line:
                    parts = line.split()
                    if len(parts) >= 2:
                        for part in parts[1:]:
                            if part and part.strip():
                                candidate = part.strip()
                                if (not candidate.startswith('KO:') and
                                    not candidate.startswith('ko:') and
                                    not candidate.replace(':', '').isdigit() and
                                    len(candidate) < 50):
                                    gene_id = candidate
                                    break
                
                # Method 3: Take everything after D and spaces
                if not gene_id:
                    candidate = line[1:].lstrip()
                    if ' ' in candidate:
                        gene_id = candidate.split()[0]
                    else:
                        gene_id = candidate
                
                if gene_id:
                    gene_id = gene_id.strip()
                    result[current_category][current_pathway].append(gene_id)
                    
                    if LOG.level <= logging.DEBUG and len(result[current_category][current_pathway]) <= 3:
                        LOG.debug(f"Line {line_num}: Added gene '{gene_id}' to '{current_pathway}'")
    
    # Calculate statistics
    total_categories = len(result)
    total_pathways = sum(len(pathways) for pathways in result.values())
    total_genes = sum(
        len(genes)
        for pathways in result.values()
        for genes in pathways.values()
    )
    
    LOG.info(f"Parsed {line_count} non-empty lines")
    LOG.info(f"Found {total_categories} categories, {total_pathways} pathways, {d_line_count} D-lines, {total_genes} genes")
    
    # Check if we actually found any genes
    if total_genes == 0:
        LOG.warning("No genes found in the file. This might be a structure-only file.")
        # Still return the structure if we found categories/pathways
        if total_categories > 0 or total_pathways > 0:
            LOG.info("Returning hierarchical structure without gene counts")
    
    return result


# ============================================================================
# PLOTTING FUNCTIONS
# ============================================================================

def create_kos_plot(
    keg_dict: Dict[str, Dict[str, List[str]]],
    output_file: str,
    width: float = 10,
    height: float = 8,
    dpi: int = 300
) -> int:
    """
    Create visualization of KEGG KO hierarchy.
    
    Args:
        keg_dict: Parsed KEGG KO data
        output_file: Output file path
        width: Figure width
        height: Figure height
        dpi: DPI for output
    """
    try:
        # Try to import matplotlib
        import matplotlib
        matplotlib.use('Agg')  # Non-interactive backend
        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches
        from matplotlib.colors import to_hex
        
        title = "KEGG Orthology Hierarchy"
        xlabel = "Number of KEGG Orthologs (KOs)"
        
        LOG.info(f"Creating KOs plot: {output_file}")
        
        # Prepare data for plotting
        categories = []
        pathways = []
        counts = []
        colors = []
        category_colors = {}
        
        # Define color palette - use tab20 for distinct colors
        color_palette = plt.cm.tab20.colors
        
        # Get all unique categories first
        all_categories = list(keg_dict.keys())
        
        # Assign colors to categories BEFORE processing
        for cat_idx, category in enumerate(all_categories):
            color_idx = cat_idx % len(color_palette)
            category_colors[category] = to_hex(color_palette[color_idx])
            LOG.debug(f"Assigned color {to_hex(color_palette[color_idx])} to category '{category}'")
        
        # Organize data
        y_positions = []
        current_y = 0
        
        # First pass: Count total KOs and check if we have data
        total_kos = 0
        has_kos = False
        
        for category, pathway_dict in keg_dict.items():
            # Count KOs in this category
            category_ko_count = 0
            for items in pathway_dict.values():
                count = len(set(items))
                category_ko_count += count
                if count > 0:
                    has_kos = True
            
            total_kos += category_ko_count
        
        if not has_kos:
            LOG.warning("No KO data found in any pathways. Plotting structure only.")
        
        # Second pass: Build plot data
        for category, pathway_dict in keg_dict.items():
            # Get the assigned color for this category
            cat_color = category_colors.get(category, '#808080')  # Default to gray if not found
            
            # Add category header
            categories.append(category)
            pathways.append("")  # Empty for category header
            
            # Calculate total KOs in this category for the category header
            category_total = sum(len(set(items)) for items in pathway_dict.values())
            counts.append(category_total)
            colors.append(cat_color)  # Use the category color for header too
            y_positions.append(current_y)
            current_y += 1
            
            # Add pathways in this category
            for pathway, items in pathway_dict.items():
                count = len(set(items))  # Unique KOs
                
                # Include pathway even if count is 0 (for structure display)
                categories.append(category)
                pathways.append(pathway)
                counts.append(count)
                colors.append(cat_color)  # Same color for all pathways in this category
                y_positions.append(current_y)
                current_y += 1
        
        if not pathways:
            LOG.error("No data to plot - empty file")
            return 1
        
        # Create figure - adjust width based on number of categories
        num_categories = len(category_colors)
        
        # Increase figure width for better legend placement
        extra_width = min(5, num_categories * 0.3)  # Max extra width of 5 inches
        fig_width = width + extra_width
        
        # Increase height based on number of rows
        num_rows = len(y_positions)
        fig_height = height + (num_rows * 0.1)  # Add extra height for many rows
        
        fig, ax = plt.subplots(figsize=(fig_width, fig_height))
        
        # Create horizontal bar chart
        bars = ax.barh(
            y_positions,
            counts,
            color=colors,
            edgecolor='black',
            linewidth=0.5
        )
        
        # Customize plot
        ax.set_xlabel(xlabel, fontsize=11, fontweight='bold')
        
        if has_kos:
            ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
        else:
            ax.set_title(f"{title} (Structure Only - No KOs)", 
                        fontsize=14, fontweight='bold', pad=20, color='red')
        
        # Remove y-axis ticks but keep labels
        ax.set_yticks(y_positions)
        
        # Create custom y-tick labels
        y_tick_labels = []
        for cat, path, count in zip(categories, pathways, counts):
            if path:  # Pathway
                # Truncate long pathway names
                if len(path) > 50:
                    label = path[:47] + '...'
                else:
                    label = path
                y_tick_labels.append(f"  {label}")
            else:  # Category
                y_tick_labels.append(f"▶ {cat}")
        
        ax.set_yticklabels(y_tick_labels, fontsize=9)
        
        # Add value labels on bars (only if count > 0)
        max_count = max(counts) if counts else 0
        for bar, count in zip(bars, counts):
            if count > 0:
                width_val = bar.get_width()
                # Adjust text color based on bar color brightness
                bar_color = bar.get_facecolor()
                # Convert to RGB and check brightness
                if isinstance(bar_color, tuple) and len(bar_color) >= 3:
                    r, g, b = bar_color[:3]
                    brightness = (0.299 * r + 0.587 * g + 0.114 * b)
                    text_color = 'black' if brightness > 0.5 else 'white'
                else:
                    text_color = 'black'
                
                ax.text(
                    width_val + max_count * 0.01 if max_count > 0 else 0.5,
                    bar.get_y() + bar.get_height() / 2,
                    f'{count}',
                    va='center',
                    ha='left',
                    fontsize=8,
                    color=text_color
                )
            elif not has_kos and pathways:  # Show "0" for structure-only display
                ax.text(
                    0.5,
                    bar.get_y() + bar.get_height() / 2,
                    '0',
                    va='center',
                    ha='left',
                    fontsize=8,
                    color='gray'
                )
        
        # Add legend for categories - only include actual categories with data
        legend_patches = []
        
        # Create a set of categories that actually have pathways with data
        categories_with_data = set()
        for category, pathway_dict in keg_dict.items():
            for items in pathway_dict.values():
                if len(set(items)) > 0:
                    categories_with_data.add(category)
        
        # If no categories have data, show all categories
        if not categories_with_data:
            categories_with_data = set(keg_dict.keys())
        
        for category in sorted(categories_with_data):
            if category in category_colors:
                patch = mpatches.Patch(
                    color=category_colors[category], 
                    label=category,
                    edgecolor='black',
                    linewidth=0.5
                )
                legend_patches.append(patch)
        
        if legend_patches:
            # Calculate optimal legend position based on number of categories
            num_categories = len(legend_patches)
            
            if num_categories <= 6:
                # Few categories: place legend outside to the right
                legend_loc = 'upper right'
                bbox_anchor = (1.15, 1)
                legend_fontsize = 9
                ncol = 1
            elif num_categories <= 12:
                # Moderate number of categories: place below or adjust position
                legend_loc = 'upper right'
                bbox_anchor = (1.25, 1)  # Move further right
                legend_fontsize = 8  # Smaller font
                ncol = 1
            else:
                # Many categories: place legend below the plot
                legend_loc = 'upper center'
                bbox_anchor = (0.5, -0.15)  # Below the plot
                legend_fontsize = 7  # Even smaller font
                ncol = 2  # Use 2 columns
                # Increase figure height to accommodate legend
                fig.set_size_inches(fig_width, fig_height + (num_categories * 0.1))
            
            legend = ax.legend(
                handles=legend_patches,
                loc=legend_loc,
                bbox_to_anchor=bbox_anchor,
                fontsize=legend_fontsize,
                title='Categories',
                title_fontsize=legend_fontsize + 1,
                ncol=ncol,
                frameon=True,
                fancybox=True,
                shadow=False,
                borderaxespad=0.5
            )
            
            # Make sure legend frame is visible
            legend.get_frame().set_facecolor('white')
            legend.get_frame().set_alpha(0.9)
            legend.get_frame().set_edgecolor('black')
            legend.get_frame().set_linewidth(0.5)
        
        # Adjust layout with more padding on the right for legend
        if num_categories <= 12:
            plt.tight_layout(rect=[0, 0, 0.85, 1])
        else:
            plt.tight_layout()
        
        # Save figure
        output_path = Path(output_file)
        plt.savefig(str(output_path), dpi=dpi, bbox_inches='tight')
        plt.close()
    
        LOG.info(f"✅ KOs plot saved: {output_path}")
        
        # Create data file
        stem = output_path.stem  # Gets filename without extension
        data_file = output_path.with_name(f"{stem}.tsv")
    
        with open(data_file, 'w') as f:
            f.write("Category\tPathway\tKO_Count\tSample_KOs\n")
            for cat, path_dict in keg_dict.items():
                for pathway, item_list in path_dict.items():
                    count = len(set(item_list))
                    sample_items = ','.join(sorted(item_list)[:5]) if item_list else "None"
                    f.write(f"{cat}\t{pathway}\t{count}\t{sample_items}\n")
    
        LOG.info(f"📊 KO data saved: {data_file}")
        
        # Also save color mapping for debugging
        color_file = output_path.with_name(f"{stem}_colors.tsv")
        with open(color_file, 'w') as f:
            f.write("Category\tColor\n")
            for category, color in category_colors.items():
                f.write(f"{category}\t{color}\n")
        LOG.debug(f"Color mapping saved: {color_file}")
    
        return 0
        
    except ImportError:
        LOG.error("❌ matplotlib not installed. Install with: pip install matplotlib")
        LOG.error("Or install gpathex with: pip install 'gpathex [viz]'")
        return 1
    
    except Exception as e:
        LOG.error(f"❌ Error creating KOs plot: {e}")
        import traceback
        traceback.print_exc()
        return 1


def create_genes_plot(
    keg_dict: Dict[str, Dict[str, List[str]]],
    output_file: str,
    width: float = 10,
    height: float = 8,
    dpi: int = 300
) -> int:
    """
    Create visualization of KEGG gene annotations.
    
    Args:
        keg_dict: Parsed KEGG gene data
        output_file: Output file path
        width: Figure width
        height: Figure height
        dpi: DPI for output
    """
    try:
        # Try to import matplotlib
        import matplotlib
        matplotlib.use('Agg')  # Non-interactive backend
        import matplotlib.pyplot as plt
        import matplotlib.patches as mpatches
        from matplotlib.colors import to_hex
        
        title = "KEGG Gene Annotation Results"
        xlabel = "Number of Genes"
        
        LOG.info(f"Creating genes plot: {output_file}")
        
        # Prepare data for plotting
        categories = []
        pathways = []
        counts = []
        colors = []
        category_colors = {}
        
        # Define color palette - use a palette with many distinct colors
        # Use tab20c colormap for more distinct colors
        color_palette = plt.cm.tab20c.colors
        
        # Get all unique categories first
        all_categories = list(keg_dict.keys())
        
        # Assign colors to categories BEFORE processing
        for cat_idx, category in enumerate(all_categories):
            color_idx = cat_idx % len(color_palette)
            category_colors[category] = to_hex(color_palette[color_idx])
            LOG.debug(f"Assigned color {to_hex(color_palette[color_idx])} to category '{category}'")
        
        # Organize data
        y_positions = []
        current_y = 0
        
        # First pass: Count total genes and check if we have data
        total_genes = 0
        has_genes = False
        
        for category, pathway_dict in keg_dict.items():
            # Count genes in this category
            category_gene_count = 0
            for items in pathway_dict.values():
                count = len(set(items))
                category_gene_count += count
                if count > 0:
                    has_genes = True
            
            total_genes += category_gene_count
        
        if not has_genes:
            LOG.warning("No gene data found in any pathways. Plotting structure only.")
        
        # Second pass: Build plot data
        for category, pathway_dict in keg_dict.items():
            # Get the assigned color for this category
            cat_color = category_colors.get(category, '#808080')  # Default to gray if not found
            
            # Add category header
            categories.append(category)
            pathways.append("")  # Empty for category header
            
            # Calculate total genes in this category for the category header
            category_total = sum(len(set(items)) for items in pathway_dict.values())
            counts.append(category_total)
            colors.append(cat_color)  # Use the category color for header too
            y_positions.append(current_y)
            current_y += 1
            
            # Add pathways in this category
            for pathway, items in pathway_dict.items():
                count = len(set(items))  # Unique genes
                
                # Include pathway even if count is 0 (for structure display)
                categories.append(category)
                pathways.append(pathway)
                counts.append(count)
                colors.append(cat_color)  # Same color for all pathways in this category
                y_positions.append(current_y)
                current_y += 1
        
        if not pathways:
            LOG.error("No data to plot - empty file")
            return 1
        
        # Create figure - adjust width based on number of categories
        num_categories = len(category_colors)
        
        # Increase figure width for better legend placement
        extra_width = min(5, num_categories * 0.3)  # Max extra width of 5 inches
        fig_width = width + extra_width
        
        # Increase height based on number of rows
        num_rows = len(y_positions)
        fig_height = height + (num_rows * 0.1)  # Add extra height for many rows
        
        fig, ax = plt.subplots(figsize=(fig_width, fig_height))
        
        # Create horizontal bar chart
        bars = ax.barh(
            y_positions,
            counts,
            color=colors,
            edgecolor='black',
            linewidth=0.5
        )
        
        # Customize plot
        ax.set_xlabel(xlabel, fontsize=11, fontweight='bold')
        
        if has_genes:
            ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
        else:
            ax.set_title(f"{title} (Structure Only - No Genes)", 
                        fontsize=14, fontweight='bold', pad=20, color='red')
        
        # Remove y-axis ticks but keep labels
        ax.set_yticks(y_positions)
        
        # Create custom y-tick labels
        y_tick_labels = []
        for cat, path, count in zip(categories, pathways, counts):
            if path:  # Pathway
                # Truncate long pathway names
                if len(path) > 50:
                    label = path[:47] + '...'
                else:
                    label = path
                y_tick_labels.append(f"  {label}")
            else:  # Category
                y_tick_labels.append(f"▶ {cat}")
        
        ax.set_yticklabels(y_tick_labels, fontsize=9)
        
        # Add value labels on bars (only if count > 0)
        max_count = max(counts) if counts else 0
        for bar, count in zip(bars, counts):
            if count > 0:
                width_val = bar.get_width()
                # Adjust text color based on bar color brightness
                bar_color = bar.get_facecolor()
                # Convert to RGB and check brightness
                if isinstance(bar_color, tuple) and len(bar_color) >= 3:
                    r, g, b = bar_color[:3]
                    brightness = (0.299 * r + 0.587 * g + 0.114 * b)
                    text_color = 'black' if brightness > 0.5 else 'white'
                else:
                    text_color = 'black'
                
                ax.text(
                    width_val + max_count * 0.01 if max_count > 0 else 0.5,
                    bar.get_y() + bar.get_height() / 2,
                    f'{count}',
                    va='center',
                    ha='left',
                    fontsize=8,
                    color=text_color
                )
            elif not has_genes and pathways:  # Show "0" for structure-only display
                ax.text(
                    0.5,
                    bar.get_y() + bar.get_height() / 2,
                    '0',
                    va='center',
                    ha='left',
                    fontsize=8,
                    color='gray'
                )
        
        # Add legend for categories - only include actual categories with data
        legend_patches = []
        legend_labels = []
        
        # Create a set of categories that actually have pathways with data
        categories_with_data = set()
        for category, pathway_dict in keg_dict.items():
            for items in pathway_dict.values():
                if len(set(items)) > 0:
                    categories_with_data.add(category)
        
        # If no categories have data, show all categories
        if not categories_with_data:
            categories_with_data = set(keg_dict.keys())
        
        for category in sorted(categories_with_data):
            if category in category_colors:
                patch = mpatches.Patch(
                    color=category_colors[category], 
                    label=category,
                    edgecolor='black',
                    linewidth=0.5
                )
                legend_patches.append(patch)
                legend_labels.append(category)
        
        if legend_patches:
            # Calculate optimal legend position based on number of categories
            num_legend_items = len(legend_patches)
            
            if num_legend_items <= 6:
                # Few categories: place legend outside to the right
                legend_loc = 'upper right'
                bbox_anchor = (1.15, 1)
                legend_fontsize = 9
                ncol = 1
            elif num_legend_items <= 12:
                # Moderate number of categories: place below or adjust position
                legend_loc = 'upper right'
                bbox_anchor = (1.25, 1)  # Move further right
                legend_fontsize = 8  # Smaller font
                ncol = 1
            else:
                # Many categories: place legend below the plot
                legend_loc = 'upper center'
                bbox_anchor = (0.5, -0.15)  # Below the plot
                legend_fontsize = 7  # Even smaller font
                ncol = 2  # Use 2 columns
                # Increase figure height to accommodate legend
                fig.set_size_inches(fig_width, fig_height + (num_legend_items * 0.1))
            
            legend = ax.legend(
                handles=legend_patches,
                loc=legend_loc,
                bbox_to_anchor=bbox_anchor,
                fontsize=legend_fontsize,
                title='Categories',
                title_fontsize=legend_fontsize + 1,
                ncol=ncol,
                frameon=True,
                fancybox=True,
                shadow=False,
                borderaxespad=0.5
            )
            
            # Make sure legend frame is visible
            legend.get_frame().set_facecolor('white')
            legend.get_frame().set_alpha(0.9)
            legend.get_frame().set_edgecolor('black')
            legend.get_frame().set_linewidth(0.5)
        
        # Adjust layout with more padding on the right for legend
        if num_legend_items <= 12:
            plt.tight_layout(rect=[0, 0, 0.85, 1])
        else:
            plt.tight_layout()
        
        # Save figure
        output_path = Path(output_file)
        plt.savefig(str(output_path), dpi=dpi, bbox_inches='tight')
        plt.close()
    
        LOG.info(f"✅ Genes plot saved: {output_path}")
        
        # Create data file
        stem = output_path.stem  # Gets filename without extension
        data_file = output_path.with_name(f"{stem}.tsv")
    
        with open(data_file, 'w') as f:
            f.write("Category\tPathway\tGene_Count\tSample_Genes\n")
            for cat, path_dict in keg_dict.items():
                for pathway, item_list in path_dict.items():
                    count = len(set(item_list))
                    sample_items = ','.join(sorted(item_list)[:5]) if item_list else "None"
                    f.write(f"{cat}\t{pathway}\t{count}\t{sample_items}\n")
    
        LOG.info(f"📊 Gene data saved: {data_file}")
        
        # Also save color mapping for debugging
        color_file = output_path.with_name(f"{stem}_colors.tsv")
        with open(color_file, 'w') as f:
            f.write("Category\tColor\n")
            for category, color in category_colors.items():
                f.write(f"{category}\t{color}\n")
        LOG.debug(f"Color mapping saved: {color_file}")
    
        return 0
        
    except ImportError:
        LOG.error("❌ matplotlib not installed. Install with: pip install matplotlib")
        LOG.error("Or install gpathex with: pip install 'gpathex [viz]'")
        return 1
    
    except Exception as e:
        LOG.error(f"❌ Error creating genes plot: {e}")
        import traceback
        traceback.print_exc()
        return 1


# ============================================================================
# COMMAND LINE FUNCTIONS
# ============================================================================

def plot_kegg_kos_results(
    keg_file: str,
    output_file: str,
    width: float = 10,
    height: float = 8,
    dpi: int = 300
) -> int:
    """
    Plot KEGG KO hierarchy results.
    """
    keg_path = Path(keg_file)
    
    if not keg_path.exists():
        LOG.error(f"KEGG file not found: {keg_path}")
        return 1
    
    LOG.info(f"🍹 Creating KEGG KO hierarchy visualization")
    LOG.info(f"📁 Input file: {keg_path}")
    LOG.info(f"📁 Output file: {output_file}")
    LOG.info(f"📏 Figure size: {width} x {height} inches")
    LOG.info(f"🖼️  DPI: {dpi}")
    
    # Parse as KO hierarchy file
    keg_dict = parse_ko_hierarchy_file(keg_path)
    if not keg_dict:
        LOG.error("No data parsed from KO hierarchy file")
        return 1
    
    # Create KO plot
    return create_kos_plot(keg_dict, output_file, width, height, dpi)


def plot_kegg_genes_results(
    keg_file: str,
    output_file: str,
    width: float = 10,
    height: float = 8,
    dpi: int = 300
) -> int:
    """
    Plot KEGG gene annotation results.
    """
    keg_path = Path(keg_file)
    
    if not keg_path.exists():
        LOG.error(f"KEGG file not found: {keg_path}")
        return 1
    
    LOG.info(f"🍹 Creating KEGG gene annotation visualization")
    LOG.info(f"📁 Input file: {keg_path}")
    LOG.info(f"📁 Output file: {output_file}")
    LOG.info(f"📏 Figure size: {width} x {height} inches")
    LOG.info(f"🖼️  DPI: {dpi}")
    
    # Parse as gene annotation file
    keg_dict = parse_gene_annotation_file(keg_path)
    if not keg_dict:
        LOG.error("No data parsed from gene annotation file")
        return 1
    
    # Create gene plot
    return create_genes_plot(keg_dict, output_file, width, height, dpi)


# ============================================================================
# MAIN FUNCTIONS FOR EACH COMMAND
# ============================================================================

def main_plot_kos():
    """Command-line entry point for plot-keg-kos."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Plot KEGG KO hierarchy (ko00001.keg format)"
    )
    parser.add_argument(
        "--keg", "-k",
        required=True,
        help="Input .keg file with KO hierarchy (e.g., ko00001.keg)"
    )
    parser.add_argument(
        "--out", "-o",
        default="kegg_kos_plot",
        help="Output file or prefix (default: kegg_kos_plot). If no extension provided, format will be added."
    )
    parser.add_argument(
        "--format",
        choices=["png", "pdf", "svg", "jpg", "jpeg"],
        default="png",
        help="Output format (default: png)"
    )
    parser.add_argument(
        "--width",
        type=float,
        default=10,
        help="Figure width in inches (default: 10)"
    )
    parser.add_argument(
        "--height",
        type=float,
        default=8,
        help="Figure height in inches (default: 8)"
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=300,
        help="Figure DPI (default: 300)"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Verbose output"
    )
    
    args = parser.parse_args()
    
    # Set up logging
    log_level = "DEBUG" if args.verbose else config.get("LOG_LEVEL")
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format=config.get("LOG_FORMAT")
    )
    
    # Handle output file name
    output_path = Path(args.out)
    
    # Check if output already has a valid image extension
    valid_extensions = ['.png', '.jpg', '.jpeg', '.pdf', '.svg']
    has_valid_extension = any(output_path.suffix.lower() == ext for ext in valid_extensions)
    
    if has_valid_extension:
        # Use the provided filename as-is
        output_file = str(output_path)
        LOG.debug(f"Output file has extension: {output_path.suffix}")
    else:
        # Add extension based on format
        extension = args.format
        if extension == 'jpeg':
            extension = 'jpg'  # Normalize jpeg to jpg
        output_file = f"{args.out}.{extension}"
        output_path = Path(output_file)
        LOG.debug(f"Added extension .{extension} to output file")
    
    # Ensure the output directory exists
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    LOG.info(f"Output plot file: {output_path}")
    
    return plot_kegg_kos_results(
        keg_file=args.keg,
        output_file=output_file,
        width=args.width,
        height=args.height,
        dpi=args.dpi
    )


def main_plot_genes():
    """Command-line entry point for plot-keg-genes."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Plot KEGG gene annotations (from make-keg command)"
    )
    parser.add_argument(
        "--keg", "-k",
        required=True,
        help="Input .keg file with gene annotations (from make-keg)"
    )
    parser.add_argument(
        "--out", "-o",
        default="kegg_genes_plot",
        help="Output file or prefix (default: kegg_genes_plot). If no extension provided, format will be added."
    )
    parser.add_argument(
        "--format",
        choices=["png", "pdf", "svg", "jpg", "jpeg"],
        default="png",
        help="Output format (default: png)"
    )
    parser.add_argument(
        "--width",
        type=float,
        default=10,
        help="Figure width in inches (default: 10)"
    )
    parser.add_argument(
        "--height",
        type=float,
        default=8,
        help="Figure height in inches (default: 8)"
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=300,
        help="Figure DPI (default: 300)"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Verbose output"
    )
    
    args = parser.parse_args()
    
    # Set up logging
    log_level = "DEBUG" if args.verbose else config.get("LOG_LEVEL")
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format=config.get("LOG_FORMAT")
    )
    
    # Handle output file name
    output_path = Path(args.out)
    
    # Check if output already has a valid image extension
    valid_extensions = ['.png', '.jpg', '.jpeg', '.pdf', '.svg']
    has_valid_extension = any(output_path.suffix.lower() == ext for ext in valid_extensions)
    
    if has_valid_extension:
        # Use the provided filename as-is
        output_file = str(output_path)
        LOG.debug(f"Output file has extension: {output_path.suffix}")
    else:
        # Add extension based on format
        extension = args.format
        if extension == 'jpeg':
            extension = 'jpg'  # Normalize jpeg to jpg
        output_file = f"{args.out}.{extension}"
        output_path = Path(output_file)
        LOG.debug(f"Added extension .{extension} to output file")
    
    # Ensure the output directory exists
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    LOG.info(f"Output plot file: {output_path}")
    
    return plot_kegg_genes_results(
        keg_file=args.keg,
        output_file=output_file,
        width=args.width,
        height=args.height,
        dpi=args.dpi
    )


# For backward compatibility
def main():
    """Legacy main function for backward compatibility."""
    # Check if we should run genes or KOs
    import sys
    if len(sys.argv) > 1 and "--help" in sys.argv:
        # Show help for both
        print("Legacy plot-keg command. Use one of:")
        print("  gpathex plot-keg-kos --help")
        print("  gpathex plot-keg-genes --help")
        return 0
    
    # Default to KOs for backward compatibility
    return main_plot_kos()


if __name__ == "__main__":
    sys.exit(main())
