"""Functions for training and running segmentation."""

import math

import os

import time

import click

import matplotlib.pyplot as plt

import numpy as np

import scipy.signal

import skimage.draw

import torch

import torchvision

import tqdm

import echonet

@click.command("segmentation")

@click.option("--data_dir", type=click.Path(exists=True, file_okay=False), default=None)

@click.option("--output", type=click.Path(file_okay=False), default=None)

@click.option("--model_name", type=click.Choice(

    sorted(name for name in torchvision.models.segmentation.__dict__

           if name.islower() and not name.startswith("__") and callable(torchvision.models.segmentation.__dict__[name]))),

    default="deeplabv3_resnet50")

@click.option("--pretrained/--random", default=False)

@click.option("--weights", type=click.Path(exists=True, dir_okay=False), default=None)

@click.option("--run_test/--skip_test", default=False)

@click.option("--save_video/--skip_video", default=False)

@click.option("--num_epochs", type=int, default=50)

@click.option("--lr", type=float, default=1e-5)

@click.option("--weight_decay", type=float, default=0)

@click.option("--lr_step_period", type=int, default=None)

@click.option("--num_train_patients", type=int, default=None)

@click.option("--num_workers", type=int, default=4)

@click.option("--batch_size", type=int, default=20)

@click.option("--device", type=str, default=None)

@click.option("--seed", type=int, default=0)

def run(

    data_dir=None,

    output=None,

    model_name="deeplabv3_resnet50",

    pretrained=False,

    weights=None,

    run_test=False,

    save_video=False,

    num_epochs=50,

    lr=1e-5,

    weight_decay=1e-5,

    lr_step_period=None,

    num_train_patients=None,

    num_workers=4,

    batch_size=20,

    device=None,

    seed=0,

):

    """Trains/tests segmentation model.

    Args:

        data_dir (str, optional): Directory containing dataset. Defaults to

            `echonet.config.DATA_DIR`.

        output (str, optional): Directory to place outputs. Defaults to

            output/segmentation/<model_name>_<pretrained/random>/.

        model_name (str, optional): Name of segmentation model. One of ``deeplabv3_resnet50'',

            ``deeplabv3_resnet101'', ``fcn_resnet50'', or ``fcn_resnet101''

            (options are torchvision.models.segmentation.<model_name>)

            Defaults to ``deeplabv3_resnet50''.

        pretrained (bool, optional): Whether to use pretrained weights for model

            Defaults to False.

        weights (str, optional): Path to checkpoint containing weights to

            initialize model. Defaults to None.

        run_test (bool, optional): Whether or not to run on test.

            Defaults to False.

        save_video (bool, optional): Whether to save videos with segmentations.

            Defaults to False.

        num_epochs (int, optional): Number of epochs during training

            Defaults to 50.

        lr (float, optional): Learning rate for SGD

            Defaults to 1e-5.

        weight_decay (float, optional): Weight decay for SGD

            Defaults to 0.

        lr_step_period (int or None, optional): Period of learning rate decay

            (learning rate is decayed by a multiplicative factor of 0.1)

            Defaults to math.inf (never decay learning rate).

        num_train_patients (int or None, optional): Number of training patients

            for ablations. Defaults to all patients.

        num_workers (int, optional): Number of subprocesses to use for data

            loading. If 0, the data will be loaded in the main process.

            Defaults to 4.

        device (str or None, optional): Name of device to run on. Options from

            https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.device

            Defaults to ``cuda'' if available, and ``cpu'' otherwise.

        batch_size (int, optional): Number of samples to load per batch

            Defaults to 20.

        seed (int, optional): Seed for random number generator. Defaults to 0.

    """

    # Seed RNGs

    np.random.seed(seed)

    torch.manual_seed(seed)

    # Set default output directory

    if output is None:

        output = os.path.join("output", "segmentation", "{}_{}".format(model_name, "pretrained" if pretrained else "random"))

    os.makedirs(output, exist_ok=True)

    # Set device for computations

    if device is None:

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Set up model

    model = torchvision.models.segmentation.__dict__[model_name](pretrained=pretrained, aux_loss=False)

    model.classifier[-1] = torch.nn.Conv2d(model.classifier[-1].in_channels, 1, kernel_size=model.classifier[-1].kernel_size)  # change number of outputs to 1

    if device.type == "cuda":

        model = torch.nn.DataParallel(model)

    model.to(device)

    if weights is not None:

        checkpoint = torch.load(weights)

        model.load_state_dict(checkpoint['state_dict'])

    # Set up optimizer

    optim = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay)

    if lr_step_period is None:

        lr_step_period = math.inf

    scheduler = torch.optim.lr_scheduler.StepLR(optim, lr_step_period)

    # Compute mean and std

    mean, std = echonet.utils.get_mean_and_std(echonet.datasets.Echo(root=data_dir, split="train"))

    tasks = ["LargeFrame", "SmallFrame", "LargeTrace", "SmallTrace"]

    kwargs = {"target_type": tasks,

              "mean": mean,

              "std": std

              }

    # Set up datasets and dataloaders

    dataset = {}

    dataset["train"] = echonet.datasets.Echo(root=data_dir, split="train", **kwargs)

    if num_train_patients is not None and len(dataset["train"]) > num_train_patients:

        # Subsample patients (used for ablation experiment)

        indices = np.random.choice(len(dataset["train"]), num_train_patients, replace=False)

        dataset["train"] = torch.utils.data.Subset(dataset["train"], indices)

    dataset["val"] = echonet.datasets.Echo(root=data_dir, split="val", **kwargs)

    # Run training and testing loops

    with open(os.path.join(output, "log.csv"), "a") as f:

        epoch_resume = 0

        bestLoss = float("inf")

        try:

            # Attempt to load checkpoint

            checkpoint = torch.load(os.path.join(output, "checkpoint.pt"))

            model.load_state_dict(checkpoint['state_dict'])

            optim.load_state_dict(checkpoint['opt_dict'])

            scheduler.load_state_dict(checkpoint['scheduler_dict'])

            epoch_resume = checkpoint["epoch"] + 1

            bestLoss = checkpoint["best_loss"]

            f.write("Resuming from epoch {}\n".format(epoch_resume))

        except FileNotFoundError:

            f.write("Starting run from scratch\n")

        for epoch in range(epoch_resume, num_epochs):

            print("Epoch #{}".format(epoch), flush=True)

            for phase in ['train', 'val']:

                start_time = time.time()

                for i in range(torch.cuda.device_count()):

                    torch.cuda.reset_peak_memory_stats(i)

                ds = dataset[phase]

                dataloader = torch.utils.data.DataLoader(

                    ds, batch_size=batch_size, num_workers=num_workers, shuffle=True, pin_memory=(device.type == "cuda"), drop_last=(phase == "train"))

                loss, large_inter, large_union, small_inter, small_union = echonet.utils.segmentation.run_epoch(model, dataloader, phase == "train", optim, device)

                overall_dice = 2 * (large_inter.sum() + small_inter.sum()) / (large_union.sum() + large_inter.sum() + small_union.sum() + small_inter.sum())

                large_dice = 2 * large_inter.sum() / (large_union.sum() + large_inter.sum())

                small_dice = 2 * small_inter.sum() / (small_union.sum() + small_inter.sum())

                f.write("{},{},{},{},{},{},{},{},{},{},{}\n".format(epoch,

                                                                    phase,

                                                                    loss,

                                                                    overall_dice,

                                                                    large_dice,

                                                                    small_dice,

                                                                    time.time() - start_time,

                                                                    large_inter.size,

                                                                    sum(torch.cuda.max_memory_allocated() for i in range(torch.cuda.device_count())),

                                                                    sum(torch.cuda.max_memory_reserved() for i in range(torch.cuda.device_count())),

                                                                    batch_size))

                f.flush()

            scheduler.step()

            # Save checkpoint

            save = {

                'epoch': epoch,

                'state_dict': model.state_dict(),

                'best_loss': bestLoss,

                'loss': loss,

                'opt_dict': optim.state_dict(),

                'scheduler_dict': scheduler.state_dict(),

            }

            torch.save(save, os.path.join(output, "checkpoint.pt"))

            if loss < bestLoss:

                torch.save(save, os.path.join(output, "best.pt"))

                bestLoss = loss

        # Load best weights

        if num_epochs != 0:

            checkpoint = torch.load(os.path.join(output, "best.pt"))

            model.load_state_dict(checkpoint['state_dict'])

            f.write("Best validation loss {} from epoch {}\n".format(checkpoint["loss"], checkpoint["epoch"]))

        if run_test:

            # Run on validation and test

            for split in ["val", "test"]:

                dataset = echonet.datasets.Echo(root=data_dir, split=split, **kwargs)

                dataloader = torch.utils.data.DataLoader(dataset,

                                                         batch_size=batch_size, num_workers=num_workers, shuffle=False, pin_memory=(device.type == "cuda"))

                loss, large_inter, large_union, small_inter, small_union = echonet.utils.segmentation.run_epoch(model, dataloader, False, None, device)

                overall_dice = 2 * (large_inter + small_inter) / (large_union + large_inter + small_union + small_inter)

                large_dice = 2 * large_inter / (large_union + large_inter)

                small_dice = 2 * small_inter / (small_union + small_inter)

                with open(os.path.join(output, "{}_dice.csv".format(split)), "w") as g:

                    g.write("Filename, Overall, Large, Small\n")

                    for (filename, overall, large, small) in zip(dataset.fnames, overall_dice, large_dice, small_dice):

                        g.write("{},{},{},{}\n".format(filename, overall, large, small))

                f.write("{} dice (overall): {:.4f} ({:.4f} - {:.4f})\n".format(split, *echonet.utils.bootstrap(np.concatenate((large_inter, small_inter)), np.concatenate((large_union, small_union)), echonet.utils.dice_similarity_coefficient)))

                f.write("{} dice (large):   {:.4f} ({:.4f} - {:.4f})\n".format(split, *echonet.utils.bootstrap(large_inter, large_union, echonet.utils.dice_similarity_coefficient)))

                f.write("{} dice (small):   {:.4f} ({:.4f} - {:.4f})\n".format(split, *echonet.utils.bootstrap(small_inter, small_union, echonet.utils.dice_similarity_coefficient)))

                f.flush()

    # Saving videos with segmentations

    dataset = echonet.datasets.Echo(root=data_dir, split="test",

                                    target_type=["Filename", "LargeIndex", "SmallIndex"],  # Need filename for saving, and human-selected frames to annotate

                                    mean=mean, std=std,  # Normalization

                                    length=None, max_length=None, period=1  # Take all frames

                                    )

    dataloader = torch.utils.data.DataLoader(dataset, batch_size=10, num_workers=num_workers, shuffle=False, pin_memory=False, collate_fn=_video_collate_fn)

    # Save videos with segmentation

    if save_video and not all(os.path.isfile(os.path.join(output, "videos", f)) for f in dataloader.dataset.fnames):

        # Only run if missing videos

        model.eval()

        os.makedirs(os.path.join(output, "videos"), exist_ok=True)

        os.makedirs(os.path.join(output, "size"), exist_ok=True)

        echonet.utils.latexify()

        with torch.no_grad():

            with open(os.path.join(output, "size.csv"), "w") as g:

                g.write("Filename,Frame,Size,HumanLarge,HumanSmall,ComputerSmall\n")

                for (x, (filenames, large_index, small_index), length) in tqdm.tqdm(dataloader):

                    # Run segmentation model on blocks of frames one-by-one

                    # The whole concatenated video may be too long to run together

                    y = np.concatenate([model(x[i:(i + batch_size), :, :, :].to(device))["out"].detach().cpu().numpy() for i in range(0, x.shape[0], batch_size)])

                    start = 0

                    x = x.numpy()

                    for (i, (filename, offset)) in enumerate(zip(filenames, length)):

                        # Extract one video and segmentation predictions

                        video = x[start:(start + offset), ...]

                        logit = y[start:(start + offset), 0, :, :]

                        # Un-normalize video

                        video *= std.reshape(1, 3, 1, 1)

                        video += mean.reshape(1, 3, 1, 1)

                        # Get frames, channels, height, and width

                        f, c, h, w = video.shape  # pylint: disable=W0612

                        assert c == 3

                        # Put two copies of the video side by side

                        video = np.concatenate((video, video), 3)

                        # If a pixel is in the segmentation, saturate blue channel

                        # Leave alone otherwise

                        video[:, 0, :, w:] = np.maximum(255. * (logit > 0), video[:, 0, :, w:])  # pylint: disable=E1111

                        # Add blank canvas under pair of videos

                        video = np.concatenate((video, np.zeros_like(video)), 2)

                        # Compute size of segmentation per frame

                        size = (logit > 0).sum((1, 2))

                        # Identify systole frames with peak detection

                        trim_min = sorted(size)[round(len(size) ** 0.05)]

                        trim_max = sorted(size)[round(len(size) ** 0.95)]

                        trim_range = trim_max - trim_min

                        systole = set(scipy.signal.find_peaks(-size, distance=20, prominence=(0.50 * trim_range))[0])

                        # Write sizes and frames to file

                        for (frame, s) in enumerate(size):

                            g.write("{},{},{},{},{},{}\n".format(filename, frame, s, 1 if frame == large_index[i] else 0, 1 if frame == small_index[i] else 0, 1 if frame in systole else 0))

                        # Plot sizes

                        fig = plt.figure(figsize=(size.shape[0] / 50 * 1.5, 3))

                        plt.scatter(np.arange(size.shape[0]) / 50, size, s=1)

                        ylim = plt.ylim()

                        for s in systole:

                            plt.plot(np.array([s, s]) / 50, ylim, linewidth=1)

                        plt.ylim(ylim)

                        plt.title(os.path.splitext(filename)[0])

                        plt.xlabel("Seconds")

                        plt.ylabel("Size (pixels)")

                        plt.tight_layout()

                        plt.savefig(os.path.join(output, "size", os.path.splitext(filename)[0] + ".pdf"))

                        plt.close(fig)

                        # Normalize size to [0, 1]

                        size -= size.min()

                        size = size / size.max()

                        size = 1 - size

                        # Iterate the frames in this video

                        for (f, s) in enumerate(size):

                            # On all frames, mark a pixel for the size of the frame

                            video[:, :, int(round(115 + 100 * s)), int(round(f / len(size) * 200 + 10))] = 255.

                            if f in systole:

                                # If frame is computer-selected systole, mark with a line

                                video[:, :, 115:224, int(round(f / len(size) * 200 + 10))] = 255.

                            def dash(start, stop, on=10, off=10):

                                buf = []

                                x = start

                                while x < stop:

                                    buf.extend(range(x, x + on))

                                    x += on

                                    x += off

                                buf = np.array(buf)

                                buf = buf[buf < stop]

                                return buf

                            d = dash(115, 224)

                            if f == large_index[i]:

                                # If frame is human-selected diastole, mark with green dashed line on all frames

                                video[:, :, d, int(round(f / len(size) * 200 + 10))] = np.array([0, 225, 0]).reshape((1, 3, 1))

                            if f == small_index[i]:

                                # If frame is human-selected systole, mark with red dashed line on all frames

                                video[:, :, d, int(round(f / len(size) * 200 + 10))] = np.array([0, 0, 225]).reshape((1, 3, 1))

                            # Get pixels for a circle centered on the pixel

                            r, c = skimage.draw.disk((int(round(115 + 100 * s)), int(round(f / len(size) * 200 + 10))), 4.1)

                            # On the frame that's being shown, put a circle over the pixel

                            video[f, :, r, c] = 255.

                        # Rearrange dimensions and save

                        video = video.transpose(1, 0, 2, 3)

                        video = video.astype(np.uint8)

                        echonet.utils.savevideo(os.path.join(output, "videos", filename), video, 50)

                        # Move to next video

                        start += offset

def run_epoch(model, dataloader, train, optim, device):

    """Run one epoch of training/evaluation for segmentation.

    Args:

        model (torch.nn.Module): Model to train/evaulate.

        dataloder (torch.utils.data.DataLoader): Dataloader for dataset.

        train (bool): Whether or not to train model.

        optim (torch.optim.Optimizer): Optimizer

        device (torch.device): Device to run on

    """

    total = 0.

    n = 0

    pos = 0

    neg = 0

    pos_pix = 0

    neg_pix = 0

    model.train(train)

    large_inter = 0

    large_union = 0

    small_inter = 0

    small_union = 0

    large_inter_list = []

    large_union_list = []

    small_inter_list = []

    small_union_list = []

    with torch.set_grad_enabled(train):

        with tqdm.tqdm(total=len(dataloader)) as pbar:

            for (_, (large_frame, small_frame, large_trace, small_trace)) in dataloader:

                # Count number of pixels in/out of human segmentation

                pos += (large_trace == 1).sum().item()

                pos += (small_trace == 1).sum().item()

                neg += (large_trace == 0).sum().item()

                neg += (small_trace == 0).sum().item()

                # Count number of pixels in/out of computer segmentation

                pos_pix += (large_trace == 1).sum(0).to("cpu").detach().numpy()

                pos_pix += (small_trace == 1).sum(0).to("cpu").detach().numpy()

                neg_pix += (large_trace == 0).sum(0).to("cpu").detach().numpy()

                neg_pix += (small_trace == 0).sum(0).to("cpu").detach().numpy()

                # Run prediction for diastolic frames and compute loss

                large_frame = large_frame.to(device)

                large_trace = large_trace.to(device)

                y_large = model(large_frame)["out"]

                loss_large = torch.nn.functional.binary_cross_entropy_with_logits(y_large[:, 0, :, :], large_trace, reduction="sum")

                # Compute pixel intersection and union between human and computer segmentations

                large_inter += np.logical_and(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum()

                large_union += np.logical_or(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum()

                large_inter_list.extend(np.logical_and(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2)))

                large_union_list.extend(np.logical_or(y_large[:, 0, :, :].detach().cpu().numpy() > 0., large_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2)))

                # Run prediction for systolic frames and compute loss

                small_frame = small_frame.to(device)

                small_trace = small_trace.to(device)

                y_small = model(small_frame)["out"]

                loss_small = torch.nn.functional.binary_cross_entropy_with_logits(y_small[:, 0, :, :], small_trace, reduction="sum")

                # Compute pixel intersection and union between human and computer segmentations

                small_inter += np.logical_and(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum()

                small_union += np.logical_or(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum()

                small_inter_list.extend(np.logical_and(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2)))

                small_union_list.extend(np.logical_or(y_small[:, 0, :, :].detach().cpu().numpy() > 0., small_trace[:, :, :].detach().cpu().numpy() > 0.).sum((1, 2)))

                # Take gradient step if training

                loss = (loss_large + loss_small) / 2

                if train:

                    optim.zero_grad()

                    loss.backward()

                    optim.step()

                # Accumulate losses and compute baselines

                total += loss.item()

                n += large_trace.size(0)

                p = pos / (pos + neg)

                p_pix = (pos_pix + 1) / (pos_pix + neg_pix + 2)

                # Show info on process bar

                pbar.set_postfix_str("{:.4f} ({:.4f}) / {:.4f} {:.4f}, {:.4f}, {:.4f}".format(total / n / 112 / 112, loss.item() / large_trace.size(0) / 112 / 112, -p * math.log(p) - (1 - p) * math.log(1 - p), (-p_pix * np.log(p_pix) - (1 - p_pix) * np.log(1 - p_pix)).mean(), 2 * large_inter / (large_union + large_inter), 2 * small_inter / (small_union + small_inter)))

                pbar.update()

    large_inter_list = np.array(large_inter_list)

    large_union_list = np.array(large_union_list)

    small_inter_list = np.array(small_inter_list)

    small_union_list = np.array(small_union_list)

    return (total / n / 112 / 112,

            large_inter_list,

            large_union_list,

            small_inter_list,

            small_union_list,

            )

def _video_collate_fn(x):

    """Collate function for Pytorch dataloader to merge multiple videos.

    This function should be used in a dataloader for a dataset that returns

    a video as the first element, along with some (non-zero) tuple of

    targets. Then, the input x is a list of tuples:

      - x[i][0] is the i-th video in the batch

      - x[i][1] are the targets for the i-th video

    This function returns a 3-tuple:

      - The first element is the videos concatenated along the frames

        dimension. This is done so that videos of different lengths can be

        processed together (tensors cannot be "jagged", so we cannot have

        a dimension for video, and another for frames).

      - The second element is contains the targets with no modification.

      - The third element is a list of the lengths of the videos in frames.

    """

    video, target = zip(*x)  # Extract the videos and targets

    # ``video'' is a tuple of length ``batch_size''

    #   Each element has shape (channels=3, frames, height, width)

    #   height and width are expected to be the same across videos, but

    #   frames can be different.

    # ``target'' is also a tuple of length ``batch_size''

    # Each element is a tuple of the targets for the item.

    i = list(map(lambda t: t.shape[1], video))  # Extract lengths of videos in frames

    # This contatenates the videos along the the frames dimension (basically

    # playing the videos one after another). The frames dimension is then

    # moved to be first.

    # Resulting shape is (total frames, channels=3, height, width)

    video = torch.as_tensor(np.swapaxes(np.concatenate(video, 1), 0, 1))

    # Swap dimensions (approximately a transpose)

    # Before: target[i][j] is the j-th target of element i

    # After:  target[i][j] is the i-th target of element j

    target = zip(*target)

    return video, target, i