import os
import shutil
import time

from nbconvert.preprocessors import ExecutePreprocessor, CellExecutionError
from nbconvert.exporters import RSTExporter
from nbconvert.writers import FilesWriter
import nbformat

from PIL import Image, ImageChops

def compile_tutorial(tutorial_name, force_recompile=False):
	print('- Compiling tutorial ' + tutorial_name + '...')

	notebook_path = 'tutorial_notebooks/' + tutorial_name + '/' + tutorial_name + '.ipynb'
	export_path = 'tutorials/' + tutorial_name + '/' + tutorial_name
	thumb_dest = os.path.dirname(export_path) + '/thumb.png'

	if not os.path.exists(os.path.dirname(export_path)):
		os.makedirs(os.path.dirname(export_path))
	
	# Read in notebook
	notebook = nbformat.read(notebook_path, 4)

	# Scrape title, description and thumbnail
	first_cell = notebook.cells[0]

	title = first_cell.source.splitlines()[0]
	if '#' in title:
		title = title.replace('#', '').strip()
	
	description = ''
	for line in first_cell.source.splitlines()[1:]:
		if line.strip():
			description = line.strip()
			break

	if not description:
		print('  Description could not be found in the notebook.')
	
	if 'thumbnail_figure_index' in notebook.metadata:
		thumbnail_figure_index = notebook.metadata['thumbnail_figure_index']
	else:
		thumbnail_figure_index = -1
	
	if 'level' in notebook.metadata:
		level = notebook.metadata['level'].capitalize()
	elif 'difficulty' in notebook.metadata:
		level = notebook.metadata['difficulty'].capitalize()
	else:
		level = 'Unknown'
	
	# Check if the tutorial was already compiled.
	if os.path.exists(export_path + '.rst'):
		if os.path.getmtime(export_path + '.rst') > os.path.getmtime(notebook_path):
			if force_recompile:
				print('  Already compiled. Recompiling anyway...')
			else:
				print('  Already compiled. Skipping...')
				return title, level, description, thumb_dest.split('/', 1)[-1]

	# Execute notebook if not already executed
	already_executed = any(c.get('outputs') or c.get('execution_count') for c in notebook.cells if c.cell_type == 'code')

	resources = {}
	
	if not already_executed:
		ep = ExecutePreprocessor(timeout=600, kernel_name='python3')
		try:
			start = time.time()

			additional_cell_1 = {
				"cell_type": "code",
				"execution_count": None,
				"metadata": {},
				"outputs": [],
				"source": r"%matplotlib inline" + '\n' +
					r"%config InlineBackend.print_figure_kwargs = {'bbox_inches': None, 'figsize': (8, 6)}"
				}
			
			additional_cell_2 = {
				"cell_type": "code",
				"execution_count": None,
				"metadata": {},
				"outputs": [],
				"source": "import matplotlib as mpl\nmpl.rcParams['figure.figsize'] = (8, 6)\nmpl.rcParams['figure.dpi'] = 150\nmpl.rcParams['savefig.dpi'] = 150"
				}
			
			notebook.cells.insert(1, nbformat.from_dict(additional_cell_1))
			notebook.cells.insert(2, nbformat.from_dict(additional_cell_2))

			km, kc = ep.start_new_kernel(cwd=os.path.abspath(os.path.dirname(notebook_path)))
			kc.allow_stdin = False

			notebook, resources = ep.preprocess(notebook, km=km)

			notebook.cells.pop(2)
			notebook.cells.pop(1)

			km.shutdown_kernel()
			
			end = time.time()
			print('  Compilation took %d seconds.' % (end - start))
		except CellExecutionError as err:
			print('  Error while processing notebook:')
			print('  ', err)
	else:
		print('  Notebook was already executed.')

	exporter = RSTExporter()
	output, resources = exporter.from_notebook_node(notebook, resources)

	writer = FilesWriter(build_directory=os.path.dirname(export_path))
	writer.write(output, resources, notebook_name=os.path.basename(export_path))

	pictures = sorted(resources['outputs'], key=output.find)

	try:
		thumbnail_source = pictures[thumbnail_figure_index]

		# Read in thumbnail source image
		img = Image.open(os.path.dirname(export_path) + '/' + thumbnail_source)

		# Trim whitespace
		bg = Image.new(img.mode, img.size, img.getpixel((0, 0)))
		diff = ImageChops.difference(img, bg)
		diff = ImageChops.add(diff, diff)
		bbox = diff.getbbox()
		if bbox:
			img = img.crop(bbox)

		# Resize image to have a width of 400px
		img.thumbnail([400, 1000])

		# Save thumbnail
		img.save(thumb_dest)
	except:
		shutil.copyfile('_static/no_thumb.png', thumb_dest)
	
	print('  Done!')

	return title, level, description, thumb_dest.split('/', 1)[-1]

index_preamble = '''
Tutorials
=========

These tutorials demonstrate the features of HCIPy in the context of a standard workflow. Tutorials are separated in three categories, depending on the required level of familiarity with HCIPy.
'''

entry_template = '''
.. only:: html

	.. container:: tutorial_item

		:doc:`{title} <{name}/{name}>`

		.. container:: tutorial_row

			.. container:: tutorial_thumbnail

				.. figure:: {thumbnail_file}

			.. container:: tutorial_description

				**Level:** {level}
				
				**Description:** {description}
'''

beginner_preamble = '''
Beginner
--------

These tutorials provide an introduction to the basic parts of HCIPy. New users should read these tutorials to get started with HCIPy.
'''

intermediate_preamble = '''
Intermediate
------------

These tutorials show the main functionality using the built-in classes of HCIPy. These tutorials focus on one aspect of high-contrast imaging.
'''

advanced_preamble = '''
Advanced
--------
These tutorials show how to use HCIPy for your own research. This includes extending HCIPy with your own optical elements and advanced use cases.
'''

expert_preamble = '''
Expert
------

These tutorials provide examples from actual published research that made heavy use of HCIPy for their optical propagations.
'''

unknown_preamble = '''
Unknown
-------

These tutorials do not have their level of difficulty rated.
'''

level_preambles = [beginner_preamble, intermediate_preamble, advanced_preamble, expert_preamble, unknown_preamble]

def compile_all_tutorials():
	print('Compling all tutorials...')

	tutorials = {}
	tutorial_names = sorted(os.listdir('tutorial_notebooks/'))
	tutorial_names = [name for name in tutorial_names if 'checkpoint' not in name]

	for name in tutorial_names:
		tutorials[name] = compile_tutorial(name)

	# Sort by level
	levels = ['Beginner', 'Intermediate', 'Advanced', 'Expert', 'Unknown']

	tutorial_names_by_level = [[] for i in range(len(levels))]
	for name in tutorial_names:
		tutorial_names_by_level[levels.index(tutorials[name][1])].append(name)

	f = open('tutorials/index.rst', 'w')
	f.write(index_preamble)

	for preamble, names in zip(level_preambles, tutorial_names_by_level):
		# Don't write this level if there are no tutorials in it.
		if len(names) == 0:
			continue

		f.write(preamble)
		
		# Write toctree
		f.write('\n.. toctree::\n    :maxdepth: 1\n    :hidden:\n\n')
		for name in names:
			f.write('    ' + name + '/' + name + '\n')
		f.write('\n\n')
		
		# Write list
		for name in names:
			title, level, desc, thumb = tutorials[name]
			f.write(entry_template.format(thumbnail_file=thumb, title=title, level=level, description=desc, name=name))
	
	f.close()
