"""
Jupyter widget functions for interactive lockdrop experimentation.

This module contains functions for creating and managing interactive widgets
for the experimental lockdrop notebook.
"""

import json
import traceback
import ipywidgets as widgets
from datetime import datetime
from IPython.display import clear_output

from .constants import NUM_STARS, NUM_GALAXIES, SCENARIOS
from .calculations import (
    calculate_dynamic_allocations, calculate_bonus_pools,
    calculate_final_allocations, generate_test_output
)
from .display import print_analysis_tables, print_table_with_borders
from .visualization import create_visualization


def create_experiment_widgets():
    """Create input widgets for the experimental notebook."""
    # Define reasonable default values
    default_participants = {
        'stars_1_years': 8000,
        'stars_2_years': 8000,
        'stars_3_years': 8000,
        'stars_4_years': 8000,
        'stars_5_years': 8000,
        'galaxies_1_years': 40,
        'galaxies_2_years': 40,
        'galaxies_3_years': 40,
        'galaxies_4_years': 40,
        'galaxies_5_years': 40
    }

    # Create input widgets
    print("="*50)

    # Star participation widgets with bounds
    stars_1_years = widgets.BoundedIntText(value=default_participants['stars_1_years'], min=0, max=NUM_STARS, description='1 Year:', style={'description_width': '80px'})
    stars_2_years = widgets.BoundedIntText(value=default_participants['stars_2_years'], min=0, max=NUM_STARS, description='2 Years:', style={'description_width': '80px'})
    stars_3_years = widgets.BoundedIntText(value=default_participants['stars_3_years'], min=0, max=NUM_STARS, description='3 Years:', style={'description_width': '80px'})
    stars_4_years = widgets.BoundedIntText(value=default_participants['stars_4_years'], min=0, max=NUM_STARS, description='4 Years:', style={'description_width': '80px'})
    stars_5_years = widgets.BoundedIntText(value=default_participants['stars_5_years'], min=0, max=NUM_STARS, description='5 Years:', style={'description_width': '80px'})

    star_controls = widgets.VBox([stars_1_years, stars_2_years, stars_3_years, stars_4_years, stars_5_years])

    galaxies_1_years = widgets.BoundedIntText(value=default_participants['galaxies_1_years'], min=0, max=NUM_GALAXIES, description='1 Year:', style={'description_width': '80px'})
    galaxies_2_years = widgets.BoundedIntText(value=default_participants['galaxies_2_years'], min=0, max=NUM_GALAXIES, description='2 Years:', style={'description_width': '80px'})
    galaxies_3_years = widgets.BoundedIntText(value=default_participants['galaxies_3_years'], min=0, max=NUM_GALAXIES, description='3 Years:', style={'description_width': '80px'})
    galaxies_4_years = widgets.BoundedIntText(value=default_participants['galaxies_4_years'], min=0, max=NUM_GALAXIES, description='4 Years:', style={'description_width': '80px'})
    galaxies_5_years = widgets.BoundedIntText(value=default_participants['galaxies_5_years'], min=0, max=NUM_GALAXIES, description='5 Years:', style={'description_width': '80px'})

    galaxy_controls = widgets.VBox([galaxies_1_years, galaxies_2_years, galaxies_3_years, galaxies_4_years, galaxies_5_years])

    # Preset scenario dropdown
    scenario_dropdown = widgets.Dropdown(
        options=[('Select Preset', '')] + [(scenario['description'], k) for k, scenario in SCENARIOS.items()],
        description='Preset:',
        style={'description_width': '60px'},
        layout={'width': '400px'}
    )

    # Calculate button
    calculate_button = widgets.Button(description='🔄 Calculate Allocations', button_style='success', layout={'width': '200px', 'margin': '10px 5px'})

    # Export button
    export_button = widgets.Button(description='💾 Export Results', button_style='info', layout={'width': '180px', 'margin': '10px 5px'})

    # Output area
    output_area = widgets.Output()

    # Set up preset selection handler
    def on_preset_change(change):
        if change['new'] and change['new'] != '':
            scenario = SCENARIOS[change['new']]

            # Update widget values
            stars_1_years.value = scenario['stars']['1_year']
            stars_2_years.value = scenario['stars']['2_year']
            stars_3_years.value = scenario['stars']['3_year']
            stars_4_years.value = scenario['stars']['4_year']
            stars_5_years.value = scenario['stars']['5_year']
            galaxies_1_years.value = scenario['galaxies']['1_year']
            galaxies_2_years.value = scenario['galaxies']['2_year']
            galaxies_3_years.value = scenario['galaxies']['3_year']
            galaxies_4_years.value = scenario['galaxies']['4_year']
            galaxies_5_years.value = scenario['galaxies']['5_year']

    scenario_dropdown.observe(on_preset_change, names='value')

    # Set up export button handler
    def on_export_click(button=None):
        """Handle export button click with timestamp."""
        # Generate timestamp for filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f'lockdrop_calculations_result_{timestamp}.json'

        # Get current values from widgets
        participation_counts = {
            'stars_1_years': stars_1_years.value,
            'stars_2_years': stars_2_years.value,
            'stars_3_years': stars_3_years.value,
            'stars_4_years': stars_4_years.value,
            'stars_5_years': stars_5_years.value,
            'galaxies_1_years': galaxies_1_years.value,
            'galaxies_2_years': galaxies_2_years.value,
            'galaxies_3_years': galaxies_3_years.value,
            'galaxies_4_years': galaxies_4_years.value,
            'galaxies_5_years': galaxies_5_years.value
        }

        try:
            # Calculate allocations
            allocation_data = calculate_dynamic_allocations(participation_counts)
            bonus_data = calculate_bonus_pools(allocation_data)
            final_data = calculate_final_allocations(allocation_data, bonus_data)

            # Generate test output
            test_output = generate_test_output(final_data)

            # Add metadata
            export_data = {
                'metadata': {
                    'source': 'lockdrop-calculations.ipynb',
                    'timestamp': timestamp,
                    'participation_counts': participation_counts
                },
                'allocations': test_output
            }

            # Save to file
            with open(filename, 'w') as f:
                json.dump(export_data, f, indent=2)

            print(f"✅ Results exported to: {filename}")

        except Exception as e:
            print(f"❌ Export failed: {str(e)}")

    export_button.on_click(on_export_click)

    # Return all widgets as a dictionary
    return {
        'star_controls': star_controls,
        'galaxy_controls': galaxy_controls,
        'scenario_dropdown': scenario_dropdown,
        'calculate_button': calculate_button,
        'export_button': export_button,
        'output_area': output_area,
        'widgets': {
            'stars_1_years': stars_1_years,
            'stars_2_years': stars_2_years,
            'stars_3_years': stars_3_years,
            'stars_4_years': stars_4_years,
            'stars_5_years': stars_5_years,
            'galaxies_1_years': galaxies_1_years,
            'galaxies_2_years': galaxies_2_years,
            'galaxies_3_years': galaxies_3_years,
            'galaxies_4_years': galaxies_4_years,
            'galaxies_5_years': galaxies_5_years
        }
    }


def create_calculate_function(widget_dict):
    """Create the calculation and display function for the experimental notebook."""
    from decimal import Decimal
    import pandas as pd

    def calculate_and_display(button=None):
        """Calculate allocations and display results based on input parameters."""

        widgets_map = widget_dict['widgets']
        output_area = widget_dict['output_area']

        with output_area:
            clear_output(wait=True)

            # Get current values from widgets
            participation_counts = {
                'stars_1_years': widgets_map['stars_1_years'].value,
                'stars_2_years': widgets_map['stars_2_years'].value,
                'stars_3_years': widgets_map['stars_3_years'].value,
                'stars_4_years': widgets_map['stars_4_years'].value,
                'stars_5_years': widgets_map['stars_5_years'].value,
                'galaxies_1_years': widgets_map['galaxies_1_years'].value,
                'galaxies_2_years': widgets_map['galaxies_2_years'].value,
                'galaxies_3_years': widgets_map['galaxies_3_years'].value,
                'galaxies_4_years': widgets_map['galaxies_4_years'].value,
                'galaxies_5_years': widgets_map['galaxies_5_years'].value
            }

            # Validate inputs
            total_stars = sum(participation_counts[key] for key in participation_counts if 'stars' in key)
            total_galaxies = sum(participation_counts[key] for key in participation_counts if 'galaxies' in key)

            if total_stars > NUM_STARS:
                print(f"⚠️  ERROR: Total stars ({total_stars:,}) exceeds maximum available ({NUM_STARS:,})")
                return

            if total_galaxies > NUM_GALAXIES:
                print(f"⚠️  ERROR: Total galaxies ({total_galaxies:,}) exceeds maximum available ({NUM_GALAXIES:,})")
                return

            if total_stars == 0 and total_galaxies == 0:
                print("⚠️  ERROR: Must have at least some participants")
                return

            try:
                # Perform calculations
                print("🔄 Calculating allocations...\n")

                allocation_data = calculate_dynamic_allocations(participation_counts)
                bonus_data = calculate_bonus_pools(allocation_data)
                final_data = calculate_final_allocations(allocation_data, bonus_data)

                # Display results
                print_analysis_tables(allocation_data, bonus_data, final_data)

                # Create visualization
                print("\n📊 Generating visualization...")
                create_visualization(allocation_data, final_data)

                # Calculate participation insights
                star_participation_rate = Decimal(total_stars) / NUM_STARS
                galaxy_participation_rate = Decimal(total_galaxies) / NUM_GALAXIES

                print("\n" + "="*80)
                print("🎯 EXPERIMENT INSIGHTS")
                print("=" * 80)

                insights_df = pd.DataFrame({
                    'Metric': [
                        'Total Participants',
                        'Star Participation Rate',
                        'Galaxy Participation Rate',
                        '5-Year Star Bonus',
                        '5-Year Galaxy Bonus',
                        'Highest Individual Allocation',
                        'Lowest Individual Allocation'
                    ],
                    'Value': [
                        f"{total_stars + total_galaxies:,}",
                        f"{star_participation_rate:.1%}",
                        f"{galaxy_participation_rate:.1%}",
                        f"{float(bonus_data['bonus_per_star_5_years']):,.2f} $Z" if participation_counts['stars_5_years'] > 0 else "N/A (no 5Y stars)",
                        f"{float(bonus_data['bonus_per_galaxy_5_years']):,.2f} $Z" if participation_counts['galaxies_5_years'] > 0 else "N/A (no 5Y galaxies)",
                        f"{max(float(final_data['final_star_allocations'][5]), float(final_data['final_galaxy_allocations'][5])):,.2f} $Z",
                        f"{min(float(final_data['final_star_allocations'][1]), float(final_data['final_galaxy_allocations'][1])):,.2f} $Z"
                    ]
                })
                print_table_with_borders(insights_df, "🎯 EXPERIMENT INSIGHTS")
                print("\n" + "="*80)

            except Exception as e:
                print(f"❌ Error during calculation: {str(e)}")
                traceback.print_exc()

    return calculate_and_display


def create_experimental_interface():
    """Create complete experimental interface with widgets and handlers."""
    from IPython.display import display

    # Create all widgets
    widget_dict = create_experiment_widgets()

    # Create calculation function
    calculate_function = create_calculate_function(widget_dict)

    # Wire up the calculate button
    widget_dict['calculate_button'].on_click(calculate_function)

    # Display the interface - preset dropdown first
    preset_label = widgets.HTML('<h3>🎮 Preset Scenarios</h3>')
    preset_layout = widgets.VBox([preset_label, widget_dict['scenario_dropdown']])
    
    # Then participation controls
    star_label = widgets.HTML('<h3>⭐ Star Participation</h3>')
    galaxy_label = widgets.HTML('<h3>🌌 Galaxy Participation</h3>')

    input_layout = widgets.HBox([
        widgets.VBox([star_label, widget_dict['star_controls']]),
        widgets.VBox([galaxy_label, widget_dict['galaxy_controls']])
    ])

    # Action buttons
    controls_layout = widgets.HBox([
        widget_dict['calculate_button'],
        widget_dict['export_button']
    ])

    # Display in order: preset first, then participation controls, then buttons
    display(preset_layout)
    display(input_layout)
    display(controls_layout)
    display(widget_dict['output_area'])

    return None