"""
Test suite for FloatSampler with comprehensive validation
"""

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

import matplotlib.pyplot as plt
import numpy as np

from samplers import FloatSampler
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 generate_validate_float_data(
    size: int = 1000, seed: Optional[int] = None
) -> List[Optional[float]]:
    """Generate test float data with different distribution patterns based on seed"""
    if seed is None:
        seed = get_validate_seed()
    random.seed(seed)
    np.random.seed(seed)

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

    if pattern == 0:  # Normal-like distribution around 0
        print(f"  🎲 Generating normal-like float distribution (seed {seed})")
        common_floats = [0.0, 1.0, -1.0, 2.0, -2.0]
        # Most data concentrated around 0
        for _ in range(size // 2):
            data.append(random.gauss(0, 5))  # Normal distribution
        for _ in range(size // 4):
            data.append(random.choice(common_floats))

    elif pattern == 1:  # Exponential-like distribution (many small values)
        print(f"  🎲 Generating exponential-like float distribution (seed {seed})")
        common_floats = [0.1, 0.5, 1.0, 2.0, 5.0]
        # Many small positive values
        for _ in range(size // 2):
            data.append(random.expovariate(0.5))  # Exponential distribution
        for _ in range(size // 4):
            data.append(random.choice(common_floats))

    elif pattern == 2:  # Large value distribution
        print(f"  🎲 Generating large-value float distribution (seed {seed})")
        common_floats = [100.0, 1000.0, 10000.0, 50000.0, 100000.0]
        # Concentrated on large values
        for _ in range(size // 2):
            data.append(random.uniform(1000, 100000))  # Large values
        for _ in range(size // 4):
            data.append(random.choice(common_floats))

    else:  # Bimodal distribution (small and large values)
        print(f"  🎲 Generating bimodal float distribution (seed {seed})")
        common_floats = [0.01, 0.1, 1000.0, 10000.0, 50000.0]
        # Two peaks: very small and very large
        for _ in range(size // 4):
            data.append(random.uniform(0.001, 1.0))  # Small peak
        for _ in range(size // 4):
            data.append(random.uniform(10000, 100000))  # Large peak
        for _ in range(size // 4):
            data.append(random.choice(common_floats))

    # Add some nulls (consistent across patterns)
    null_count = size // 20
    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.gauss(0, 10))
        elif pattern == 1:
            data.append(random.expovariate(0.2))
        elif pattern == 2:
            data.append(random.uniform(100, 10000))
        else:
            if random.random() < 0.5:
                data.append(random.uniform(0.001, 1.0))
            else:
                data.append(random.uniform(1000, 50000))

    # Shuffle the data
    random.shuffle(data)
    return data[:size]


def create_validate_float_sampler_from_data(
    float_data: List[Optional[float]], **kwargs
):
    """Helper function to create FloatSampler from float data using TiDB types for testing"""
    num_buckets = kwargs.get("num_buckets", 10)
    top_n_size = kwargs.get("top_n_size", 5)
    precision = kwargs.get("precision", None)
    seed = kwargs.get("seed", None)
    if seed is None:
        seed = get_validate_seed()

    # Process data
    numeric_data = []
    null_count = 0

    for value in float_data:
        if value is None:
            null_count += 1
        else:
            numeric_data.append(float(value))

    # Generate TiDB histogram from numeric data
    tidb_buckets, tidb_top_n = _generate_tidb_histogram_from_numeric_data(
        numeric_data, num_buckets, top_n_size, precision
    )

    # Create FloatSampler with TiDB types
    return FloatSampler(
        buckets=tidb_buckets,
        top_n=tidb_top_n,
        null_count=null_count,
        precision=precision,
        is_cumulative_count=False,
        seed=seed,
    )


def _generate_tidb_histogram_from_numeric_data(
    numeric_data: List[float],
    num_buckets: int,
    top_n_size: int,
    precision: Optional[int] = None,
):
    """Generate TiDB histogram buckets and TopN from numeric data (for testing)"""
    if not numeric_data:
        return [], []

    # Find TopN most frequent values
    value_counts = Counter(numeric_data)
    most_common = value_counts.most_common(top_n_size)

    # Create TiDB numeric TopN items
    tidb_top_n = [
        TiDBNumericTopNItem(
            data_type="FLOAT" if precision is None else "DECIMAL",
            encoded=val,  # Store the actual float value
            count=count,
        )
        for val, count in most_common
    ]

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

    if not histogram_data:
        return [], tidb_top_n

    # Create histogram buckets
    min_val = min(histogram_data)
    max_val = max(histogram_data)

    if min_val == max_val:
        # All values are the same
        tidb_bucket = TiDBNumericBucket(
            data_type="FLOAT" if precision is None else "DECIMAL",
            lower_bound=min_val,
            upper_bound=max_val,
            count=len(histogram_data),
            repeats=len(histogram_data),
            ndv=1,
            precision=precision,
        )
        return [tidb_bucket], tidb_top_n

    # Create bucket boundaries
    bucket_edges = np.linspace(min_val, max_val, num_buckets + 1)
    tidb_buckets = []

    for i in range(num_buckets):
        lower = bucket_edges[i]
        upper = bucket_edges[i + 1]

        # Count values in this bucket
        if i == num_buckets - 1:  # Last bucket includes upper bound
            bucket_values = [val for val in histogram_data if lower <= val <= upper]
        else:
            bucket_values = [val for val in histogram_data if lower <= val < upper]

        if bucket_values:
            # Count repeats of upper bound
            repeats = sum(1 for val in bucket_values if val == upper)

            tidb_bucket = TiDBNumericBucket(
                data_type="FLOAT" if precision is None else "DECIMAL",
                lower_bound=lower,
                upper_bound=upper,
                count=len(bucket_values),
                repeats=repeats,
                ndv=len(set(bucket_values)),
                precision=precision,
            )
            tidb_buckets.append(tidb_bucket)

    return tidb_buckets, tidb_top_n


def validate_float_sampler_basic():
    """Test basic FloatSampler functionality"""
    print("=== Test Basic FloatSampler Functionality ===")

    # Generate test data
    seed = get_validate_seed()
    validate_data = generate_validate_float_data(1000, seed=seed)

    print(f"Generated {len(validate_data)} test floats")
    print(f"Null count: {sum(1 for x in validate_data if x is None)}")
    print(f"Unique floats: {len(set(x for x in validate_data if x is not None))}")

    # Create FloatSampler using test helper
    sampler = create_validate_float_sampler_from_data(
        validate_data, num_buckets=10, top_n_size=5, seed=seed
    )

    # Get statistics from underlying ITS
    stats = sampler.its.get_statistics()
    print("\nFloatSampler Statistics:")
    print(f"  Number of buckets: {stats['num_buckets']}")
    print(f"  Number of TopN: {stats['num_topn']}")
    print(f"  Total count: {stats['total_count']}")
    print(f"  Null probability: {stats['null_probability']:.3f}")
    print(f"  Precision: {sampler.precision}")

    # Sample floats
    samples = sampler.sample_batch(100)
    non_null_samples = [s for s in samples if s is not None]

    print("\nSampling Results:")
    print(f"  Generated {len(samples)} samples")
    print(f"  Non-null samples: {len(non_null_samples)}")
    print(f"  Sample examples: {non_null_samples[:10]}")

    # Verify all samples are floats
    all_floats = all(isinstance(s, float) for s in non_null_samples)
    print(f"  All samples are floats: {all_floats}")

    # Check for reasonable value ranges
    if non_null_samples:
        min_sample = min(non_null_samples)
        max_sample = max(non_null_samples)
        print(f"  Sample range: {min_sample:.6f} to {max_sample:.6f}")

    return validate_data, sampler, samples


def validate_float_sampler_precision():
    """Test FloatSampler with different precision settings (for DECIMAL types)"""
    print("\n=== Test FloatSampler with Precision (DECIMAL) ===")

    # Generate decimal-like data
    decimal_data = [
        12.34,
        56.78,
        90.12,
        34.56,
        78.90,
        123.456,
        789.012,
        345.678,
        901.234,
        1.5,
        2.25,
        3.125,
        4.0625,
        5.03125,
    ] * 10  # Repeat for more data

    seed = get_validate_seed()
    precision_tests = [2, 3, 4]

    for precision in precision_tests:
        print(f"\n--- Testing precision: {precision} ---")

        sampler = create_validate_float_sampler_from_data(
            decimal_data, num_buckets=5, top_n_size=3, precision=precision, seed=seed
        )

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

        print(f"Generated {len(non_null_samples)} samples")
        print(f"Sample examples: {non_null_samples[:8]}")

        # Check precision
        if non_null_samples:
            # Check if all samples have the correct decimal places
            correct_precision = True
            for sample in non_null_samples:
                # Convert to string and count decimal places
                str_sample = (
                    f"{sample:.{precision + 2}f}"  # Add extra precision to check
                )
                if "." in str_sample:
                    decimal_part = str_sample.split(".")[1].rstrip("0")
                    if len(decimal_part) > precision:
                        correct_precision = False
                        break

            print(
                f"Correct precision ({precision} decimal places): {correct_precision}"
            )

            # Show some detailed examples
            for i, sample in enumerate(non_null_samples[:3]):
                print(f"  Sample {i + 1}: {sample} -> {sample:.{precision}f}")


def validate_float_sampler_from_tidb_stats():
    """Test FloatSampler with mock TiDB statistics"""
    print("\n=== Test FloatSampler from TiDB Stats ===")

    seed = get_validate_seed()
    # Mock TiDB float statistics using TiDB types
    tidb_buckets = [
        TiDBNumericBucket(
            data_type="FLOAT",
            lower_bound=0.0,
            upper_bound=1.0,
            count=150,
            repeats=15,
            ndv=50,
        ),
        TiDBNumericBucket(
            data_type="FLOAT",
            lower_bound=1.0,
            upper_bound=10.0,
            count=200,
            repeats=20,
            ndv=80,
        ),
        TiDBNumericBucket(
            data_type="FLOAT",
            lower_bound=10.0,
            upper_bound=100.0,
            count=300,
            repeats=30,
            ndv=120,
        ),
        TiDBNumericBucket(
            data_type="FLOAT",
            lower_bound=100.0,
            upper_bound=1000.0,
            count=250,
            repeats=25,
            ndv=100,
        ),
        TiDBNumericBucket(
            data_type="FLOAT",
            lower_bound=1000.0,
            upper_bound=10000.0,
            count=100,
            repeats=10,
            ndv=40,
        ),
    ]

    tidb_top_n = [
        TiDBNumericTopNItem(data_type="FLOAT", encoded=3.14159, count=80),
        TiDBNumericTopNItem(data_type="FLOAT", encoded=2.71828, count=60),
        TiDBNumericTopNItem(data_type="FLOAT", encoded=1.41421, count=50),
        TiDBNumericTopNItem(data_type="FLOAT", encoded=0.57721, count=40),
    ]

    # Create sampler from TiDB stats
    tidb_sampler = FloatSampler(
        buckets=tidb_buckets,
        top_n=tidb_top_n,
        null_count=30,
        is_cumulative_count=False,
        seed=seed,
    )

    # Get statistics
    stats = tidb_sampler.its.get_statistics()
    print("TiDB FloatSampler Statistics:")
    print(f"  Number of buckets: {stats['num_buckets']}")
    print(f"  Number of TopN: {stats['num_topn']}")
    print(f"  Total count: {stats['total_count']}")

    # Sample floats
    samples = tidb_sampler.sample_batch(100)
    non_null_samples = [s for s in samples if s is not None]

    print("\nSampling Results:")
    print(f"  Generated {len(samples)} samples")
    print(f"  Non-null samples: {len(non_null_samples)}")
    print(f"  Sample examples: {non_null_samples[:15]}")

    # Count TopN occurrences
    topn_values = {3.14159, 2.71828, 1.41421, 0.57721}
    topn_count = sum(1 for s in non_null_samples if s in topn_values)
    print(
        f"  TopN values in samples: {topn_count} / {len(non_null_samples)} ({topn_count / len(non_null_samples) * 100:.1f}%)"
    )

    # Verify all are floats
    all_floats = all(isinstance(s, float) for s in non_null_samples)
    print(f"  All samples are floats: {all_floats}")

    return tidb_sampler, samples


def visualize_float_distributions(
    original_data: List[Optional[float]],
    generated_samples: List[Optional[float]],
    sampler: FloatSampler,
    output_dir: str = "tests",
):
    """Create visualization comparing float distributions"""
    print("\n=== Creating Float Distribution Visualizations ===")

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

    # Filter out nulls
    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 figure with subplots
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle("FloatSampler - Distribution Comparison", fontsize=16)

    # 1. Histogram comparison
    ax1 = axes[0, 0]
    bins = np.linspace(
        min(min(original_clean), min(generated_clean)),
        max(max(original_clean), max(generated_clean)),
        50,
    )

    ax1.hist(
        original_clean,
        bins=bins,
        alpha=0.7,
        label="Original",
        density=True,
        color="blue",
    )
    ax1.hist(
        generated_clean,
        bins=bins,
        alpha=0.7,
        label="Generated",
        density=True,
        color="red",
    )
    ax1.set_xlabel("Float Value")
    ax1.set_ylabel("Density")
    ax1.set_title("Float Distribution Comparison")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # 2. Log-scale histogram (for wide range data)
    ax2 = axes[0, 1]

    # Filter positive values for log scale
    orig_positive = [x for x in original_clean if x > 0]
    gen_positive = [x for x in generated_clean if x > 0]

    if orig_positive and gen_positive:
        log_bins = np.logspace(
            np.log10(min(min(orig_positive), min(gen_positive))),
            np.log10(max(max(orig_positive), max(gen_positive))),
            50,
        )

        ax2.hist(
            orig_positive,
            bins=log_bins,
            alpha=0.7,
            label="Original",
            density=True,
            color="blue",
        )
        ax2.hist(
            gen_positive,
            bins=log_bins,
            alpha=0.7,
            label="Generated",
            density=True,
            color="red",
        )
        ax2.set_xscale("log")
        ax2.set_xlabel("Float Value (log scale)")
        ax2.set_ylabel("Density")
        ax2.set_title("Log-Scale Distribution")
        ax2.legend()
        ax2.grid(True, alpha=0.3)

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

    # Define ranges
    ranges = [
        (-float("inf"), -100, "< -100"),
        (-100, -1, "-100 to -1"),
        (-1, 0, "-1 to 0"),
        (0, 1, "0 to 1"),
        (1, 10, "1 to 10"),
        (10, 100, "10 to 100"),
        (100, 1000, "100 to 1K"),
        (1000, float("inf"), "> 1K"),
    ]

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

    for min_r, max_r, label in ranges:
        orig_count = sum(1 for x in original_clean if min_r <= x < max_r)
        gen_count = sum(1 for x in generated_clean if min_r <= x < max_r)
        if orig_count > 0 or gen_count > 0:  # Only include ranges with data
            orig_range_counts.append(orig_count)
            gen_range_counts.append(gen_count)
            range_labels.append(label)

    if range_labels:
        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, rotation=45, ha="right")
        ax3.legend()
        ax3.grid(True, alpha=0.3)

    # 4. Sampler buckets visualization
    ax4 = axes[1, 1]
    its_stats = sampler.its.get_statistics()

    if its_stats["bucket_ranges"]:
        bucket_centers = []
        bucket_counts = []
        bucket_widths = []

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

        ax4.bar(
            bucket_centers,
            bucket_counts,
            width=bucket_widths,
            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, "float_sampler_comparison.png")
    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    print(f"Visualization saved to: {output_path}")

    # Print statistics
    print("\nDistribution Comparison:")
    print(f"  Original floats: {len(original_clean)}")
    print(f"  Generated floats: {len(generated_clean)}")
    print(f"  Original range: {min(original_clean):.6f} - {max(original_clean):.6f}")
    print(f"  Generated range: {min(generated_clean):.6f} - {max(generated_clean):.6f}")
    print(f"  Original mean: {np.mean(original_clean):.6f}")
    print(f"  Generated mean: {np.mean(generated_clean):.6f}")

    return output_path


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

    seed = get_validate_seed()
    # Test with empty data
    try:
        empty_sampler = FloatSampler(
            [], [], null_count=0, is_cumulative_count=False, seed=seed
        )
        empty_samples = empty_sampler.sample_batch(5)
        print(f"Empty data samples: {empty_samples}")
    except Exception as e:
        print(f"Empty data error: {e}")

    # Test with only nulls
    try:
        null_sampler = FloatSampler(
            [], [], null_count=100, is_cumulative_count=False, seed=seed
        )
        null_samples = null_sampler.sample_batch(5)
        print(f"Only nulls samples: {null_samples}")
    except Exception as e:
        print(f"Only nulls error: {e}")

    # Test with very small floats using TiDB types
    try:
        small_buckets = [
            TiDBNumericBucket(
                data_type="FLOAT",
                lower_bound=1e-10,
                upper_bound=1e-8,
                count=50,
                repeats=5,
                ndv=20,
            )
        ]
        small_topn = [TiDBNumericTopNItem(data_type="FLOAT", encoded=1e-9, count=30)]

        small_sampler = FloatSampler(
            small_buckets,
            small_topn,
            null_count=5,
            is_cumulative_count=False,
            seed=seed,
        )
        small_samples = small_sampler.sample_batch(10)
        non_null_small = [s for s in small_samples if s is not None]
        print(f"Very small float samples: {non_null_small}")

    except Exception as e:
        print(f"Very small float error: {e}")

    # Test with very large floats using TiDB types
    try:
        large_buckets = [
            TiDBNumericBucket(
                data_type="FLOAT",
                lower_bound=1e6,
                upper_bound=1e8,
                count=50,
                repeats=5,
                ndv=30,
            )
        ]
        large_topn = [TiDBNumericTopNItem(data_type="FLOAT", encoded=5e7, count=30)]

        large_sampler = FloatSampler(
            large_buckets,
            large_topn,
            null_count=5,
            is_cumulative_count=False,
            seed=seed,
        )
        large_samples = large_sampler.sample_batch(10)
        non_null_large = [s for s in large_samples if s is not None]
        print(f"Very large float samples: {non_null_large}")

    except Exception as e:
        print(f"Very large float error: {e}")

    # Test with negative floats using TiDB types
    try:
        negative_buckets = [
            TiDBNumericBucket(
                data_type="FLOAT",
                lower_bound=-1000.0,
                upper_bound=-100.0,
                count=40,
                repeats=4,
                ndv=25,
            ),
            TiDBNumericBucket(
                data_type="FLOAT",
                lower_bound=-100.0,
                upper_bound=0.0,
                count=60,
                repeats=6,
                ndv=35,
            ),
        ]
        negative_topn = [
            TiDBNumericTopNItem(data_type="FLOAT", encoded=-273.15, count=20)
        ]  # Absolute zero

        negative_sampler = FloatSampler(
            negative_buckets,
            negative_topn,
            null_count=5,
            is_cumulative_count=False,
            seed=seed,
        )
        negative_samples = negative_sampler.sample_batch(15)
        non_null_negative = [s for s in negative_samples if s is not None]
        print(f"Negative float samples: {non_null_negative}")

        # Verify they're floats and negative/zero
        all_floats = all(isinstance(s, float) for s in non_null_negative)
        all_negative_or_zero = all(s <= 0 for s in non_null_negative)
        print(f"Negative values are floats: {all_floats}")
        print(f"All values <= 0: {all_negative_or_zero}")

    except Exception as e:
        print(f"Negative float error: {e}")

    # Test precision edge case using TiDB types
    try:
        precision_buckets = [
            TiDBNumericBucket(
                data_type="DECIMAL",
                lower_bound=0.1,
                upper_bound=0.9,
                count=100,
                repeats=10,
                ndv=50,
                precision=8,
            )
        ]
        precision_sampler = FloatSampler(
            precision_buckets,
            [],
            null_count=0,
            precision=8,
            is_cumulative_count=False,
            seed=seed,
        )
        precision_samples = precision_sampler.sample_batch(5)
        non_null_precision = [s for s in precision_samples if s is not None]
        print(f"High precision samples (8 decimals): {non_null_precision}")

    except Exception as e:
        print(f"High precision error: {e}")


def run_comprehensive_float_sampler_test():
    """Run comprehensive test suite for FloatSampler"""
    print("🚀 Starting Comprehensive FloatSampler Test Suite")
    print("=" * 60)

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

    # 1. Basic functionality test
    original_data, sampler, generated_samples = validate_float_sampler_basic()

    # 2. Precision test
    validate_float_sampler_precision()

    # 3. TiDB stats test
    tidb_sampler, tidb_samples = validate_float_sampler_from_tidb_stats()

    # 4. Create visualizations
    try:
        visualize_float_distributions(original_data, generated_samples, sampler)
        print("✅ Visualizations created successfully")
    except Exception as e:
        print(f"❌ Visualization failed: {e}")
        import traceback

        traceback.print_exc()

    # 5. Test edge cases
    validate_edge_cases()

    print("\n" + "=" * 60)
    print("🎉 Comprehensive FloatSampler Test Suite Completed!")


if __name__ == "__main__":
    run_comprehensive_float_sampler_test()