from typing import Dict
from typing import Any
from typing import Union


def printHam(ham: Dict[str, Any], digits: int = 4) -> None:
    """
    Display the basic information of the Hamiltonian.

    :param ham: the Hamiltonian dictionary
    :param digits: the precision of the numbers printed
    :return: None
    """

    qubitNum = ham['circuit']['qubits']

    print(f"\n====================\n1. Basic information\n====================\n")
    print(f"Title: `{ham['file']['title']}`")
    print(f"Qubits: {qubitNum}")
    print(f"System energy level: {ham['circuit']['sys_level']}")
    print(f"Sampling interval: {ham['circuit']['dt']} ns")
    print(f"Circuit duration: {ham['circuit']['max_time_ns']} ns")
    print(f"Calculation steps: {ham['circuit']['max_time_dt']}")

    # Obtain the max length name
    maxNameLengthDrift = 0 if len(ham["drift"]) == 0 else max([len(key) for key in ham["drift"]])
    maxNameLengthControl = 0 if len(ham["control"]) == 0 else max([len(key) for key in ham["control"]])
    maxNameLength = str(max(max(maxNameLengthDrift, maxNameLengthControl), 10))

    # Print abstract of operator
    print(f"\n============\n2. Operators\n============\n")
    qubitFormat = "{0: <5}  {1: <7}  {2: <" + maxNameLength + "}  {3: <9}  {4: <6} {5: <6}"
    print(qubitFormat.format('-' * 5, '-' * 7, '-' * 10, '-' * 9, '-' * 6, '-' * 6))
    print(qubitFormat.format("Qubit", "Type", "Name", "On qubits", "Pulses", "Amp"))
    for qubit in range(qubitNum):
        print(qubitFormat.format('-' * 5, '-' * 7, '-' * 10, '-' * 9, '-' * 6, '-' * 6))
        for key in ham["drift"]:
            drifts = ham["drift"][key]
            if qubit in drifts["on_qubits"]:
                print(qubitFormat.format(qubit, "Drift", key, f"{drifts['on_qubits']}", 0, f"{drifts['amp']}"))
        for key in ham["control"]:
            ctrls = ham["control"][key]
            if qubit in ctrls["on_qubits"]:
                print(qubitFormat.format(qubit, "Control", key, f"{ctrls['on_qubits']}",
                                         len(ctrls['waveforms']), "-"))

    # Print abstract of waveforms
    def paraRound(para: Dict[str, Any]) -> Union[Dict[str, Any], None]:
        """ Reduce the length of pulse Parameters """
        if para is None:
            return None
        else:
            for key in para:
                para[key] = round(para[key], digits)
            return para

    print(f"\n============\n3. Waveforms\n============\n")
    qubitFormat = "{0: <9}  {1: <" + maxNameLength + "}  {2: <20} {3: <5}  {4: <7}  {5: <45}"
    print(qubitFormat.format('-' * 9, '-' * 10, '-' * 15, '-' * 5, '-' * 7, '-' * 45))
    print(qubitFormat.format("On qubits", "Control", "Waveform", "Start", "Duration", "Params (Sequences)"))
    for key in ham["control"]:
        ctrls = ham["control"][key]
        if len(ctrls['waveforms']) > 0:
            print(qubitFormat.format('-' * 9, '-' * 10, '-' * 15, '-' * 5, '-' * 7, '-' * 45))
        for wave in ctrls['waveforms']:
            waveName = ""
            wavePara = ""
            if wave['func'] is None:
                waveName = "Manual Sequence"
                wavePara = f"Sequence contains {len(wave['sequence'])} pieces"
            elif callable(wave['func']):
                waveName = "Manual Wave"
                wavePara = f"{paraRound(wave['para'])}"
            elif isinstance(wave['func'], str):
                waveName = wave['func']
                wavePara = f"{paraRound(wave['para'])}"
            print(qubitFormat.format(f"{ctrls['on_qubits']}", key, waveName, wave['insert_ns'],
                                     wave['duration_ns'], wavePara))