#! /usr/bin/env python3

import os
import pathlib
import re
import subprocess
import sys
import multiprocessing

assert sys.version_info.major == 3 and sys.version_info.minor >= 7, sys.version_info

if len(sys.argv) != 2:
    print(f"usage: {sys.argv[0]} INPUT_DIR\n\n"
          f"E.g.:  {sys.argv[0]} output_nse_D1Q3_SRT_4_4_symbolic", file=sys.stderr)
    sys.exit(1)

# input_dir: directory with subs_list files and EPDE coeffs (output of lbmat)
input_dir = pathlib.Path(sys.argv[1])
assert input_dir.exists(), input_dir

# subs_dir: directory with the "subs_list*" files
subs_dir = input_dir / "subs"
assert subs_dir.exists(), subs_dir

# target_files: "coef_*" files where the substitutions will be expanded
# (the glob pattern should match all PDEs)
target_files = sorted(input_dir.glob("pde*_subs/coef_*"))
if not target_files:
    print(f"no coef_* files found in {input_dir}", file=sys.stderr)
    sys.exit(1)

def natural_sorted(l, *, reverse=False):
    """ Sort the given list in the way that humans expect.
        For example: ["img_1", "img_2", "img_3", "img_20"]
        instead of ["img_1", "img_2", "img_20", "img_3"]

        :param bool reverse: if True, numbers are extracted in reverse order
    """
    convert = lambda text: int(text) if text.isdigit() else text
    if reverse is True:
        alphanum_key = lambda key: [convert(c) for c in reversed(re.split('([0-9]+)', str(key)))]
    else:
        alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', str(key))]
    return sorted(l, key=alphanum_key)

def simplify(expr):
    try:
        if isinstance(expr, str):
            expressions = expr
        else:
            expressions = "\n".join(expr)
        stdout = subprocess.run("./simplify", text=True, input=expressions, capture_output=True, check=True).stdout
    except subprocess.CalledProcessError as e:
        print(f"expressions: {expressions}")
        print(f"GiNaC error: {e.stderr}")
        raise
    results = [line.strip() for line in stdout.splitlines()]
    if isinstance(expr, str):
        # only 1 expression
        return results[0].strip()
    else:
        return results

def simplify_par(expressions, pool, nproc):
    chunks = [expressions[i:i + nproc] for i in range(0, len(expressions), nproc)]
    results = pool.map(simplify, chunks)
    # flatten list of lists
    results = [item for sublist in results for item in sublist]
    return results

# map of all substitutions
subs_map = {}

def substitute_coeffs(expression):
    def repl(match):
        coeff = match.group(0)
        if coeff in subs_map:
            return "(" + subs_map[coeff] + ")"
        else:
            print(f"ERROR: coeff {coeff} not found in the map, cannot be substituted", flush=True)
            return coeff
    coeff_pattern = r"a\d+|c\d+_\d+"
    expression = re.sub(coeff_pattern, repl, expression)
    return expression

def parse_subs_list(f_in):
    print(f"Parsing file {f_in}", flush=True)
    for line in open(f_in, "r").readlines():
        # parse the line: "coeff==expression"
        coeff, expression = line.strip().split("==")
        # generate results
        yield coeff, expression

def parse_expr_list(f_in):
    print(f"Parsing file {f_in}", flush=True)
    for line in open(f_in, "r").readlines():
        # expected format: one expression per line
        expression = line.strip()
        # correct the expression: replace functions
        #sed 's/(t,x,y,z)//g' $inputfile > $outputfile
        #sed -i 's/fu/u/g' $outputfile
        #sed -i 's/fv/v/g' $outputfile
        #sed -i 's/fw/w/g' $outputfile
        #sed -i 's/frho/rho/g' $outputfile
        expression = expression.replace("(t,x,y,z)", "").replace("fu", "u").replace("fv", "v").replace("fw", "w").replace("frho", "rho")
        # generate results
        yield expression

def find_all_coeffs(expression, running_set):
    """
    :param running_set: a growing set of the coefficients that were already found
    """
    coeff_pattern = r"a\d+|c\d+_\d+"
    for match in re.finditer(coeff_pattern, expression):
        coeff = match.group(0)
        if coeff not in running_set:
            yield from find_all_coeffs(subs_map[coeff], running_set)
            yield coeff

# read subs lists, do not actually subsittute the coeffs
print("STEP 1: reading substitution lists", flush=True)
for f_in in sorted(subs_dir.glob("subs_list*")):
    for coeff, expression in parse_subs_list(f_in):
        # save the expression in memory
        subs_map[coeff] = expression

print("STEP 2: reading files with target expressions", flush=True)
target_expressions = []
for f_in in target_files:
    for expression in parse_expr_list(f_in):
        target_expressions.append(expression)

print("STEP 3: finding all coefficients to be substituted", flush=True)
target_coeffs = set()
for expression in target_expressions:
    for coeff in find_all_coeffs(expression, target_coeffs):
        target_coeffs.add(coeff)
print(f"{len(target_coeffs)} coeffs needed for substitutions")

print("STEP 4: substituting the selected coefficients", flush=True)
buf_coeffs = []
buf_exprs = []
# GOTCHA: coeffs need to be sorted properly
for coeff in natural_sorted(target_coeffs):
    # get the expression from memory
    expression = subs_map[coeff]
    # substitute all coeffs in the expression
    expression = substitute_coeffs(expression)
    # save the expression in memory
    subs_map[coeff] = expression
    # add to the buffer
    buf_coeffs.append(coeff)
    buf_exprs.append(expression)
# simplify the expressions
nproc = multiprocessing.cpu_count()
with multiprocessing.Pool(nproc) as pool:
    buf_exprs = simplify_par(buf_exprs, pool, nproc)
for coeff, expression in zip(buf_coeffs, buf_exprs):
    # update the simplified expression in memory
    subs_map[coeff] = expression

print("STEP 5: substituting coefficients in the target directory", flush=True)
# the target expressions were already read, so they can be simplified at once (after the coefficients are substituted)
target_expressions_new = []
for expression in target_expressions:
    # substitute all coeffs in the expression
    expression = substitute_coeffs(expression)
    # save the substituted expression
    target_expressions_new.append(expression)
with multiprocessing.Pool(nproc) as pool:
    target_expressions_new = simplify_par(target_expressions_new, pool, nproc)
target_expressions = dict( (k, v) for k, v in zip(target_expressions, target_expressions_new) )

for f_in in target_files:
    pde_dir = str(f_in.parent.relative_to(input_dir))
    assert pde_dir.endswith("_subs"), pde_dir
    pde_dir = pde_dir[:-len("_subs")]
    fname_out = input_dir / pde_dir / f_in.name
    assert fname_out != f_in
    # create the output directory
    os.makedirs(fname_out.parent, exist_ok=True)
    with open(fname_out, "w") as f_out:
        for expression in parse_expr_list(f_in):
            # search the updated expression
            expression = target_expressions[expression]
            # save the expression on disk
            print(f"{expression}", file=f_out)
