"""
Comprehensive test suite for InverseTransformSampler
Combines functionality tests, performance tests, and visualizations
"""

import os
import random
import time

import matplotlib.pyplot as plt
import numpy as np

from samplers import HistogramBucket, InverseTransformSampler, TopNItem


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


def create_validate_distribution(seed=None):
    """Create different test distributions based on seed to test various distribution shapes"""
    if seed is None:
        seed = get_validate_seed()

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

    if pattern == 0:  # Normal-like distribution (concentrated in middle)
        buckets = [
            HistogramBucket(0, 10, 50, 5),
            HistogramBucket(10, 20, 100, 10),
            HistogramBucket(20, 30, 200, 20),
            HistogramBucket(30, 40, 500, 50),
            HistogramBucket(40, 45, 800, 80),  # Peak
            HistogramBucket(45, 50, 1200, 120),  # Peak
            HistogramBucket(50, 55, 1500, 150),  # Peak
            HistogramBucket(55, 60, 1000, 100),
            HistogramBucket(60, 65, 600, 60),
            HistogramBucket(65, 70, 400, 40),
            HistogramBucket(70, 80, 200, 20),
            HistogramBucket(80, 90, 100, 10),
            HistogramBucket(90, 100, 50, 5),
        ]
        top_n = [
            TopNItem(42.5, 300),
            TopNItem(47.5, 400),
            TopNItem(52.5, 500),
            TopNItem(57.5, 350),
        ]

    elif pattern == 1:  # Left-skewed distribution (peak at start)
        buckets = [
            HistogramBucket(0, 10, 2000, 200),  # Peak at start
            HistogramBucket(10, 20, 1500, 150),
            HistogramBucket(20, 30, 1000, 100),
            HistogramBucket(30, 40, 800, 80),
            HistogramBucket(40, 50, 600, 60),
            HistogramBucket(50, 60, 400, 40),
            HistogramBucket(60, 70, 200, 20),
            HistogramBucket(70, 80, 100, 10),
            HistogramBucket(80, 90, 50, 5),
            HistogramBucket(90, 100, 25, 2),
        ]
        top_n = [TopNItem(5.0, 500), TopNItem(15.0, 300), TopNItem(25.0, 200)]

    elif pattern == 2:  # Right-skewed distribution (peak at end)
        buckets = [
            HistogramBucket(0, 10, 25, 2),
            HistogramBucket(10, 20, 50, 5),
            HistogramBucket(20, 30, 100, 10),
            HistogramBucket(30, 40, 200, 20),
            HistogramBucket(40, 50, 400, 40),
            HistogramBucket(50, 60, 600, 60),
            HistogramBucket(60, 70, 800, 80),
            HistogramBucket(70, 80, 1000, 100),
            HistogramBucket(80, 90, 1500, 150),  # Peak at end
            HistogramBucket(90, 100, 2000, 200),
        ]
        top_n = [TopNItem(85.0, 400), TopNItem(92.5, 300), TopNItem(97.0, 500)]

    else:  # Bimodal distribution (two peaks)
        buckets = [
            HistogramBucket(0, 10, 100, 10),
            HistogramBucket(10, 20, 800, 80),  # First peak
            HistogramBucket(20, 30, 1200, 120),  # First peak
            HistogramBucket(30, 40, 400, 40),
            HistogramBucket(40, 50, 200, 20),  # Valley
            HistogramBucket(50, 60, 150, 15),  # Valley
            HistogramBucket(60, 70, 300, 30),
            HistogramBucket(70, 80, 1000, 100),  # Second peak
            HistogramBucket(80, 90, 1500, 150),  # Second peak
            HistogramBucket(90, 100, 600, 60),
        ]
        top_n = [
            TopNItem(20.0, 250),
            TopNItem(25.0, 200),
            TopNItem(75.0, 300),
            TopNItem(85.0, 350),
        ]

    return buckets, top_n


def create_large_validate_distribution():
    """Create a large realistic distribution for performance testing"""
    # Create many buckets to simulate real TiDB statistics
    buckets = []
    for i in range(100):  # 100 buckets
        lower = i * 10.0
        upper = (i + 1) * 10.0
        count = random.randint(50, 500)
        repeats = random.randint(0, count // 10)
        buckets.append(HistogramBucket(lower, upper, count, repeats))

    # Create many TopN items
    top_n = []
    for i in range(50):  # 50 TopN items
        value = random.uniform(0, 1000)
        count = random.randint(100, 1000)
        top_n.append(TopNItem(value, count))

    return buckets, top_n


def generate_original_samples(buckets, top_n, n_samples=10000):
    """Generate samples from the original distribution for comparison"""
    samples = []

    # Calculate total count
    total_count = sum(b.count for b in buckets) + sum(t.count for t in top_n)

    # Generate samples
    for _ in range(n_samples):
        u = random.random() * total_count
        cumulative = 0

        # Check TopN first
        for item in top_n:
            cumulative += item.count
            if u <= cumulative:
                samples.append(item.value)
                break
        else:
            # Check buckets
            for bucket in buckets:
                cumulative += bucket.count
                if u <= cumulative:
                    # Sample within bucket
                    if bucket.repeats > 0 and random.random() < (
                        bucket.repeats / bucket.count
                    ):
                        samples.append(bucket.upper_bound)
                    else:
                        samples.append(
                            random.uniform(bucket.lower_bound, bucket.upper_bound)
                        )
                    break

    return samples


def validate_inverse_transform_sampler():
    """Test the inverse transform sampler with visualization"""
    print("=== Testing Inverse Transform Sampler ===")

    # Create test distribution based on seed
    seed = get_validate_seed()
    print(f"🎲 Using test seed: {seed}")
    buckets, top_n = create_validate_distribution(seed)

    # Create sampler with absolute counts (not cumulative)
    sampler = InverseTransformSampler(
        buckets=buckets,
        top_n=top_n,
        null_count=100,
        is_cumulative_count=False,  # Test with absolute counts
        seed=seed,
    )

    # Print statistics
    stats = sampler.get_statistics()
    print("\nSampler Statistics:")
    print(f"  Number of buckets: {stats['num_buckets']}")
    print(f"  Number of TopN values: {stats['num_topn']}")
    print(f"  Total count: {stats['total_count']}")
    print(f"  Null probability: {stats['null_probability']:.3f}")

    # Generate samples
    n_samples = 10000
    sample_results = sampler.sample_batch(n_samples)

    # Count nulls
    null_count = sum(1 for s in sample_results if s.value is None)
    non_null_samples = [s.value for s in sample_results if s.value is not None]

    print("\nSampling Results:")
    print(f"  Total samples: {n_samples}")
    print(f"  Null samples: {null_count} ({null_count / n_samples * 100:.1f}%)")
    print(f"  Non-null samples: {len(non_null_samples)}")

    # Generate original samples for comparison
    original_samples = generate_original_samples(buckets, top_n, n_samples)

    # Create visualization
    plt.figure(figsize=(15, 10))

    # Plot 1: Histogram comparison
    plt.subplot(2, 2, 1)
    bins = np.linspace(0, 100, 50)
    plt.hist(
        original_samples,
        bins=bins,
        alpha=0.5,
        density=True,
        label="Original",
        color="blue",
    )
    plt.hist(
        non_null_samples,
        bins=bins,
        alpha=0.5,
        density=True,
        label="Generated",
        color="red",
    )
    plt.xlabel("Value")
    plt.ylabel("Density")
    plt.title("Histogram Comparison")
    plt.legend()
    plt.grid(True, alpha=0.3)

    # Plot 2: Cumulative distribution
    plt.subplot(2, 2, 2)
    orig_sorted = np.sort(original_samples)
    gen_sorted = np.sort(non_null_samples)
    orig_cdf = np.arange(1, len(orig_sorted) + 1) / len(orig_sorted)
    gen_cdf = np.arange(1, len(gen_sorted) + 1) / len(gen_sorted)
    plt.plot(orig_sorted, orig_cdf, label="Original", linewidth=2)
    plt.plot(gen_sorted, gen_cdf, label="Generated", linewidth=2, linestyle="--")
    plt.xlabel("Value")
    plt.ylabel("Cumulative Probability")
    plt.title("Cumulative Distribution Function")
    plt.legend()
    plt.grid(True, alpha=0.3)

    # Plot 3: Bucket visualization
    plt.subplot(2, 2, 3)
    bucket_centers = []
    bucket_counts = []
    bucket_widths = []

    for bucket in buckets:
        center = (bucket.lower_bound + bucket.upper_bound) / 2
        bucket_centers.append(center)
        bucket_counts.append(bucket.count)
        bucket_widths.append(bucket.upper_bound - bucket.lower_bound)

    if bucket_centers:
        plt.bar(
            bucket_centers,
            bucket_counts,
            width=bucket_widths,
            alpha=0.7,
            color="green",
            edgecolor="black",
        )
        plt.xlabel("Value")
        plt.ylabel("Count")
        plt.title("Histogram Buckets")
        plt.grid(True, alpha=0.3)

    # Plot 4: Statistics
    plt.subplot(2, 2, 4)
    stats_text = f"""Statistics Comparison:

Original Distribution:
  Mean: {np.mean(original_samples):.2f}
  Std: {np.std(original_samples):.2f}
  Min: {np.min(original_samples):.2f}
  Max: {np.max(original_samples):.2f}

Generated Distribution:
  Mean: {np.mean(non_null_samples):.2f}
  Std: {np.std(non_null_samples):.2f}
  Min: {np.min(non_null_samples):.2f}
  Max: {np.max(non_null_samples):.2f}

Sampler Info:
  Buckets: {stats['num_buckets']}
  TopN: {stats['num_topn']}
  Total Count: {stats['total_count']:,}
  Null Rate: {stats['null_probability']:.3f}
"""

    plt.text(
        0.05,
        0.95,
        stats_text,
        fontsize=9,
        verticalalignment="top",
        fontfamily="monospace",
        transform=plt.gca().transAxes,
    )
    plt.xlim(0, 1)
    plt.ylim(0, 1)
    plt.axis("off")
    plt.title("Statistics Summary")

    plt.tight_layout()
    plt.savefig("tests/its_comprehensive_test.png", dpi=300, bbox_inches="tight")
    print("  📊 Visualization saved to: tests/its_comprehensive_test.png")

    # Calculate statistics
    print("\nDistribution Statistics:")
    print(
        f"  Original - Mean: {np.mean(original_samples):.2f}, Std: {np.std(original_samples):.2f}"
    )
    print(
        f"  Generated - Mean: {np.mean(non_null_samples):.2f}, Std: {np.std(non_null_samples):.2f}"
    )

    # Test TopN accuracy
    topn_values = [t.value for t in top_n]
    topn_in_samples = sum(1 for s in non_null_samples if s in topn_values)
    print("\nTopN Values:")
    print(
        f"  Expected TopN ratio: {sum(t.count for t in top_n) / stats['total_count']:.3f}"
    )
    print(f"  Actual TopN ratio: {topn_in_samples / len(non_null_samples):.3f}")

    plt.show()
    return sampler


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

    # Test empty distribution
    sampler = InverseTransformSampler(
        buckets=[], top_n=[], null_count=0, is_cumulative_count=False
    )
    samples = sampler.sample_batch(10)
    print(f"Empty distribution samples: {samples}")

    # Test only nulls
    sampler = InverseTransformSampler(
        buckets=[], top_n=[], null_count=100, is_cumulative_count=False
    )
    samples = sampler.sample_batch(10)
    print(f"Only nulls samples: {samples}")

    # Test only TopN
    top_n = [TopNItem(10.0, 50), TopNItem(20.0, 50)]
    sampler = InverseTransformSampler(
        buckets=[], top_n=top_n, null_count=0, is_cumulative_count=False
    )
    samples = sampler.sample_batch(20)
    print(f"Only TopN samples: {samples}")

    # Test single bucket
    buckets = [HistogramBucket(0, 100, 100, 10)]
    sampler = InverseTransformSampler(
        buckets=buckets, top_n=[], null_count=0, is_cumulative_count=False
    )
    samples = sampler.sample_batch(20)
    print(f"Single bucket samples: {samples[:10]}...")


def benchmark_single_sampling(
    sampler: InverseTransformSampler, num_samples: int = 100000
):
    """Benchmark single sample() calls - simulates row-by-row database insertion"""
    print(f"  Benchmarking {num_samples:,} individual sample() calls...")

    start_time = time.time()
    samples = []

    # Simulate row-by-row sampling like in prepare.py
    for _ in range(num_samples):
        sample_result = sampler.sample()
        samples.append(sample_result.value)

    end_time = time.time()
    elapsed = end_time - start_time

    non_null_count = sum(1 for s in samples if s is not None)
    samples_per_second = num_samples / elapsed

    print(f"    Time: {elapsed:.4f} seconds")
    print(f"    Rate: {samples_per_second:,.0f} samples/second")
    print(f"    Non-null samples: {non_null_count:,}/{num_samples:,}")

    return elapsed, samples_per_second, samples


def benchmark_batch_sampling(
    sampler: InverseTransformSampler, num_samples: int = 100000
):
    """Benchmark batch sample_batch() calls for comparison"""
    print(f"  Benchmarking batch sampling of {num_samples:,} samples...")

    start_time = time.time()
    samples = sampler.sample_batch(num_samples)
    end_time = time.time()

    elapsed = end_time - start_time
    non_null_count = sum(1 for s in samples if s is not None)
    samples_per_second = num_samples / elapsed

    print(f"    Time: {elapsed:.4f} seconds")
    print(f"    Rate: {samples_per_second:,.0f} samples/second")
    print(f"    Non-null samples: {non_null_count:,}/{num_samples:,}")

    return elapsed, samples_per_second, samples


def validate_performance():
    """Test performance with large distribution"""
    print("\n=== Testing Performance ===")

    # Create large distribution
    buckets, top_n = create_large_validate_distribution()
    sampler = InverseTransformSampler(
        buckets=buckets,
        top_n=top_n,
        null_count=1000,
        is_cumulative_count=False,  # Use absolute counts
        seed=42,
    )

    stats = sampler.get_statistics()
    optimization_info = stats["optimization_info"]

    print("  Distribution size:")
    print(f"    Buckets: {stats['num_buckets']}")
    print(f"    TopN items: {stats['num_topn']}")
    print(f"    Total count: {stats['total_count']:,}")

    print("  Optimized data structures:")
    print(f"    CDF entries: {optimization_info['cdf_entries']}")
    print(
        f"    Pre-computed repeat probabilities: {optimization_info['pre_computed_repeat_probs']}"
    )

    # Benchmark single sampling
    print("\n  Single Sample Performance:")
    elapsed, sps, samples = benchmark_single_sampling(sampler, 100000)

    # Benchmark batch sampling
    print("\n  Batch Sample Performance:")
    batch_elapsed, batch_sps, _ = benchmark_batch_sampling(sampler, 100000)

    print("\n  Performance Comparison:")
    print(f"    Single sampling: {sps:,.0f} samples/sec")
    print(f"    Batch sampling: {batch_sps:,.0f} samples/sec")
    print(f"    Single vs Batch ratio: {sps / batch_sps:.2f}")

    return sampler, sps


def validate_row_by_row_simulation():
    """Simulate the exact usage pattern from prepare.py"""
    print("\n=== Row-by-Row Database Insertion Simulation ===")

    # Create multiple samplers like in metadata/schema_manager.py column_sampler_dict
    samplers = {}

    # Simulate different column types
    for i in range(5):
        buckets, top_n = create_large_validate_distribution()
        sampler = InverseTransformSampler(
            buckets=buckets,
            top_n=top_n,
            null_count=100,
            is_cumulative_count=False,
            seed=42 + i,
        )
        samplers[f"table.column_{i}"] = sampler

    num_rows = 10000
    print(
        f"  Simulating {num_rows:,} database rows with {len(samplers)} columns each..."
    )

    start_time = time.time()

    # Simulate the exact pattern from prepare.py lines 144-152
    for row_idx in range(num_rows):
        row_values = []
        for column_name, sampler in samplers.items():
            # This mimics: column_sampler_dict[(table_schema, table_name, cn)].sample()
            sample_result = sampler.sample()
            value = sample_result.value
            row_values.append(value)

        # Simulate database insertion (just append to simulate I/O time)
        # In real prepare.py this would be: insert_data(conn, sql_template, row_values)

    end_time = time.time()
    elapsed = end_time - start_time

    total_samples = num_rows * len(samplers)
    rows_per_second = num_rows / elapsed
    samples_per_second = total_samples / elapsed

    print(f"  Total time: {elapsed:.4f} seconds")
    print(f"  Rows per second: {rows_per_second:,.0f}")
    print(f"  Total samples: {total_samples:,}")
    print(f"  Samples per second: {samples_per_second:,.0f}")

    print("  ✅ Row-by-row pattern simulation completed")
    return elapsed, rows_per_second


def validate_cumulative_count_feature():
    """Test the cumulative count conversion feature"""
    print("\n=== Testing Cumulative Count Conversion ===")

    # Create test data with TiDB-style cumulative counts
    cumulative_buckets = [
        HistogramBucket(0.0, 10.0, 100, 0),  # 100 cumulative
        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 - filtered)
    ]

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

    # 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, seed=42
    )

    # Compare statistics
    stats_cum = sampler_cumulative.get_statistics()
    stats_abs = sampler_absolute.get_statistics()

    print("  Cumulative conversion:")
    print(f"    Original buckets: 4 -> Converted buckets: {stats_cum['num_buckets']}")
    print(f"    Total count: {stats_cum['total_count']}")
    print(f"    Bucket counts: {[count for _, _, count in stats_cum['bucket_ranges']]}")

    print("  Absolute (no conversion):")
    print(f"    Buckets: {stats_abs['num_buckets']}")
    print(f"    Total count: {stats_abs['total_count']}")
    print(f"    Bucket counts: {[count for _, _, count in stats_abs['bucket_ranges']]}")

    # Verify they produce the same result
    if stats_cum["bucket_ranges"] == stats_abs["bucket_ranges"]:
        print("    ✅ Both methods produce identical distributions")
    else:
        print("    ❌ Methods produce different distributions")

    return True


def run_comprehensive_its_test():
    """Run comprehensive test suite combining functionality and performance"""
    print("🚀 Starting Comprehensive ITS Test Suite")
    print("=" * 80)

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

    try:
        # 1. Basic functionality test with visualization
        sampler = validate_inverse_transform_sampler()

        # 2. Edge cases
        validate_edge_cases()

        # 3. Cumulative count conversion test
        validate_cumulative_count_feature()

        # 4. Performance tests
        perf_sampler, sps = validate_performance()

        # 5. Row-by-row simulation
        db_elapsed, rows_per_sec = validate_row_by_row_simulation()

        print("\n" + "=" * 80)
        print("🎉 Comprehensive ITS Test Suite Completed!")

        print("\n📊 Summary:")
        print("  ✅ Basic functionality: Distribution preservation verified")
        print("  ✅ Edge cases: All handled correctly")
        print("  ✅ Cumulative count conversion: Working properly")
        print(f"  ✅ Performance: {sps:,.0f} samples/second")
        print(f"  ✅ Database simulation: {rows_per_sec:,.0f} rows/second")
        print("  ✅ Visualizations: Generated for analysis")

        print("\n🎯 Key Features Validated:")
        print("  • Inverse transform sampling accuracy")
        print("  • TiDB cumulative count support")
        print("  • Row-by-row sampling efficiency")
        print("  • Multiple input format support")
        print("  • Optimized data structures")
        print("  • Statistical distribution preservation")

        return True

    except ImportError as e:
        if "matplotlib" in str(e):
            print("⚠️  matplotlib not available - skipping visual tests")
            print("   Install with: pip install matplotlib")
            # Run non-visual tests
            validate_edge_cases()
            validate_cumulative_count_feature()
            buckets, top_n = create_large_validate_distribution()
            sampler = InverseTransformSampler(
                buckets=buckets, top_n=top_n, null_count=1000, is_cumulative_count=False
            )
            benchmark_single_sampling(sampler, 50000)
            return True
        else:
            raise e

    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_comprehensive_its_test()
    exit(0 if success else 1)