"""
Layout Engine for SaltShaker
Pure layout logic with simple radial stacking

Key features:
- Simple radial stacking from outer to inner
- Single-pass algorithm
- Group-aware positioning with consistent spacing
- No dynamic expansion or complex budgets
"""
from __future__ import annotations
from typing import Dict, List, Tuple, Optional, Any, TypedDict
from dataclasses import dataclass
import pandas as pd
import numpy as np
import logging
from math import ceil, log

from .types import (
    LayoutResult,
    SectorBudget,  # Keep for compatibility but won't use sectors
    GroupBandLayout,
    SingleEventLayout
)
from ..config import PlotConfig

logger = logging.getLogger(__name__)


# Type definition for layout return dictionary
class RingLayoutResult(TypedDict):
    """Type definition for ring layout return value"""
    group_bands: List[GroupBandLayout]
    single_events: List[SingleEventLayout]


class LayoutEngine:
    """
    Simplified layout engine with radial stacking
    
    Algorithm:
    1. Determine which type (del/dup) goes outer based on median size
    2. For each type, stack groups from outer to inner radius
    3. Multi-event groups get consistent spacing within bands
    4. Single events share radius levels when possible
    """
    
    def __init__(self, config: PlotConfig, genome_length: int = 16569):
        """
        Initialize layout engine
        
        Args:
            config: Plot configuration (must be PlotConfig instance)
            genome_length: Mitochondrial genome length (bp)
            
        Raises:
            TypeError: If config is not a PlotConfig instance
            ValueError: If genome_length is invalid
        """
        if not isinstance(config, PlotConfig):
            raise TypeError(f"config must be PlotConfig instance, got {type(config)}")
        if genome_length <= 0:
            raise ValueError(f"genome_length must be positive, got {genome_length}")
            
        self.config = config
        self.layout_config = config.layout
        self.genome_length = genome_length
        
        logger.info(f"LayoutEngine initialized (simplified radial stacking)")
    
    def _calculate_required_space(self, events: pd.DataFrame, event_type: str) -> float:
        """
        Calculate actual space required for events in a ring
        
        This calculates space based on:
        - Number and size of multi-event groups (bands)
        - Number of single events (shared levels)
        - Gaps between groups
        
        Args:
            events: Events to calculate space for
            event_type: 'del' or 'dup' for logging
            
        Returns:
            Required space in pixels
        """
        if events.empty:
            return 0.0
        
        # Separate multi-event groups from single events
        group_counts = events['group'].value_counts().to_dict()
        
        # Multi-event groups: 2+ events OR any BL group
        multi_groups = [g for g, c in group_counts.items() if c > 1 or g.startswith('BL')]
        single_groups = [g for g, c in group_counts.items() if c == 1 and not g.startswith('BL')]
        
        total_space = 0.0
        band_details = []
        
        # Space for multi-event group bands
        for group_id in multi_groups:
            n_events = group_counts[group_id]
            band_size = self._calculate_band_size(n_events)
            # Ensure minimum viable height
            band_size = max(band_size, self.layout_config.min_viable_band_height)
            total_space += band_size
            total_space += self.layout_config.group_gap  # Gap after each group
            band_details.append(f"{group_id}={band_size:.1f}px")
        
        # Space for single events (shared levels)
        single_space = 0.0
        if single_groups:
            # More realistic estimate: assume we can pack singles based on angular spacing
            # With 360° and min 45° spacing, max ~8 events per level
            # But be conservative: assume 4-6 per level on average
            max_per_level = self.layout_config.single_event_max_per_level
            n_singles = len(single_groups)
            n_levels = max(1, (n_singles + max_per_level - 1) // max_per_level)  # Ceiling division
            
            # Each level needs minimum height for a band
            space_per_level = max(
                self.layout_config.min_viable_band_height,
                self.layout_config.group_gap * 2  # Conservative: 2 gaps worth
            )
            single_space = n_levels * space_per_level
            
            # SMART VISIBILITY: If ONLY singles (no multi-event groups), use larger minimum
            # This ensures single-event-only plots are visible and don't look tiny
            if not multi_groups:
                min_visible_radius = 150.0  # Minimum for good visibility
                single_space = max(single_space, min_visible_radius)
                logger.debug(f"  Single-only layout: using minimum {min_visible_radius:.1f}px for visibility")
            
            total_space += single_space
        
        logger.info(f"Required space for {event_type}: {len(multi_groups)} groups + "
                   f"{len(single_groups)} singles = {total_space:.1f}px")
        logger.debug(f"  Multi-event bands: {', '.join(band_details) if band_details else 'none'}")
        logger.debug(f"  Singles: {len(single_groups)} events, {n_levels if single_groups else 0} levels, "
                    f"{single_space:.1f}px")
        
        return total_space
    
    def calculate_layout(
        self,
        del_events: pd.DataFrame,
        dup_events: pd.DataFrame,
        requested_radius: float
    ) -> LayoutResult:
        """
        Calculate layout for all events with automatic radius expansion
        
        This method ensures ALL groups get valid bands by expanding total_radius
        if the requested radius is insufficient. NO fallback labels, NO invalid bands.
        
        NOTE: This method modifies del_events and dup_events in-place by adding
        a 'radius' column. If you need to preserve original data, pass copies.
        
        Args:
            del_events: DataFrame with deletion events (modified in-place)
            dup_events: DataFrame with duplication events (modified in-place)
            requested_radius: Requested total radius (will be expanded if needed)
        
        Returns:
            LayoutResult with all positions calculated and actual total_radius used
        """
        logger.info(f"Calculating layout for {len(del_events)} dels, {len(dup_events)} dups")
        
        # Ensure events have required columns (modifying in-place)
        for df in [del_events, dup_events]:
            if not df.empty and 'radius' not in df.columns:
                df['radius'] = 0.0
        
        # Step 1: Determine ring assignment (outer vs inner)
        # Put events with SHORTER arcs on OUTER ring for better visibility
        del_median_size = del_events['delsize'].median() if not del_events.empty else 0
        dup_median_size = dup_events['delsize'].median() if not dup_events.empty else 0
        
        # Reverse logic: smaller arcs go to outer ring
        if dup_median_size < del_median_size:
            outer_type, outer_events = 'dup', dup_events
            inner_type, inner_events = 'del', del_events
        else:
            outer_type, outer_events = 'del', del_events
            inner_type, inner_events = 'dup', dup_events
        
        logger.info(f"Ring assignment: {outer_type} outer ({len(outer_events)} events), "
                   f"{inner_type} inner ({len(inner_events)} events)")
        
        # Step 2: Calculate separator position and required space
        total_events = len(outer_events) + len(inner_events)
        if total_events == 0:
            return self._create_empty_layout(requested_radius)
        
        # Calculate space requirements for each ring
        outer_required = self._calculate_required_space(outer_events, outer_type)
        inner_required = self._calculate_required_space(inner_events, inner_type)
        
        # Calculate total radius needed INCLUDING separator and buffers
        buffer_zone = self.layout_config.buffer_zone
        
        # The separator is a FRACTION of the total radius (default 15%)
        # So if separator_fraction = 0.15, then:
        # - 15% of radius is separator
        # - 85% of radius is for events
        # Therefore: events_space = total_radius * (1 - separator_fraction)
        # Solving for total_radius: total_radius = events_space / (1 - separator_fraction)
        
        events_space_needed = outer_required + inner_required
        
        # Calculate total radius to accommodate events + separator as fraction
        # Add buffers: 3x buffer (inner edge, both sides of separator, outer edge)
        radius_for_events_and_separator = events_space_needed / (1.0 - self.layout_config.separator_fraction)
        total_radius_needed = radius_for_events_and_separator + (3 * buffer_zone)
        
        # Use the larger of requested or required
        actual_total_radius = max(requested_radius, total_radius_needed)
        
        if actual_total_radius > requested_radius:
            expansion_factor = actual_total_radius / requested_radius
            logger.warning(f"Expanding radius: requested={requested_radius:.1f}px, "
                         f"needed={total_radius_needed:.1f}px, using={actual_total_radius:.1f}px "
                         f"(expansion={expansion_factor:.2f}x)")
        else:
            logger.info(f"Sufficient space: requested={requested_radius:.1f}px, "
                       f"needed={total_radius_needed:.1f}px")
        
        # Now calculate actual separator size and ring allocations from actual_total_radius
        available_for_events_and_separator = actual_total_radius - (3 * buffer_zone)
        separator_size = available_for_events_and_separator * self.layout_config.separator_fraction
        available_for_events = available_for_events_and_separator - separator_size
        
        # CRITICAL: Use the REQUIRED space directly, don't scale it down
        # The whole point of expanding the radius was to give events their required space!
        outer_radius = outer_required
        inner_radius = inner_required
        
        # Verify we actually have enough space (should always be true after expansion)
        if outer_radius + inner_radius > available_for_events:
            logger.error(f"SPACE CALCULATION BUG: Required {outer_radius + inner_radius:.1f}px "
                        f"but only have {available_for_events:.1f}px available. "
                        f"This should never happen after radius expansion!")
        
        logger.info(f"Space allocation: {outer_type}={outer_radius:.1f}px, "
                   f"{inner_type}={inner_radius:.1f}px, separator={separator_size:.1f}px")
        
        
        # Calculate ring ranges with proper space allocation
        # Layout from bottom to top:
        # [buffer] [inner_ring] [buffer] [separator] [buffer] [outer_ring] [buffer]
        
        buffer_zone = self.layout_config.buffer_zone
        
        # Calculate positions from bottom up
        current_pos = buffer_zone  # Start after bottom buffer
        
        # Inner ring
        inner_start = current_pos
        inner_end = inner_start + inner_radius
        current_pos = inner_end + buffer_zone  # Add buffer after inner ring
        
        # Separator (centered in its allocated space)
        separator_start = current_pos
        separator_circle_radius = separator_start + (separator_size / 2)  # Middle of separator
        current_pos = separator_start + separator_size + buffer_zone  # Add buffer after separator
        
        # Outer ring
        outer_start = current_pos
        outer_end = outer_start + outer_radius
        
        # Define ranges (min, max) for layout
        inner_range = (inner_start, inner_end)
        outer_range = (outer_start, outer_end)
        
        # Blacklist radius IS the separator circle (middle of separator band)
        blacklist_radius = separator_circle_radius
        
        logger.info(f"Outer ring ({outer_type}): {outer_range[0]:.1f}-{outer_range[1]:.1f} px (buffered)")
        logger.info(f"Inner ring ({inner_type}): {inner_range[0]:.1f}-{inner_range[1]:.1f} px (buffered)")
        logger.info(f"Blacklist separator radius: {blacklist_radius:.1f} px (buffer={buffer_zone}px)")
        logger.info(f"  -> Outer events: {outer_range[0]:.1f}px to {outer_range[1]:.1f}px")
        logger.info(f"  -> Separator at: {blacklist_radius:.1f}px")
        logger.info(f"  -> Inner events: {inner_range[0]:.1f}px to {inner_range[1]:.1f}px")
        
        # Step 3: Layout each ring with guaranteed sufficient space
        outer_layout = self._layout_events_in_ring(outer_events, outer_range, outer_type)
        inner_layout = self._layout_events_in_ring(inner_events, inner_range, inner_type)
        
        
        # Step 4: Combine results
        all_events = pd.concat([outer_events, inner_events], ignore_index=False)
        
        # Combine group bands and single events
        all_group_bands = outer_layout['group_bands'] + inner_layout['group_bands']
        all_single_events = outer_layout['single_events'] + inner_layout['single_events']
        
        # Determine radius ranges
        if outer_type == 'del':
            del_radius_range = outer_range
            dup_radius_range = inner_range
        else:
            del_radius_range = inner_range
            dup_radius_range = outer_range
        
        # Calculate actual used radius
        total_radius_used = all_events['radius'].max() if not all_events.empty else 0
        
        return LayoutResult(
            events=all_events,
            sector_budgets={},  # No sectors in simplified version
            group_bands=all_group_bands,
            single_events=all_single_events,
            total_radius_used=actual_total_radius,  # Use actual radius (may be expanded)
            del_radius_range=del_radius_range,
            dup_radius_range=dup_radius_range,
            blacklist_radius=blacklist_radius,
            layout_algorithm='simplified_radial_with_expansion',
            layout_stats={
                'outer_type': outer_type,
                'inner_type': inner_type,
                'n_group_bands': len(all_group_bands),
                'n_single_events': len(all_single_events),
                'requested_radius': requested_radius,
                'actual_radius': actual_total_radius,
                'radius_expanded': actual_total_radius > requested_radius
            }
        )
    
    def _layout_events_in_ring(
        self,
        events: pd.DataFrame,
        radius_range: Tuple[float, float],
        event_type: str
    ) -> RingLayoutResult:
        """
        Layout all events within a ring using simple stacking
        
        Args:
            events: Events to layout (modified in place)
            radius_range: (min_radius, max_radius) for this ring
            event_type: 'del' or 'dup' for logging
        
        Returns:
            RingLayoutResult with 'group_bands' and 'single_events' lists
        """
        if events.empty:
            return {'group_bands': [], 'single_events': []}
        
        min_radius, max_radius = radius_range
        available_space = max_radius - min_radius
        
        # Separate multi-event groups from single events
        # SPECIAL CASE: BL groups (blacklist-crossing) should ALWAYS get group bands
        # even if they have only 1 event, for visibility
        group_counts = events['group'].value_counts().to_dict()
        
        # Critical diagnostic: log ALL groups in this ring
        logger.info(f"=== {event_type.upper()} RING GROUPS ===")
        for group_id, count in sorted(group_counts.items()):
            logger.info(f"  {group_id}: {count} event(s)")
        
        # Multi-event groups: 2+ events OR any BL group (even with 1 event)
        multi_groups = sorted(
            [g for g, c in group_counts.items() if c > 1 or g.startswith('BL')],
            key=lambda g: group_counts[g]
        )  # Smallest first
        
        # Single events: 1 event AND not a BL group
        single_groups = [g for g, c in group_counts.items() if c == 1 and not g.startswith('BL')]
        
        logger.info(f"Ring for {event_type}: {len(multi_groups)} multi-event groups "
                   f"(including BL groups), {len(single_groups)} single events")
        logger.info(f"  Multi-event groups: {multi_groups}")
        logger.info(f"  Single-event groups: {single_groups}")
        logger.info(f"  Available space: {available_space:.1f}px ({min_radius:.1f} to {max_radius:.1f})")
        
        group_bands = []
        current_radius = max_radius  # Start from outer edge
        
        # Layout multi-event groups
        for group_id in multi_groups:
            group_mask = events['group'] == group_id
            group_events = events[group_mask]
            n_events = len(group_events)
            
            # Calculate band size
            band_size = self._calculate_band_size(n_events)
            
            # Ensure band meets minimum viable height
            min_viable = self.layout_config.min_viable_band_height
            if band_size < min_viable:
                band_size = min_viable
                logger.debug(f"Group {group_id}: increased band to minimum viable {min_viable:.1f}px")
            
            # Calculate band bounds
            band_top = current_radius
            band_bottom = band_top - band_size
            
            # Validate band - this should NEVER fail with proper space allocation
            if band_bottom < min_radius:
                logger.error(f"LAYOUT FAILURE: Group {group_id} exceeds available space. "
                           f"Required band bottom={band_bottom:.1f}px < min_radius={min_radius:.1f}px. "
                           f"This indicates a bug in space calculation.")
                # Clamp to minimum
                band_bottom = min_radius
            
            if band_top <= band_bottom:
                logger.error(f"LAYOUT FAILURE: Group {group_id} invalid band: "
                           f"top={band_top:.1f}px <= bottom={band_bottom:.1f}px. "
                           f"This should never happen with proper space allocation.")
                # Skip this group - it will have no layout
                continue
            
            # Layout events within band
            band_layouts = self._layout_group_band(
                events, 
                group_events.index.tolist(),
                group_id,
                band_top,
                band_bottom
            )
            group_bands.extend(band_layouts)
            
            # Move to next band with gap
            current_radius = band_bottom - self.layout_config.group_gap
        
        # Layout single events on shared levels
        single_event_indices = [events[events['group'] == g].index[0] for g in single_groups]
        single_layouts = self._layout_single_events(
            events,
            single_event_indices,
            current_radius,
            min_radius
        )
        
        # Diagnostic: report what got laid out
        laid_out_groups = set([b.group_id for b in group_bands])
        laid_out_singles = set([events.loc[s.event_index, 'group'] for s in single_layouts])
        all_groups_in_ring = set(group_counts.keys())
        missing_groups = all_groups_in_ring - laid_out_groups - laid_out_singles
        
        logger.info(f"  Groups with bands: {sorted(laid_out_groups)}")
        logger.info(f"  Groups as singles: {sorted(laid_out_singles)}")
        if missing_groups:
            logger.error(f"  MISSING from layout: {sorted(missing_groups)} ← THESE WILL BECOME FALLBACKS!")
        
        return {
            'group_bands': group_bands,
            'single_events': single_layouts
        }
    
    def _calculate_band_size(self, n_events: int) -> float:
        """
        Calculate band size for a group
        
        Simple formula: base size + scale factor * number of events
        More conservative scaling to fit large groups like G2 (196 events)
        """
        base = self.layout_config.base_band_size
        scale = self.layout_config.band_size_scale_factor
        max_size = self.layout_config.max_band_size
        
        # Apply logarithmic scaling for very large groups to conserve space
        # This prevents a single large group (like G2 with 196 events) from taking all space
        if n_events > 50:
            # Use log scaling for large groups: base + scale * (50 + log(n-50))
            size = base + (scale * (50 + (10 * log(n_events - 49, 2))))
        else:
            size = base + (n_events * scale)
        
        return min(size, max_size)
    
    def _layout_group_band(
        self,
        events: pd.DataFrame,
        group_indices: List[int],
        group_id: str,
        band_top: float,
        band_bottom: float
    ) -> List[GroupBandLayout]:
        """
        Layout events within a group band with consistent spacing
        
        Args:
            events: Full DataFrame (modified in place)
            group_indices: Indices of events in this group
            group_id: Group identifier
            band_top: Outer radius of band
            band_bottom: Inner radius of band
        
        Returns:
            List with single GroupBandLayout (no splitting in simplified version)
        """
        n_events = len(group_indices)
        band_height = band_top - band_bottom
        
        # Band validation should be done by caller - this should never happen
        assert band_height > 0, f"Invalid band: top={band_top}, bottom={band_bottom}"
        
        # Calculate spacing - UNIFORM across all groups for consistent visual appearance
        min_spacing = self.layout_config.min_event_spacing
        
        if self.layout_config.uniform_within_group_spacing:
            # Use TARGET spacing for consistent dense appearance across all groups
            # This makes small groups (G3, G4) look as dense as large groups (G2, G5)
            spacing = self.layout_config.target_event_spacing
            
            # Check if we have enough space with target spacing
            required_height = n_events * spacing
            if required_height > band_height:
                # Not enough space - scale down proportionally
                spacing = band_height / n_events
                logger.warning(f"Group {group_id}: target spacing {self.layout_config.target_event_spacing:.1f} "
                             f"too large, using {spacing:.1f} px")
            
            # Still respect minimum
            if spacing < min_spacing:
                logger.warning(f"Group {group_id}: spacing {spacing:.1f} < min {min_spacing}")
                spacing = min_spacing
        else:
            # Original proportional spacing (varies by group size)
            if n_events > 1:
                spacing = band_height / (2 * n_events - 1)
                if spacing < min_spacing:
                    logger.warning(f"Group {group_id}: calculated spacing {spacing:.1f} < min {min_spacing}")
                    spacing = min_spacing
            else:
                spacing = band_height

        group_df = events.loc[group_indices].copy()
        group_df = group_df.sort_values('deg1')
        sorted_indices = group_df.index.tolist()
        
        # Assign radii with consistent spacing
        for i, idx in enumerate(sorted_indices):
            # Center events in their slots
            radius = band_top - (i * spacing) - (spacing / 2)
            # Ensure within band bounds
            radius = max(band_bottom, min(band_top, radius))
            events.loc[idx, 'radius'] = radius

        # Create layout descriptor
        return [GroupBandLayout(
            group_id=group_id,
            sub_band_index=0,  # No splitting in simplified version
            band_top=band_top,
            band_bottom=band_bottom,
            n_events=n_events,
            event_indices=sorted_indices,
            within_spacing=spacing,
            sector_id=0  # No sectors
        )]
    
    def _layout_single_events(
        self,
        events: pd.DataFrame,
        single_indices: List[int],
        start_radius: float,
        min_radius: float
    ) -> List[SingleEventLayout]:
        """
        Layout single events on shared radius levels
        
        Args:
            events: Full DataFrame (modified in place)
            single_indices: Indices of single events
            start_radius: Starting radius for first level
            min_radius: Minimum allowed radius
        
        Returns:
            List of SingleEventLayout objects
        """
        if not single_indices:
            return []
        
        logger.info(f"  Laying out {len(single_indices)} single events, "
                   f"start_radius={start_radius:.1f}, min_radius={min_radius:.1f}")
        
        if start_radius < min_radius:
            logger.error(f"  SINGLES LAYOUT FAILURE: start_radius ({start_radius:.1f}) < min_radius ({min_radius:.1f})! "
                       f"No space left for singles after multi-event bands.")
            # Can't layout singles - no space
            return []
        
        layouts = []
        shared_levels = []
        current_radius = start_radius
        gap = self.layout_config.group_gap
        
        for idx in single_indices:
            event_deg = events.loc[idx, 'deg1']
            group_id = events.loc[idx, 'group']
            
            # Try to place on existing level
            placed = False
            for level in shared_levels:
                # Check angular spacing with other events on this level
                can_share = all(
                    self._angular_distance(event_deg, deg) >= self.layout_config.single_event_min_angular_spacing
                    for deg in level['degrees']
                )
                
                if can_share and len(level['events']) < self.layout_config.single_event_max_per_level:
                    level['events'].append(idx)
                    level['degrees'].append(event_deg)
                    events.loc[idx, 'radius'] = level['radius']
                    placed = True
                    logger.debug(f"    {group_id}: shared level {len(shared_levels)-1} at {level['radius']:.1f}px")
                    break
            
            if not placed:
                # Create new level
                new_radius = current_radius - gap
                if new_radius < min_radius:
                    logger.warning(f"    {group_id}: Cannot create new level - would be below min_radius "
                                 f"(need {new_radius:.1f}, min is {min_radius:.1f}). SKIPPING!")
                    continue  # Skip this single event
                    
                shared_levels.append({
                    'radius': new_radius,
                    'events': [idx],
                    'degrees': [event_deg]
                })
                events.loc[idx, 'radius'] = new_radius
                current_radius = new_radius
                logger.debug(f"    {group_id}: new level {len(shared_levels)-1} at {new_radius:.1f}px")
        
        # Create SingleEventLayout objects
        for level_id, level in enumerate(shared_levels):
            for idx in level['events']:
                layouts.append(SingleEventLayout(
                    event_index=idx,
                    radius=level['radius'],
                    sector_id=0,  # No sectors
                    shared_level_id=level_id,
                    n_events_on_level=len(level['events'])
                ))
        
        logger.info(f"  Successfully laid out {len(layouts)} of {len(single_indices)} singles on {len(shared_levels)} levels")
        
        return layouts
    
    def _angular_distance(self, deg1: float, deg2: float) -> float:
        """Calculate minimum angular distance between two angles"""
        diff = abs(deg1 - deg2)
        return min(diff, 360 - diff)
    
    def _create_empty_layout(self, total_radius: float) -> LayoutResult:
        """Create empty layout for when there are no events"""
        return LayoutResult(
            events=pd.DataFrame(),
            sector_budgets={},
            group_bands=[],
            single_events=[],
            total_radius_used=0,
            del_radius_range=(0, 0),
            dup_radius_range=(0, 0),
            blacklist_radius=total_radius / 2,
            layout_algorithm='simplified_radial',
            layout_stats={}
        )