#!/usr/bin/env python3

"""Code to generate plots for Extended Data Fig. 4."""

import os

import matplotlib

import matplotlib.pyplot as plt

import numpy as np

import echonet

def main(root=os.path.join("timing", "video"),

         fig_root=os.path.join("figure", "complexity"),

         FRAMES=(1, 8, 16, 32, 64, 96),

         pretrained=True):

    """Generate plots for Extended Data Fig. 4."""

    echonet.utils.latexify()

    os.makedirs(fig_root, exist_ok=True)

    fig = plt.figure(figsize=(6.50, 2.50))

    gs = matplotlib.gridspec.GridSpec(1, 3, width_ratios=[2.5, 2.5, 1.50])

    ax = (plt.subplot(gs[0]), plt.subplot(gs[1]), plt.subplot(gs[2]))

    # Create legend

    for (model, color) in zip(["EchoNet-Dynamic (EF)", "R3D", "MC3"], matplotlib.colors.TABLEAU_COLORS):

        ax[2].plot([float("nan")], [float("nan")], "-", color=color, label=model)

    ax[2].set_title("")

    ax[2].axis("off")

    ax[2].legend(loc="center")

    for (model, color) in zip(["r2plus1d_18", "r3d_18", "mc3_18"], matplotlib.colors.TABLEAU_COLORS):

        for split in ["val"]:  # ["val", "train"]:

            print(model, split)

            data = [load(root, model, frames, 1, pretrained, split) for frames in FRAMES]

            time = np.array(list(map(lambda x: x[0], data)))

            n = np.array(list(map(lambda x: x[1], data)))

            mem_allocated = np.array(list(map(lambda x: x[2], data)))

            # mem_cached = np.array(list(map(lambda x: x[3], data)))

            batch_size = np.array(list(map(lambda x: x[4], data)))

            # Plot Time (panel a)

            ax[0].plot(FRAMES, time / n, "-" if pretrained else "--", marker=".", color=color, linewidth=(1 if split == "train" else None))

            print("Time:\n" + "\n".join(map(lambda x: "{:8d}: {:f}".format(*x), zip(FRAMES, time / n))))

            # Plot Memory (panel b)

            ax[1].plot(FRAMES, mem_allocated / batch_size / 1e9, "-" if pretrained else "--", marker=".", color=color, linewidth=(1 if split == "train" else None))

            print("Memory:\n" + "\n".join(map(lambda x: "{:8d}: {:f}".format(*x), zip(FRAMES, mem_allocated / batch_size / 1e9))))

            print()

    # Labels for panel a

    ax[0].set_xticks(FRAMES)

    ax[0].text(-0.05, 1.10, "(a)", transform=ax[0].transAxes)

    ax[0].set_xlabel("Clip length (frames)")

    ax[0].set_ylabel("Time Per Clip (seconds)")

    # Labels for panel b

    ax[1].set_xticks(FRAMES)

    ax[1].text(-0.05, 1.10, "(b)", transform=ax[1].transAxes)

    ax[1].set_xlabel("Clip length (frames)")

    ax[1].set_ylabel("Memory Per Clip (GB)")

    # Save figure

    plt.tight_layout()

    plt.savefig(os.path.join(fig_root, "complexity.pdf"))

    plt.savefig(os.path.join(fig_root, "complexity.eps"))

    plt.close(fig)

def load(root, model, frames, period, pretrained, split):

    """Loads runtime and memory usage for specified hyperparameter choice."""

    with open(os.path.join(root, "{}_{}_{}_{}".format(model, frames, period, "pretrained" if pretrained else "random"), "log.csv"), "r") as f:

        for line in f:

            line = line.split(",")

            if len(line) < 4:

                # Skip lines that are not csv (these lines log information)

                continue

            if line[1] == split:

                *_, time, n, mem_allocated, mem_cached, batch_size = line

                time = float(time)

                n = int(n)

                mem_allocated = int(mem_allocated)

                mem_cached = int(mem_cached)

                batch_size = int(batch_size)

                return time, n, mem_allocated, mem_cached, batch_size

    raise ValueError("File missing information.")

if __name__ == "__main__":

    main()