#!/usr/bin/env python3
# Copyright 2020 The Skywater PDK Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import json
import os
import pprint
import re
import shutil
import sys

from collections import defaultdict


RE_EXPAND_NUMBERS = re.compile(r'( [ \-]?(([0-9]+\.[0-9]+)|(0))\b)')
RE_FIX_VERSION = re.compile(r'(?P<v>VERSION)\s*(?P<a>[0-9])\.(?P<b>[0-9])0+', re.I)
RE_NEWLINE = re.compile(r'\n\n')


RE_MACRO = re.compile('MACRO (?P<cellname>[^\\n]*?)$\\n(?P<content>.*?\\n?)END (?P=cellname)\\n', re.DOTALL|re.MULTILINE)
RE_FOREIGN = re.compile('^   FOREIGN (?P<cellname>.*?) ;$', re.MULTILINE)
RE_PIN = re.compile('  PIN (?P<pinname>[^\\n]*)\\n(?P<content>.*?\\n)  END (?P=pinname)\\n', re.DOTALL|re.MULTILINE)
RE_PIN_TYPE = re.compile("USE ([^;]*);")
RE_PIN_DIR = re.compile("DIRECTION ([^;]*);")

RE_PORT = re.compile('    PORT\n.*?    END\n', re.DOTALL)
RE_OBS  = re.compile('  OBS\n.*?  END\n',  re.DOTALL)

RE_RECT = re.compile(
        '^(\s+RECT)\s+'
        '(-?[0-9]+(?:\.[0-9]+)?)\s+'
        '(-?[0-9]+(?:\.[0-9]+)?)\s+'
        '(-?[0-9]+(?:\.[0-9]+)?)\s+'
        '(-?[0-9]+(?:\.[0-9]+)?)\s+'
        '(;\n)', re.MULTILINE)
RE_POLYGON = re.compile(
        '^(\s+POLYGON)\s+'
        '((?:(?:-?[0-9]+(?:\.[0-9]+)?)\s+)+)'
        '(;\n)', re.MULTILINE)

RE_ANTENNA = re.compile('((?P<a>\s+ANTENNA([^\s]*)AREA\s+[0-9.]+)\s+LAYER\s+[^\s]+\s*;)+', re.DOTALL|re.MULTILINE)

RE_BUS_PIN = re.compile('(?P<n>[^\[]*)(\[(?P<i>[0-9]*)\])?')


def rewrite_antenna(contents):
    contents = RE_ANTENNA.sub('\g<a> ;\n', contents)
    while RE_NEWLINE.search(contents):
        contents = RE_NEWLINE.sub('\n', contents)
    return contents


def r_key(m):
    o = []
    for i in m:
        try:
            i = float(i)
        except ValueError:
            pass
        o.append(i)

    return o


def pin_key(pn):
    """
    >>> pin_key('D')
    ('D', 0)

    >>> pin_key('D[0]')
    ('D', 0)

    >>> pin_key('D[220]')
    ('D', 220)

    >>> pin_key('X123')
    ('X123', 0)

    >>> pin_key('X123XY2')
    ('X123XY2', 0)
    """
    m = RE_BUS_PIN.match(pn)

    n = m.group('n')
    i = m.group('i')
    if i is None:
        i = 0
    return (n, int(i))


def sort_geo(p):
    p = p.splitlines(True)
    assert p[-1].strip() == 'END', '\n'+pprint.pformat(p)

    start = p.pop(0)
    end   = p.pop(-1)

    layers = []
    for l in p:
        s = l.strip()
        if s.startswith('LAYER'):
            layers.append((l, []))
        else:
            layers[-1][1].append(l)

    for l, r in layers:
        o1 = []
        for a in r:
            m = RE_RECT.match(a)
            if m:
                o1.append(list(m.groups()))
                continue
            m = RE_POLYGON.match(a)
            if m:
                a, b, c = m.groups()
                o1.append([a]+b.split()+[c])
                continue
            assert False, a

        positions = 0
        for b in o1:
            positions = max(positions, len(b))
            for i in range(0, len(b)-2):
                if '.' in b[i+1]:
                    a1, a2 = b[i+1].split('.', 1)
                else:
                    a1 = b[i+1]
                    a2 = ''
                while len(a2) < 6:
                    a2 += '0'
                b[i+1] = a1+'.'+a2

        o1.sort(key=r_key)
        #pprint.pprint(o1)

        maxlen = [0,]*(positions-2)
        for b in o1:
            assert b[0].strip() == 'RECT' or b[0].strip() == 'POLYGON', b
            assert b[-1].strip() == ';', b
            for i in range(0, len(b)-2):
                maxlen[i] = max(maxlen[i], len(b[i+1]))

        for b in o1:
            for i in range(0, len(b)-2):
                b[i+1] = b[i+1].rjust(maxlen[i])

        r[:] = o1

    layers.sort()

    for i, (l, r) in enumerate(layers):
        layers[i] = (l, [' '.join(a) for a in r])

    return ''.join([start]+[a+''.join(b) for a, b in layers]+[end])


def canonicalize_lef(contents):
    # Expand all the numbers into '-0.000000'
    def expand(m):
        before = contents[m.start(0):m.start(1)]
        after  = contents[m.end(1):m.end(0)]
        num = m.group(1)
        if '.' not in num:
            assert num == ' 0', num
            num = num + '.'
        if num[1] != '-' and num[1] != ' ':
            num = ' '+num
        while len(num) < len(' -0.000000'):
            num += '0'
        while len(num) > len(' -0.000000') and num[-1] == '0':
            num = num[:-1]
        #print('%20s'%repr(m.group(1)), repr(num))
        return before+num+after

    contents = RE_EXPAND_NUMBERS.sub(expand, contents)
    contents = RE_FIX_VERSION.sub('\g<v> \g<a>.\g<b>', contents)

    # Remove extra antenna constructions
    contents = rewrite_antenna(contents)

    m = RE_MACRO.search(contents)
    cellname = m.group('cellname')
    before = contents[:m.start(0)]
    body = m.group(0)
    end = contents[m.end(0):]

    # FIXME Hack!
    assert body.count('CLASS CORE ;') <= 1, "%s %s\n%s" % (body.count('CLASS CORE ;'), cellname, body)
    if 'tap' in cellname:
        if body.count('CLASS CORE ;') == 1:
            print("Rewriting class for TAP cell")
            assert 'CLASS CORE WELLTAP ;' not in body, body
            body = body.replace('CLASS CORE ;', 'CLASS CORE WELLTAP ;')
    elif 'fill' in cellname:
        if body.count('CLASS CORE ;') == 1:
            print("Rewriting class for FILL cell")
            assert 'CLASS CORE SPACER ;' not in body, body
            body = body.replace('CLASS CORE ;', 'CLASS CORE SPACER ;')
    #elif 'conb' in cellname:
    #    if body.count('CLASS CORE ;') == 1:
    #        print("Rewriting class for TIE cell")
    #        assert 'CLASS CORE TIELOW ;' not in body, body
    #        body = body.replace('CLASS CORE ;', 'CLASS CORE TIELOW ;\n  CLASS CORE TIEHIGH ;')

    # Sort the pins in the output
    output = [before]
    output_sig_pins = {}
    output_pwr_pins = {}

    last_pin_endpos = 0
    for pin in RE_PIN.finditer(body):
        between = body[last_pin_endpos:pin.start(0)]
        if between.strip():
            output.append(between)
        last_pin_endpos = pin.end(0)

        pn = pin.group('pinname')
        pc = pin.group('content')

        pc_other = []
        pc_ports = []

        last_port_endpos = 0
        for port in RE_PORT.finditer(pc):
            between = pc[last_port_endpos:port.start(0)]
            if between.strip():
                pc_other.extend(between.splitlines(True))
            last_port_endpos = port.end(0)

            pc_ports.append(port.group(0))

        for i, p in enumerate(pc_ports):
            po = sort_geo(p)
            pc_ports[i] = po

        pc_ports.sort()
        pc_other.extend(pc[last_port_endpos:].splitlines(True))
        pc_other.sort()

        pc = ''.join([f'  PIN {pn}\n']+pc_other+pc_ports+[f'  END {pn}\n'])

        pin_type = RE_PIN_TYPE.search(pc)
        if pin_type:
            pin_type = pin_type.group(1).strip()
        else:
            pin_type = 'SIGNAL'

        if pin_type in ('SIGNAL',):
            output_pins = output_sig_pins
        else:
            output_pins = output_pwr_pins

        assert pn not in output_pins, pn
        output_pins[pn] = pc

    for pn in sorted(output_sig_pins, key=pin_key):
        output.append(output_sig_pins[pn])

    for pn in sorted(output_pwr_pins):
        output.append(output_pwr_pins[pn])

    after = body[last_pin_endpos:]

    last_obs_endpos = 0
    for obs in RE_OBS.finditer(after):
        between = after[last_obs_endpos:obs.start(0)]
        if between.strip():
            output.append(between)
        last_obs_endpos = obs.end(0)

        oc = obs.group(0)
        oc = oc.replace('        RECT', '      RECT')
        oc = oc.replace('        POLYGON', '      POLYGON')
        oc = oc.replace('      LAYER', '    LAYER')
        oo = sort_geo(oc)
        output.append(oo)

    after = after[last_obs_endpos:]
    output.append(after)

    output.append(end)
    contents = ''.join(output)
    if not contents.endswith('END LIBRARY\n'):
        contents += 'END LIBRARY\n'

    contents = contents.replace(
        "# SPDX-License-Identifier: Apache-2.0\nVERSION",
        "# SPDX-License-Identifier: Apache-2.0\n\nVERSION")

    return contents


def rewrite(input_path, output_path):
    with open(input_path, 'r') as f:
        contents = f.read()

    contents = canonicalize_lef(contents)

    if output_path == '/dev/stdout':
        sys.stdout.write(contents)
    else:
        with open(output_path, 'w') as f:
            f.write(contents)


def main(files):
    for input_path, output_path in files:
        try:
            rewrite(input_path, output_path)
        except Exception as e:
            print(f'Failure {input_path} -> {output_path}')
            raise


if __name__ == "__main__":
    import doctest
    fails, _ = doctest.testmod()
    if fails>0:
        sys.exit(1)
    args = list(sys.argv[1:])
    if args[0] == '--inplace':
        files = []
        for a in args[1:]:
            files.append((a, a))
    else:
        files = []
        for a in args:
            files.append((a, '/dev/stdout'))

    sys.exit(main(files))