"""
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.solver import LinearSolver
from CBD.depGraph import DepGraph, createDepGraph, gvDepGraph
from CBD.converters.latexify import CBD2Latex, _BLOCK_MAP
import CBD.naivelog as naivelog

from jinja2 import Template
from shutil import copyfile
import os

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)

	@staticmethod
	def flatten(t):
		"""
		Flattens a list of sublists.

		Args:
			t (iter):   The list to flatten.
		"""
		return [item for sublist in t for item in sublist]

	def obtain_data_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.
		"""
		i = "0" if curIt == 0 else "i"
		conf = {"path_sep": '_', "ignore_path": False, "escape_nonlatex": False, "time_variable": "time"}
		norder = []
		for c in order:
			nc = []
			if len(c) == 1 and c[0] not in [c[0].getBlockConnectedToInput(x).block for x in c[0].getInputPortNames()]:
				n = c[0]
				bd = CBD2Latex.get_block_data(n, self.model, conf)
				if bd is None:
					if n.getBlockType() == "WireBlock":
						bl = n.getBlockConnectedToInput("IN1")
						nc = ([n.getPath('_')], bl.block.getPath('_') + '_' + bl.output_port + '[' + i + ']')
				else:
					kv = {x: n.getBlockConnectedToInput(x) for x in n.getInputPortNames()}
					for k, v in kv.items():
						bd[1].apply(n.getPath('_') + '_' + k, v.block.getPath('_') + '_' + v.output_port)
					# if not isinstance(bd[1], str):
					# 	bd[1].args = ["(double)%s" % str(x) for x in bd[1].args]
					if n.getBlockType() == "DelayBlock":
						if curIt == 0:
							bd[1].apply_time(time=-1, t=i, fmt="[{time}]")
							bd = bd[0], bd[1].args[1]
						else:
							bd[1].apply_time(t=i, fmt="[{time}]")
							bd = bd[0], bd[1].args[0]
					try:
						bd[1].apply_time(t=i, fmt="[{time}]")
					except: pass
					nc = ([bd[0]], str(bd[1]))
				if n.getBlockType() != "WireBlock":
					n.compute(curIt)
			else:
				self.solver.checkValidity(self.model.getPath("."), c)

				# Get matrix representation
				m1, m2 = self.solver.constructInput(c, curIt)
				mat = []
				for r in range(m1.rows):
					row = []
					for col in range(m1.cols):
						row.append(float(m1[r, col]))
					row.append(float(m2[r]))
					mat.append(row)

				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 len(mat) == 1:
					nc = ([deps[c[mat[0].index(1)]]], mat[0][-1])
				else:
					dlist = []
					for b in c:
						for d in deps[b]:
							if d not in dlist:
								dlist.append(d)

					nc = (dlist, str(mat).replace("[", "{").replace("]", "}"))
			norder.append(nc)
		return norder

	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)
		copyfile(os.path.join(path, "lsolve.c"), os.path.join(tpath, "lsolve.c"))
		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('_'))
				continue
			for port in block.getSignals().keys():
				variables.append(block.getPath('_') + '_' + port)

		# Obtain the data
		comp0 = self.get_order(0)
		comp = self.get_order(1)
		self.model.addFixedRateClock()
		comp0 = self.obtain_data_from_order(comp0, 0)
		comp = self.obtain_data_from_order(comp, 1)

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

if __name__ == '__main__':
	from CBD.Core import CBD
	from CBD.lib.std import *
	class Test(CBD):
		def __init__(self, name):
			super().__init__(name, [], ["x"])
			self.addBlock(AdderBlock("A"))
			# self.addBlock(ProductBlock("B"))
			# self.addConnection("A", "A")
			# self.addConnection("A", "B")
			# self.addConnection("B", "A")
			# self.addConnection("B", "B")
			# self.addBlock(ProductBlock("B"))
			self.addBlock(ConstantBlock("C", 3.0))
			self.addBlock(ConstantBlock("D", 6.0))
			self.addConnection("C", "A")
			# self.addConnection("B", "A")
			# self.addConnection("A", "B")
			self.addConnection("D", "A")
			self.addConnection("A", "x")

	test = Test("test")
	gen = CBD2C(test, 3000)
	gen.generate("example.c")

