#!/usr/bin/env python3
"""Accumulate values along a chromosome-length array.

    python3 genomic_arrays.py                     # runs the built-in example
    python3 genomic_arrays.py intervals_example.bed

Allocates a zeroed NumPy array the length of the chromosome and adds a value
across a range at a time. NumPy handles the slice arithmetic, so accumulation
stays fast even when the array is a quarter of a billion positions long.

Requires NumPy.
"""

import sys

import numpy as np

# Human chromosome 1, GRCh37.
CHR1_LENGTH = 247_249_719


class GenomicArray:
    """A dense numeric track over one chromosome.

    Intervals are half-open, [start, end), matching BED and Python slicing:
    `chunk_add(3, 5, 10)` touches positions 3 and 4, not 5. If your
    annotations came from a one-based, closed-interval format such as GFF,
    convert them before you get here rather than after you have plotted the
    result.
    """

    def __init__(self, size, dtype=np.float64):
        self.size = size
        self.store = np.zeros(size, dtype=dtype)

    def chunk_add(self, start, end, value):
        """Add `value` to every position in [start, end)."""
        if start < 0 or end > self.size:
            raise IndexError(
                f"interval [{start}, {end}) outside array of length {self.size}")
        if end < start:
            raise ValueError(f"end {end} precedes start {start}")
        self.store[start:end] += value
        return self.store

    def __len__(self):
        return self.size


def load_bed(path, array):
    """Add every interval in a 4-column BED file: chrom, start, end, value."""
    count = 0
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line or line.startswith(("#", "track", "browser")):
                continue
            fields = line.split()
            start, end = int(fields[1]), int(fields[2])
            value = float(fields[3]) if len(fields) > 3 else 1.0
            array.chunk_add(start, end, value)
            count += 1
    return count


def main(argv=None):
    argv = sys.argv[1:] if argv is None else argv

    if argv:
        # Sized to the example file rather than a whole chromosome, so the
        # demo does not allocate two gigabytes to show ten intervals.
        array = GenomicArray(1000)
        added = load_bed(argv[0], array)
        print(f"added {added} intervals from {argv[0]}")
        print("positions 0-20:", array.store[0:20])
        print("max coverage:  ", array.store.max())
        print("covered bases: ", int((array.store > 0).sum()))
    else:
        array = GenomicArray(CHR1_LENGTH)
        array.chunk_add(3, 5, 10)
        array.chunk_add(4, 10, 20)
        print("positions 0-20:", array.store[0:20])

    return 0


if __name__ == "__main__":
    raise SystemExit(main())