"""
Fast and efficient IntSampler test with visualization
Replaces the slow validate_int_sampler.py with a more focused approach
"""

import os
import random
from collections import Counter
from typing import List, Optional

from samplers import IntSampler
from samplers.tidb_types import TiDBNumericBucket, TiDBNumericTopNItem


def get_validate_seed():
    """Get seed from environment or use default"""
    return int(os.environ.get("ITS_VALIDATION_SEED", 42))


def create_focused_validate_data(
    size: int = 1000, seed: Optional[int] = None
) -> List[Optional[int]]:
    """Generate focused test data with different distribution patterns based on seed"""
    if seed is None:
        seed = get_validate_seed()
    random.seed(seed)

    # Use seed to create different distribution patterns
    pattern = seed % 4  # Create 4 different distribution patterns
    data = []

    if pattern == 0:  # Small values distribution (0-100)
        print(f"  🎲 Generating small-value integer distribution (seed {seed})")
        common_values = [1, 2, 5, 10, 20]
        # Focus on small integers
        for _ in range(int(size * 0.4)):
            data.append(random.choice(common_values))
        for _ in range(int(size * 0.5)):
            data.append(random.randint(0, 100))

    elif pattern == 1:  # Medium values distribution (100-1000)
        print(f"  🎲 Generating medium-value integer distribution (seed {seed})")
        common_values = [100, 200, 500, 750, 1000]
        # Focus on medium integers
        for _ in range(int(size * 0.4)):
            data.append(random.choice(common_values))
        for _ in range(int(size * 0.5)):
            data.append(random.randint(100, 1000))

    elif pattern == 2:  # Large values distribution (1000-10000)
        print(f"  🎲 Generating large-value integer distribution (seed {seed})")
        common_values = [1000, 2000, 5000, 7500, 10000]
        # Focus on large integers
        for _ in range(int(size * 0.4)):
            data.append(random.choice(common_values))
        for _ in range(int(size * 0.5)):
            data.append(random.randint(1000, 10000))

    else:  # Mixed distribution (all ranges)
        print(f"  🎲 Generating mixed-range integer distribution (seed {seed})")
        common_values = [1, 50, 500, 2000, 8000]
        # Mix of all ranges
        for _ in range(int(size * 0.2)):
            data.append(random.choice(common_values))
        for _ in range(int(size * 0.2)):
            data.append(random.randint(0, 100))  # Small
        for _ in range(int(size * 0.2)):
            data.append(random.randint(100, 1000))  # Medium
        for _ in range(int(size * 0.15)):
            data.append(random.randint(1000, 10000))  # Large

    # Add nulls (consistent across patterns) - 2% of data
    null_count = int(size * 0.02)
    for _ in range(null_count):
        data.append(None)

    # Fill remaining with pattern-specific data
    remaining = size - len(data)
    for _ in range(remaining):
        if pattern == 0:
            data.append(random.randint(0, 100))
        elif pattern == 1:
            data.append(random.randint(100, 1000))
        elif pattern == 2:
            data.append(random.randint(1000, 10000))
        else:
            # Mixed: choose range randomly
            choice = random.randint(0, 2)
            if choice == 0:
                data.append(random.randint(0, 100))
            elif choice == 1:
                data.append(random.randint(100, 1000))
            else:
                data.append(random.randint(1000, 10000))

    random.shuffle(data)
    return data[:size]


def create_int_sampler_from_data(int_data: List[Optional[int]]) -> IntSampler:
    """Create IntSampler from integer data using TiDB types (simplified version)"""
    # Process data
    numeric_data = [x for x in int_data if x is not None]
    null_count = sum(1 for x in int_data if x is None)

    if not numeric_data:
        return IntSampler([], [], null_count=null_count)

    # Find TopN values (top 5 most frequent)
    value_counts = Counter(numeric_data)
    most_common = value_counts.most_common(5)
    # Create TiDB numeric TopN items
    tidb_top_n = [
        TiDBNumericTopNItem(
            data_type="INT", encoded=val, count=count  # Store the actual integer value
        )
        for val, count in most_common
    ]

    # Remove TopN values for histogram
    topn_values = {val for val, _ in most_common}
    histogram_data = [val for val in numeric_data if val not in topn_values]

    # Create TiDB numeric buckets (fewer buckets for speed)
    tidb_buckets = []
    if histogram_data:
        min_val = min(histogram_data)
        max_val = max(histogram_data)

        if min_val != max_val:
            # Create 5 buckets for simplicity
            bucket_size = (max_val - min_val) / 5
            for i in range(5):
                lower = min_val + i * bucket_size
                upper = min_val + (i + 1) * bucket_size

                # Count values in bucket
                if i == 4:  # Last bucket includes upper bound
                    bucket_values = [v for v in histogram_data if lower <= v <= upper]
                else:
                    bucket_values = [v for v in histogram_data if lower <= v < upper]

                if bucket_values:
                    tidb_buckets.append(
                        TiDBNumericBucket(
                            data_type="INT",
                            lower_bound=lower,
                            upper_bound=upper,
                            count=len(bucket_values),
                            repeats=0,
                        )
                    )

    return IntSampler(
        buckets=tidb_buckets,
        top_n=tidb_top_n,
        null_count=null_count,
        is_cumulative_count=False,  # Critical: use absolute counts for test data
        seed=get_validate_seed(),
    )


def validate_int_sampler_distribution():
    """Test IntSampler distribution accuracy"""
    print("=== Testing IntSampler Distribution Accuracy ===")

    # Generate focused test data
    seed = get_validate_seed()
    original_data = create_focused_validate_data(1000, seed=seed)
    print(f"Generated {len(original_data)} original integers")

    original_clean = [x for x in original_data if x is not None]
    print(
        f"Non-null: {len(original_clean)}, Range: {min(original_clean)}-{max(original_clean)}"
    )

    # Create sampler
    sampler = create_int_sampler_from_data(original_data)

    # Get sampler stats
    stats = sampler.its.get_statistics()
    print(f"Sampler: {stats['num_buckets']} buckets, {stats['num_topn']} TopN values")

    # Generate samples
    n_samples = 2000
    generated_samples = [sampler.sample() for _ in range(n_samples)]
    generated_clean = [x for x in generated_samples if x is not None]

    print(
        f"Generated {len(generated_clean)} samples, Range: {min(generated_clean)}-{max(generated_clean)}"
    )

    # Compare TopN frequencies
    original_counter = Counter(original_clean)
    generated_counter = Counter(generated_clean)

    print("\nTopN Value Comparison:")
    total_original = len(original_clean)
    total_generated = len(generated_clean)

    for item in stats["topn_values"][:5]:  # Top 5
        value, expected_count = item
        value = int(value)

        orig_count = original_counter[value]
        gen_count = generated_counter[value]

        orig_ratio = orig_count / total_original
        gen_ratio = gen_count / total_generated

        print(
            f"  Value {value}: Original {orig_count} ({orig_ratio:.3f}) -> Generated {gen_count} ({gen_ratio:.3f})"
        )

    # Compare range distributions
    print("\nRange Distribution Comparison:")
    ranges = [
        (0, 50, "0-50"),
        (50, 200, "50-200"),
        (200, 1000, "200-1K"),
        (1000, 10000, "1K-10K"),
        (10000, float("inf"), "10K+"),
    ]

    for lower, upper, label in ranges:
        orig_count = sum(1 for x in original_clean if lower <= x < upper)
        gen_count = sum(1 for x in generated_clean if lower <= x < upper)

        orig_ratio = orig_count / total_original
        gen_ratio = gen_count / total_generated

        print(
            f"  {label}: Original {orig_count} ({orig_ratio:.3f}) -> Generated {gen_count} ({gen_ratio:.3f})"
        )

    return original_data, generated_samples, sampler


def create_simple_visualization(
    original_data, generated_samples, sampler, output_dir="tests"
):
    """Create simple text-based visualization (no matplotlib dependency)"""
    print("\n=== Creating Simple Distribution Analysis ===")

    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)

    original_clean = [x for x in original_data if x is not None]
    generated_clean = [x for x in generated_samples if x is not None]

    # Create text report
    report_path = os.path.join(output_dir, "int_sampler_analysis.txt")

    with open(report_path, "w") as f:
        f.write("IntSampler Distribution Analysis Report\n")
        f.write("=" * 50 + "\n\n")

        f.write("Original Data:\n")
        f.write(f"  Total samples: {len(original_data)}\n")
        f.write(f"  Non-null: {len(original_clean)}\n")
        f.write(f"  Range: {min(original_clean)} - {max(original_clean)}\n")
        f.write(f"  Mean: {sum(original_clean) / len(original_clean):.2f}\n\n")

        f.write("Generated Data:\n")
        f.write(f"  Total samples: {len(generated_samples)}\n")
        f.write(f"  Non-null: {len(generated_clean)}\n")
        f.write(f"  Range: {min(generated_clean)} - {max(generated_clean)}\n")
        f.write(f"  Mean: {sum(generated_clean) / len(generated_clean):.2f}\n\n")

        # Top values comparison
        original_counter = Counter(original_clean)
        generated_counter = Counter(generated_clean)

        f.write("Top Values Comparison:\n")
        f.write("Value\tOriginal\tGenerated\tOrig%\tGen%\n")

        all_top_values = set(dict(original_counter.most_common(10)).keys()) | set(
            dict(generated_counter.most_common(10)).keys()
        )

        for value in sorted(all_top_values)[:15]:
            orig_count = original_counter[value]
            gen_count = generated_counter[value]
            orig_pct = orig_count / len(original_clean) * 100
            gen_pct = gen_count / len(generated_clean) * 100
            f.write(
                f"{value}\t{orig_count}\t{gen_count}\t{orig_pct:.1f}%\t{gen_pct:.1f}%\n"
            )

        f.write("\nSampler Statistics:\n")
        stats = sampler.its.get_statistics()
        f.write(f"  Buckets: {stats['num_buckets']}\n")
        f.write(f"  TopN values: {stats['num_topn']}\n")
        f.write(f"  Total count: {stats['total_count']}\n")
        f.write(f"  Null probability: {stats['null_probability']:.3f}\n")

    print(f"Analysis report saved to: {report_path}")

    # Try to create matplotlib visualization if available
    try:
        pass

        create_matplotlib_visualization(
            original_clean, generated_clean, sampler, output_dir
        )
    except ImportError:
        print("matplotlib not available - skipping visual plots")
        print("Text-based analysis complete!")


def create_matplotlib_visualization(
    original_clean, generated_clean, sampler, output_dir
):
    """Create matplotlib visualization if numpy/matplotlib available"""
    import matplotlib.pyplot as plt
    import numpy as np

    print("Creating matplotlib visualization...")

    # Create figure
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle("IntSampler - Distribution Comparison (Fixed)", fontsize=14)

    # 1. Histogram comparison (properly binned for actual data range)
    ax1 = axes[0, 0]

    # Create appropriate bins based on actual data range
    min_val = min(min(original_clean), min(generated_clean))
    max_val = max(max(original_clean), max(generated_clean))

    # Use 30 bins across the actual data range for good resolution
    bins = np.linspace(min_val, max_val, 31)

    ax1.hist(
        original_clean,
        bins=bins,
        alpha=0.6,
        label="Original",
        density=True,
        color="blue",
        edgecolor="black",
        linewidth=0.5,
    )
    ax1.hist(
        generated_clean,
        bins=bins,
        alpha=0.6,
        label="Generated",
        density=True,
        color="red",
        edgecolor="black",
        linewidth=0.5,
    )
    ax1.set_xlabel("Integer Value")
    ax1.set_ylabel("Density")
    ax1.set_title("Integer Distribution Comparison")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # 2. Top values comparison
    ax2 = axes[0, 1]
    original_counter = Counter(original_clean)
    generated_counter = Counter(generated_clean)

    top_values = sorted(
        set(dict(original_counter.most_common(10)).keys())
        | set(dict(generated_counter.most_common(10)).keys())
    )[:12]

    if top_values:
        x_pos = np.arange(len(top_values))
        orig_counts = [original_counter[v] for v in top_values]
        gen_counts = [generated_counter[v] for v in top_values]

        width = 0.35
        ax2.bar(
            x_pos - width / 2,
            orig_counts,
            width,
            label="Original",
            alpha=0.7,
            color="blue",
        )
        ax2.bar(
            x_pos + width / 2,
            gen_counts,
            width,
            label="Generated",
            alpha=0.7,
            color="red",
        )

        ax2.set_xlabel("Top Integer Values")
        ax2.set_ylabel("Count")
        ax2.set_title("Top Values Comparison")
        ax2.set_xticks(x_pos)
        ax2.set_xticklabels([str(v) for v in top_values], rotation=45)
        ax2.legend()
        ax2.grid(True, alpha=0.3)

    # 3. Range distribution
    ax3 = axes[1, 0]
    ranges = [
        (0, 50, "0-50"),
        (50, 200, "50-200"),
        (200, 1000, "200-1K"),
        (1000, float("inf"), "1K+"),
    ]

    orig_range_counts = []
    gen_range_counts = []
    range_labels = []

    for lower, upper, label in ranges:
        orig_count = sum(1 for x in original_clean if lower <= x < upper)
        gen_count = sum(1 for x in generated_clean if lower <= x < upper)
        orig_range_counts.append(orig_count)
        gen_range_counts.append(gen_count)
        range_labels.append(label)

    x_pos = np.arange(len(range_labels))
    width = 0.35

    ax3.bar(
        x_pos - width / 2,
        orig_range_counts,
        width,
        label="Original",
        alpha=0.7,
        color="blue",
    )
    ax3.bar(
        x_pos + width / 2,
        gen_range_counts,
        width,
        label="Generated",
        alpha=0.7,
        color="red",
    )

    ax3.set_xlabel("Value Ranges")
    ax3.set_ylabel("Count")
    ax3.set_title("Range Distribution Comparison")
    ax3.set_xticks(x_pos)
    ax3.set_xticklabels(range_labels)
    ax3.legend()
    ax3.grid(True, alpha=0.3)

    # 4. Histogram buckets
    ax4 = axes[1, 1]
    stats = sampler.its.get_statistics()

    if stats["bucket_ranges"]:
        bucket_centers = []
        bucket_counts = []

        for lower, upper, count in stats["bucket_ranges"]:
            center = (lower + upper) / 2
            bucket_centers.append(center)
            bucket_counts.append(count)

        ax4.bar(
            bucket_centers, bucket_counts, alpha=0.7, color="green", edgecolor="black"
        )
        ax4.set_xlabel("Bucket Center")
        ax4.set_ylabel("Count")
        ax4.set_title("Histogram Buckets")
        ax4.grid(True, alpha=0.3)

    plt.tight_layout()

    # Save the plot
    output_path = os.path.join(output_dir, "int_sampler_comparison.png")
    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    print(f"Visualization saved to: {output_path}")
    plt.close()


def validate_edge_cases():
    """Test IntSampler edge cases"""
    print("\n=== Testing Edge Cases ===")

    # Test empty data
    empty_sampler = IntSampler([], [], null_count=0, is_cumulative_count=False)
    empty_samples = empty_sampler.sample_batch(5)
    print(f"Empty data: {empty_samples}")

    # Test only nulls
    null_sampler = IntSampler([], [], null_count=100, is_cumulative_count=False)
    null_samples = null_sampler.sample_batch(5)
    print(f"Only nulls: {null_samples}")

    # Test single value using TiDB bucket
    single_bucket = [
        TiDBNumericBucket(
            data_type="INT",
            lower_bound=42,
            upper_bound=42,
            count=100,
            repeats=100,
            ndv=1,
        )
    ]
    single_sampler = IntSampler(single_bucket, [], is_cumulative_count=False)
    single_samples = single_sampler.sample_batch(5)
    print(f"Single value: {single_samples}")


def run_fast_int_sampler_test():
    """Run fast IntSampler test suite"""
    print("🚀 Fast IntSampler Test Suite")
    print("=" * 50)

    # Set seeds for reproducibility
    seed = get_validate_seed()
    print(f"🎲 Using test seed: {seed}")
    random.seed(seed)

    try:
        # Main distribution test
        original_data, generated_samples, sampler = validate_int_sampler_distribution()

        # Create visualization/analysis
        create_simple_visualization(original_data, generated_samples, sampler)

        # Test edge cases
        validate_edge_cases()

        print("\n" + "=" * 50)
        print("✅ Fast IntSampler Test Suite Completed Successfully!")

        return True

    except Exception as e:
        print(f"\n❌ Test suite failed: {e}")
        import traceback

        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = run_fast_int_sampler_test()
    exit(0 if success else 1)