
import argparse
import sys
import math

try:
    import tskit  # noqa: F401
except Exception as e:
    # We allow the script to be inspected without tskit, but bail at runtime.
    pass


def clamp(v, lo, hi):
    return max(lo, min(hi, v))


def write_track_line(out, name, description, visibility=2, use_score=True, item_rgb=True):
    parts = [f'track name="{name}"']
    if description:
        parts.append(f'description="{description}"')
    parts.append(f'visibility={int(visibility)}')
    if use_score:
        parts.append("useScore=1")
    if item_rgb:
        parts.append('itemRgb="On"')
    out.write(" ".join(parts) + "\n")


def tree_stat(tree, stat):
    """
    Return a per-tree statistic used for score/color.
    Currently supported:
      - root_time: time of the (single) root; if multiple roots, max root time
      - total_branch_length: sum of branch lengths (in time units)
      - n_roots: number of roots (integer)
    """
    if stat == "root_time":
        # If there's a single root, use that; if multiple, take the max time.
        if hasattr(tree, "roots"):
            roots = list(tree.roots)
        else:
            # older tskit: use tree.root
            roots = [tree.root] if tree.root != tskit.NULL else []
        if not roots:
            return 0.0
        return max(tree.time(r) for r in roots)

    if stat == "total_branch_length":
        # Sum over edges in this tree: sum_t (time(parent) - time(child))
        total = 0.0
        for u in tree.nodes():
            p = tree.parent(u)
            if p != tskit.NULL:
                total += tree.time(p) - tree.time(u)
        return max(0.0, total)

    if stat == "n_roots":
        return float(len(list(tree.roots)))

    raise ValueError(f"Unsupported stat: {stat}")


def main():
    p = argparse.ArgumentParser(
        description="Create a BED file from a tskit tree sequence."
    )
    p.add_argument("input", help="Input .trees (tskit tree sequence)")
    p.add_argument("output", help="Output .bed path, or '-' for stdout")
    p.add_argument("--chrom", default="chr1", help="Chromosome/contig name for BED (default: chr1)")
    p.add_argument("--mode", choices=["trees", "sites"], default="trees",
                   help="Emit intervals per tree or per site (default: trees)")
    p.add_argument("--stat", choices=["root_time", "total_branch_length", "n_roots"],
                   default="root_time", help="Per-tree stat for score/color (trees mode)")
    p.add_argument("--ucsc-track", action="store_true",
                   help="Prepend a UCSC track line to the BED output")
    p.add_argument("--track-name", default="tskit-BED",
                   help="Name to use in UCSC track line (if --ucsc-track)")
    p.add_argument("--track-description", default="Generated by tskit_to_bed.py",
                   help="Description to use in UCSC track line")
    p.add_argument("--no-bed9", action="store_true",
                   help="Write simple 3-column BED (chrom, start, end). Default is 9-column BED with score/color.")
    p.add_argument("--precision", type=int, default=0,
                   help="Round interval coordinates to this many decimals before casting to int (default: 0)")

    args = p.parse_args()

    ts = tskit.load(args.input)

    # Choose output
    if args.output == "-":
        out = sys.stdout
        close_out = False
    else:
        out = open(args.output, "w", encoding="utf-8")
        close_out = True

    try:
        if args.ucsc_track:
            # continue
            write_track_line(
                out,
                name=args.track_name,
                description=args.track_description,
                visibility=2,
                use_score=not args.no_bed9,
                item_rgb=not args.no_bed9,
            )

        chrom = args.chrom

        if args.mode == "sites":
            # Emit per-site 3-col or 6/9-col BED. We keep it simple: 3-col unless bed9 is desired.
            for var in ts.variants():
                # tskit positions are float; BED wants ints. Round then cast.
                start = int(round(var.position, args.precision))
                end = start + 1
                if args.no_bed9:
                    out.write(f"{chrom}\t{start}\t{end}\n")
                else:
                    name = getattr(var, "id", None)
                    if name is None:
                        # fallback to pos + first alt allele
                        alleles = [a for a in var.alleles if a is not None]
                        name = f"pos{start}:{'/'.join(alleles[:2])}" if alleles else f"pos{start}"
                    score = 0
                    strand = "."
                    thickStart = start
                    thickEnd = end
                    rgb = "0,0,0"
                    out.write(f"{chrom}\t{start}\t{end}\t{name}\t{score}\t{strand}\t{thickStart}\t{thickEnd}\t{rgb}\n")
            return

        # trees mode
        # First pass: collect stats and intervals
        intervals = []
        stat_values = []
        for tree in ts.trees(tracked_samples=None):
            left = int(round(tree.interval.left, args.precision))
            right = int(round(tree.interval.right, args.precision))
            if right <= left:
                continue
            val = tree_stat(tree, args.stat)
            intervals.append((left, right, tree.index, val))
            stat_values.append(val)

        if not intervals:
            return

        # Determine scaling for score/color
        vmax = max(stat_values) if stat_values else 1.0
        if vmax <= 0:
            vmax = 1.0

        for left, right, tidx, val in intervals:
            if args.no_bed9:
                out.write(f"{chrom}\t{left}\t{right}\n")
                continue

            # name, score, strand, thickStart, thickEnd, itemRgb
            name = f"tree_{tidx}"
            score = int(round(1000.0 * (val / vmax)))
            score = clamp(score, 0, 1000)
            strand = "."
            thickStart = left
            thickEnd = right

            # Map val to grayscale 0..255
            if vmax > 0:
                g = int(round(255.0 * (val / vmax)))
            else:
                g = 0
            g = clamp(g, 0, 255)
            rgb = f"{g},{g},{g}"

            out.write(f"{chrom}\t{left}\t{right}\t{name}\t{score}\t{strand}\t{thickStart}\t{thickEnd}\t{rgb}\n")

    finally:
        if close_out:
            out.close()


if __name__ == "__main__":
    main()
