#!/usr/bin/env python3
# poner-gnd: la masa del circuito se llama GND (el nodo 0 ya no se admite).
#
# Uso: poner-gnd [--probar] [ARCHIVO.d|.ds|.ds2|DIRECTORIO ...]
#
# Sin argumentos recorre todo el corpus. En las ramas de las secciones ELEMENTS
# del circuito cambia el nodo 0 por GND; si el circuito no tiene ni 0 ni GND, el
# nodo mas chico pasa a ser GND (el menor numero, o si son todos nombres el
# primero en orden alfabetico). Las ramas de los modelos (MODEL
# DESCRIPTION y parecidas) no se tocan: sus nodos son las patas del modelo, no
# nodos del circuito. Las ramas comentadas (# R1, 1-2 = ...) tambien cambian,
# para que al descomentarlas sigan valiendo.
#
# --probar dice que cambiaria sin escribir nada.

import os, re, sys

CORPUS = os.path.dirname(os.path.abspath(__file__))

# encabezados de seccion, como los reconoce slex/d2datdinc/dato.c
SECCIONES = ['CIRCUIT DESCRIPTION', 'ELEMENT DESCRIPTION', 'ELEMENT DEFINITION',
             'ELEMENTS', 'OUTPUTS', 'DEFINED PARAMETERS', 'INITIAL CONDITIONS',
             'RUN CONTROLS', 'CONTROL EQUATIONS', 'EQUATIONS', 'VARIABLES',
             'CONSTANTS', 'FUNCTIONS', 'FORCED VARIABLES', 'DATA', 'DATABLOB',
             'MODELS', 'MODEL DESCRIPTION', 'MULTIELEMENT DEFINITION', 'RESULTADOS',
             'GRAFICOS', 'ACEPTAR GRAFICOS', 'RERUN DESCRIPTION', 'CONTINUE',
             'RE-OUTPUT', 'PARAPLOT', 'CONTROLS', 'ECUACIONES', 'SENSITIVITY',
             'MONTE CARLO', 'WORST CASE', 'WORST-CASE', 'OPTIMIZATION', 'RERUN',
             'EXECUTE', 'CLEAR', 'END']
DE_MODELOS = {'MODELS', 'MODEL DESCRIPTION', 'MULTIELEMENT DEFINITION'}

# NOMBRE o TIPO-NOMBRE, coma, nodos separados por '-', '=' (o ',' en el formato
# viejo de exprtk_varios: E1, 0-1 , VE1 - 1 , ...)
RAMA = re.compile(r'^(\s*#?\s*[A-Za-z][A-Za-z0-9_]*(?:-[A-Za-z0-9_]+)?\s*,\s*)'
                  r'([A-Za-z0-9]+(?:\s*-\s*[A-Za-z0-9]+)+)(\s*[=,].*)$', re.S)
# las que siguen en la misma linea, separadas por coma como permite SCEPTRE
# (C1, 2-0 = 1, R2, 2-0 = 12(11,13)); estas si piden el '='
OTRA = re.compile(r',\s*[A-Za-z][A-Za-z0-9_]*(?:-[A-Za-z0-9_]+)?\s*,\s*'
                  r'([A-Za-z0-9]+(?:\s*-\s*[A-Za-z0-9]+)+)(?=\s*=)')


def tramos(linea):
    """(desde, hasta) de los nodos de cada rama de la linea"""
    m = RAMA.match(linea)
    if not m:
        return []
    r = [m.span(2)]
    if m.group(3).lstrip().startswith('='):
        r += [o.span(1) for o in OTRA.finditer(linea, m.end(2))]
    return r


def seccion(linea):
    t = re.sub(r'[\s_]+', ' ', linea.split('#')[0]).strip().upper()
    if not t:
        return None
    for s in sorted(SECCIONES, key=len, reverse=True):
        if t == s or (s != 'END' and t.startswith(s + ' ') and s == 'ELEMENTS'):
            return s
    return None


def ramas(lineas):
    """indices de las lineas que son ramas del circuito, con sus nodos"""
    en_elementos, en_modelo = False, False
    for i, l in enumerate(lineas):
        s = seccion(l)
        if s:
            if s in DE_MODELOS:
                en_modelo = True
            elif s == 'CIRCUIT DESCRIPTION':
                en_modelo = False
            en_elementos = s == 'ELEMENTS' and not en_modelo
            continue
        if en_elementos:
            t = tramos(l)
            if t:
                yield i, t, [n.strip() for a, b in t for n in l[a:b].split('-')]


def convertir(ruta, probar):
    texto = open(ruta, encoding='utf-8', errors='surrogateescape').read()
    lineas = texto.splitlines(keepends=True)
    halladas = list(ramas(lineas))
    if not halladas:
        return None
    nodos = {n for _, _, ns in halladas for n in ns}
    if 'GND' in nodos and '0' not in nodos:
        return None
    if '0' in nodos:
        viejo = '0'
    else:
        # el mas chico: el menor numero, o si son todos nombres el primero en
        # orden alfabetico (como elige la masa slex/din/din.c)
        numeros = sorted((n for n in nodos if n.isdigit()), key=int)
        viejo = numeros[0] if numeros else sorted(nodos)[0]
    cambiadas = 0
    patron = r'(?<![A-Za-z0-9])%s(?![A-Za-z0-9])' % re.escape(viejo)
    for i, t, _ in halladas:
        l = lineas[i]
        for a, b in reversed(t):
            nuevos = re.sub(patron, 'GND', l[a:b])
            if nuevos != l[a:b]:
                l = l[:a] + nuevos + l[b:]
                cambiadas += 1
        lineas[i] = l
    if not probar:
        with open(ruta, 'w', encoding='utf-8', errors='surrogateescape') as f:
            f.write(''.join(lineas))
    return 'nodo %s -> GND en %d ramas' % (viejo, cambiadas)


def archivos(args):
    for a in args or [CORPUS]:
        if os.path.isdir(a):
            for d, _, fs in sorted(os.walk(a)):
                for f in sorted(fs):
                    if f.endswith(('.d', '.ds', '.ds2')):
                        yield os.path.join(d, f)
        else:
            yield a


def main():
    args = sys.argv[1:]
    probar = '--probar' in args
    args = [a for a in args if a != '--probar']
    cuenta = {}
    for ruta in archivos(args):
        r = convertir(ruta, probar)
        if r:
            print('%s: %s' % (os.path.relpath(ruta, CORPUS), r))
            clave = r.split(' en ')[0] if '->' in r else r
            cuenta[clave] = cuenta.get(clave, 0) + 1
    for k, v in sorted(cuenta.items(), key=lambda x: -x[1]):
        print('# %4d  %s' % (v, k))


if __name__ == '__main__':
    main()
