"""
This module performs schematic drawing using circuitikz from a netlist.

>>> from lcapy import Schematic
>>> sch = Schematic()
>>> sch.add('P1 1 0.1; down')
>>> sch.add('R1 3 1; right')
>>> sch.add('L1 2 3; right')
>>> sch.add('C1 3 0; down')
>>> sch.add('P2 2 0.2; down')
>>> sch.add('W 0 0.1; right')
>>> sch.add('W 0.2 0.2; right')
>>> sch.draw()

Copyright 2014--2021 Michael Hayes, UCECE
"""

# Strings starting with ;; are schematic options.  They are parsed in
# netfile.py and added to the opts attribute of the netlist.  They get
# passed to the draw method of the Schematic class.


from __future__ import print_function
from .latex import latex_format_label, latex_format_node_label
from .expr import Expr
from . import schemcpts
import sympy as sym
from .engformatter import EngFormatter
from .schemmisc import Pos
from .schemplacer import schemplacer
from .opts import Opts
from .netfile import NetfileMixin
from .system import LatexRunner, PDFConverter, tmpfilename
from os import path, remove
from collections import OrderedDict
from warnings import warn
import math


__all__ = ('Schematic', )


def display_matplotlib(filename, dpi=150, scale=2):

    from matplotlib.pyplot import figure, figaspect
    from matplotlib.image import imread

    img = imread(filename)
    h, w, d = img.shape
    width = scale * w / dpi
    height = scale * h / dpi
    fig = figure(figsize=(width, height))
    ax = fig.add_subplot(111)
    ax.imshow(img)
    ax.axis('off')
    fig.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)


def png_image_size(filename):

    import struct

    with open(filename, 'rb') as f:
        header = f.read(25)

    if (header[:8] != b'\211PNG\r\n\032\n' and (header[12:16] != b'IHDR')):
        raise Exception('%s not a png image' % filename)
    w, h = struct.unpack('>LL', header[16:24])
    width = int(w)
    height = int(h)
    return width, height


class SchematicOpts(Opts):

    def __init__(self):

        super (SchematicOpts, self).__init__(
            {'draw_nodes': 'primary',
             'label_values': True,
             'label_ids': True,
             'annotate_values': False,
             'label_nodes': 'primary',
             'scale' : 1.0,
             'dpi' : 150,
             'cpt_size' : 1.5,
             'node_spacing' : 2.0,
             'help_lines' : 0.0,
             'style' : 'american',
             'voltage dir' : 'RP'})


class Node(object):

    def __init__(self, name):

        self.name = name
        self._port = False
        self._count = 0
        parts = name.split('_')
        # Primary 1, 2, a, a_3.  Not primary _1, _2, _a, _a_3, 0_3, a_b_c
        self.primary = (name[0] != '_' and len(parts) <= 2) and not (name[0].isdigit() and len(parts) != 1)

        # See MX1.sch for example.
        if '._' in name:
            self.primary = False

        self.elt_list = []
        self.pos = 'unknown'
        # Sanitised name
        self.s = name.replace('.', '@')
        self.label = latex_format_node_label(self.name)
        self.labelpos = None
        self.pin = False
        self.pinlabel = ''
        self.namepos = None
        self.pinname = ''
        self.pinpos = None
        self.clock = False
        self.auxiliary = None
        # Reference to node that defines the pin
        self.ref = None
        # Node is connected to a ground, etc. (an implicit power
        # supply connection)
        self.implicit = False

    @property
    def basename(self):
        fields = self.name.split('.')
        return fields[-1]

    @property
    def cptname(self):
        fields = self.name.split('.')
        if len(fields) < 2:
            return None
        return fields[-2]

    def __repr__(self):
        return '%s @ (%s)' % (self.name, self.pos)

    def append(self, elt):
        """Add new element to the node"""

        if elt.type == 'P':
            self._port = True

        self.elt_list.append(elt)
        if elt.type not in ('A', 'O'):
            self._count += 1

    @property
    def count(self):
        """Number of elements (including wires but not open-circuits)
        connected to the node"""

        return self._count

    def belongs(self, cpt_name):
        return self.cptname == cpt_name

    def visible(self, draw_nodes):
        """Return true if node drawn.
        `draw_nodes' can be `all', 'none', 'connections', 'primary', None,
        True, or False."""

        if self.auxiliary:
            return False

        if draw_nodes in ('all', True):
            return True

        if self.pin:
            return False

        if self.implicit:
            return False

        if self._port:
            return True

        if draw_nodes in ('none', None, False):
            return False

        # Implied port
        if self.count == 1:
            return True

        if draw_nodes in ('connections', 'connected'):
            return self.count > 2

        if draw_nodes == 'primary':
            return self.primary

        raise ValueError('Unknown argument %s for draw_nodes' % draw_nodes)

    @property
    def port(self):
        """Return true if node is a port"""

        return self._port or self.count == 1

    def show_label(self, label_nodes):

        if self.label == '':
            return False

        name = self.basename

        # pins is for backward compatibility
        if label_nodes in ('none', 'pins', 'false', False):
            return False
        elif label_nodes in ('all', 'true', True):
            return True
        elif label_nodes == 'alpha':
            return self.primary and name[0].isalpha()
        elif label_nodes == 'primary':
            return self.primary

        # handle label_nodes = '{1, 2}' etc.
        if label_nodes[0] == '{' and label_nodes[-1] == '}':
            label_nodes = label_nodes[1:-1]
        labels = [foo.strip() for foo in label_nodes.split(',')]
        return self.label in labels

    def debug(self):
        print(' %s @ (%s), count=%d, pin=%s' % (self.name, self.pos,
                                                self.count, self.pin))


class Schematic(NetfileMixin):

    def __init__(self, filename=None, allow_anon=False, **kwargs):

        self.elements = OrderedDict()
        self.nodes = {}
        self.hints = False
        self._init_parser(schemcpts, allow_anon)
        self.cpt_size = 1.2
        self.node_spacing = 2.0
        self.scale = 1.0
        self.dummy_node = 0
        self.context = None
        self.debug = 0

        if filename is not None:
            self.netfile_add(filename)

    def __repr__(self):

        return self.netlist()

    def __getitem__(self, name):
        """Return component by name"""
        try:
            return self.elements[name]
        except KeyError:
            raise AttributeError('Unknown component %s' % name)

    def netfile_add(self, filename):
        """Add the nets from file with specified filename"""

        self._netfile_add(filename)

    def add(self, string):
        """Add a component to the netlist.
        The general form is: 'Name Np Nm args'
        where Np is the positive node and Nm is the negative node.

        A positive current is defined to flow from the positive node
        to the negative node.
        """

        self._add(string)

    def netlist(self):
        """Return the current netlist"""

        return '\n'.join([elt.__str__() for elt in self.elements.values()])

    def _node_add(self, nodename, elt, auxiliary=False):

        if nodename not in self.nodes:
            self.nodes[nodename] = Node(nodename)
        node = self.nodes[nodename]

        node.append(elt)

        if node.auxiliary is None:
            node.auxiliary = auxiliary
        elif node.auxiliary and not auxiliary:
            # An auxiliary node can be made a proper node.
            node.auxiliary = False

        return node

    def _cpt_add(self, cpt):

        if cpt.offset != 0 and len(cpt.node_names) == 2:
            # Change component and add wires.

            n1 = cpt.node_names[0]
            n2 = cpt.node_names[1]
            # If have only two parallel components, can used fixed node names.
            self.dummy_node += 1
            on1 = n1 + 'off%d' % self.dummy_node
            self.dummy_node += 1
            on2 = n2 + 'off%d' % self.dummy_node

            angle = 90
            size = cpt.offset
            if size < 0:
                size = -size
                angle = -angle

            w1 = 'W %s %s; rotate=%s, size=%s' % (n1, on1, cpt.angle + angle, size)
            w2 = 'W %s %s; rotate=%s, size=%s' % (n2, on2, cpt.angle + angle, size)

            self.add(w1)
            self.add(w2)

            # Add open-circuit to ensure alignment.
            o = 'O %s %s; rotate=%s, size=%s' % (n1, n2, cpt.angle, cpt.size)
            self.add(o)

            # Rename nodes
            parts = cpt.net.split(' ')
            parts[1] = on1
            parts[2] = on2
            net = ' '.join(parts)
            cpt.opts.strip('offset')
            self.add('%s; %s' % (net, cpt.opts))
            return

        def tex_name(name, subscript=None):

            if subscript is None:
                subscript = ''

            if len(name) > 1:
                name = r'\mathrm{%s}' % name
            if len(subscript) > 1:
                subscript = r'\mathrm{%s}' % subscript
            if len(subscript) == 0:
                return name

            return '%s_{%s}' % (name, subscript)

        # There are two possible labels for a component:
        # 1. Component identifier, e.g., R1
        # 2. Component value, expression, or symbol
        id_label = tex_name(cpt.type, cpt.id)
        value_label = None

        if cpt.type in ('O', 'P', 'W') or id_label.find('#') != -1:
            id_label = None

        if cpt.type in ('S', 'SW', 'U'):
            value_label = ''

        if cpt.args != ():

            # TODO, extend for mechanical and acoustical components.
            units_map = {'V': 'V', 'I': 'A', 'R': '$\Omega$',
                         'C': 'F', 'L': 'H'}

            expr = cpt.args[0]
            if cpt.classname in ('Vstep', 'Istep'):
                expr = '(%s) * Heaviside(t)' % expr
                value_label = Expr(expr, cache=False).latex_math()
            elif cpt.classname in ('Vs', 'Is'):
                value_label = Expr(expr, cache=False).latex_math()
            elif cpt.classname == 'TF':
                expr = sym.sympify(expr)
                if expr.is_Pow and expr.args[1] == -1:
                    value_label = '%s:1' % (1 / expr)
                else:
                    value_label = '1:%s' % expr
            elif cpt.type in ('F', 'H') and len(cpt.args) > 1:
                # This is hard to give a reasonable label since the
                # control current is specified by a voltage source.
                # The user will have to override manually.
                expr = cpt.args[1]
                value_label = Expr(expr, cache=False).latex_math()
            elif cpt.classname not in ('TP',):
                try:
                    # Handle things like 9/1000 that can occur
                    # when substituting cpt values.
                    value = float(sym.Rational(expr))
                    if cpt.type in units_map:
                        value_label = EngFormatter().latex_math(value, units_map[cpt.type])
                    else:
                        value_label = Expr(expr, cache=False).latex_math()

                except (ValueError, TypeError):
                    # This catches non numeric arg.
                    value_label = Expr(expr, cache=False).latex_math()

        # Currently, we only annnotated the component with the value,
        # expression, or symbol.  If this is not specified, it
        # defaults to the component identifier.  Note, some objects
        # we do not want to label, such as wires and ports.
        cpt.id_label = '' if id_label is None else latex_format_label(id_label)
        cpt.value_label = cpt.id_label if value_label is None else latex_format_label(value_label)

        if cpt.opts_string != '':
            self.hints = True

        if cpt.name in self.elements:
            warn('Overriding component %s' % cpt.name)
            # Need to search lists and update component.

        self.elements[cpt.name] = cpt

        for node in cpt.auxiliary_node_names:
            self._node_add(node, cpt, auxiliary=True)

        # Note, an auxiliary node can be trumped...
        for node in cpt.required_node_names:
            self._node_add(node, cpt, auxiliary=False)

    def _setup(self):
        # This is called before node positions are assigned.

        for elt in self.elements.values():
            elt.setup()

        for nodename, node in self.nodes.items():
            if node.ref is not None:
                continue
            if node.cptname is not None and not node.implicit:
                raise ValueError('Unreferenced pin connection %s for %s' % (node.name, node.elt_list))

    def make_graphs(self, debug=0):

        placer = schemplacer(self.elements, self.nodes, 'graph', debug)
        placer._make_graphs()
        return placer.xgraph, placer.ygraph

    def _positions_calculate(self, method='graph', debug=False):

        placer = schemplacer(self.elements, self.nodes, method, debug)
        self.width, self.height = placer.solve(self.node_spacing)

    def _tikz_draw(self, style_args='', **kwargs):

        method = kwargs.pop('method', 'graph')

        self._setup()

        self._positions_calculate(method, self.debug)

        # Note, scale does not scale the font size.
        opts = ['scale=%.2f' % self.scale,
                'transform shape',
                '/tikz/circuitikz/bipoles/length=%.2fcm' % self.cpt_size]
        opts.append(style_args)
        if 'options' in kwargs:
            opts.append(kwargs.pop('options'))

        if 'font' in kwargs:
            font = kwargs.pop('font')
            opts.append('font=' + font)

        if 'voltage dir' in kwargs:
            voltage_dir = kwargs.pop('voltage dir')
            opts.append('voltage dir=' + voltage_dir)

        s = r'\begin{tikzpicture}[%s]''\n' % ', '.join(opts)

        # Add preamble
        if 'preamble' in kwargs:
            s += '  ' + kwargs.pop('preamble') + '\n'

        help_lines = float(kwargs.pop('help_lines', 0))
        color = kwargs.pop('color', 'blue')
        if help_lines != 0:
            start = Pos(-0.5, -0.5) * self.node_spacing
            stop = Pos(self.width + 0.5, self.height + 0.5) * self.node_spacing

            s += r'  \draw[help lines, %s] (%s) grid [xstep=%s, ystep=%s] (%s);''\n' % (
                color, start, help_lines, help_lines, stop)

        # Write coordinates.  TODO, not all coordinates are needed
        # so those can be weeded out to simplify the generated file.
        for n in self.nodes.values():
            s += r'  \coordinate (%s) at (%s);''\n' % (n.s, n.pos)

        # Keyword args for second pass
        kwargs2 = kwargs.copy()

        # Pass 1: Draw components
        for elt in self.elements.values():
            if elt.ignore:
                continue
            if elt.directive:
                for key, val in elt.opts.items():
                    # Update kwargs
                    kwargs[key] = val

            s += elt.draw(**kwargs)
            s += elt.draw_nodes(**kwargs)
            s += elt.draw_pins()

        # Pass 2: Add the node labels
        for elt in self.elements.values():
            if elt.ignore:
                continue
            if elt.directive:
                for key, val in elt.opts.items():
                    kwargs2[key] = val
            s += elt.draw_node_labels(**kwargs2)

        # Add postamble
        if 'postamble' in kwargs:
            s += '  ' + kwargs.pop('postamble') + '\n'

        s += r'\end{tikzpicture}''\n'

        return s

    def tikz_draw(self, filename, **kwargs):

        root, ext = path.splitext(filename)

        self.debug = kwargs.pop('debug', 0)
        style = kwargs.pop('style', 'american')
        self.dpi = float(kwargs.pop('dpi', 150))
        self.cpt_size = float(kwargs.pop('cpt_size', 1.2))
        self.node_spacing = float(kwargs.pop('node_spacing', 2.0))
        self.scale = float(kwargs.pop('scale', 1.0))

        include = ''
        if 'include' in kwargs:
            include += kwargs.pop('include')
        if 'includefile' in kwargs:
            file = open(kwargs.pop('includefile'))
            include += file.read()
            file.close()

        if style == 'american':
            style_args = 'american currents, american voltages'
        elif style == 'british':
            style_args = 'american currents, european voltages'
        elif style == 'european':
            style_args = ('european currents, european voltages'
                          ', european inductors, european resistors')
        else:
            raise ValueError('Unknown style %s' % style)

        # For debugging when do not want to write to file
        nosave = kwargs.pop('nosave', False)

        latexrunner = LatexRunner(self.debug & 2)

        self.circuitikz_date, self.circuitikz_version = latexrunner.circuitikz_version()
        if self.circuitikz_date is None:
            raise RuntimeError('circuitikz is not installed')

        content = self._tikz_draw(style_args=style_args, **kwargs)

        if nosave:
            return

        if self.debug & 1:
            print('circuitikz version %s (%s)' % (self.circuitikz_version,
                                                  self.circuitikz_date))
            print('width=%d, height=%d, dpi=%d, cpt_size=%.2f, node_spacing=%.2f, scale=%.2f'
                  % (self.width, self.height, self.dpi,
                     self.cpt_size, self.node_spacing, self.scale))
            print('Nodes:')
            for node in self.nodes.values():
                node.debug()

        if ext in ('.pytex', '.schtex', '.pgf'):
            if include != '':
                warn('Include option ignored for non-standalone file')
            open(filename, 'w').write(content)
            return

        # Need amsmath for operatorname
        template = ('\\documentclass[a4paper]{standalone}\n'
                    '\\usepackage{amsmath}\n'
                    '\\usepackage{circuitikz}\n'
                    '\\usetikzlibrary{fit, shapes, arrows, patterns, decorations.text, decorations.markings}\n'
                    '%s\n'
                    '\\begin{document}\n%s\\end{document}')
        content = template % (include, content)
        tex_filename = filename.replace(ext, '.tex')
        open(tex_filename, 'w').write(content)

        if ext == '.tex':
            return

        pdf_filename = tex_filename.replace('.tex', '.pdf')
        latexrunner.run(tex_filename)
        if not (self.debug & 1):
            latexrunner.cleanup(tex_filename, pdf_filename)

        if not path.exists(pdf_filename):
            raise RuntimeError('Could not generate %s with pdflatex' %
                               pdf_filename)

        if ext == '.pdf':
            return

        if ext == '.svg':
            pdfconverter = PDFConverter(self.debug & 2)
            pdfconverter.to_svg(pdf_filename, root + '.svg')
            if not (self.debug & 1):
                remove(pdf_filename)
            return

        if ext == '.png':
            pdfconverter = PDFConverter(self.debug & 2)
            pdfconverter.to_png(pdf_filename, root + '.png', self.dpi)
            if not (self.debug & 1):
                remove(pdf_filename)
            return

        raise RuntimeError('Cannot create file of type %s' % ext)

    def draw(self, filename=None, **kwargs):
        """
        filename specifies the name of the file to produce.  If None,
        the schematic is displayed on the screen.

        Note, if using Jupyter, then need to first issue command %matplotlib inline

        kwargs include:
           'label_ids': True to show component ids
           'label_values': True to display component values
           'annotate_values': True to display component values as separate label
           'draw_nodes': True to show all nodes,
             False or 'none' to show no nodes,
             'primary' to show primary nodes,
             'connections' to show nodes that connect more than two components,
             'all' to show all nodes
           'label_nodes': True to label all nodes,
             False or 'none' to label no nodes,
             'primary' to label primary nodes,
             'alpha' to label nodes starting with a letter,
             'pins' to label nodes that are pins on a chip,
             'all' to label all nodes,
             'none' to label no nodes
           'include': name of file to include before \\begin{document}
           'style': 'american', 'british', or 'european'
           'scale': schematic scale factor, default 1.0
           'node_spacing': spacing between component nodes, default 2.0
           'cpt_size': size of a component, default 1.5
           'dpi': dots per inch for png files
           'help_lines': distance between lines in grid, default 0.0 (disabled)
           'debug': Non-zero to display debug information
        """

        # None means don't care, so remove.
        for key in list(kwargs.keys()):
            if kwargs[key] is None:
                kwargs.pop(key)

        # Remove options that may be overridden
        for elt in self.elements.values():
            for key in list(elt.opts.keys()):
                if key in kwargs:
                    elt.opts.remove(key)

        # Default options
        opts = SchematicOpts()
        for key, val in opts.items():
            if key not in kwargs:
                kwargs[key] = val

        # Global options (usually at end of netlist for historical reasons)
        # These can be overwritten
        for eltname in self.elements:
            elt = self.elements[eltname]
            if not elt.directive:
                continue
            for key, val in elt.opts.items():
                # val is a str
                kwargs[key] = val

        def in_ipynb():
            try:
                ip = get_ipython()
                cfg = ip.config

                kernapp = cfg['IPKernelApp']

                # Check if processing ipynb file for Jupyter notebook.
                if 'connection_file' in kernapp:
                    return True
                elif kernapp and kernapp['parent_appname'] == 'ipython-notebook':
                    return True
                else:
                    return False
            except (NameError, KeyError):
                return False

        if not self.hints:
            raise RuntimeWarning('No schematic drawing hints provided!')

        if in_ipynb() and filename is None:
            png = 'png' in kwargs and kwargs.pop('png')
            svg = 'svg' in kwargs and kwargs.pop('svg')

            if not png and not svg:
                svg = False

            if svg:
                try:
                    from IPython.display import SVG, display_svg

                    svgfilename = tmpfilename('.svg')
                    self.tikz_draw(svgfilename, **kwargs)

                    # Create and display SVG image object.
                    # Note, there is a problem displaying multiple SVG
                    # files since the later ones inherit the namespace of
                    # the first ones.
                    display_svg(SVG(filename=svgfilename))
                    return
                except:
                    pass

            from IPython.display import Image, display_png

            png_filename = tmpfilename('.png')
            self.tikz_draw(png_filename, **kwargs)

            # Create and display PNG image object.
            # There are two problems:
            # 1. The image metadata (width, height) is ignored
            #    when the ipynb file is loaded.
            # 2. The image metadata (width, height) is not stored
            #    when the ipynb file is written non-interactively.
            width, height = png_image_size(png_filename)
            # width, height specify the image dimension in pixels
            display_png(Image(filename=png_filename,
                              width=width, height=height))
            if not (self.debug & 1):
                remove(png_filename)
            return

        if filename is None:
            filename = tmpfilename('.png')
            # Thicken up lines to reduce aliasing causing them to
            # disappear, especially when using pdftoppm.
            self.tikz_draw(filename=filename,
                           options='/tikz/circuitikz/bipoles/thickness=2',
                           **kwargs)
            display_matplotlib(filename, self.dpi)
            if not (self.debug & 1):
                remove(filename)
            return

        self.tikz_draw(filename=filename, **kwargs)

def test():

    sch = Schematic()

    sch.add('P1 1 0.1')
    sch.add('R1 1 3; right')
    sch.add('L1 3 2; right')
    sch.add('C1 3 0; down')
    sch.add('P2 2 0.2')
    sch.add('W 0.1 0; right')
    sch.add('W 0 0.2; right')

    sch.draw()
    return sch
