from qiskit import QuantumCircuit, transpile
import ipywidgets as widgets
from IPython.display import HTML
import logging

class WidgetLayout:
    timeline_layout = { 
        'border': '1px inset #eee',
        'padding': '2px',
        'height': '275px',
        'overflow': 'auto',
        'width': '70%'
    }

    accordion_layout = {
        'border': '1px inset #eee',
        'padding': '0',
        'height': '275px',
        'width': '30%'
    }
    
    tabular_data_layout = {
        'padding': '5px',
        'grid_template_columns': 'repeat(2, 50%)'
    }

timeline_panel = widgets.VBox([], layout = { 'width': '100%' })
timeline_wpr = widgets.Box([timeline_panel], layout = WidgetLayout.timeline_layout)

stats_labels = [
    widgets.Label('1Q ops'), widgets.Label(''),
    widgets.Label('2Q ops'), widgets.Label(''),
    widgets.Label('Depth'), widgets.Label(''),
    widgets.Label('Size'), widgets.Label(''),
    widgets.Label('Width'), widgets.Label(''),
    
]

stats_panel = widgets.GridBox(stats_labels, layout = WidgetLayout.tabular_data_layout)

properties_panel = widgets.GridBox([], layout = WidgetLayout.tabular_data_layout)

accordion = widgets.Accordion(children = [stats_panel, properties_panel], layout = WidgetLayout.accordion_layout)
accordion.set_title(0, 'Circuit Stats')
accordion.set_title(1, 'Property Set')

main_view = widgets.HBox([timeline_wpr, accordion], layout = { 'width': '100%' })

modal = widgets.Output(layout = { 'width': '100%' })

tp_widget = widgets.VBox([
    main_view,
    modal
], layout = { 'width': '100%' })

tp_widget.add_class('tp-widget')

import sys
from types import ModuleType, FunctionType
from gc import get_referents

# Custom objects know their class.
# Function objects seem to know way too much, including modules.
# Exclude modules as well.
BLACKLIST = type, ModuleType, FunctionType

def getsize(obj):
    """sum size of object & members."""
    #if isinstance(obj, BLACKLIST):
    #    raise TypeError('getsize() does not take argument of type: '+ str(type(obj)))
    seen_ids = set()
    size = 0
    objects = [obj]
    while objects:
        need_referents = []
        for obj in objects:
            if not isinstance(obj, BLACKLIST) and id(obj) not in seen_ids:
            # if id(obj) not in seen_ids:
                seen_ids.add(id(obj))
                size += sys.getsizeof(obj)
                need_referents.append(obj)
        objects = get_referents(*need_referents)
    return size

from enum import Enum

class PassType(Enum):
    ANALYSIS = 'A'
    TRANSFORMATION = 'T'

class Property:
    LARGE_VALUE_THRESHOLD = 2 ** 20

    def __init__(self, name, type, value) -> None:
        self.name = name
        self.type = type
        
        if getsize(value) > self.LARGE_VALUE_THRESHOLD:
            self.value = 'LARGE_VALUE'
        else:
            self.value = value

    def __repr__(self) -> str:
        return f"{self.name} ({self.type.__name__}) : {self.value}"

class CircuitStats:
    def __init__(self) -> None:
        self.width = None
        self.size = None
        self.depth = None
        self.ops_1q = None
        self.ops_2q = None

    def __eq__(self, other):
        return self.width == other.width and self.size == other.size and self.depth == other.depth and self.ops_1q == other.ops_1q and self.ops_2q == other.ops_2q

    def __repr__(self) -> str:
        return f"CircuitStats(width={self.width}, size={self.size}, depth={self.depth}, 1Q ops={self.ops_1q}, 2Q ops={self.ops_2q})"

from time import time

class LogEntry:
    def __init__(self, levelname, msg, args) -> None:
        self.levelname = levelname
        self.msg = msg
        self.args = args
        self.time = time()

    def __repr__(self) -> str:
        return f"[{self.levelname}] {self.msg}"

class TranspilationStep:
    def __init__(self, name, type) -> None:
        self.name = name
        self.type = type
        self.docs = ''
        self.run_docs = ''
        self.duration = 0
        self.circuit_stats = CircuitStats()
        self.property_set = {}
        self.logs = []
        self.dag = None

    def __repr__(self) -> str:
        return f"(name={self.name}, type={self.type})"

class ButtonWithValue(widgets.Button):
    def __init__(self, *args, **kwargs):
        self.value = kwargs['value']
        kwargs.pop('value', None)
        super(ButtonWithValue, self).__init__(*args, **kwargs)

from datetime import datetime
import html

%matplotlib inline

def on_logs(btn):
    step = transpilation_sequence.steps[int(btn.value)]
    logs = step.logs

    #TEMP:
    if len(logs) > 22:
        logs[7].levelname = 'WARNING'
        logs[10].levelname = 'ERROR'
        logs[13].levelname = 'CRITICAL'
        logs[15].levelname = 'INFO'
        logs[16].levelname = 'INFO'
        logs[20].levelname = 'ERROR'
        logs[21].levelname = 'CRITICAL'
    #END-TEMP

    html_str = '<div class="logs-wpr">'
    for entry in logs:
        html_str = html_str + "<pre class='date'>{0}</pre><pre class='level {1}'>[{1}]</pre><pre class='log-entry {1}'>{2}</pre>".format(datetime.fromtimestamp(entry.time).strftime("%H:%M:%S.%f")[:-3], entry.levelname, entry.msg % entry.args)
    html_str = html_str + '</div>'
    
    show_modal(step.name + ' Logs', html_str, '')

def on_circ(btn):
    from qiskit.converters import dag_to_circuit
    from io import BytesIO
    from binascii import b2a_base64
    
    circ = dag_to_circuit(transpilation_sequence.steps[int(btn.value)].dag)
    fig = circ.draw('mpl', idle_wires = False, with_layout = False, scale = 0.8, fold = 20)

    bio = BytesIO()
    #fig.canvas.print_png(bio)
    fig.savefig(bio, format = 'png', bbox_inches = 'tight')

    img_data = b2a_base64(bio.getvalue()).decode()
    img_html = '<img src="data:image/png;base64,{}&#10;">'.format(img_data)

    show_modal('Circuit Plot', img_html, '840px')

def on_info(btn):
    step = transpilation_sequence.steps[int(btn.value)]

    html_str = '<pre class="help">' + step.docs + '</pre>'
    html_str = html_str + '<div class="help-header"><span style="color: #e83e8c;">' + step.name + '</span>.run(<span style="color: #0072c3;">dag</span>)</div>'
    html_str = html_str + '<pre class="help">' + step.run_docs + '</pre>'

    show_modal(step.name, html_str, '')

def on_property(btn):
    step_index, property_name = btn.value.split(',')
    property_set = transpilation_sequence.steps[int(step_index)].property_set
    property = property_set[property_name]

    html_str = '<table>'
    for v in property.value:
        ###
        # ??? if type(v) == tuple:
        if property_name == 'block_list':
            html_str = html_str + '<tr><td><pre>' + html.escape(str([node.name for node in v])) + '</pre></td></tr>'
        else:
            html_str = html_str + '<tr><td><pre>' + html.escape(str(v)) + '</pre></td></tr>'
    html_str = html_str + '</table>'
    show_modal('Property: ' + property_name, html_str, '')

def show_modal(title, content, width):
    modal.remove_class('show')
    html_str = """
    <script>
        // Fix z-index issue:
        var stackingContextRoot = document.querySelector(".tp-widget").closest("div.cell");
        if(stackingContextRoot) {{
            stackingContextRoot.classList.add("stacking-context");
        }}

        var closeButton = document.getElementsByClassName("close-wdgt-modal")[0];
        closeButton.addEventListener("click", function(ev) {{
            document.querySelectorAll(".output_wrapper").forEach(function(div) {{
                div.classList.remove("show");
            }});
        }});
    </script>
    <div class="wdgt-modal-background"></div>
    <div class="wdgt-modal" style="width: {width}">
        <div class="wdgt-modal-header">
            <h3>{title}</h3>
            <span class="close-wdgt-modal">
                <img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAMAAAAoLQ9TAAAAdVBMVEUAAABNTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU1NTU0N3NIOAAAAJnRSTlMAAQIDBAUGBwgRFRYZGiEjQ3l7hYaqtLm8vsDFx87a4uvv8fP1+bbY9ZEAAAB8SURBVBhXXY5LFoJAAMOCIP4VBRXEv5j7H9HFDOizu2TRFljedgCQHeocWHVaAWStXnKyl2oVWI+kd1XLvFV1D7Ng3qrWKYMZ+MdEhk3gbhw59KvlH0eTnf2mgiRwvQ7NW6aqNmncukKhnvo/zzlQ2PR/HgsAJkncH6XwAcr0FUY5BVeFAAAAAElFTkSuQmCC" width="16" height="16" alt="">
            </span>
        </div>
        <div class="wdgt-modal-content">{content}</div></div>
    """.format(title = title, content = content, width = width)
    modal.outputs = []
    modal.append_display_data(HTML(html_str))
    modal.add_class('show')

class TranspilationSequence:
    def __init__(self) -> None:
        self.steps = []
        self._collected_logs = {}
        self._prop_labels_map = {}

    def add_step(self, step) -> None:
        if step.name in self._collected_logs:
            step.logs = self._collected_logs[step.name]
            self._collected_logs.pop(step.name, None)

        ######
        if len(transpilation_sequence.steps) == 0:
            prev_stats = CircuitStats()
        else:
            prev_stats = self.steps[-1].circuit_stats

        if prev_stats.ops_1q != step.circuit_stats.ops_1q:
            stats_labels[1].value = str(step.circuit_stats.ops_1q)
            stats_labels[1].add_class('highlight')
        else:
            stats_labels[1].remove_class('highlight')

        if prev_stats.ops_2q != step.circuit_stats.ops_2q:
            stats_labels[3].value = str(step.circuit_stats.ops_2q)
            stats_labels[3].add_class('highlight')
        else:
            stats_labels[3].remove_class('highlight')

        if prev_stats.depth != step.circuit_stats.depth:
            stats_labels[5].value = str(step.circuit_stats.depth)
            stats_labels[5].add_class('highlight')
        else:
            stats_labels[5].remove_class('highlight')
            
        if prev_stats.size != step.circuit_stats.size:
            stats_labels[7].value = str(step.circuit_stats.size)
            stats_labels[7].add_class('highlight')
        else:
            stats_labels[7].remove_class('highlight')

        if prev_stats.width != step.circuit_stats.width:
            stats_labels[9].value = str(step.circuit_stats.width)
            stats_labels[9].add_class('highlight')
        else:
            stats_labels[9].remove_class('highlight')
        ######

        for prop_name in step.property_set:
            property = step.property_set[prop_name]
            if property.name in self._prop_labels_map:
                index = self._prop_labels_map[property.name]
            else:
                ####
                u = widgets.Label(value = property.name)
                if type(property.value) == list:
                    prop_label = widgets.Label('(list)', layout = { 'width': '80%' })
                    prop_button = ButtonWithValue(value = str(len(self.steps)) + ',' + property.name, description = '...', layout = { 'width': '20%' })
                    prop_button.on_click(on_property)
                    v = widgets.HBox([prop_label, prop_button], layout = { 'width': '100%' })
                else:
                    v = widgets.Label(value = str(property.value))
                index = len(properties_panel.children)
                properties_panel.children = properties_panel.children + (u, v)
                ####
                
                self._prop_labels_map[property.name] = index

            u = widgets.Label(value = property.name)
            if type(property.value) != list:
                properties_panel.children[index + 1].value = str(property.value)
        ######
            
        a = widgets.Label(step.type.value)
        a.add_class(step.type.name.lower())
        
        b = widgets.Label(step.name)

        from math import log10
        if step.duration > 0:
            duration_font_size = 10
            duration_font_size = 10 + round(log10(step.duration))
            t = widgets.Label(str(round(step.duration, 1)) + ' ms')
            t.add_class('fs' + str(duration_font_size))
        else:
            t = widgets.Label('')

        if type(step.logs) == list and len(step.logs) > 0:
            c = ButtonWithValue(value = str(len(self.steps)), tooltip = str(len(step.logs)) + ' logs entries', icon = 'align-justify', layout = { 'width': '32px' })
            c.on_click(on_logs)
        else:
            c = widgets.Label('', layout = { 'width': '32px' })

        # Due to a bug in DAGCircuit.__eq__, we can not use ``step.dag != None``
        from qiskit.dagcircuit import DAGCircuit
        if isinstance(step.dag, DAGCircuit):
            d = ButtonWithValue(value = str(len(self.steps)), icon = 'sliders', layout = { 'width': '32px' })
            d.on_click(on_circ)
        else:
            d = widgets.Label('', layout = { 'width': '32px' })

        e = ButtonWithValue(value = str(len(self.steps)), icon = 'info-circle', layout = { 'width': '32px' })
        e.on_click(on_info)
        
        item = widgets.GridBox([a, b, t, c, d, e], layout = { 'width': '100%', 'grid_template_columns': '20px auto 70px 35px 35px 35px', 'min_height': '35px' })
        item.add_class('transpilation-step')
        timeline_panel.children = timeline_panel.children + (item, )
        ######

        self.steps.append(step)

    def add_log_entry(self, pass_name, log_entry) -> None:
        if not pass_name in self._collected_logs:
            self._collected_logs[pass_name] = []

        self._collected_logs[pass_name].append(log_entry)

transpilation_sequence = TranspilationSequence()

class TranspilerLoggingHandler(logging.Handler):
    def __init__(self, *args, **kwargs):
        self.loggers_map = kwargs['loggers_map']
        kwargs.pop('loggers_map', None)
        super(TranspilerLoggingHandler, self).__init__(*args, **kwargs)

    def emit(self, record):
        log_entry = LogEntry(record.levelname, record.msg, record.args)
        transpilation_sequence.add_log_entry(self.loggers_map[record.name], log_entry)

from qiskit.transpiler.basepasses import AnalysisPass, TransformationPass

all_loggers = logging.Logger.manager.loggerDict
passes_loggers = {key:value for (key, value) in all_loggers.items() if key.startswith('qiskit.transpiler.passes.')}

loggers_map = {}
for _pass in AnalysisPass.__subclasses__():
    if _pass.__module__ in passes_loggers.keys():
        loggers_map[_pass.__module__] = _pass.__name__

for _pass in TransformationPass.__subclasses__():
    if _pass.__module__ in passes_loggers.keys():
        loggers_map[_pass.__module__] = _pass.__name__

handler = TranspilerLoggingHandler(loggers_map = loggers_map)
logger = logging.getLogger('qiskit.transpiler.passes')
logger.setLevel(logging.DEBUG)
logger.addHandler(handler)


_properties = {}

def transpiler_callback(**kwargs):
    pass_ = kwargs['pass_']

    pass_name = pass_.name()
    pass_type = ''
    if pass_.is_analysis_pass:
        pass_type = PassType.ANALYSIS
    elif pass_.is_transformation_pass:
        pass_type = PassType.TRANSFORMATION

    transpilation_step = TranspilationStep(pass_name, pass_type)
    transpilation_step.docs = pass_.__doc__
    transpilation_step.run_docs = getattr(pass_, 'run').__doc__

    duration_value = round(1000 * kwargs['time'], 2)
    transpilation_step.duration = duration_value

    # Properties
    _props_updated = False
    property_set = kwargs['property_set']
    for key in property_set:
        if key not in _properties.keys():
            _props_updated = True
            _properties[key] = property_set[key]

    if _props_updated:

        from collections import defaultdict
        for property_name in property_set:
            #if property_name == 'commutation_set':
            #    print(property_set[property_name])
            property_type = type(property_set[property_name])
            property_value = None
            if property_type in (int, float, bool, str):
                property_value = property_set[property_name]
            elif property_type == list:
                property_value = []
                for value in property_set[property_name]:
                    property_value.append(value)
            elif property_type == defaultdict:
                #print('defaultdict')
                property_value = []
                for key, value in property_set[property_name].items():
                    #print(key, value)
                    #print('------------------------')
                    property_value.append(str(key) + ': ' + str(value))
            else:
                property_value = '(' + property_type.__name__ + ')'

            transpilation_step.property_set[property_name] = Property(property_name, property_type, property_value)

    # circuit stats:
    dag = kwargs['dag']
    transpilation_step.circuit_stats.width = dag.width()
    transpilation_step.circuit_stats.size = dag.size()
    transpilation_step.circuit_stats.depth = dag.depth()

    nodes = dag.op_nodes(include_directives = False)
    circ_ops_1q = 0
    circ_ops_2q = 0
    for node in nodes:
        if len(node.qargs) > 1:
            circ_ops_2q += 1
        else:
            circ_ops_1q += 1
    
    transpilation_step.circuit_stats.ops_1q = circ_ops_1q
    transpilation_step.circuit_stats.ops_2q = circ_ops_2q

    # ???
    if transpilation_step.type == PassType.TRANSFORMATION and transpilation_step.circuit_stats.depth < 200:
        # if len(transpilation_sequence.steps) == 0 or transpilation_sequence.steps[-1].circuit_stats != transpilation_step.circuit_stats:
        transpilation_step.dag = dag

    transpilation_sequence.add_step(transpilation_step)

%%html
<style>
.tp-widget { border:1px solid #aaa; }
.p-TabPanel-tabContents { padding: 10px !important; }
.p-Collapse-header { padding: 1px 5px; background: #eee; }
.p-Collapse-open > .p-Collapse-header { background: #ddd; }
.p-Collapse-contents { padding-top: 0; padding-left: 0; padding-bottom: 0;
    padding-right: 0; height: 220px; background: #f5f5f5; }
.p-Collapse-contents button {
    width: 20%;
    background: #fff;
    text-align: center;
    padding: 0;
    font-weight: bold; }

.widget-gridbox { background: #f5f5f5; }
.widget-checkbox { padding-left: 40px; }
div.output_scroll { box-shadow: none }

.p-Accordion .widget-gridbox .widget-label { background-color: #fff; padding: 0 3px;
    font-family: 'Roboto Mono', monospace; font-size: 14px; }
                                            
.highlight { background-color: #ffc !important; font-weight: bold; }

.transpilation-step { background: #fff; border-bottom: 1px solid #ccc; }
.transpilation-step button { background: #fff; }
.transpilation-step .transformation { color: red; }
.transpilation-step .analysis { color: green; }

.transpilation-step div.fs10 { font-size: 10px; }
.transpilation-step div.fs11 { font-size: 11px; }
.transpilation-step div.fs12 { font-size: 12px; }
.transpilation-step div.fs13 { font-size: 13px; }
.transpilation-step div.fs14 { font-size: 14px; }
.transpilation-step div.fs15 { font-size: 15px; }
                            
.transpilation-step > :nth-child(1) { text-align: center; font-weight: bold; font-size: 14px; }
.transpilation-step > :nth-child(2) { font-family: 'Roboto Mono', monospace; font-size: 15px; }
.transpilation-step > :nth-child(3) { font-family: 'Roboto Mono', monospace; font-size: 10px; color: #900;}
.transpilation-step > :nth-child(4) { color: #ff9d85; }
.transpilation-step > :nth-child(5) { color: #b587f7; }
.transpilation-step > :nth-child(6) { color: #6ea2c9; }

.logs-wpr { display: grid; grid-template-columns: 70px 60px auto; }
.logs-wpr pre.date { font-size: 10px; }
.logs-wpr pre.level { font-size: 10px; text-align: right; padding-right: 5px; }
.logs-wpr pre.log-entry { font-size: 12px; }
.logs-wpr pre.DEBUG { color: #000000; }
.logs-wpr pre.INFO { color: #1c84a2; }
.logs-wpr pre.WARNING { color: #ed7723; }
.logs-wpr pre.ERROR { color: #d64e4a; }
.logs-wpr pre.CRITICAL { color: white; background: #d64e4a; }

div.output_area pre.help { font-family: Helvetica,Arial,sans-serif; font-size: 13px;
    border: 1px solid #ccc; padding: 10px;}
div.help-header {
    font-family: 'Roboto Mono', monospace;
    font-size: 12px;
    border: 1px solid #ccc;
    border-bottom: none;
    margin-top: 4px;
    padding: 5px 10px;
    font-weight: bold;
    background: #f5f5f5;
}

.stacking-context { position: relative; z-index: 99997; }

.wdgt-modal-background {
    width: 100%;
    height: 100%;
    background-color: rgba(0,0,0,0.5);
    position: fixed;
    top: 0;
    left: 0;
    display: none;
    z-index: 99998;
}

.wdgt-modal {
    position: fixed;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    margin: auto;
    display: none;
    width: 700px;
    height: 300px;
    background-color: #fff;
    box-sizing: border-box;
    border: 1px solid #eee;
    border-radius: 5px;
    overflow:hidden;
    z-index: 99999;
}

.wdgt-modal-header {
    background-color: #ccc;
    border-bottom: 1px solid #dddddd;
    box-sizing: border-box;
    height: 50px;
    padding: 0;
}
.wdgt-modal-header h3 {
    margin: 0 !important;
    box-sizing: border-box;
    padding-left: 15px;
    line-height: 50px;
    color: #4d4d4d;
    font-size: 16px;
    display: inline-block;
}
.wdgt-modal-header span.close-wdgt-modal {
    box-sizing: border-box;
    border-left: 1px solid #dddddd;
    float: right;
    line-height: 50px;
    margin-top: 18px;
    padding: 0 15px 0 15px;
    cursor: pointer;
}
.wdgt-modal-header span.close-wdgt-modal:hover img {
    opacity: 0.6;
}

.wdgt-modal-content { overflow-y: auto; height: 250px; margin: 0; padding: 10px; }
.wdgt-modal-content pre { background: #fff; text-align: left; }
.show .wdgt-modal, .show .wdgt-modal-background { display: block; }

div.output_area .wdgt-modal-content img { max-width: none; /*height: 100%;*/ }
</style>

from qiskit.circuit.random import random_circuit
from qiskit.test.mock import FakeAlmaden

backend = FakeAlmaden()

transpilation_sequence = TranspilationSequence()
circ = random_circuit(5, 20, measure = True)

tp_widget

start_time = time()
circ_tr = transpile(circ, backend, optimization_level = 3, callback = transpiler_callback)
print(time() - start_time)

