#!/usr/bin/env python3

import csv
from pathlib import Path
import argparse
from collections import defaultdict
import shlex
import re
from termcolor import colored
import sys
from prettytable import PrettyTable
import json
import clean_spice_files

sourcetodests = defaultdict(list)

RE_SUBCKT = re.compile(r'\.subckt\s+(?P<subcktname>[^ ]+)\s?(?P<nodenames>.*)')
# RE_MODEL = re.compile(r'^[\.]*model\s+(?P<modelname>[^ ]+)\s+(?P<modeltype>[^ ]+)')

# RE_MODEL_END = re.compile(r'^[\s\n]*ends(?P<end>.*)(\n|$)')
# RE_SUBCKT_END = re.compile(r'^.ends')
RE_GROUP = re.compile(r'[*\s]*?(?P<group>[^*\n]+?)[*\s]*(\n|$)')
RE_MODE_LINE = re.compile(r'(?P<modeid>\d+)\s*:\s*type=\s*(?P<modetype>.*)(\n|$)')
RE_INLINE_SUBCKT_MODEL = re.compile(r'^(?P<subcktname>[^ ]+)\s*\((?P<nodenames>.*)\)\s*(?P<modelname>[^ ]+)\s*(?P<parameters>.*)($|\n)')

RE_MODEL = re.compile(r'[\.]?model')

def order_parameters(paramname):
    priority = {
        'subckt': 0,
        'inlinesubckt': 1,
        'nodes': 2,
        'model': 3,
        'type': 4,
        'mode': 5,
        'modetype': 6,
        'lmin': 7,
        'lmax': 8,
        'wmin': 9,
        'wmax': 10,
        'level': 11,
        'tnom': 12,
        'version': 13,
        'tox': 14,
        'toxm': 15,
    }
    if paramname in priority:
        return (priority[paramname], paramname)
    return (1000, paramname)


def csv_to_pm3(parameters, csventry):
    # we have following orders
    # inlinesubckt nodes model type
    # inlinesubckt nodes model type mode modetype
    # model type
    # subckt nodes
    # subckt nodes model type
    result = []
    if parameters[0] == 'model':
        for entry in csventry:
            assert 'N/A' not in entry[:2], entry
            result.append(f'.model {entry[0]} {entry[1]}')
            for parameter, value in zip(parameters[2:], entry[2:]):
                if value != 'N/A':
                    result.append(f'+{parameter}= {value}')
    elif parameters[0] == 'subckt':
        lastsubckt = None
        firstsubckt = True
        for entry in csventry:
            assert 'N/A' not in entry[:2], entry
            if lastsubckt != entry[0]:
                if not firstsubckt:
                    result.append('.ends')
                firstsubckt = False
                result.append(f'.subckt {entry[0]} {" ".join(entry[1].split(","))}')
                lastsubckt = entry[0]
            startpoint = 2
            if entry[2] == 'model':
                assert 'N/A' not in entry[2:4], entry
                result.append(f'.model {entry[2]} {entry[3]}')
                startpoint = 4
            for parameter, value in zip(parameters[startpoint:], entry[startpoint:]):
                if value != 'N/A':
                    result.append(f'+{parameter}= {value}')
        if not firstsubckt:
            result.append('.ends')
    elif parameters[0] == 'inlinesubckt':
        lastsubcktdef = None
        lastmodel = None
        firstsubcktdef = True
        firstmodel = True
        for entry in csventry:
            assert 'N/A' not in entry[:4], entry
            if not entry[0].endswith('_dummy') and entry[0] != lastsubcktdef:
                if not firstsubcktdef:
                    result.append(f'ends {lastsubcktdef}')
                lastsubcktdef = entry[0]
                firstsubcktdef = False
                result.append(f'inline subckt {entry[0]} ({" ".join(entry[1].split(","))})')
            if 'mode' in set(parameters):
                if entry[2] != lastmodel:
                    if not firstmodel:
                        result.append('}')
                    firstmodel = False
                    result.append(f'model {entry[2]} {entry[3]} {{')
                    lastmodel = entry[2]
                result.append(f'{entry[4]}: type={entry[5]}')
                for parameter, value in zip(parameters[6:], entry[6:]):
                    result.append(f'+{parameter}= {value}')
            else:
                result.append(f'{entry[0]} ({" ".join(entry[1].split(","))}) {entry[2]}')
                for parameter, value in zip(parameters[3:], entry[3:]):
                    if value != 'N/A':
                        result.append(f'+{parameter}= {value}')
        if not firstmodel:
            result.append('}')
        if not firstsubcktdef:
            result.append(f'ends {lastsubcktdef}')
    result = clean_spice_files.basic_cleaning(result)
    result = clean_spice_files.makefixedwidthcolumn(result)
    return result

def pm3_to_csv(lines):
    result = []
    params = {}
    params['subckt'] = {}
    params['model'] = {}
    params['mode'] = {}
    scopetype = None
    leftovers = []
    def produce_line():
        finparams = {}
        finparams.update(params['subckt'])
        finparams.update(params['model'])
        finparams.update(params['mode'])
        if len(finparams) > 0:
            result.append(finparams)
        scopetype = None
    for line in lines:
        try:
            if not line.strip():
                continue
            origline = line
            line = line.strip()
            elements = shlex.split(line, posix=False)
            if line.startswith('*'):
                continue
            m = RE_MODE_LINE.match(line)
            if m:
                if scopetype == 'mode':
                    produce_line()
                params['mode'] = {}
                params['mode']['mode'] = m.group('modeid')
                params['mode']['modetype'] = m.group('modetype')
                scopetype = 'mode'
                continue
            ind = 0
            def parse_parameters():
                if scopetype is None:
                    return
                ind = 0
                prevkey = None
                while ind < len(elements):
                    if not elements[ind].strip():
                        ind += 1
                    elif elements[ind].endswith('='):
                        if scopetype is None:
                            print(colored(line, 'magenta'), file=sys.stderr)
                            ind += 2
                            continue
                        params[scopetype][elements[ind][:-1]] = elements[ind+1]
                        prevkey = elements[ind][:-1]
                        ind += 2
                    elif '=' in elements[ind]:
                        res = elements[ind].split('=')
                        params[scopetype][res[0]] = res[1]
                        prevkey = res[0]
                        ind += 1
                    elif prevkey is not None:
                        params[scopetype][prevkey] += elements[ind]
                        ind += 1
                    else:
                        print(colored(line, 'cyan'), file=sys.stderr)
                        raise Exception(line)
            if elements[0] == '.subckt':
                produce_line()
                params['subckt'] = {}
                params['model'] = {}
                params['mode'] = {}
                params['subckt']['subckt'] = elements[1]
                nodelist = []
                ind = 2
                for element in elements[2:]:
                    if '=' in element:
                        break
                    element = element.replace('(', '')
                    element = element.replace(')', '')
                    nodelist.append(element)
                    ind += 1
                params['subckt']['nodes'] = ','.join(nodelist)
                scopetype = 'subckt'
                elements = elements[ind:]
                parse_parameters()
            elif RE_MODEL.match(elements[0]):
                if scopetype != 'subckt':
                    produce_line()
                params['model'] = {}
                params['mode'] = {}
                params['model']['model'] = elements[1]
                params['model']['type'] = elements[2]
                scopetype = 'model'
                if len(elements) >=4 and elements[3] == '{':
                    continue
                elements = elements[3:]
                parse_parameters()
            elif elements[0].startswith('+'):
                elements[0] = elements[0][1:]
                parse_parameters()
            elif elements[0] == 'parameters':
                elements = elements[1:]
                parse_parameters()
            else:
                m = RE_INLINE_SUBCKT_MODEL.match(line)
                if m:
                    if scopetype is not None:
                        produce_line()
                    params['subckt'] = {}
                    params['model'] = {}
                    params['mode'] = {}
                    scopetype = 'subckt'
                    params['subckt']['inlinesubckt'] = m.group('subcktname')
                    params['subckt']['nodes'] = ','.join(m.group('nodenames').split(' '))
                    params['model']['model'] = m.group('modelname')
                    params['model']['type'] = params['subckt']['inlinesubckt']
                    elements = m.group('parameters').split()
                    parse_parameters()
                else:
                    if not line.strip().startswith('ends') and not line.strip().startswith('inline subckt'):
                        if scopetype is None or not line.strip() in ['{', '}']:
                            leftovers.append(origline.rstrip())
                            print(colored(line, 'yellow'), file=sys.stderr)
        except:
            print(colored(line, 'red'), file=sys.stderr)
            raise
    produce_line()
    parameters = set()
    for entry in result:
        parameters.update(entry.keys())

    parameters = sorted(parameters, key=order_parameters)

    csvcontent = [parameters]

    for entry in result:
        row = []
        for param in parameters:
            if param in entry:
                row.append(entry[param])
            else:
                row.append('N/A')
        csvcontent.append(row)

    return csvcontent, leftovers

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "input_dir",
        help="The path to the directory containing skywater-pdk",
        type=Path)
    parser.add_argument(
        "--sourcetodests",
        help="Mapping from source files to destination files",
        type=Path)
    parser.add_argument(
        "--dry-run",
        help="Do not perform actions on filesystem, only print",
        action="store_true")

    args = parser.parse_args()

    files = sorted(args.input_dir.rglob('*.pm3'))

    allfilesnum = len(files)

    parameters = []
    # file model parameter value
    filesdata = defaultdict(lambda: defaultdict(dict))

    for num, filename in enumerate(files):
        with open(filename, 'r') as data:
            print(f'[{num:05d}/{allfilesnum:05d}]: {filename}')
            print(f'[{num:05d}/{allfilesnum:05d}]: {filename}', file=sys.stderr)
            # parameters across many entries
            lines = []
            for line in data.readlines():
                lines.append(line)

            result, leftovers = pm3_to_csv(lines)
            if len(result) > 1:
                with open(filename.with_suffix('.csv'), 'w', newline='') as csvfile:
                    writer = csv.writer(csvfile, delimiter=';')
                    writer.writerows(result)
                    sourcetodests[str(filename)].append(str(Path(filename).with_suffix('.csv')))

                with open(filename.with_suffix('.table'), 'w') as tablefile:
                    t = PrettyTable(result[0])
                    t.align = 'r'
                    t.border = False
                    for line in result[1:]:
                        t.add_row(line)
                    tablefile.write(str(t))
                    sourcetodests[str(filename)].append(str(Path(filename).with_suffix('.table')))
                # print(colored('\n'.join(csv_to_pm3(result[0], result[1:])), 'green'))
                binnedspice = csv_to_pm3(result[0], result[1:])
                csv2bsimname = filename.with_suffix(f'.csv2bsim{filename.suffix}')
                with open(csv2bsimname, 'w') as csv2bsimfile:
                    csv2bsimfile.write('\n'.join(binnedspice))
                with open(filename.with_suffix(f'.csvwrap{filename.suffix}'), 'w') as csvwrapfile:
                    leftovers.append(f'.include "{csv2bsimname}"')
                    csvwrapfile.write('\n'.join(leftovers))
    with open(args.sourcetodests, 'w') as srctodst:
        json.dump(sourcetodests, srctodst, indent=2)