"""
Program: Figure 2 - Movie Image Plots
Date: 24 November 2018
Author: Lebo Molefe
"""

# IMPORTANT: For this file to run, change the directory within the `get_file_path` function to the `movies` folder.

import numpy as np
from mraw import mraw
import matplotlib.pyplot as plt

# Gets the .cih file path from a movie number.
def get_file_path(movie):
    """
    Return the file path for the .cih file with the movie number specified.
    :param movie: movie number
    :return: A string file path of the form 'directory/movie####/a.cih'
    """

    directory = 'E:/data/Internships/2018_Lebo/movies/'
    # Write the movie number as a string such as '0004'.
    movie_number = str(movie).zfill(4)
    filepath = directory + 'movie{0}/a.cih'.format(movie_number)

    return filepath


# Settings
gridsize = (4, 2)
fig = plt.figure(figsize=(4, 8))
plt.suptitle('')
fontsize = 18
fontcolor = 'w'
letter = ['a', 'b', 'c', 'd',
          'e', 'f', 'g', 'h']
# Calculated times in ms, based on frame rate.
t1 = 0.01
t2 = 0.07
t3 = 0.21
t4 = 0.01
t5 = 0.05
t6 = 0.25
label = ['$t$ = {0} ms'.format(t1), '$t$ = {0} ms'.format(t2), '$t$ = {0} ms'.format(t3),
         '$t$ = {0} ms'.format(t4), '$t$ = {0} ms'.format(t5), '$t$ = {0} ms'.format(t6)]

# Extracting data from `[movie, frame1, frame2, frame3].`
# The second and third movie frames are the first and second peak bubble areas.
for i, movie_frames in enumerate([[635, 8117, 8123, 8137], [595, 2517, 2521, 2541]]):
    movie = movie_frames[0]
    frame1 = movie_frames[1]
    frame2 = movie_frames[2]
    frame3 = movie_frames[3]

    ax1 = plt.subplot2grid(gridsize, (0, i), rowspan=1, colspan=1)
    ax2 = plt.subplot2grid(gridsize, (1, i), rowspan=1, colspan=1)
    ax3 = plt.subplot2grid(gridsize, (2, i), rowspan=1, colspan=1)
    ax4 = plt.subplot2grid(gridsize, (3, i), rowspan=1, colspan=1)

    f1 = mraw(get_file_path(movie)).get_frame(frame1)
    f2 = mraw(get_file_path(movie)).get_frame(frame2)
    f3 = mraw(get_file_path(movie)).get_frame(frame3)

    # Subtract the background.
    b = mraw(get_file_path(movie)).get_frame(0)

    px = 30
    py = 70

    # Plot a scale bar that is about 33.2 for the triangle and 33.9 pixels for the square (equal to one mm).
    start_px = 30
    if i == 0:
        x = np.arange(start_px, start_px + 33 + 1, 1)
        y = [230] * len(x)
    elif i == 1:
        x = np.arange(start_px, start_px + 34 + 1, 1)
        y = [230] * len(x)

    # First plot is a schematic of the recording area.
    ax1.imshow(f1 * 0 + 4095, cmap=plt.cm.gray, vmin=0, vmax=4095)
    if i == 0:
        ax1.plot((117, 117, 267, 267, 117), (57, 207, 207, 57, 57), linestyle='-', color='k')
        ax1.plot((180, 180, 280, 280, 180), (45, 123, 123, 45, 45), linestyle='-', color='k')  # A frame of 100 x 78
        ax1.plot(235, 85, marker='o', markersize=6, color='k')
        ax1.fill_between(x=(180, 280), y1=45, y2=123, color='k', alpha=0.4)
    elif i == 1:
        ax1.plot((267, 267, 130, 267), (57, 207, 132, 57), linestyle='-', color='k')
        ax1.plot((180, 180, 280, 280, 180), (70, 148, 148, 70, 70), linestyle='-', color='k')  # A frame of 100 x 78
        ax1.plot(245, 95, marker='o', markersize=6, color='k')
        ax1.fill_between(x=(180, 280), y1=70, y2=148, color='k', alpha=0.4)
    ax1.set_xticks([])
    ax1.set_yticks([])
    ax1.text(px, py, letter[0 + 1 * i], fontsize=fontsize, color='k')

    # Multiplying by 1.5 increases the contrast.
    ax2.imshow(f1 * 1.5, cmap=plt.cm.gray, vmin=0, vmax=4095)
    ax2.plot(x, y, linestyle='-', color='w')
    ax2.set_xticks([])
    ax2.set_yticks([])
    ax2.set_xlabel(label[0 + 3 * i], fontsize=fontsize)
    ax2.text(px, py, letter[2 + 1 * i], fontsize=fontsize, color=fontcolor)
    ax3.imshow(f2 * 1.5, cmap=plt.cm.gray, vmin=0, vmax=4095)
    ax3.set_xticks([])
    ax3.set_yticks([])
    ax3.set_xlabel(label[1 + 3 * i], fontsize=fontsize)
    ax3.text(px, py, letter[4 + 1 * i], fontsize=fontsize, color=fontcolor)
    ax4.imshow(f3 * 1.5, cmap=plt.cm.gray, vmin=0, vmax=4095)
    ax4.set_xticks([])
    ax4.set_yticks([])
    ax4.set_xlabel(label[2 + 3 * i], fontsize=fontsize)
    ax4.text(px, py, letter[6 + 1 * i], fontsize=fontsize, color=fontcolor)

#plt.tight_layout(h_pad=-10)
plt.tight_layout(h_pad=-10)
plt.savefig('figure02.pdf', bbox_inches='tight', pad_inches=0)
plt.show()