"""
Test cumulative count conversion functionality
"""

import os

from samplers import HistogramBucket, IntSampler, InverseTransformSampler
from samplers.tidb_types import TiDBNumericBucket

try:
    import matplotlib.pyplot as plt
    import numpy as np

    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False


def validate_cumulative_count_conversion():
    """Test that cumulative counts are correctly converted to absolute counts"""
    print("=== Testing Cumulative Count Conversion ===")

    # Create test data with TiDB-style cumulative counts
    # Simulating: bucket0=100, bucket1=200, bucket2=250 cumulative
    # Should convert to: bucket0=100, bucket1=100, bucket2=50 absolute
    cumulative_buckets = [
        HistogramBucket(0.0, 10.0, 100, 0),  # 100 items in first bucket
        HistogramBucket(10.0, 20.0, 200, 0),  # 200 cumulative (100 in this bucket)
        HistogramBucket(20.0, 30.0, 250, 0),  # 250 cumulative (50 in this bucket)
        HistogramBucket(
            30.0, 40.0, 250, 0
        ),  # 250 cumulative (0 in this bucket - should be filtered)
    ]

    # Test with is_cumulative_count=True (TiDB style)
    sampler_cumulative = InverseTransformSampler(
        buckets=cumulative_buckets, is_cumulative_count=True, seed=42
    )

    stats = sampler_cumulative.get_statistics()
    print("  Cumulative conversion:")
    print("    Original buckets: 4")
    print(f"    Converted buckets: {stats['num_buckets']}")
    print(f"    Total count: {stats['total_count']}")
    print(f"    Bucket ranges: {stats['bucket_ranges']}")

    # Verify the conversion worked correctly
    expected_counts = [100, 100, 50]  # Absolute counts after conversion
    actual_counts = [count for _, _, count in stats["bucket_ranges"]]

    assert len(actual_counts) == 3, f"Expected 3 buckets, got {len(actual_counts)}"
    assert (
        actual_counts == expected_counts
    ), f"Expected {expected_counts}, got {actual_counts}"
    assert (
        stats["total_count"] == 250
    ), f"Expected total 250, got {stats['total_count']}"

    print(f"    ✅ Conversion correct: {actual_counts}")

    # Test with is_cumulative_count=False (standard histogram style)
    absolute_buckets = [
        HistogramBucket(0.0, 10.0, 100, 0),  # 100 items (absolute)
        HistogramBucket(10.0, 20.0, 100, 0),  # 100 items (absolute)
        HistogramBucket(20.0, 30.0, 50, 0),  # 50 items (absolute)
    ]

    sampler_absolute = InverseTransformSampler(
        buckets=absolute_buckets,
        is_cumulative_count=False,  # Standard histogram format
        seed=42,
    )

    stats_abs = sampler_absolute.get_statistics()
    print("  Absolute count (no conversion):")
    print(f"    Buckets: {stats_abs['num_buckets']}")
    print(f"    Total count: {stats_abs['total_count']}")
    print(f"    Bucket ranges: {stats_abs['bucket_ranges']}")

    # Both should produce the same final distribution
    assert (
        stats["bucket_ranges"] == stats_abs["bucket_ranges"]
    ), "Cumulative and absolute should produce same result"

    print("    ✅ Both methods produce identical results")
    return True


def validate_sampling_distribution_correctness():
    """Test that sampling distribution is correct after cumulative conversion"""
    print("\n=== Testing Sampling Distribution Correctness ===")

    # Create cumulative data:
    # bucket0: [0-10] with 1000 items
    # bucket1: [10-20] with 500 items
    # bucket2: [20-30] with 200 items
    cumulative_buckets = [
        HistogramBucket(0.0, 10.0, 1000, 0),  # 1000 cumulative
        HistogramBucket(10.0, 20.0, 1500, 0),  # 1500 cumulative (500 in this bucket)
        HistogramBucket(20.0, 30.0, 1700, 0),  # 1700 cumulative (200 in this bucket)
    ]

    sampler = InverseTransformSampler(
        buckets=cumulative_buckets, is_cumulative_count=True, seed=42
    )

    # Generate samples
    sample_results = sampler.sample_batch(10000)
    non_null_samples = [s.value for s in sample_results if s.value is not None]

    # Count samples in each range
    range_counts = [0, 0, 0]
    for sample in non_null_samples:
        if 0 <= sample < 10:
            range_counts[0] += 1
        elif 10 <= sample < 20:
            range_counts[1] += 1
        elif 20 <= sample < 30:
            range_counts[2] += 1

    total_samples = sum(range_counts)
    percentages = [count / total_samples * 100 for count in range_counts]

    print(f"  Sample distribution (n={total_samples}):")
    print(
        f"    Range [0-10]:  {range_counts[0]} samples ({percentages[0]:.1f}%) - expected ~58.8%"
    )
    print(
        f"    Range [10-20]: {range_counts[1]} samples ({percentages[1]:.1f}%) - expected ~29.4%"
    )
    print(
        f"    Range [20-30]: {range_counts[2]} samples ({percentages[2]:.1f}%) - expected ~11.8%"
    )

    # Verify approximate distribution (with some tolerance)
    expected_percentages = [58.8, 29.4, 11.8]  # 1000/1700, 500/1700, 200/1700

    for i, (actual, expected) in enumerate(zip(percentages, expected_percentages)):
        diff = abs(actual - expected)
        assert (
            diff < 3.0
        ), f"Range {i} percentage {actual:.1f}% differs from expected {expected:.1f}% by {diff:.1f}%"

    print("    ✅ Distribution matches expected proportions")
    return True


def validate_tuple_format_with_cumulative():
    """Test tuple format with cumulative count conversion"""
    print("\n=== Testing Tuple Format with Cumulative Counts ===")

    # Test with TiDB format (like from prepare.py)
    cumulative_buckets = [
        TiDBNumericBucket(
            data_type="INT",
            lower_bound=0,
            upper_bound=10,
            count=500,
            repeats=0,
            ndv=100,
        ),  # 500 cumulative
        TiDBNumericBucket(
            data_type="INT",
            lower_bound=10,
            upper_bound=20,
            count=800,
            repeats=0,
            ndv=100,
        ),  # 800 cumulative (300 in this bucket)
        TiDBNumericBucket(
            data_type="INT",
            lower_bound=20,
            upper_bound=30,
            count=1000,
            repeats=0,
            ndv=100,
        ),  # 1000 cumulative (200 in this bucket)
    ]

    # Test IntSampler with cumulative counts (default behavior)
    int_sampler = IntSampler(
        cumulative_buckets,
        is_cumulative_count=True,
    )  # Should use is_cumulative_count=True by default

    samples = int_sampler.sample_batch(1000)
    non_null_samples = [s for s in samples if s is not None]

    print("  IntSampler with tuple format:")
    print(f"    Generated {len(non_null_samples)} samples")
    print(f"    Value range: {min(non_null_samples)} to {max(non_null_samples)}")

    # Check distribution
    range_counts = [0, 0, 0]
    for sample in non_null_samples:
        if 0 <= sample < 10:
            range_counts[0] += 1
        elif 10 <= sample < 20:
            range_counts[1] += 1
        elif 20 <= sample < 30:
            range_counts[2] += 1

    total = sum(range_counts)
    percentages = [count / total * 100 if total > 0 else 0 for count in range_counts]

    print(
        f"    Distribution: {percentages[0]:.1f}% | {percentages[1]:.1f}% | {percentages[2]:.1f}%"
    )
    print("    Expected:     50.0% | 30.0% | 20.0%")

    # Verify reasonable distribution
    assert (
        percentages[0] > percentages[1] > percentages[2]
    ), "First bucket should have most samples"

    print("    ✅ Tuple format with cumulative conversion works")
    return True


def validate_edge_cases():
    """Test edge cases for cumulative count conversion"""
    print("\n=== Testing Edge Cases ===")

    # Test empty buckets
    print("  Testing empty buckets...")
    empty_sampler = InverseTransformSampler(buckets=[], is_cumulative_count=True)
    sample_results = empty_sampler.sample_batch(10)
    samples = [s.value for s in sample_results]
    assert all(s is None for s in samples), "Empty buckets should return all nulls"
    print("    ✅ Empty buckets handled correctly")

    # Test single bucket
    print("  Testing single bucket...")
    single_bucket = [HistogramBucket(0.0, 10.0, 100, 0)]
    single_sampler = InverseTransformSampler(
        buckets=single_bucket, is_cumulative_count=True
    )
    sample_results = single_sampler.sample_batch(20)
    samples = [s.value for s in sample_results]
    non_null_samples = [s for s in samples if s is not None]
    assert all(
        0 <= s <= 10 for s in non_null_samples
    ), "Samples should be in range [0, 10]"
    print(f"    ✅ Single bucket: {len(non_null_samples)} samples in correct range")

    # Test all zero increases (pathological TiDB case)
    print("  Testing all same cumulative count (zero increases)...")
    zero_increase_buckets = [
        HistogramBucket(0.0, 10.0, 100, 0),
        HistogramBucket(10.0, 20.0, 100, 0),  # Same count = 0 increase
        HistogramBucket(20.0, 30.0, 100, 0),  # Same count = 0 increase
    ]
    zero_sampler = InverseTransformSampler(
        buckets=zero_increase_buckets, is_cumulative_count=True
    )
    stats = zero_sampler.get_statistics()

    print(f"    Original buckets: 3, Converted buckets: {stats['num_buckets']}")
    assert (
        stats["num_buckets"] == 1
    ), "Should only keep first bucket with positive count"

    sample_results = zero_sampler.sample_batch(10)
    samples = [s.value for s in sample_results]
    non_null_samples = [s for s in samples if s is not None]
    assert all(
        0 <= s <= 10 for s in non_null_samples
    ), "All samples should be from first bucket"
    print("    ✅ Zero increases handled: only first bucket kept")

    return True


def visualize_cumulative_conversion(output_dir: str = "tests"):
    """Create visualization showing cumulative count conversion effect"""
    print("\n=== Creating Cumulative Count Conversion Visualization ===")

    if not HAS_MATPLOTLIB:
        print("matplotlib not available - creating text-based analysis only")
        create_text_analysis(output_dir)
        return

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

    # Create test data showing the conversion effect
    cumulative_buckets = [
        HistogramBucket(0.0, 10.0, 300, 0),  # 300 cumulative
        HistogramBucket(10.0, 20.0, 600, 0),  # 600 cumulative (300 in this bucket)
        HistogramBucket(20.0, 30.0, 800, 0),  # 800 cumulative (200 in this bucket)
        HistogramBucket(30.0, 40.0, 850, 0),  # 850 cumulative (50 in this bucket)
    ]

    # Test both conversion methods
    sampler_cumulative = InverseTransformSampler(
        buckets=cumulative_buckets, is_cumulative_count=True, seed=42  # TiDB style
    )

    # Generate samples from both
    sample_results_cumulative = sampler_cumulative.sample_batch(5000)
    samples_cumulative = [s.value for s in sample_results_cumulative]
    samples_cumulative = [s for s in samples_cumulative if s is not None]

    # Create figure
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    fig.suptitle("Cumulative Count Conversion - Before vs After", fontsize=16)

    # 1. Original cumulative counts (bar chart)
    ax1 = axes[0, 0]
    bucket_centers = [5, 15, 25, 35]
    cumulative_counts = [300, 600, 800, 850]

    ax1.bar(
        bucket_centers,
        cumulative_counts,
        width=8,
        alpha=0.7,
        color="orange",
        edgecolor="black",
    )
    ax1.set_xlabel("Bucket Center")
    ax1.set_ylabel("Cumulative Count")
    ax1.set_title("Original TiDB Cumulative Counts")
    ax1.grid(True, alpha=0.3)

    # Add count labels on bars
    for center, count in zip(bucket_centers, cumulative_counts):
        ax1.text(center, count + 10, str(count), ha="center", va="bottom")

    # 2. Converted absolute counts
    ax2 = axes[0, 1]
    absolute_counts = [300, 300, 200, 50]  # After conversion

    ax2.bar(
        bucket_centers,
        absolute_counts,
        width=8,
        alpha=0.7,
        color="green",
        edgecolor="black",
    )
    ax2.set_xlabel("Bucket Center")
    ax2.set_ylabel("Absolute Count")
    ax2.set_title("Converted Absolute Counts")
    ax2.grid(True, alpha=0.3)

    # Add count labels on bars
    for center, count in zip(bucket_centers, absolute_counts):
        ax2.text(center, count + 5, str(count), ha="center", va="bottom")

    # 3. Sample distribution histogram
    ax3 = axes[1, 0]

    bins = np.arange(0, 45, 5)
    ax3.hist(samples_cumulative, bins=bins, alpha=0.7, color="blue", edgecolor="black")
    ax3.set_xlabel("Sample Value")
    ax3.set_ylabel("Sample Count")
    ax3.set_title("Generated Sample Distribution")
    ax3.grid(True, alpha=0.3)

    # 4. Conversion process illustration
    ax4 = axes[1, 1]

    # Show the conversion process
    bucket_labels = ["[0-10]", "[10-20]", "[20-30]", "[30-40]"]
    x_pos = np.arange(len(bucket_labels))

    width = 0.35
    ax4.bar(
        x_pos - width / 2,
        cumulative_counts,
        width,
        label="Cumulative",
        alpha=0.7,
        color="orange",
    )
    ax4.bar(
        x_pos + width / 2,
        absolute_counts,
        width,
        label="Absolute",
        alpha=0.7,
        color="green",
    )

    ax4.set_xlabel("Bucket Range")
    ax4.set_ylabel("Count")
    ax4.set_title("Conversion Comparison")
    ax4.set_xticks(x_pos)
    ax4.set_xticklabels(bucket_labels)
    ax4.legend()
    ax4.grid(True, alpha=0.3)

    plt.tight_layout()

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

    # Print analysis
    print("\nConversion Analysis:")
    print(f"  Original cumulative: {cumulative_counts}")
    print(f"  Converted absolute:  {absolute_counts}")
    print(f"  Total samples generated: {len(samples_cumulative)}")

    # Verify distribution matches expected proportions
    range_counts = [0, 0, 0, 0]
    for sample in samples_cumulative:
        if 0 <= sample < 10:
            range_counts[0] += 1
        elif 10 <= sample < 20:
            range_counts[1] += 1
        elif 20 <= sample < 30:
            range_counts[2] += 1
        elif 30 <= sample < 40:
            range_counts[3] += 1

    total_expected = sum(absolute_counts)
    expected_ratios = [count / total_expected for count in absolute_counts]
    actual_ratios = [count / len(samples_cumulative) for count in range_counts]

    print("\nDistribution Verification:")
    for i, (bucket, expected, actual) in enumerate(
        zip(bucket_labels, expected_ratios, actual_ratios)
    ):
        print(
            f"  {bucket}: Expected {expected:.3f}, Actual {actual:.3f} (diff: {abs(expected - actual):.3f})"
        )


def create_text_analysis(output_dir: str):
    """Create text-based analysis when matplotlib is not available"""
    report_path = os.path.join(output_dir, "cumulative_count_analysis.txt")

    with open(report_path, "w") as f:
        f.write("Cumulative Count Conversion Analysis\n")
        f.write("=" * 50 + "\n\n")

        f.write("TiDB uses cumulative counts in histogram buckets:\n")
        f.write("  Bucket 1: count=300 (300 items in this bucket)\n")
        f.write("  Bucket 2: count=600 (300 items in this bucket, 600 total)\n")
        f.write("  Bucket 3: count=800 (200 items in this bucket, 800 total)\n")
        f.write("  Bucket 4: count=850 (50 items in this bucket, 850 total)\n\n")

        f.write("After conversion to absolute counts:\n")
        f.write("  Bucket 1: count=300 (300 items)\n")
        f.write("  Bucket 2: count=300 (300 items)\n")
        f.write("  Bucket 3: count=200 (200 items)\n")
        f.write("  Bucket 4: count=50 (50 items)\n\n")

        f.write("This conversion ensures correct probability distribution\n")
        f.write("when sampling from the histogram buckets.\n")

    print(f"Text analysis saved to: {report_path}")


def run_cumulative_count_tests():
    """Run all cumulative count tests"""
    print("🚀 Starting Cumulative Count Test Suite")
    print("=" * 60)

    try:
        validate_cumulative_count_conversion()
        validate_sampling_distribution_correctness()
        validate_tuple_format_with_cumulative()
        validate_edge_cases()

        # Create visualization
        visualize_cumulative_conversion()

        print("\n" + "=" * 60)
        print("🎉 All Cumulative Count Tests Passed!")

        print("\n📊 Test Summary:")
        print("  ✅ Cumulative to absolute count conversion")
        print("  ✅ Sampling distribution correctness")
        print("  ✅ Tuple format compatibility")
        print("  ✅ Edge case handling")
        print("  ✅ TiDB format support with is_cumulative_count=True (default)")
        print("  ✅ Standard histogram support with is_cumulative_count=False")

        return True

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

        traceback.print_exc()
        return False


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