#!/usr/bin/env python3

import qckt.Canvas as cnv
import qckt.Gates as gts
import qckt.gatesutils as gutils
import qckt.noisemodel as ns
from qckt.qException import QCktException
import numpy as np


class GateWrapper:
	def __init__(self, qckt, gateCls):
		self.qckt = qckt
		self.GateCls = gateCls
	def __call__(self, *args, **kwargs):
		gateObj = self.GateCls(*args, **kwargs)
		self.qckt.circuit.append(gateObj)
		if gateObj.check_qbit_args(self.qckt.nqubits) == False:
			raise QCktException(f"Error: qubit arguments incorrect. {gateObj.name}{str(gateObj.qbits)}")
		return gateObj
	def set_noise_on_all(self, noise_chan):
		# don't need to check NoiseChannel or NoiseChannelSequence, consolidate_gate_noise() checks that and converts
		self.qckt.noise_profile_gates[self.GateCls] = noise_chan


class QCkt:

	def __init__(self, nqubits, nclbits=0, name=None, noise_profile=None):
		self.nqubits = nqubits
		self.nclbits = nclbits
		self.circuit = []
		self.name = name
		self.canvas = cnv.Canvas(self)
		self.noise_profile_gates = {}  # this stores a map with gateCls as key, and noise spec as value
		self.noise_profile = None  # set_noise_profile() below updates it
		self.set_noise_profile(noise_profile)

		for gclass in gts.GatesList:
			setattr(self, gclass.__name__, GateWrapper(self,gclass))

		self.idx = 0 # for iterations

	def __iter__(self):
		self.it = iter(self.circuit)
		return self
	def __next__(self):
		return next(self.it)

	def get_gates_list(self):
		raise QCktException('ERROR: get_gates_list() is available under qckt package. See README.md. Usage: qckt.get_gates_list()')

	def append(self, otherckt, name=None):
		# otherckt - the ckt to be appended
		nq = max(self.nqubits, otherckt.nqubits)
		nc = max(self.nclbits, otherckt.nclbits)
		if name is None:
			name = self.name
		newckt = QCkt(nq,nc,name=name)

		# copy over noise_profile_gates list into new circuit
		newckt.__copyover_noise_profile_gates__(srcckt=self)
		newckt.__copyover_noise_profile_gates__(srcckt=otherckt)

		# copy circuits over to the new circuit
		for g in self.circuit:
			newckt.circuit.append(g)
		for g in otherckt.circuit:
			newckt.circuit.append(g)

		return newckt

	# change the qubits order to a different order
	def realign(self, newnq, newnc, inpqubits, name=None):
		# newq and newc are the new sizes of the qubits register and clbits register
		# inpqubits gives the new positions of the qubits
		# See README.md for details
		if self.nqubits != len(inpqubits):
			errmsg = "Error: error aligning qubits, number of qubits do not match"
			raise QCktException(errmsg)
		if name is None:
			name = self.name
		newckt = QCkt(newnq, newnc, name=name)

		# copy over noise_profile_gates list into new circuit
		newckt.__copyover_noise_profile_gates__(srcckt=self)

		for g in self.circuit:
			galigned = g.realign(inpqubits)
			newckt.circuit.append(galigned)
		return newckt

	def __copyover_noise_profile_gates__(self, srcckt):
		# copy over noise_profile_gates from source circuit
		for gcls in srcckt.noise_profile_gates:
			self.noise_profile_gates[gcls] = srcckt.noise_profile_gates[gcls]



	def assemble(self):
		assembled = []
		noise_wrapper = self.NOISE
		noise_cls = noise_wrapper.GateCls
		if self.noise_profile is not None and self.noise_profile.noise_chan_init is not None:
			noise_gate = noise_cls(self.noise_profile.noise_chan_init, list(range(self.nqubits)))
			cktstep = noise_gate.assemble(self.noise_profile, self.noise_profile_gates)
			cktstep['tag'] = 'INIT'
			assembled.append(cktstep)
		for gt in self.circuit:
			cktstep = gt.assemble(self.noise_profile, self.noise_profile_gates)
			cktstep['tag'] = 'NS' if cktstep['op'] == 'noise' else gt.name
			assembled.append(cktstep)
			if self.noise_profile is not None and self.noise_profile.noise_chan_allsteps is not None and gt.is_noise_step():
				for kop,qbt in self.noise_profile.noise_chan_allsteps:
					noise_gate = noise_cls(kop, qbt)
					cktstep = noise_gate.assemble(self.noise_profile, self.noise_profile_gates)
					cktstep['tag'] = 'AS'
					assembled.append(cktstep)
		return assembled

	def to_opMatrix(self):
		oplist = []
		for q in self.circuit:
			op = q.to_fullmatrix(self.nqubits)
			if op is not None: ## Border, Probe gates return None
				oplist.append(op)
		if len(oplist) == 0:
			return np.matrix(np.eye(2**self.nqubits),dtype=complex)
		opmat = gutils.combine_opmatrices_seq(oplist)
		return opmat

	def set_noise_profile(self, noise_profile=None):
		fixed_noise_profile = noise_profile
		if noise_profile is not None:
			# validate the keys in the dict for noise_profile argument
			if type(noise_profile) is not ns.NoiseProfile:
				raise QCktException(f'ERROR: noise_profile expected to be NoiseProfile object or None')

			# validate the noise_chan_init field
			opseq_init = noise_profile.noise_chan_init
			if opseq_init is not None:
				if type(opseq_init) is ns.NoiseChannelSequence:
					pass
				elif type(opseq_init) is ns.NoiseChannel:
					opseq_init = ns.NoiseChannelSequence(opseq_init)
				else:
					raise QCktException('ERROR: noise_profile.noise_chan_init must be NoiseChannelSequence or NoiseChannel object.')
				# if noise_chan_init is specified, it must have only 1-qubit noise channel
				for op in opseq_init:
					if op.nqubits != 1:
						raise QCktException(f'ERROR: noise_opseq_init must use 1-qubit noise channel')

			# validate the noise_chan_allgates field
			opseq_allgates = noise_profile.noise_chan_allgates
			if opseq_allgates is not None:
				if type(opseq_allgates) is ns.NoiseChannelSequence:
					pass
				elif type(opseq_allgates) is ns.NoiseChannel:
					opseq_allgates = ns.NoiseChannelSequence(opseq_allgates)
				else:
					raise QCktException('ERROR: noise_profile.noise_chan_allgates must be NoiseChannel or NoiseChannel object.')

			# validate the noise_chan_qubits field
			opseq_qubits = noise_profile.noise_chan_qubits
			if opseq_qubits is not None:
				if type(opseq_qubits) is not ns.NoiseChannelApplierSequense:
					raise QCktException('ERROR: noise_profile.noise_chan_qubits must be NoiseChannelApplierSequense object.')

			# validate the noise_chan_allsteps field
			opseq_allsteps = noise_profile.noise_chan_allsteps
			if opseq_allsteps is not None:
				if type(opseq_allsteps) is not ns.NoiseChannelApplierSequense:
					raise QCktException('ERROR: noise_profile.noise_chan_allsteps must be NoiseChannelApplierSequense object.')

			fixed_noise_profile = ns.NoiseProfile(noise_chan_init=opseq_init, noise_chan_allgates=opseq_allgates, noise_chan_qubits=opseq_qubits, noise_chan_allsteps=opseq_allsteps)
		# save the noise_profile
		self.noise_profile = fixed_noise_profile

	def get_size(self):
		return self.nqubits, self.nclbits

	def draw(self,show_noise=True):
		assembled_circuit = self.assemble()
		self.canvas.draw(assembled_circuit, show_noise=show_noise)

	def list(self, show_noise=True):
		assembled_circuit = self.assemble()
		if self.name is not None:
			print(self.name)
		for cktstep in assembled_circuit:
			if cktstep['op'] == 'gate':
				prline = f"{cktstep['name']}"
				gateparam_list = []
				for p in cktstep['gateObj'].gateparams:
					if type(p) is int:
						pstr = f'{p:d}'
					elif type(p) is float:
						pstr = f'{p:.4f}'
					else:
						pstr = str(p)
					gateparam_list.append(f'{pstr}')
				if len(gateparam_list) > 0:
					gateparam_str = ','.join(gateparam_list)
					prline = f'{prline}({gateparam_str})'
				prline = f"{prline}{cktstep['qubits']}"
				if cktstep['gateObj'].cbit_cond is not None:
					prline = f"{prline}.ifcbit{cktstep['gateObj'].cbit_cond}"
				print(prline)
				if show_noise:
					noise_ch_appl_seq = cktstep['gateObj'].form_gatenoise_ch_appl_seq(self.noise_profile, self.noise_profile_gates)
					if len(noise_ch_appl_seq.name) > 0:
						prline = f"{cktstep['name']}:{noise_ch_appl_seq.name}"
						print(f'* {prline}')
			elif cktstep['op'] == 'measure':
				prline = f"{cktstep['name']}"
				prline = f"{prline}{cktstep['qubits']}"
				if cktstep['gateObj'].cbits is not None:
					prline = f"{prline}{cktstep['gateObj'].cbits}"
				print(prline)
			elif cktstep['op'] == 'probe':
				prline = 'PROBE>>>>>>>>>>>>>>'
				print(prline)
			elif cktstep['op'] == 'noop':
				prline = 'BORDER#############'
				print(prline)
			elif cktstep['op'] == 'noise':
				if show_noise:
					prline = f"{cktstep['tag']}:{cktstep['name']}{cktstep['qubits']}"
					print(f'* {prline}')
			else:
				raise QCktException(f"INTERNAL ERROR: Unknown op {cktstep['op']}")

	def custom_gate(self, cgate_name, opMatrix):
		raise QCktException('ERROR: custom_gate() is discontinued, use qckt.define_gate() instead. See README.md.')


if __name__ == "__main__":
	pass
