import pickle
import sys
import os
import io
import subprocess
import json
from functools import lru_cache
from typing import List, Tuple, Optional, Any, Dict

cache = lru_cache(None)

__all__ = ["format_flamegraph", "segments", "memory", "compare"]

def _frame_fmt(f, full_filename=False):
    i = f['line']
    fname = f['filename']
    if not full_filename:
        fname = fname.split('/')[-1]
    func = f['name']
    return f'{fname}:{i}:{func}'

@cache
def _frame_filter(name, filename):
    omit_functions = [
        "unwind::unwind",
        "StackContext::gather",
        "StackContext::_gather",
        "_start",
        "__libc_start_main"
    ]
    omit_filenames = [
        "core/boxing",
        "/Register",
        "/Redispatch",
        "pythonrun.c",
        "Modules/main.c",
        "Objects/call.c",
        "Objects/methodobject.c",
        "pycore_ceval.h",
        "ceval.c",
        "cpython/abstract.h",
    ]
    for of in omit_functions:
        if of in name:
            return False
    for of in omit_filenames:
        if of in filename:
            return False
    return True

def _frames_fmt(frames, full_filename=False, reverse=False):
    if reverse:
        frames = reversed(frames)
    return [_frame_fmt(f, full_filename) for f in frames if _frame_filter(f['name'], f['filename'])]

def format_flamegraph(flamegraph_lines, flamegraph_script=None):
    if flamegraph_script is None:
        flamegraph_script = f'/tmp/{os.getuid()}_flamegraph.pl'
    if not os.path.exists(flamegraph_script):
        import urllib.request
        print(f"Downloading flamegraph.pl to: {flamegraph_script}")
        urllib.request.urlretrieve(
            'https://raw.githubusercontent.com/brendangregg/FlameGraph/master/flamegraph.pl', flamegraph_script)
        subprocess.run(['chmod', '+x', flamegraph_script])
    args = [flamegraph_script, '--countname', 'bytes']
    p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, encoding='utf-8')
    assert p.stdin is not None
    assert p.stdout is not None
    p.stdin.write(flamegraph_lines)
    p.stdin.close()
    result = p.stdout.read()
    p.stdout.close()
    p.wait()
    assert p.wait() == 0
    return result

def _write_blocks(f, prefix, blocks):
    for b in blocks:
        if 'history' not in b:
            f.write(f'{prefix};{b["state"]} {b["size"]}\n')
            continue
        accounted_for_size = 0
        for h in b['history']:
            sz = h['real_size']
            accounted_for_size += sz
            if 'frames' in h:
                frames = h['frames']
                if frames:
                    frame_s = ';'.join(_frames_fmt(frames, reverse=True))
                else:
                    frame_s = "<non-python>"
                f.write(f'{prefix};{b["state"]};{frame_s} {sz}\n')
            else:
                f.write(f'{prefix};{b["state"]};<no-context> {sz}\n')
        gaps = b['size'] - accounted_for_size
        if gaps:
            f.write(f'{prefix};{b["state"]};<gaps> {gaps}\n')

def segments(snapshot, format_flamegraph=format_flamegraph):
    f = io.StringIO()
    for seg in snapshot['segments']:
        prefix = f'stream_{seg["stream"]};seg_{seg["address"]}'
        _write_blocks(f, prefix, seg['blocks'])
    return format_flamegraph(f.getvalue())

def memory(snapshot, format_flamegraph=format_flamegraph):
    f = io.StringIO()
    for seg in snapshot['segments']:
        prefix = f'stream_{seg["stream"]}'
        _write_blocks(f, prefix, seg['blocks'])
    return format_flamegraph(f.getvalue())

def compare(before, after, format_flamegraph=format_flamegraph):
    def _seg_key(seg):
        return (seg['address'], seg['total_size'])

    def _seg_info(seg):
        return f'stream_{seg["stream"]};seg_{seg["address"]}'

    f = io.StringIO()

    before_segs = {_seg_key(seg) for seg in before}
    after_segs = {_seg_key(seg) for seg in after}

    print(f'only_before = {[a for a,_ in (before_segs - after_segs)]}')
    print(f'only_after = {[a for a,_ in (after_segs - before_segs)]}')

    for seg in before:
        if _seg_key(seg) not in after_segs:
            _write_blocks(f, f'only_before;{_seg_info(seg)}', seg['blocks'])

    for seg in after:
        if _seg_key(seg) not in before_segs:
            _write_blocks(f, f'only_after;{_seg_info(seg)}', seg['blocks'])

    return format_flamegraph(f.getvalue())

def _format_size(num):
    # https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size
    for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
        if abs(num) < 1024.0:
            return f"{num:3.1f}{unit}B"
        num /= 1024.0
    return f"{num:.1f}YiB"

class Bytes:
    def __init__(self, value):
        self.value = value

    def __add__(self, rhs):
        return Bytes(self.value + rhs)

    def __repr__(self):
        return _format_size(self.value)

def calc_active(seg):
    return sum(b['size'] for b in seg['blocks'] if b['state'] == 'active_allocated')

def _report_free(free_external, free_internal):
    total = free_external + free_internal
    suffix = ''
    if total != 0:
        pct = (free_internal / total) * 100
        suffix = f' ({pct:.1f}% internal)'
    return f'{Bytes(total)}{suffix}'

PAGE_SIZE = 1024 * 1024 * 20
legend = f"""\

Legend:
    [a     ] - a segment in the allocator
     ^-- a page {Bytes(PAGE_SIZE)} of memory in the segment
    a-z: pages filled with a single block's content
    ' ': page is completely free
    *: page if completely full with multiple blocks
    0-9: page is partially full with tensors of multiple blocks (9 == 90% full)
    (X% internal) - of the free memory, X% is free because we rounded the size of the allocation.
"""

def segsum(data):
    """" Visually reports how the allocator has filled its segments. This printout can help debug fragmentation issues
    since free fragments will appear as gaps in this printout.  The amount of free space is reported for each segment.
    We distinguish between internal free memory which occurs because the allocator rounds the allocation size, and
    external free memory, which are the gaps between allocations in a segment.
    Args:
        data: snapshot dictionary created from _snapshot()

    """
    segments = []
    out = io.StringIO()
    out.write(f"Summary of segments >= {Bytes(PAGE_SIZE)} in size\n")
    total_reserved = 0
    total_allocated = 0
    free_external = 0
    free_internal = 0
    for seg in sorted(data['segments'], key=lambda x: (x['total_size'], calc_active(x))):
        total_reserved += seg['total_size']

        seg_free_external = 0
        seg_free_internal = 0
        seg_allocated = 0
        all_ranges = []
        boffset = 0
        for b in seg['blocks']:
            active = b['state'] == 'active_allocated'
            if 'history' in b:
                # use the more accureate real_size to account for internal fragmenetation if we have it
                for h in b['history']:
                    if active:
                        all_ranges.append((h['addr'] - seg['address'], h['real_size'], active))
                        seg_allocated += h['real_size']
                        assert len(b['history']) == 1
                        seg_free_internal += b['size'] - h['real_size']
            else:
                if active:
                    all_ranges.append((boffset, b['size'], True))
                    seg_allocated += b['size']
            if not active:
                seg_free_external += b['size']

            boffset += b['size']

        total_allocated += seg_allocated
        free_external += seg_free_external
        free_internal += seg_free_internal

        nseg = (seg['total_size'] - 1) // PAGE_SIZE + 1
        occupied = [' ' for _ in range(nseg)]
        frac = [0.0 for _ in range(nseg)]
        active_size = 0
        for i, (start_, size, active) in enumerate(all_ranges):
            active_size += size
            finish_ = (start_ + size)
            start = start_ // PAGE_SIZE
            finish = (finish_ - 1) // PAGE_SIZE + 1
            m = chr((ord('a' if active else 'A') + (i % 26)))
            for j in range(start, finish):
                s = max(start_, j * PAGE_SIZE)
                e = min(finish_, (j + 1) * PAGE_SIZE)
                frac[j] += (e - s) / PAGE_SIZE
                if occupied[j] != ' ':
                    occupied[j] = '0123456789*'[int(frac[j] * 10)]
                else:
                    occupied[j] = m
        stream = '' if seg['stream'] == 0 else f', stream_{seg["stream"]}'
        body = ''.join(occupied)
        assert seg_free_external + seg_free_internal + seg_allocated == seg['total_size']
        stream = f' stream_{seg["stream"]}' if seg['stream'] != 0 else ''
        if seg['total_size'] >= PAGE_SIZE:
            out.write(f'[{body}] {Bytes(seg["total_size"])} allocated, '
                      f'{_report_free(seg_free_external, seg_free_internal)} free{stream}\n')
    out.write(f'segments: {len(data["segments"])}\n')
    out.write(f'total_reserved: {Bytes(total_reserved)}\n')
    out.write(f'total_allocated: {Bytes(total_allocated)}\n')
    internal_external = f' ({Bytes(free_internal)} internal + {Bytes(free_external)} external)' if free_internal else ''
    out.write(f'total_free: {_report_free(free_external, free_internal)}\n')
    out.write(legend)
    assert free_internal + free_external + total_allocated == total_reserved
    return out.getvalue()

def trace(data):
    out = io.StringIO()

    def format(entries):
        segment_intervals : list = []
        segment_addr_to_name = {}
        allocation_addr_to_name = {}

        free_names : list = []
        next_name = 0

        def _name():
            nonlocal next_name
            if free_names:
                return free_names.pop()
            r, m = next_name // 26, next_name % 26
            next_name += 1
            return f'{chr(ord("a") + m)}{"" if r == 0 else r}'

        def find_segment(addr):
            for name, saddr, size in segment_intervals:
                if addr >= saddr and addr < saddr + size:
                    return name, saddr
            for i, seg in enumerate(data['segments']):
                saddr = seg['address']
                size = seg['allocated_size']
                if addr >= saddr and addr < saddr + size:
                    return f'seg_{i}', saddr
            return None, None
        count = 0
        out.write(f'{len(entries)} entries\n')


        total_reserved = 0
        for seg in data['segments']:
            total_reserved += seg['total_size']

        for count, e in enumerate(entries):
            if e['action'] == 'alloc':
                addr, size = e['addr'], e['size']
                n = _name()
                seg_name, seg_addr = find_segment(addr)
                if seg_name is None:
                    seg_name = "MEM"
                    offset = addr
                else:
                    offset = addr - seg_addr
                out.write(f'{n} = {seg_name}[{offset}:{Bytes(size)}]\n')
                allocation_addr_to_name[addr] = (n, size, count)
                count += size
            elif e['action'] == 'free_requested':
                addr, size = e['addr'], e['size']
                name, _, _ = allocation_addr_to_name.get(addr, (addr, None, None))
                out.write(f'del {name} # {Bytes(size)}\n')
            elif e['action'] == 'free_completed':
                addr, size = e['addr'], e['size']
                count -= size
                name, _, _ = allocation_addr_to_name.get(addr, (addr, None, None))
                out.write(f'# free completed for {name} {Bytes(size)}\n')
                if name in allocation_addr_to_name:
                    free_names.append(name)
                    del allocation_addr_to_name[name]
            elif e['action'] == 'segment_alloc':
                addr, size = e['addr'], e['size']
                name = _name()
                out.write(f'{name} = cudaMalloc({addr}, {Bytes(size)})\n')
                segment_intervals.append((name, addr, size))
                segment_addr_to_name[addr] = name
            elif e['action'] == 'segment_free':
                addr, size = e['addr'], e['size']
                name = segment_addr_to_name.get(addr, addr)
                out.write(f'cudaFree({name}) # {Bytes(size)}\n')
                if name in segment_addr_to_name:
                    free_names.append(name)
                    del segment_addr_to_name[name]
            elif e['action'] == 'oom':
                size = e['size']
                free = e['device_free']
                out.write(f'raise OutOfMemoryError() # {Bytes(size)} requested, {Bytes(free)} free in CUDA\n')
            else:
                out.write(f'{e}\n')
        out.write(f"TOTAL MEM: {Bytes(count)}")
    for i, d in enumerate(data['device_traces']):
        if d:
            out.write(f'Device {i} ----------------\n')
            format(d)
    return out.getvalue()

class PlotWriter:
    def __init__(self, categories: List[str] = None):
        string_table: List[str] = []

        # compresses lists of strings that have common suffixes
        # such as stack traces with the most recent frame first.
        # (reference to string table, another suffix table entry)
        suffix_table: List[Tuple[int, Optional[int]]] = []

        elements_size = []

        # indexes into the suffix_table
        elements_info = []

        # indexes into category table
        elements_category: Optional[List[int]] = None if categories is None else []

        # indexes into the elements_ tables
        actions: List[int] = []

        # indexes into the elements_ tables
        initially_allocated: List[int] = []

        @cache
        def intern_str(s):
            string_table.append(s)
            return len(string_table) - 1

        @cache
        def intern_suffix(sid, restid):
            suffix_table.append((sid, restid))
            return len(suffix_table) - 1

        def intern_stack(frames):
            next_id = None
            for f in reversed(frames):
                next_id = intern_suffix(intern_str(f), next_id)
            return next_id

        def add_element(size, lines, category=None):
            nonlocal elements_category
            # note: struct of arrays format to info about elements
            # avoids a lot of repeated string keys when serialized
            elements_size.append(size)
            elements_info.append(intern_stack(lines))

            # lazily create since we will not always have categories
            if categories is not None:
                assert category >= 0 and category < len(categories)
                assert elements_category is not None
                elements_category.append(category)

            return len(elements_size) - 1

        def to_html():
            r = {
                'actions': actions,
                'elements_size': elements_size,
                'elements_info': elements_info,
                'elements_category': elements_category,
                'suffix_table': suffix_table,
                'string_table': string_table,
                'initially_allocated': list(reversed(initially_allocated)),
                'categories': categories,
            }
            plot_data = json.dumps(r)
            return _memory_over_time_template.replace('$PLOT_DATA', plot_data)

        self.add_element = add_element
        self.allocate = actions.append
        self.free = actions.append
        self.initially_allocated = initially_allocated.append
        self.to_html = to_html
        self.categories = categories

def _choose_device(data, device):
    if device is None:
        for i, t in enumerate(data['device_traces']):
            if len(t) > 0:
                if device is not None:
                    raise ValueError(f'Both device {device} and {i} have traces, use --device to specify which trace.')
                device = i
    return device

def trace_plot(data, device=None, plot_segments=False):
    """Generate a visualization over time of the memory usage recorded by the trace as an html file.

    Args:
        data: Memory snapshot as generated from torch.cuda.memory._snapshot()
        device (torch.device, optional): Generate the trace for this device, needed if multiple devices have allocations.
        plot_segments (bool, optional): Plots memory returned from cudaMalloc, rather than individual allocations.
                                        Defaults to False.

    Returns:
        str: HTML of visualization
    """
    w = PlotWriter()
    addr_to_alloc = {}
    device = _choose_device(data, device)
    if device is None:
        raise ValueError('No trace information was recorded.')

    trace = data['device_traces'][device]

    if plot_segments:
        addr_prefix = 's'
        alloc = 'segment_alloc'
        free = 'segment_free'
    else:
        addr_prefix = 'b'
        alloc = 'alloc'
        free = 'free_completed'

    addr_versions: Dict[int, int] = {}

    def add_element(addr, size, frames, extra=()):
        next_version = addr_versions[addr] = addr_versions.get(addr, 0) + 1
        frames = [f"{addr_prefix}{addr:x}_{next_version - 1} {_format_size(size)} allocation ({size} bytes)",
                  *extra,
                  *_frames_fmt(frames, full_filename=True)]
        return w.add_element(size, frames)

    for i, e in enumerate(trace):
        if e['action'] == alloc:
            elemid = add_element(e['addr'], e['size'], e.get('frames', []))
            addr_to_alloc[e['addr']] = elemid
            w.allocate(elemid)
        elif e['action'] == free:
            idx = addr_to_alloc.pop(e['addr'], None)
            if idx is None:
                idx = add_element(e['addr'], e['size'], e.get('frames', []), extra=('alloc not recorded, stack trace for free:',))
                w.initially_allocated(idx)
            w.free(idx)
    return w.to_html()

def profile_plot(profile, device=None):
    """Generate a visualization over time of the memory usage recorded by kineto memory profiling as an html file.

    Args:
        profile: profile as genererated by `torch.profiler.profile(profile_memory=True)`
        device (torch.device, optional): Generate the trace for this device, needed if multiple devices have allocations.

    Returns:
        str: HTML of visualization
    """
    import torch
    from torch.profiler._memory_profiler import Action, TensorKey, Category
    from torch._C._profiler import _EventType
    memory_profile = profile._memory_profile()

    if device is None:
        if torch.cuda.is_available():
            device = torch.device('cuda', torch.cuda.current_device())
        else:
            device = torch.device('cpu')

    w = PlotWriter(categories=["unknown", *(c.name.lower() for c in Category)])


    allocation_stacks = {}
    for event in memory_profile._op_tree.sorted_nodes:
        if event.tag == _EventType.Allocation:
            parent = event.parent
            python_parents = []
            while parent:
                if parent.tag in (_EventType.PyCall, _EventType.PyCCall):
                    python_parents.append(parent)
                parent = parent.parent
            key = TensorKey.from_allocation(event.extra_fields)

            # Corner case: If allocation doesn't have an ID (can't prove it was used as a Tensor)
            #              key will be None. I should add some way to identify these, I just haven't yet.
            if key and event.extra_fields.alloc_size > 0:
                allocation_stacks[key] = python_parents

    def add_element(size, tensor_key, version):
        category = memory_profile._categories.get(tensor_key, version)
        category = category.value if category is not None else 0
        stack = allocation_stacks.get(tensor_key, ())
        assert w.categories is not None
        return w.add_element(size,
                             [f"{_format_size(size)} ({size} bytes) allocation ({w.categories[category]})",
                              *(p.name for p in stack)],
                             category)

    kv_to_elem = {}
    for time, action, (tensor_key, version), size in memory_profile.timeline:
        if tensor_key.device != device:
            continue
        if action == Action.CREATE:
            kv_to_elem[(tensor_key, version)] = elemid = add_element(size, tensor_key, version)
            w.allocate(elemid)
        elif action == Action.DESTROY:
            w.free(kv_to_elem.pop((tensor_key, version)))
        elif action == Action.INCREMENT_VERSION:
            w.free(kv_to_elem.pop((tensor_key, version)))
            kv_to_elem[(tensor_key, version + 1)] = elemid = add_element(size, tensor_key, version + 1)
            w.allocate(elemid)
        elif action == Action.PREEXISTING:
            kv_to_elem[(tensor_key, version)] = elemid = add_element(size, tensor_key, version)
            w.initially_allocated(elemid)
    return w.to_html()

# note: this template should eventually move to its own file,
# however, we first need to package _memory_viz.py so that it can be
# pip-installed separately from pytorch so it is easy to run e.g.
# on a laptop with downloaded snapshots. Currently this is
# accomplished by downloading _memory_viz.py so the template
# needs to be included
_memory_over_time_template = r"""
<!DOCTYPE html>
<html>
<head></head>
<body>
<script type="module">
import * as d3 from "https://cdn.jsdelivr.net/npm/d3@7.7.0/+esm";
import {schemeTableau10} from "https://cdn.skypack.dev/d3-scale-chromatic@3";
import {axisLeft} from "https://cdn.skypack.dev/d3-axis@3";
import {scaleLinear} from "https://cdn.skypack.dev/d3-scale@4";
import {zoom, zoomIdentity} from "https://cdn.skypack.dev/d3-zoom@3";
import {brushX} from "https://cdn.skypack.dev/d3-brush@3";

let alloc_data = $PLOT_DATA

function process_alloc_data(fraction_of_memory_reported=1) {
    let current = []
    let current_data = []
    let data = []
    let max_size = 0

    let total_mem = 0
    let timestep = 0

    let max_at_time = []
    function advance(n, max) {
        timestep += n
        for (let i = 0; i < n; i++) {
            max_at_time.push(max)
        }
    }

    let mini_points = []

    let sizes = alloc_data.elements_size.map(x => x).sort((x, y) => y - x)
    let total_size = sizes.reduce((x, y) => x + y)
    const memory_threshold = fraction_of_memory_reported * total_size
    let total_seen = 0
    let memory_threshold_size = 0

    for (const [i, size] of sizes.entries()) {
        total_seen += size
        if (total_seen > memory_threshold) {
            memory_threshold_size = size
            break
        }
    }

    function add_allocation(elem) {
        let size = alloc_data.elements_size[elem]
        current.push(elem)
        let color = elem
        if (alloc_data.elements_category !== null) {
            color = alloc_data.elements_category[elem]
        }
        let e = {elem: elem, timesteps: [timestep], offsets: [total_mem], size: size, color: color}
        current_data.push(e)
        data.push(e)
        total_mem += size
    }

    for (const elem of alloc_data.initially_allocated) {
        add_allocation(elem)
    }

    for (const action of alloc_data.actions) {
        const elem = action
        const idx = current.findIndex(x => x === elem)
        const size = alloc_data.elements_size[elem]
        if (size < memory_threshold_size) {
            continue
        }
        // first time we see an action we add it
        // second time we remove it
        if (idx == -1) {
            add_allocation(elem)
            advance(1, total_mem)
        } else {
            advance(1, total_mem)
            const removed = current_data[idx]
            removed.timesteps.push(timestep)
            removed.offsets.push(removed.offsets.at(-1))
            current.splice(idx, 1)
            current_data.splice(idx, 1)

            if (idx < current.length) {
                for (let j = idx; j < current.length; j++) {
                    const e = current_data[j]
                    e.timesteps.push(timestep)
                    e.offsets.push(e.offsets.at(-1))
                    e.timesteps.push(timestep + 3)
                    e.offsets.push(e.offsets.at(-1) - size)
                }
                advance(3, total_mem)
            }
            total_mem -= size
        }
        max_size = Math.max(total_mem, max_size)
    }

    for (const elem of current_data) {
        elem.timesteps.push(timestep)
        elem.offsets.push(elem.offsets.at(-1))
    }
    return {
        max_size: max_size,
        allocations_over_time: data,
        max_at_time: max_at_time,
        context_for_id:  (elem) => {
            let strings = []
            let id = alloc_data.elements_info[elem]
            while (id !== null) {
                const [sid, next_id] = alloc_data.suffix_table[id]
                strings.push(alloc_data.string_table[sid])
                id = next_id
            }
            return `${strings.join('\n')}\n`
        }
    }
}

function MemoryPlot(svg, data, left_pad, colors=schemeTableau10) {
    function format_points(d) {
        const size = d.size
        const xs = d.timesteps.map(t => xscale(t))
        const bottom = d.offsets.map(t => yscale(t))
        const top = d.offsets.map(t => yscale(t + size))

        const p0 = xs.map((x, i) => `${x},${bottom[i]}`)
        const p1 = xs.map((x, i) => `${x},${top[i]}`).reverse()

        return `${p0.join(' ')} ${p1.join(' ')}`
    }

    let max_timestep = data.max_at_time.length
    let max_size = data.max_size

    let width = svg.attr('width')
    let height = svg.attr('height')
    let plot_width = width - left_pad
    let plot_height = height

    let yscale = scaleLinear().domain([0, max_size]).range([plot_height, 0]);
    let heightscale = scaleLinear().domain([0, max_size]).range([0, plot_height]);
    let yaxis = axisLeft(yscale).tickFormat(d3.format("~s"))
    let xscale = scaleLinear().domain([0, max_timestep]).range([0, plot_width])
    let plot_coordinate_space = svg.append("g").attr("transform", `translate(${left_pad}, ${0})`)
    let plot_outer = plot_coordinate_space.append('g')

    function view_rect(a) {
        return a.append('rect').attr('x', 0).attr('y', 0)
                .attr('width', plot_width).attr('height', plot_height)
                .attr('fill', 'white')
    }

    view_rect(plot_outer)

    let cp = svg.append("clipPath").attr("id", "clip")
    view_rect(cp)
    plot_outer.attr('clip-path', "url(#clip)")


    let zoom_group = plot_outer.append("g")
    let scrub_group = zoom_group.append('g')

    let plot = scrub_group.selectAll("polygon")
    .data(data.allocations_over_time)
    .enter()
    .append("polygon")
    .attr('points', format_points)
    .attr('fill', d => colors[d.color % colors.length])

    let axis = plot_coordinate_space.append('g').call(yaxis)


    let scale_mini = 0
    let translate_mini = 0
    function handleZoom(e) {
        const t = e.transform
        zoom_group.attr("transform", t)
        axis.call(yaxis.scale(e.transform.rescaleY(yscale)))
    }

    const thezoom = zoom().on('zoom', handleZoom)
    plot_outer.call(thezoom)

    return {
        select_window: (stepbegin, stepend, max) => {
            let begin = xscale(stepbegin)
            let size = xscale(stepend) - xscale(stepbegin);
            let scale = plot_width / size
            let translate = -begin
            let yscale =  max_size/max
            scrub_group.attr("transform", `scale(${scale/yscale}, 1) translate(${translate}, 0)`)
            plot_outer.call(thezoom.transform, zoomIdentity.scale(yscale).translate(0, -(plot_height - plot_height/yscale)))
        },
        set_delegate: (delegate) => {
            plot.on('mouseover', function (e, d) { delegate.set_selected(d3.select(this)) } )
            .on('mousedown', function(e, d) { delegate.default_selected = d3.select(this)})
            .on('mouseleave', function (e, d) { delegate.set_selected(delegate.default_selected) } )
        }
    }
}

function ContextViewer(text, data) {
    let current_selected = null

    return {
        default_selected: null,
        set_selected: (d) => {
            if (current_selected !== null) {
                current_selected.attr('stroke', null).attr('stroke-width', null);
            }
            if (d === null) {
                text.text("")
            } else {
                const dd = d.datum()
                text.text(`${dd.elem} ${data.context_for_id(dd.elem)}`)
                d.attr('stroke', 'black').attr('stroke-width', 1).attr('vector-effect', 'non-scaling-stroke')
            }
            current_selected = d
        }
    }
}


function MiniMap(mini_svg, plot, data, left_pad, height=70) {
    let max_at_time = data.max_at_time
    let width = mini_svg.attr('width')
    let plot_width = width - left_pad
    let yscale = scaleLinear().domain([0, data.max_size]).range([height, 0]);
    let minixscale = scaleLinear().domain([0, max_at_time.length]).range([left_pad, width])

    let mini_points = [[max_at_time.length, 0], [0, 0]]

    for (const [i, m] of max_at_time.entries()) {
        let [lastx, lasty] = mini_points[mini_points.length - 1]
        if (m !== lasty) {
            mini_points.push([i, lasty])
            mini_points.push([i, m])
        } else if (i === max_at_time.length - 1) {
            mini_points.push([i, m])
        }
    }


    let points = mini_points.map(([t, o]) => `${minixscale(t)}, ${yscale(o)}`)
    points = points.join(' ')
    mini_svg.append('polygon').attr('points', points).attr('fill', schemeTableau10[0])

    let xscale = scaleLinear().domain([0, max_at_time.length]).range([0, plot_width])


    const brush = brushX()
    brush.extent([[left_pad, 0], [width, height]])
    brush.on('brush', function({selection}) {
        let [begin, end] = selection.map(x => x - left_pad)

        let stepbegin = Math.floor(xscale.invert(begin))
        let stepend = Math.floor(xscale.invert(end))
        let max = 0
        for (let i = stepbegin; i < stepend; i++) {
            max = Math.max(max, max_at_time[i])
        }
        plot.select_window(stepbegin, stepend, max)
    })
    mini_svg.call(brush)
    return {}
}

function Legend(svg, categories) {
    let xstart = width - 100
    let ystart = 30
    plot_svg.append('g').selectAll('rect')
    .data(categories)
    .enter()
    .append('rect')
    .attr('x', (c, i) => xstart)
    .attr('y', (c, i) => ystart + i*15)
    .attr('width', 10)
    .attr('height', 10)
    .attr('fill', (c, i) => schemeTableau10[i % schemeTableau10.length])
    plot_svg.append('g').selectAll('text')
    .data(categories)
    .enter()
    .append('text')
    .attr('x', (c, i) => xstart + 20)
    .attr('y', (c, i) => ystart + i*15 + 8)
    .attr("font-family", "helvetica")
    .attr('font-size', 10)
    .text((c) => c)
    return {}
}

let left_pad = 70
let width = 1024
let height = 768
let data = process_alloc_data()
let body = d3.select("body")

let plot_svg = body.append("svg").attr('width', width).attr('height', height).attr('display', 'block')
let plot = MemoryPlot(plot_svg, data, left_pad)

if (alloc_data.categories !== null) {
    Legend(plot_svg.append('g'), alloc_data.categories)
}

MiniMap(body.append("svg").attr('width', width).attr('height', 80).attr('display', 'block'), plot, data, left_pad)
let delegate = ContextViewer(body.append("div").append("pre").text('none'), data)
plot.set_delegate(delegate)

</script>
</body>
</html>
"""

def segment_plot(data: Any, device=None):
    device = _choose_device(data, device)
    if device is None:
        trace = []
    else:
        trace = data['device_traces'][device]

    string_table: List[str] = []
    suffix_table: List[Tuple[int, Optional[int]]] = []

    @cache
    def intern_str(s):
        string_table.append(s)
        return len(string_table) - 1

    @cache
    def intern_suffix(sid, restid):
        suffix_table.append((sid, restid))
        return len(suffix_table) - 1

    def intern_stack(frames):
        next_id = None
        for f in reversed(frames):
            next_id = intern_suffix(intern_str(f), next_id)
        return next_id

    def format_frames(frames):
        return intern_stack(_frames_fmt(frames, full_filename=True))

    result: Any = {
        'string_table': string_table,
        'suffix_table': suffix_table,
        'events': {
            'action': [],  # reference to string table
            'addr': [],  # for OOM, this will hold device_free value
            'size': [],
            'stream': [],
            'frames': []  # reference to suffix_table
        },
        'segments': {
            'addr': [],
            'size': [],
            'stream': []
        },
        'blocks': {
            'addr': [],
            'size': [],
            'real_size': [],
            'frames': [],  # reference to string table
            'pending_free': [],
        }
    }

    def fold_free(ts):
        # turn a free_requested/free_completed pair into a single free event
        i = 0
        while i < len(ts):
            t = ts[i]
            if i + 1 < len(ts):
                tnext = ts[i + 1]
                if t['action'] == 'free_requested' and tnext['action'] == 'free_completed' and t['addr'] == tnext['addr']:
                    yield {**t, 'action': 'free'}
                    i += 2
                    continue
            if t['action'] == 'oom':
                yield {**t, 'addr': t['device_free']}
            else:
                yield t
            i += 1

    preproc: Any = {
        'action': intern_str,
        'frames': format_frames,
    }

    events: Any = result['events']
    for event in fold_free(trace):
        for k in events.keys():
            # stack frames not recorded on event
            # happens for snapshot even when
            # frames are recorded for other things.
            if k == 'frames' and k not in event:
                events[k].append(None)
                continue
            events[k].append(preproc.get(k, lambda x: x)(event[k]))

    segments = result['segments']
    blocks = result['blocks']

    segment_names = {
        'addr': 'address',
        'size': 'total_size',
    }

    for seg in data['segments']:
        for k in segments.keys():
            sk = segment_names.get(k, k)
            segments[k].append(preproc.get(k, lambda x: x)(seg[sk]))
        addr = seg['address']
        for b in seg['blocks']:
            if b['state'] in ('active_pending_free', 'active_allocated'):
                if 'history' in b:
                    frames = b['history'][0].get('frames', [])
                    real_size = b['history'][0]['real_size']
                else:
                    real_size = b['size']
                    frames = []
                blocks['addr'].append(addr)
                blocks['size'].append(b['size'])
                blocks['real_size'].append(real_size)
                blocks['frames'].append(format_frames(frames))
                blocks['pending_free'].append(1 if b['state'] == 'active_pending_free' else 0)
            addr += b['size']

    plot_data = json.dumps(result)
    return _events_template.replace('$PLOT_DATA', plot_data)

_events_template = r"""
<!DOCTYPE html>
<html>
<head>
<style>
pre {
    margin: 0px;
}
html, body {
    height: 100%;
}
</style>
</head>
<body>
<script type="module">
import * as d3 from "https://cdn.jsdelivr.net/npm/d3@7.7.0/+esm";
import {schemeTableau10} from "https://cdn.skypack.dev/d3-scale-chromatic@3";
import {axisLeft} from "https://cdn.skypack.dev/d3-axis@3";
import {scaleLinear} from "https://cdn.skypack.dev/d3-scale@4";
import {zoom, zoomIdentity} from "https://cdn.skypack.dev/d3-zoom@3";
import {brushX} from "https://cdn.skypack.dev/d3-brush@3";

let trace_data = $PLOT_DATA

function frames_for(id) {
    let strings = []
    while (id !== null) {
        const [sid, next_id] = trace_data.suffix_table[id]
        strings.push(trace_data.string_table[sid])
        id = next_id
    }
    return `${strings.join('\n')}\n`
}

let events_data = trace_data.events
let string_table = trace_data.string_table

class Event {
    constructor(id) {
        this.idx = id
    }
    get size() { return events_data.size[this.idx] }
    get action() { return string_table[events_data.action[this.idx]] }
    get addr() { return events_data.addr[this.idx] }
    get frames() { return frames_for(this.frames_idx) }
    get frames_idx() { return events_data.frames[this.idx]}
    get stream() { return events_data.stream[this.idx] }
    get device_free() { return events_data.addr[this.idx] }
}

function createEvents() {
    let events = trace_data.events.action.map( (_, i) => new Event(i))

    function version_space() {
        let version = {}
        return (addr, increment) => {
            if (!(addr in version)) {
                version[addr] = 0
            }
            let r = version[addr]
            if (increment) {
                version[addr]++
            }
            return r
        }
    }

    let segment_version = version_space()
    let block_version = version_space()
    for (let t of events) {
        // set unique version for each time an address is used
        // so that ctrl-f can be used to search for the beginning
        // and end of allocations and segments
        switch (t.action) {
            case 'free':
            case 'free_completed':
                t.version = block_version(t.addr, true)
                break
            case 'free_requested':
            case 'alloc':
                t.version = block_version(t.addr, false)
                break
            case 'segment_free':
                t.version = segment_version(t.addr, true)
                break;
            case 'segment_alloc':
                t.version = segment_version(t.addr, false)
                break
            default:
                break
        }
    }
    trace_data.segment_version = segment_version
    trace_data.block_version = block_version
    return events
}

function Segment(addr, size, stream, frame_idx, version) {
    return {addr, size, stream, version, get frames() { return frames_for(frame_idx)}}
}
function Block(addr, size, real_size, frame_idx, free_requested, version) {
    console.assert(frame_idx !== undefined)
    return {addr, size, real_size, get frames() { return frames_for(frame_idx)}, free_requested, version}
}

function EventSelector(body, outer, events, stack_info, memory_view) {
    let events_div = outer
    .append("div")
    .attr('style', 'grid-column: 1; grid-row: 1; overflow: auto; font-family: monospace')

    let events_selection = events_div
    .selectAll("pre")
    .data(events)
    .enter()
    .append("pre")
    .text((e) => formatEvent(e))
    .attr('style', '')

    let selected_event_idx = null

    let es = {
        select(idx) {
            if (selected_event_idx !== null) {
                let selected_event = d3.select(events_div.node().children[selected_event_idx])
                selected_event.attr('style', '')
            }
            if (idx !== null) {
                let div = d3.select(events_div.node().children[idx])
                div.attr('style', `background-color: ${schemeTableau10[5]}`)
                let [reserved, allocated] = memory_view.draw(idx)
                let enter = () => eventStack(div.datum(), allocated, reserved)
                stack_info.highlight(enter)
                div.node().scrollIntoViewIfNeeded(false)
            } else {
                memory_view.draw(0)
            }
            selected_event_idx = idx
        }
    }
    body.on('keydown', (e) => {
        let actions = {ArrowDown: 1, ArrowUp: -1}
        if (selected_event_idx !== null && e.key in actions) {
            let new_idx = selected_event_idx + actions[e.key]
            es.select(Math.max(0, Math.min(new_idx, events.length - 1)))
            e.preventDefault()
        }
    })

    stack_info.register(events_selection, (t) => eventStack(t.datum()), (t) => {}, (d) => es.select(d.datum().idx))

    return es
}

function formatSize(num) {
    let orig = num
    // https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size
    const units = ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"];
    for (const unit of units) {
        if (Math.abs(num) < 1024.0) {
        return `${num.toFixed(1)}${unit}B (${orig} bytes)`;
        }
        num /= 1024.0;
    }
    return `${num.toFixed(1)}YiB`;
}

function formatEvent(event) {
    function formatAddr(event) {
        let version = event.version == 0 ? "" : `_${event.version}`
        let prefix = (event.action == "segment_free" || event.action == "segment_alloc") ? "s" : "b"
        return `${prefix}${event.addr.toString(16)}_${event.version}`
    }
    let stream = event.stream == 0 ? "" : `\n              (stream ${event.stream})`
    switch (event.action) {
        case 'oom':
            return `OOM (requested ${formatSize(event.size)}, CUDA has ${formatSize(event.device_free)} memory free)${stream}`
        case 'snapshot':
            return "snapshot"
        default:
            return `${event.action.padEnd(14)} ${formatAddr(event).padEnd(18)} ${formatSize(event.size)}${stream}`
    }
}

function eventStack(e, allocated, reserved) {
    let event = formatEvent(e)
    if (reserved !== undefined) {
        event = `(${formatSize(allocated)} allocated / ${formatSize(reserved)} reserved)\n${event}`
    }
    return event + "\n" + e.frames
}

function hashCode(num) {
  const numStr = num.toString();
  let hash = 0;
  for (let i = 0; i < numStr.length; i++) {
    const charCode = numStr.charCodeAt(i);
    hash = ((hash << 5) - hash) + charCode;
    hash = hash & hash; // Convert to 32-bit integer
  }
  return hash;
}

function addStroke(d) {
    d.attr('stroke', 'red')
    .attr('stroke-width', '2')
    .attr('vector-effect', 'non-scaling-stroke')
}

function removeStroke(d) {
    d.attr('stroke', '')
}

function MemoryView(outer, stack_info, trace_data, events) {
    let svg = outer.append('svg')
    .attr('style', 'grid-column: 2; grid-row: 1; width: 100%; height: 100%;')
    .attr('viewBox', '0 0 200 100')
    .attr('preserveAspectRatio', 'xMinYMin meet')
    let g = svg.append('g')
    let seg_zoom = zoom()
    seg_zoom.on('zoom', ({transform}) => {
        g.attr('transform', transform)
    })
    svg.call(seg_zoom)

    let segments_map = {}
    let block_map = {}

    let segments_data = trace_data.segments
    for (let [i, addr] of trace_data.segments.addr.entries()) {
        segments_map[addr] = Segment(addr, segments_data.size[i], segments_data.stream[i],
                                     null, trace_data.segment_version(addr, false))
    }

    let blocks_data = trace_data.blocks
    for (let [i, addr] of trace_data.blocks.addr.entries()) {
        block_map[addr] = Block(addr, blocks_data.size[i], blocks_data.real_size[i],
                                blocks_data.frames[i], blocks_data.pending_free[i],
                                trace_data.block_version(addr, false))
    }

    function simulate_memory(idx) {
        let l_segment_map = {...segments_map}
        let l_block_map = {...block_map}
        for (let i = events.length - 1; i > idx; i--) {
            let event = events[i]
            switch (event.action) {
                case 'free':
                    l_block_map[event.addr] = Block(event.addr, event.size, event.size, event.frames_idx, false, event.version)
                    break
                case 'free_requested':
                    l_block_map[event.addr].free_requested = false
                    break
                case 'free_completed':
                    l_block_map[event.addr] = Block(event.addr, event.size, event.size, event.frames_idx, true, event.version)
                    break
                case 'alloc':
                    delete l_block_map[event.addr]
                    break
                case 'segment_free':
                    l_segment_map[event.addr] = Segment(event.addr, event.size, event.stream, event.frames_idx, event.version)
                    break
                case 'segment_alloc':
                    delete l_segment_map[event.addr]
                    break
                case 'oom':
                    break
                default:
                    console.log(`unknown event: ${event.action}`)
                    break
            }
        }
        let new_segments = Object.values(l_segment_map)
        let new_blocks = Object.values(l_block_map)
        return [new_segments, new_blocks]
    }

    return {
        draw(idx) {
            let [segments_unsorted, blocks] = simulate_memory(idx)
            g.selectAll('g').remove()

            let segment_d = g.append('g')
            let block_g = g.append('g')
            let block_r = g.append('g')

            segment_d.selectAll('rect').remove()
            block_g.selectAll('rect').remove()
            block_r.selectAll('rect').remove()
            let segments = [...segments_unsorted].sort((x, y) => x.size == y.size ? (x.addr - y.addr) : (x.size - y.size))

            let segments_by_addr = [...segments].sort((x, y) => x.addr - y.addr)

            let max_size = segments.at(-1).size

            let xScale = scaleLinear([0, max_size], [0, 200])
            let padding = xScale.invert(1)

            let cur_row = 0
            let cur_row_size = 0
            for (let seg of segments) {
                seg.occupied = 0
                seg.internal_free = 0
                if (cur_row_size + seg.size > max_size) {
                    cur_row_size = 0
                    cur_row += 1
                }
                seg.offset = cur_row_size
                seg.row = cur_row
                cur_row_size += seg.size + padding
            }

            let num_rows = cur_row + 1

            let yScale = scaleLinear([0, num_rows], [0, 100])

            let segments_selection = segment_d.selectAll('rect').data(segments).enter()
            .append('rect')
            .attr('x', (x) => xScale(x.offset))
            .attr('y', (x) => yScale(x.row))
            .attr('width', (x) => xScale(x.size))
            .attr('height', yScale(4/5))
            .attr('stroke', 'black')
            .attr('stroke-width', '1')
            .attr('vector-effect', 'non-scaling-stroke')
            .attr('fill', 'white')

            stack_info.register(segments_selection,
                (d) => {
                    addStroke(d)
                    let t = d.datum()
                    let free = t.size - t.occupied
                    let internal = ""
                    if (t.internal_free > 0) {
                        internal = ` (${t.internal_free/free*100}% internal)`
                    }
                    return `s${t.addr.toString(16)}_${t.version}: segment ${formatSize(t.size)} allocated, ` +
                           `${formatSize(free)} free${internal} (stream ${t.stream})\n${t.frames}`
                },
                (d) => {
                    d.attr('stroke', 'black')
                    .attr('stroke-width', '1')
                    .attr('vector-effect', 'non-scaling-stroke')
                }
            )

            function find_segment(addr) {
                let left = 0;
                let right = segments_by_addr.length - 1;
                while (left <= right) {
                    let mid = Math.floor((left + right) / 2);
                    if (addr < segments_by_addr[mid].addr) {
                        right = mid - 1;
                    } else if (addr >= segments_by_addr[mid].addr + segments_by_addr[mid].size) {
                        left = mid + 1;
                    } else {
                        return segments_by_addr[mid];
                    }
                }
                return null;
            }

            for (let b of blocks) {
                b.segment = find_segment(b.addr)
                b.segment.occupied += b.real_size
                b.segment.internal_free += (b.size - b.real_size)
            }

            let block_selection = block_g.selectAll('rect').data(blocks).enter()
            .append('rect')
            .attr('x', (x) => xScale(x.segment.offset + (x.addr - x.segment.addr)))
            .attr('y', (x) => yScale(x.segment.row))
            .attr('width', (x) => xScale(x.real_size > padding/4 ? x.real_size - padding/4 : x.real_size))
            .attr('height', yScale(4/5))
            .attr('fill', (x, i) => x.free_requested ? 'red' : schemeTableau10[Math.abs(hashCode(x.addr)) % schemeTableau10.length])

            stack_info.register(block_selection, (d) => {
                addStroke(d)
                let t = d.datum()
                let requested = ""
                if (t.free_requested) {
                    requested = " (block freed but waiting due to record_stream)"
                }
                return `b${t.addr.toString(16)}_${t.version} ` +
                       `${formatSize(t.real_size)} allocation${requested} (stream ${t.segment.stream})\n` + t.frames
            }, removeStroke)

            let free_selection = block_r.selectAll('rect').data(blocks).enter()
            .append('rect')
            .attr('x', (x) => xScale(x.segment.offset + (x.addr - x.segment.addr) + x.real_size))
            .attr('y', (x) => yScale(x.segment.row))
            .attr('width', (x) => xScale(x.size - x.real_size))
            .attr('height', yScale(4/5))
            .attr('fill', (x, i) => 'red')

            stack_info.register(free_selection, (d) => {
                addStroke(d)
                let t = d.datum()
                return `Free space lost due to rounding ${formatSize(t.size - t.real_size)}` +
                       ` (stream ${t.segment.stream})\n` + t.frames
            }, removeStroke)

            let reserved = segments.reduce((x, y) => x + y.size, 0)
            let allocated = blocks.reduce((x, y) => x + y.real_size, 0)
            return [reserved, allocated]
        }
    }
}

function StackInfo(outer) {
    let stack_trace = outer
    .append('pre')
    .attr("style", "grid-column: 1 / 3; grid-row: 2; overflow: auto")
    let selected = {
        enter: () => { stack_trace.text("") },
        leave: () => {},
    }
    return {
        register(dom, enter, leave = (e) => {}, select = (e) => {}) {
            dom.on('mouseover', (e) => {
                selected.leave()
                stack_trace.text(enter(d3.select(e.target)))
            })
            .on('mousedown', (e) => {
                let obj = d3.select(e.target)
                selected = {enter: () => stack_trace.text(enter(obj)), leave: () => leave(obj)}
                select(obj)
            })
            .on('mouseleave', (e) =>  {
                leave(d3.select(e.target))
                selected.enter()
            })
        },
        highlight(enter, leave = () => {}) {
            selected = {enter: () => stack_trace.text(enter()), leave: leave}
            selected.enter()
        }
    }
}

function main() {
    // add some supplement information to trace_data
    let events = createEvents()

    let body = d3.select("body")

    let outer = body.append("div")
    .attr('style', "display: grid; grid-template-columns: 1fr 2fr; grid-template-rows: 2fr 1fr; height: 100%; gap: 10px")

    let stack_info = StackInfo(outer)
    let memory_view = MemoryView(outer, stack_info, trace_data, events)
    let event_selector = EventSelector(body, outer, events, stack_info, memory_view)

    window.addEventListener('load', function() {
        event_selector.select(events.length > 0 ? events.length - 1 : null)
    });
}
main()

</script>
</body>
</html>
"""

if __name__ == "__main__":
    import os.path
    thedir = os.path.realpath(os.path.dirname(__file__))
    if thedir in sys.path:
        # otherwise we find cuda/random.py as random...
        sys.path.remove(thedir)
    import argparse

    fn_name = 'torch.cuda.memory._snapshot()'
    pickled = f'pickled memory statistics from {fn_name}'
    parser = argparse.ArgumentParser(description=f'Visualize memory dumps produced by {fn_name}')

    subparsers = parser.add_subparsers(dest='action')

    def _output(p):
        p.add_argument('-o', '--output', default='output.svg', help='flamegraph svg (default: output.svg)')

    description = 'Prints overall allocation statistics and a visualization of how the allocators segments are currently filled.'
    stats_a = subparsers.add_parser('stats', description=description)
    stats_a.add_argument('input', help=pickled)

    description = 'Prints buffer of the most recent allocation events embedded in the snapshot in a Pythonic style.'
    trace_a = subparsers.add_parser('trace', description=description)
    trace_a.add_argument('input', help=pickled)

    description = 'Generate a flamegraph that visualizes what memory is stored in each allocator segment (aka block)'
    segments_a = subparsers.add_parser('segments', description=description)
    segments_a.add_argument('input', help=pickled)
    _output(segments_a)

    description = "Generate a flamegraph the program locations contributing to CUDA memory usage."
    memory_a = subparsers.add_parser('memory', description=description)
    memory_a.add_argument('input', help=pickled)
    _output(memory_a)

    description = 'Generate a flamegraph that shows segments (aka blocks) that have been added ' \
        'or removed between two different memorys snapshots.'
    compare_a = subparsers.add_parser('compare', description=description)
    compare_a.add_argument('before', help=pickled)
    compare_a.add_argument('after', help=pickled)
    _output(compare_a)

    plots = (
        ("trace_plot", "Generate a visualization over time of the memory usage recorded by the trace as an html file."),
        ("segment_plot", "Visualize how allocations are packed into allocator segments at each point in a trace as an html file.")
    )
    for cmd, description in plots:
        trace_plot_a = subparsers.add_parser(cmd, description=description)
        trace_plot_a.add_argument('input', help=pickled)
        help = 'visualize trace from this device (default: chooses the only device with trace info or errors)'
        trace_plot_a.add_argument('-d', '--device', type=int, default=None, help=help)
        help = 'path to save the visualization(default: output.html)'
        trace_plot_a.add_argument('-o', '--output', default='output.html', help=help)
        if cmd == "trace_plot":
            help = 'visualize change to segments rather than individual allocations'
            trace_plot_a.add_argument('-s', '--segments', action='store_true', help=help)


    args = parser.parse_args()

    def _read(name):
        if name == '-':
            f = sys.stdin.buffer
        else:
            f = open(name, 'rb')
        data = pickle.load(f)
        if isinstance(data, list):  # segments only...
            data = {'segments': data, 'traces': []}
        return data

    def _write(name, data):
        with open(name, 'w') as f:
            f.write(data)

    if args.action == 'segments':
        data = _read(args.input)
        _write(args.output, segments(data))
    elif args.action == 'memory':
        data = _read(args.input)
        _write(args.output, memory(data))
    elif args.action == 'stats':
        data = _read(args.input)
        print(segsum(data))
    elif args.action == 'trace':
        data = _read(args.input)
        print(trace(data))
    elif args.action == 'compare':
        before = _read(args.before)
        after = _read(args.after)
        _write(args.output, compare(before, after))
    elif args.action == 'trace_plot':
        data = _read(args.input)
        _write(args.output, trace_plot(data, device=args.device, plot_segments=args.segments))
    elif args.action == 'segment_plot':
        data = _read(args.input)
        _write(args.output, segment_plot(data, device=args.device))