import random
from anc.data.anc_sampler import AncSampler


class SimpleDataset:
    def __init__(self, sub_ds_count, sub_ds_size):
        self.sub_ds_count = sub_ds_count
        self.sub_ds_size = sub_ds_size

    def get_sub_lengths(self):
        return [self.sub_ds_size] * self.sub_ds_count

    def __len__(self):
        return self.sub_ds_count * self.sub_ds_size


def check_sampler_indices(shuffle, batch_size, resume_step):
    # create a fake dataset with 3 files and each file contains 100 elements
    sub_ds_size = 40
    sub_ds_count = 3
    ds = SimpleDataset(sub_ds_count, sub_ds_size)
    world = 4
    rank = 0
    num_workers = 2
    drop_last = False
    num_per_rank = (sub_ds_size * sub_ds_count + (world - 1) * int(not drop_last)) // world
    # with the above parameters (if shuffle = False), the data idxs that rank 0 would generate are [0, 30), totally 30 elements
    # since batch size = 4, the batch number would be 8, with last batch only containing 2 elements
    # thus worker 0 of rank 0 would generate [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
    #      worker 1 of rank 0 would generate [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
    sampler = AncSampler(ds, batch_size, world, rank, num_workers, shuffle=False, drop_last=drop_last, resume_step=resume_step)
    assert len(sampler) == (num_per_rank + batch_size - 1) // batch_size
    indices = []
    for item in sampler:
        cur_index, is_last_batch = item
        if not is_last_batch:
            # cur_index would be [] if resume_step % num_workers != 0
            assert len(cur_index) == batch_size or len(cur_index) == 0
        if is_last_batch:
            # each worker's last batch would set is_last_batch to True
            # so we have num_workers is_last_batch
            assert sampler.step >= len(sampler) - num_workers
        indices += cur_index
    if not shuffle and resume_step == 0:
        assert indices[:16] == [0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23]
        assert indices[-6:] == [12, 13, 14, 15, 28, 29]
    return indices


def test_resume():
    batch_size = 4
    from_beginning_indices = check_sampler_indices(shuffle=False, batch_size=batch_size, resume_step=0)
    check_sampler_indices(shuffle=True, batch_size=batch_size, resume_step=0)
    resume_step = 3
    resumed_indices = check_sampler_indices(shuffle=False, batch_size=batch_size, resume_step=resume_step)
    assert resumed_indices == from_beginning_indices[batch_size * resume_step:]


def test_get_indices_with_ratio():
    # Test setup
    ds = SimpleDataset(1, 10)  # Single dataset with 10 elements
    sampler = AncSampler(ds, batch_size=2, shuffle=False)
    base_indices = list(range(10))  # [0,1,2,3,4,5,6,7,8,9]
    
    # Test case 1: Integer ratio without shuffle
    result = sampler._get_indices_with_ratio(base_indices, 2.0)
    assert len(result) == 2, "Should have 2 copies of indices"
    assert result[0] == base_indices, "First copy should match original"
    assert result[1] == base_indices, "Second copy should match original"
    
    # Test case 2: Fractional ratio without shuffle
    result = sampler._get_indices_with_ratio(base_indices, 1.5)
    assert len(result) == 2, "Should have 2 lists (1 full + 1 partial)"
    assert result[0] == base_indices, "First copy should be complete"
    assert len(result[1]) == 5, "Second copy should be half length"
    assert result[1] == base_indices[:5], "Partial copy should be first half"
    
    # Test case 3: With chunk granularity
    sampler.chunk_granularity = 3
    result = sampler._get_indices_with_ratio(base_indices, 1.0)
    assert len(result) == 4, "Should split into 4 chunks of size 3 (last chunk size 1)"
    assert result[0] == [0,1,2], "First chunk"
    assert result[1] == [3,4,5], "Second chunk"
    assert result[2] == [6,7,8], "Third chunk"
    assert result[3] == [9], "Last chunk"

    # Test case 4: With shuffle
    sampler.shuffle = True
    sampler.chunk_granularity = -1
    random.seed(42)  # For reproducibility
    result = sampler._get_indices_with_ratio(base_indices, 1.0)
    assert len(result) == 1, "Should have 1 shuffled copy"
    assert len(result[0]) == len(base_indices), "Should contain all indices"
    assert result[0] != base_indices, "Should be shuffled"
    assert sorted(result[0]) == base_indices, "Should contain same elements"


if __name__ == "__main__":
    test_resume()
    test_get_indices_with_ratio()