"""
A module that allows the creation of a C file, in which the block's computations are inlined.
Additional library files will be created to allow for linking:

- lsolve.c
- lsolve.h


Requires `Jinja2 <https://jinja.palletsprojects.com/en/3.0.x/>`_.
The :code:`template.c` file is the jinja template file.

Note:
	To compile the generated code, execute :code:`gcc <filename>.c lsolve.c -lm`.
"""
from CBD.scheduling import TopologicalScheduler
from CBD.loopsolvers.linearsolver import LinearSolver
from CBD.depGraph import createDepGraph
from CBD.converters.latexify import CBD2Latex, Time
import CBD.naivelog as naivelog
from CBD.simulator import Simulator

from jinja2 import Template
from shutil import copyfile
import os, re


class CBD2C:
	"""
	Main conversion class. Generates the new C file.
	The C model will be executed in fixed-time steps of size
	:code:`delta`, for a total of :code:`itcnt` iterations.
	In total, the simulation will thus execute until timestep
	:code:`itcnt * delta`.

	Note:
		This class does **not** change the original model.

	Warning:
		This class only works if the dependency graph is fixed
		after the first iteration. I.e., iteration 0 can be
		different, but all others must be the same.

	Args:
		model (CBD):    CBD model that will be converted to C.
		itcnt (int):    The amount of iterations to execute the
						c model for.
		delta (float):  The stepsize of the C code.
	"""
	def __init__(self, model, itcnt, delta=1.0):
		self.itcnt = itcnt
		self.delta = delta
		self.model = model.clone()
		self.model.flatten(psep='_')
		for block in self.model.getBlocks():
			block.setBlockName(block.getBlockName().replace("-", "_"))
		self.scheduler = TopologicalScheduler()
		self.solver = LinearSolver(naivelog.getLogger("CBD"))

	def get_blocks(self):
		"""
		Gets all the blocks from the flattened model.
		"""
		return self.model.getBlocks()

	def get_order(self, curIt):
		"""
		Gets the topological order of the blocks at a specific iteration.

		Args:
			curIt (int):
		"""
		depGraph = createDepGraph(self.model, curIt)
		return self.scheduler.schedule(depGraph, curIt, 0.0)

	def obtain_eqs_from_order(self, order, curIt):
		"""
		Given the order in which the blocks must be executed,
		this function will use the :class:`CBD2Latex` class to
		construct the C-representation for the functions.

		Args:
			order:          List of components, topologically ordered.
			curIt (int):    The current iteration.

		Returns:
			A list of tuples like :code:`[([LHS, ...], RHS), ...]`;
			where the :code:`LHS` identifies an ordered list of
			left-hand sides (when there are multiple, it means an
			algebraic loop); and :code:`RHS` the string-representation
			of the results.
		"""
		conf = {
			"path_sep": '_',
			"ignore_path": False,
			"escape_nonlatex": False,
			"time_variable": "i",
			"time_format": "[{time}]"
		}
		ltx = CBD2Latex(self.model, **conf)
		eqs = [re.sub(r"-(?!\s*\d)", "_", x) for x in ltx.eq(Time(0, curIt == "i")).split("\n")]
		if eqs[-1] == '':
			eqs.pop()
		res_order = []
		# reverse the order to allow scheduling at end as well
		for c in reversed(order):
			c_order = []
			if len(c) == 1:  # NO ALGEBRAIC LOOP
				block = c[0]

				# Assumes no special blocks that change the deps exist
				deps = block.getDependencies(1 if curIt == "i" else curIt)
				path = block.getPath("_")
				for out in block.getSignals().keys():
					p = path + "_" + out + "[%s]" % curIt
					for e in eqs:
						if e.startswith(p):
							c_order.insert(0, e)
							break
				for inp in block.getInputPortNames():
					p = path + "_" + inp + "[%s]" % curIt
					in_deps = block.getBlockConnectedToInput(inp).block in deps
					for e in eqs:
						if e.startswith(p):
							if in_deps:
								c_order.insert(0, e)
							else:
								c_order.append(e)
							break
				if block.getBlockType() == "WireBlock":
					for e in eqs:
						if e.startswith(path + "[%s]" % curIt):
							c_order.append(e)
							break
			else:  # ALGEBRAIC LOOP!
				self.solver.checkValidity(self.model.getPath("."), c)
				other = []
				known = {}
				for block in c:
					path = block.getPath("_")
					for inp, b in block.getLinksIn().items():
						var = path + "_" + inp
						p = var + "[%s]" % curIt
						prev = block.getBlockConnectedToInput(inp).block
						for e in eqs:
							if e.startswith(p):
								if prev not in c:
									c_order.insert(0, e)
									known[var] = p

					for inp in block.getInputPortNames():
						var = path + "_" + inp
						if var in known: continue
						p = var + "[%s]" % curIt
						for e in eqs:
							if e.startswith(p):
								other.insert(0, e)
								break

				# Get matrix representation
				m1, m2 = self.solver.get_matrix(c, '_', known)
				mat = m1.concat(m2)

				deps = {}
				for b in c:
					D = [b.getBlockConnectedToInput(x) for x in b.getInputPortNames()]
					deps[b] = [b.block.getPath('_') + '_' + b.output_port for b in D if b.block in c]

				if mat.rows == 1:
					# TODO: fix this?
					raise NotImplementedError("Suspected singular matrix. However, currently cannot find an example "
					                          "that enters this branch. Please contact the repository owner with this "
					                          "error and all associated files.")
					# nc = ([deps[c[mat[0].index(1)]]], mat[0][-1])
					# print("NC:", nc)
				else:
					dlist = []
					for b in c:
						for d in deps[b]:
							if d not in dlist:
								dlist.append(d)
					sm = "{" + mat.format(', ', '{}', "%f", ', ').replace("{,", "{").replace(", }", " }") + "}"
					nc = (dlist, sm)

				c_order.append(nc)
				c_order += other
			res_order = c_order + res_order
		return res_order

	def generate(self, fname):
		"""
		Generates the C file and copies the sources for the
		lsolve library in the same folder.

		To compile the generated code, execute :code:`gcc <filename>.c lsolve.c -lm`.

		Args:
			fname (str):    The filename of the C-code.
		"""
		path = os.path.dirname(os.path.realpath(__file__))
		tpath = os.path.dirname(os.path.realpath(fname))
		filename = os.path.join(path, "template.c")
		with open(filename, 'r') as file:
			template = Template(file.read(), trim_blocks=True, lstrip_blocks=True)
		if not os.path.exists(os.path.join(tpath, "lsolve.c")):
			copyfile(os.path.join(path, "lsolve.c"), os.path.join(tpath, "lsolve.c"))
		if not os.path.exists(os.path.join(tpath, "lsolve.h")):
			copyfile(os.path.join(path, "lsolve.h"), os.path.join(tpath, "lsolve.h"))
		variables = []
		for block in self.get_blocks():
			if block.getBlockType() == "WireBlock":
				variables.append(block.getPath('_'))
			for port in block.getSignals().keys():
				variables.append(block.getPath('_') + '_' + port)
			for port in block.getInputPortNames():
				variables.append(block.getPath('_') + '_' + port)
		variables = [v.replace("-", "_") for v in variables]

		# Obtain the data
		comp0 = self.get_order(0)
		comp = self.get_order(1)

		eqs0 = self.obtain_eqs_from_order(comp0, 0)
		eqs = self.obtain_eqs_from_order(comp, "i")

		contents = template.render(itcnt = self.itcnt,
		                           filename = fname,
		                           delta = self.delta,
		                           variables = variables,
		                           components0 = eqs0,
		                           components = eqs )
		with open(fname, 'w') as file:
			file.write(contents)

