"""
WARNING: spaghetti code.
"""

import numpy as np
import pickle
import os

def parser(model):
	"""
	Read the .cfg file to extract layers into `layers`
	as well as model-specific parameters into `meta`
	"""
	def _parse(l, i = 1):
		return l.split('=')[i].strip()

	with open(model, 'rb') as f:
		lines = f.readlines()

	lines = [line.decode() for line in lines]	
	
	meta = dict(); layers = list() # will contains layers' info
	h, w, c = [int()] * 3; layer = dict()
	for line in lines:
		line = line.strip()
		line = line.split('#')[0]
		if '[' in line:
			if layer != dict(): 
				if layer['type'] == '[net]': 
					h = layer['height']
					w = layer['width']
					c = layer['channels']
					meta['net'] = layer
				else:
					if layer['type'] == '[crop]':
						h = layer['crop_height']
						w = layer['crop_width']
					layers += [layer]				
			layer = {'type': line}
		else:
			try:
				i = float(_parse(line))
				if i == int(i): i = int(i)
				layer[line.split('=')[0].strip()] = i
			except:
				try:
					key = _parse(line, 0)
					val = _parse(line, 1)
					layer[key] = val
				except:
					'banana ninja yadayada'

	meta.update(layer) # last layer contains meta info
	if 'anchors' in meta:
		splits = meta['anchors'].split(',')
		anchors = [float(x.strip()) for x in splits]
		meta['anchors'] = anchors
	meta['model'] = model # path to cfg, not model name
	meta['inp_size'] = [h, w, c]
	return layers, meta

def cfg_yielder(model, binary):
	"""
	yielding each layer information to initialize `layer`
	"""
	layers, meta = parser(model); yield meta;
	h, w, c = meta['inp_size']; l = w * h * c

	# Start yielding
	flat = False # flag for 1st dense layer
	conv = '.conv.' in model
	for i, d in enumerate(layers):
		#-----------------------------------------------------
		if d['type'] == '[crop]':
			yield ['crop', i]
		#-----------------------------------------------------
		elif d['type'] == '[local]':
			n = d.get('filters', 1)
			size = d.get('size', 1)
			stride = d.get('stride', 1)
			pad = d.get('pad', 0)
			activation = d.get('activation', 'logistic')
			w_ = (w - 1 - (1 - pad) * (size - 1)) // stride + 1
			h_ = (h - 1 - (1 - pad) * (size - 1)) // stride + 1
			yield ['local', i, size, c, n, stride, 
					pad, w_, h_, activation]
			if activation != 'linear': yield [activation, i]
			w, h, c = w_, h_, n
			l = w * h * c
		#-----------------------------------------------------
		elif d['type'] == '[convolutional]':
			n = d.get('filters', 1)
			size = d.get('size', 1)
			stride = d.get('stride', 1)
			pad = d.get('pad', 0)
			padding = d.get('padding', 0)
			if pad: padding = size // 2
			activation = d.get('activation', 'logistic')
			batch_norm = d.get('batch_normalize', 0) or conv
			yield ['convolutional', i, size, c, n, 
				   stride, padding, batch_norm,
				   activation]
			if activation != 'linear': yield [activation, i]
			w_ = (w + 2 * padding - size) // stride + 1
			h_ = (h + 2 * padding - size) // stride + 1
			w, h, c = w_, h_, n
			l = w * h * c
		#-----------------------------------------------------
		elif d['type'] == '[maxpool]':
			stride = d.get('stride', 1)
			size = d.get('size', stride)
			padding = d.get('padding', (size-1) // 2)
			yield ['maxpool', i, size, stride, padding]
			w_ = (w + 2*padding) // d['stride'] 
			h_ = (h + 2*padding) // d['stride']
			w, h = w_, h_
			l = w * h * c
		#-----------------------------------------------------
		elif d['type'] == '[avgpool]':
			flat = True; l = c
			yield ['avgpool', i]
		#-----------------------------------------------------
		elif d['type'] == '[softmax]':
			yield ['softmax', i, d['groups']]
		#-----------------------------------------------------
		elif d['type'] == '[connected]':
			if not flat:
				yield ['flatten', i]
				flat = True
			activation = d.get('activation', 'logistic')
			yield ['connected', i, l, d['output'], activation]
			if activation != 'linear': yield [activation, i]
			l = d['output']
		#-----------------------------------------------------
		elif d['type'] == '[dropout]': 
			yield ['dropout', i, d['probability']]
		#-----------------------------------------------------
		elif d['type'] == '[select]':
			if not flat:
				yield ['flatten', i]
				flat = True
			inp = d.get('input', None)
			if type(inp) is str:
				file = inp.split(',')[0]
				layer_num = int(inp.split(',')[1])
				with open(file, 'rb') as f:
					profiles = pickle.load(f, encoding = 'latin1')[0]
				layer = profiles[layer_num]
			else: layer = inp
			activation = d.get('activation', 'logistic')
			d['keep'] = d['keep'].split('/')
			classes = int(d['keep'][-1])
			keep = [int(c) for c in d['keep'][0].split(',')]
			keep_n = len(keep)
			train_from = classes * d['bins']
			for count in range(d['bins']-1):
				for num in keep[-keep_n:]:
					keep += [num + classes]
			k = 1
			while layers[i-k]['type'] not in ['[connected]', '[extract]']:
				k += 1
				if i-k < 0:
					break
			if i-k < 0: l_ = l
			elif layers[i-k]['type'] == 'connected':
				l_ = layers[i-k]['output']
			else:
				l_ = layers[i-k].get('old',[l])[-1]
			yield ['select', i, l_, d['old_output'],
				   activation, layer, d['output'], 
				   keep, train_from]
			if activation != 'linear': yield [activation, i]
			l = d['output']
		#-----------------------------------------------------
		elif d['type'] == '[conv-select]':
			n = d.get('filters', 1)
			size = d.get('size', 1)
			stride = d.get('stride', 1)
			pad = d.get('pad', 0)
			padding = d.get('padding', 0)
			if pad: padding = size // 2
			activation = d.get('activation', 'logistic')
			batch_norm = d.get('batch_normalize', 0) or conv
			d['keep'] = d['keep'].split('/')
			classes = int(d['keep'][-1])
			keep = [int(x) for x in d['keep'][0].split(',')]

			segment = classes + 5
			assert n % segment == 0, \
			'conv-select: segment failed'
			bins = n // segment
			keep_idx = list()
			for j in range(bins):
				offset = j * segment
				for k in range(5):
					keep_idx += [offset + k]
				for k in keep:
					keep_idx += [offset + 5 + k]
			w_ = (w + 2 * padding - size) // stride + 1
			h_ = (h + 2 * padding - size) // stride + 1
			c_ = len(keep_idx)
			yield ['conv-select', i, size, c, n, 
				   stride, padding, batch_norm,
				   activation, keep_idx, c_]
			w, h, c = w_, h_, c_
			l = w * h * c
		#-----------------------------------------------------
		elif d['type'] == '[conv-extract]':
			file = d['profile']
			with open(file, 'rb') as f:
				profiles = pickle.load(f, encoding = 'latin1')[0]
			inp_layer = None
			inp = d['input']
			out = d['output']
			inp_layer = None
			if inp >= 0:
				inp_layer = profiles[inp]
			if inp_layer is not None:
				assert len(inp_layer) == c, \
				'Conv-extract does not match input dimension'
			out_layer = profiles[out]

			n = d.get('filters', 1)
			size = d.get('size', 1)
			stride = d.get('stride', 1)
			pad = d.get('pad', 0)
			padding = d.get('padding', 0)
			if pad: padding = size // 2
			activation = d.get('activation', 'logistic')
			batch_norm = d.get('batch_normalize', 0) or conv
			
			k = 1
			find = ['[convolutional]','[conv-extract]']
			while layers[i-k]['type'] not in find:
				k += 1
				if i-k < 0: break
			if i-k >= 0:
				previous_layer = layers[i-k]
				c_ = previous_layer['filters']
			else:
				c_ = c
			
			yield ['conv-extract', i, size, c_, n, 
				   stride, padding, batch_norm,
				   activation, inp_layer, out_layer]
			if activation != 'linear': yield [activation, i]
			w_ = (w + 2 * padding - size) // stride + 1
			h_ = (h + 2 * padding - size) // stride + 1
			w, h, c = w_, h_, len(out_layer)
			l = w * h * c
		#-----------------------------------------------------
		elif d['type'] == '[extract]':
			if not flat:
				yield['flatten', i]
				flat = True
			activation = d.get('activation', 'logistic')
			file = d['profile']
			with open(file, 'rb') as f:
				profiles = pickle.load(f, encoding = 'latin1')[0]
			inp_layer = None
			inp = d['input']
			out = d['output']
			if inp >= 0:
				inp_layer = profiles[inp]
			out_layer = profiles[out]
			old = d['old']
			old = [int(x) for x in old.split(',')]
			if inp_layer is not None:
				if len(old) > 2: 
					h_, w_, c_, n_ = old
					new_inp = list()
					for p in range(c_):
						for q in range(h_):
							for r in range(w_):
								if p not in inp_layer:
									continue
								new_inp += [r + w*(q + h*p)]
					inp_layer = new_inp
					old = [h_ * w_ * c_, n_]
				assert len(inp_layer) == l, \
				'Extract does not match input dimension'
			d['old'] = old
			yield ['extract', i] + old + [activation] + [inp_layer, out_layer]
			if activation != 'linear': yield [activation, i]
			l = len(out_layer)
		#-----------------------------------------------------
		elif d['type'] == '[route]': # add new layer here
			routes = d['layers']
			if type(routes) is int:
				routes = [routes]
			else:
				routes = [int(x.strip()) for x in routes.split(',')]
			routes = [i + x if x < 0 else x for x in routes]
			for j, x in enumerate(routes):
				lx = layers[x]; 
				xtype = lx['type']
				_size = lx['_size'][:3]
				if j == 0:
					h, w, c = _size
				else: 
					h_, w_, c_ = _size
					assert w_ == w and h_ == h, \
					'Routing incompatible conv sizes'
					c += c_
			yield ['route', i, routes]
			l = w * h * c
		#-----------------------------------------------------
		elif d['type'] == '[reorg]':
			stride = d.get('stride', 1)
			yield ['reorg', i, stride]
			w = w // stride; h = h // stride; 
			c = c * (stride ** 2)
			l = w * h * c
		#-----------------------------------------------------
		else:
			exit('Layer {} not implemented'.format(d['type']))

		d['_size'] = list([h, w, c, l, flat])

	if not flat: meta['out_size'] = [h, w, c]
	else: meta['out_size'] = l