#-------------------------------------------------------------------------

# train.py

# Train NeuSomatic network

#-------------------------------------------------------------------------

import os

import traceback

import argparse

import datetime

import logging

import numpy as np

import torch

from torch.autograd import Variable

import torch.nn as nn

import torch.nn.functional as F

import torch.optim as optim

from torchvision import transforms

import torchvision

from random import shuffle

import pickle

from network import NeuSomaticNet

from dataloader import NeuSomaticDataset, matrix_transform

from merge_tsvs import merge_tsvs

type_class_dict = {"DEL": 0, "INS": 1, "NONE": 2, "SNP": 3}

vartype_classes = ['DEL', 'INS', 'NONE', 'SNP']

import torch._utils

try:

    torch._utils._rebuild_tensor_v2

except AttributeError:

    def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):

        tensor = torch._utils._rebuild_tensor(

            storage, storage_offset, size, stride)

        tensor.requires_grad = requires_grad

        tensor._backward_hooks = backward_hooks

        return tensor

    torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2

def make_weights_for_balanced_classes(count_class_t, count_class_l, nclasses_t, nclasses_l,

                                      none_count=None):

    logger = logging.getLogger(make_weights_for_balanced_classes.__name__)

    w_t = [0] * nclasses_t

    w_l = [0] * nclasses_l

    count_class_t = list(count_class_t)

    count_class_l = list(count_class_l)

    if none_count:

        count_class_t[type_class_dict["NONE"]] = none_count

        count_class_l[0] = none_count

    logger.info("count type classes: {}".format(

        list(zip(vartype_classes, count_class_t))))

    N = float(sum(count_class_t))

    for i in range(nclasses_t):

        w_t[i] = (1 - (float(count_class_t[i]) / float(N))) / float(nclasses_t)

    w_t = np.array(w_t)

    logger.info("weight type classes: {}".format(

        list(zip(vartype_classes, w_t))))

    logger.info("count length classes: {}".format(list(

        zip(range(nclasses_l), count_class_l))))

    N = float(sum(count_class_l))

    for i in range(nclasses_l):

        w_l[i] = (1 - (float(count_class_l[i]) / float(N))) / float(nclasses_l)

    w_l = np.array(w_l)

    logger.info("weight length classes: {}".format(list(

        zip(range(nclasses_l), w_l))))

    return w_t, w_l

def test(net, epoch, validation_loader, use_cuda):

    logger = logging.getLogger(test.__name__)

    net.eval()

    nclasses = len(vartype_classes)

    class_correct = list(0. for i in range(nclasses))

    class_total = list(0. for i in range(nclasses))

    class_p_total = list(0. for i in range(nclasses))

    len_class_correct = list(0. for i in range(4))

    len_class_total = list(0. for i in range(4))

    len_class_p_total = list(0. for i in range(4))

    falses = []

    for data in validation_loader:

        (matrices, labels, _, var_len_s, _), (paths) = data

        matrices = Variable(matrices)

        if use_cuda:

            matrices = matrices.cuda()

        outputs, _ = net(matrices)

        [outputs1, outputs2, outputs3] = outputs

        _, predicted = torch.max(outputs1.data.cpu(), 1)

        pos_pred = outputs2.data.cpu().numpy()

        _, len_pred = torch.max(outputs3.data.cpu(), 1)

        preds = {}

        for i, _ in enumerate(paths[0]):

            preds[i] = [vartype_classes[predicted[i]], pos_pred[i], len_pred[i]]

        if labels.size()[0] > 1:

            compare_labels = (predicted == labels).squeeze()

        else:

            compare_labels = (predicted == labels)

        false_preds = np.where(compare_labels.numpy() == 0)[0]

        if len(false_preds) > 0:

            for i in false_preds:

                falses.append([paths[0][i], vartype_classes[predicted[i]], pos_pred[i], len_pred[i],

                               list(

                                   np.round(F.softmax(outputs1[i, :], 0).data.cpu().numpy(), 4)),

                               list(

                                   np.round(F.softmax(outputs3[i, :], 0).data.cpu().numpy(), 4))])

        for i in range(len(labels)):

            label = labels[i]

            class_correct[label] += compare_labels[i].data.cpu().numpy()

            class_total[label] += 1

        for i in range(len(predicted)):

            label = predicted[i]

            class_p_total[label] += 1

        if var_len_s.size()[0] > 1:

            compare_len = (len_pred == var_len_s).squeeze()

        else:

            compare_len = (len_pred == var_len_s)

        for i in range(len(var_len_s)):

            len_ = var_len_s[i]

            len_class_correct[len_] += compare_len[i].data.cpu().numpy()

            len_class_total[len_] += 1

        for i in range(len(len_pred)):

            len_ = len_pred[i]

            len_class_p_total[len_] += 1

    for i in range(nclasses):

        SN = 100 * class_correct[i] / (class_total[i] + 0.0001)

        PR = 100 * class_correct[i] / (class_p_total[i] + 0.0001)

        F1 = 2 * PR * SN / (PR + SN + 0.0001)

        logger.info('Epoch {}: Type Accuracy of {:>5} ({}) : {:.2f}  {:.2f} {:.2f}'.format(

            epoch,

            vartype_classes[i], class_total[i],

            SN, PR, F1))

    logger.info('Epoch {}: Type Accuracy of the network on the {} test candidates: {:.4f} %'.format(

        epoch, sum(class_total), (

            100 * sum(class_correct) / float(sum(class_total)))))

    for i in range(4):

        SN = 100 * len_class_correct[i] / (len_class_total[i] + 0.0001)

        PR = 100 * len_class_correct[i] / (len_class_p_total[i] + 0.0001)

        F1 = 2 * PR * SN / (PR + SN + 0.0001)

        logger.info('Epoch {}: Length Accuracy of {:>5} ({}) : {:.2f}  {:.2f} {:.2f}'.format(

            epoch, i, len_class_total[i],

            SN, PR, F1))

    logger.info('Epoch {}: Length Accuracy of the network on the {} test candidates: {:.4f} %'.format(

        epoch, sum(len_class_total), (

            100 * sum(len_class_correct) / float(sum(len_class_total)))))

    net.train()

class SubsetNoneSampler(torch.utils.data.sampler.Sampler):

    def __init__(self, none_indices, var_indices, none_count):

        self.none_indices = none_indices

        self.var_indices = var_indices

        self.none_count = none_count

        self.current_none_id = 0

    def __iter__(self):

        logger = logging.getLogger(SubsetNoneSampler.__iter__.__name__)

        if self.current_none_id > (len(self.none_indices) - self.none_count):

            this_round_nones = self.none_indices[self.current_none_id:]

            self.none_indices = list(map(lambda i: self.none_indices[i],

                                         torch.randperm(len(self.none_indices)).tolist()))

            self.current_none_id = self.none_count - len(this_round_nones)

            this_round_nones += self.none_indices[0:self.current_none_id]

        else:

            this_round_nones = self.none_indices[

                self.current_none_id:self.current_none_id + self.none_count]

            self.current_none_id += self.none_count

        current_indices = this_round_nones + self.var_indices

        ret = iter(map(lambda i: current_indices[i],

                       torch.randperm(len(current_indices))))

        return ret

    def __len__(self):

        return len(self.var_indices) + self.none_count

def train_neusomatic(candidates_tsv, validation_candidates_tsv, out_dir, checkpoint,

                     num_threads, batch_size, max_epochs, learning_rate, lr_drop_epochs,

                     lr_drop_ratio, momentum, boost_none, none_count_scale,

                     max_load_candidates, coverage_thr, save_freq, ensemble,

                     merged_candidates_per_tsv, merged_max_num_tsvs, overwrite_merged_tsvs,

                     train_split_len,

                     normalize_channels,

                     use_cuda):

    logger = logging.getLogger(train_neusomatic.__name__)

    logger.info("----------------Train NeuSomatic Network-------------------")

    logger.info("PyTorch Version: {}".format(torch.__version__))

    logger.info("Torchvision Version: {}".format(torchvision.__version__))

    if not os.path.exists(out_dir):

        os.mkdir(out_dir)

    if not use_cuda:

        torch.set_num_threads(num_threads)

    data_transform = matrix_transform((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))

    num_channels = 119 if ensemble else 26

    net = NeuSomaticNet(num_channels)

    if use_cuda:

        logger.info("GPU training!")

        net.cuda()

    else:

        logger.info("CPU training!")

    if torch.cuda.device_count() > 1:

        logger.info("We use {} GPUs!".format(torch.cuda.device_count()))

        net = nn.DataParallel(net)

    if not os.path.exists("{}/models/".format(out_dir)):

        os.mkdir("{}/models/".format(out_dir))

    if checkpoint:

        logger.info(

            "Load pretrained model from checkpoint {}".format(checkpoint))

        pretrained_dict = torch.load(

            checkpoint, map_location=lambda storage, loc: storage)

        pretrained_state_dict = pretrained_dict["state_dict"]

        tag = pretrained_dict["tag"]

        sofar_epochs = pretrained_dict["epoch"]

        logger.info(

            "sofar_epochs from pretrained checkpoint: {}".format(sofar_epochs))

        coverage_thr = pretrained_dict["coverage_thr"]

        logger.info(

            "Override coverage_thr from pretrained checkpoint: {}".format(coverage_thr))

        if "normalize_channels" in pretrained_dict:

            normalize_channels = pretrained_dict["normalize_channels"]

        else:

            normalize_channels = False

        logger.info(

            "Override normalize_channels from pretrained checkpoint: {}".format(normalize_channels))

        prev_epochs = sofar_epochs + 1

        model_dict = net.state_dict()

        # 1. filter out unnecessary keys

        # pretrained_state_dict = {

        # k: v for k, v in pretrained_state_dict.items() if k in model_dict}

        if "module." in list(pretrained_state_dict.keys())[0] and "module." not in list(model_dict.keys())[0]:

            pretrained_state_dict = {k.split("module.")[1]: v for k, v in pretrained_state_dict.items(

            ) if k.split("module.")[1] in model_dict}

        elif "module." not in list(pretrained_state_dict.keys())[0] and "module." in list(model_dict.keys())[0]:

            pretrained_state_dict = {

                ("module." + k): v for k, v in pretrained_state_dict.items() if ("module." + k) in model_dict}

        else:

            pretrained_state_dict = {k: v for k,

                                     v in pretrained_state_dict.items() if k in model_dict}

        # 2. overwrite entries in the existing state dict

        model_dict.update(pretrained_state_dict)

        # 3. load the new state dict

        net.load_state_dict(pretrained_state_dict)

    else:

        prev_epochs = 0

        time_now = datetime.datetime.now().strftime("%y-%m-%d-%H-%M-%S")

        tag = "neusomatic_{}".format(time_now)

    logger.info("tag: {}".format(tag))

    shuffle(candidates_tsv)

    if len(candidates_tsv) > merged_max_num_tsvs:

        candidates_tsv = merge_tsvs(input_tsvs=candidates_tsv, out=out_dir,

                                    candidates_per_tsv=merged_candidates_per_tsv,

                                    max_num_tsvs=merged_max_num_tsvs,

                                    overwrite_merged_tsvs=overwrite_merged_tsvs,

                                    keep_none_types=True)

    Ls = []

    for tsv in candidates_tsv:

        idx = pickle.load(open(tsv + ".idx", "rb"))

        Ls.append(len(idx) - 1)

    Ls, candidates_tsv = list(zip(

        *sorted(zip(Ls, candidates_tsv), key=lambda x: x[0], reverse=True)))

    train_split_tsvs = []

    current_L = 0

    current_split_tsvs = []

    for i, (L, tsv) in enumerate(zip(Ls, candidates_tsv)):

        current_L += L

        current_split_tsvs.append(tsv)

        if current_L >= train_split_len or (i == len(candidates_tsv) - 1 and current_L > 0):

            logger.info("tsvs in split {}: {}".format(

                len(train_split_tsvs), current_split_tsvs))

            train_split_tsvs.append(current_split_tsvs)

            current_L = 0

            current_split_tsvs = []

    assert sum(map(lambda x: len(x), train_split_tsvs)) == len(candidates_tsv)

    train_sets = []

    none_counts = []

    var_counts = []

    none_indices_ = []

    var_indices_ = []

    samplers = []

    for split_i, tsvs in enumerate(train_split_tsvs):

        train_set = NeuSomaticDataset(roots=tsvs,

                                      max_load_candidates=int(

                                          max_load_candidates * len(tsvs) / float(len(candidates_tsv))),

                                      transform=data_transform, is_test=False,

                                      num_threads=num_threads, coverage_thr=coverage_thr,

                                      normalize_channels=normalize_channels)

        train_sets.append(train_set)

        none_indices = train_set.get_none_indices()

        var_indices = train_set.get_var_indices()

        if none_indices:

            none_indices = list(map(lambda i: none_indices[i],

                                    torch.randperm(len(none_indices)).tolist()))

        logger.info(

            "Non-somatic candidates in split {}: {}".format(split_i, len(none_indices)))

        if var_indices:

            var_indices = list(map(lambda i: var_indices[i],

                                   torch.randperm(len(var_indices)).tolist()))

        logger.info("Somatic candidates in split {}: {}".format(

            split_i, len(var_indices)))

        none_count = max(min(len(none_indices), len(

            var_indices) * none_count_scale), 1)

        logger.info(

            "Non-somatic considered in each epoch of split {}: {}".format(split_i, none_count))

        sampler = SubsetNoneSampler(none_indices, var_indices, none_count)

        samplers.append(sampler)

        none_counts.append(none_count)

        var_counts.append(len(var_indices))

        var_indices_.append(var_indices)

        none_indices_.append(none_indices)

    logger.info("# Total Train cadidates: {}".format(

        sum(map(lambda x: len(x), train_sets))))

    if validation_candidates_tsv:

        validation_set = NeuSomaticDataset(roots=validation_candidates_tsv,

                                           max_load_candidates=max_load_candidates,

                                           transform=data_transform, is_test=True,

                                           num_threads=num_threads, coverage_thr=coverage_thr,

                                           normalize_channels=normalize_channels)

        validation_loader = torch.utils.data.DataLoader(validation_set,

                                                        batch_size=batch_size, shuffle=True,

                                                        num_workers=num_threads, pin_memory=True)

        logger.info("#Validation candidates: {}".format(len(validation_set)))

    count_class_t = [0] * 4

    count_class_l = [0] * 4

    for train_set in train_sets:

        for i in range(4):

            count_class_t[i] += train_set.count_class_t[i]

            count_class_l[i] += train_set.count_class_l[i]

    weights_type, weights_length = make_weights_for_balanced_classes(

        count_class_t, count_class_l, 4, 4, sum(none_counts))

    weights_type[2] *= boost_none

    weights_length[0] *= boost_none

    logger.info("weights_type:{}, weights_length:{}".format(

        weights_type, weights_length))

    loss_s = []

    gradients = torch.FloatTensor(weights_type)

    gradients2 = torch.FloatTensor(weights_length)

    if use_cuda:

        gradients = gradients.cuda()

        gradients2 = gradients2.cuda()

    criterion_crossentropy = nn.CrossEntropyLoss(gradients)

    criterion_crossentropy2 = nn.CrossEntropyLoss(gradients2)

    criterion_smoothl1 = nn.SmoothL1Loss()

    optimizer = optim.SGD(

        net.parameters(), lr=learning_rate, momentum=momentum)

    net.train()

    len_train_set = sum(none_counts) + sum(var_counts)

    logger.info("Number of candidater per epoch: {}".format(len_train_set))

    print_freq = max(1, int(len_train_set / float(batch_size) / 4.0))

    curr_epoch = prev_epochs

    torch.save({"state_dict": net.state_dict(),

                "tag": tag,

                "epoch": curr_epoch,

                "coverage_thr": coverage_thr,

                "normalize_channels": normalize_channels},

               '{}/models/checkpoint_{}_epoch{}.pth'.format(out_dir, tag, curr_epoch))

    if len(train_sets) == 1:

        train_sets[0].open_candidate_tsvs()

        train_loader = torch.utils.data.DataLoader(train_sets[0],

                                                   batch_size=batch_size,

                                                   num_workers=num_threads, pin_memory=True,

                                                   sampler=samplers[0])

    # loop over the dataset multiple times

    n_epoch = 0

    for epoch in range(max_epochs - prev_epochs):

        n_epoch += 1

        running_loss = 0.0

        i_ = 0

        for split_i, train_set in enumerate(train_sets):

            if len(train_sets) > 1:

                train_set.open_candidate_tsvs()

                train_loader = torch.utils.data.DataLoader(train_set,

                                                           batch_size=batch_size,

                                                           num_workers=num_threads, pin_memory=True,

                                                           sampler=samplers[split_i])

            for i, data in enumerate(train_loader, 0):

                i_ += 1

                # get the inputs

                (inputs, labels, var_pos_s, var_len_s, _), _ = data

                # wrap them in Variable

                inputs, labels, var_pos_s, var_len_s = Variable(inputs), Variable(

                    labels), Variable(var_pos_s), Variable(var_len_s)

                if use_cuda:

                    inputs, labels, var_pos_s, var_len_s = inputs.cuda(

                    ), labels.cuda(), var_pos_s.cuda(), var_len_s.cuda()

                # zero the parameter gradients

                optimizer.zero_grad()

                outputs, _ = net(inputs)

                [outputs_classification, outputs_pos, outputs_len] = outputs

                var_len_labels = Variable(

                    torch.LongTensor(var_len_s.cpu().data.numpy()))

                if use_cuda:

                    var_len_labels = var_len_labels.cuda()

                loss = criterion_crossentropy(outputs_classification, labels) + 1 * criterion_smoothl1(

                    outputs_pos.squeeze(1), var_pos_s[:, 1]

                ) + 1 * criterion_crossentropy2(outputs_len, var_len_labels)

                loss.backward()

                optimizer.step()

                loss_s.append(loss.data)

                running_loss += loss.data

                if i_ % print_freq == print_freq - 1:

                    logger.info('epoch: {}, iter: {:>7}, lr: {}, loss: {:.5f}'.format(

                                n_epoch + prev_epochs, len(loss_s),

                                learning_rate, running_loss / print_freq))

                    running_loss = 0.0

            if len(train_sets) > 1:

                train_set.close_candidate_tsvs()

        curr_epoch = n_epoch + prev_epochs

        if curr_epoch % save_freq == 0:

            torch.save({"state_dict": net.state_dict(),

                        "tag": tag,

                        "epoch": curr_epoch,

                        "coverage_thr": coverage_thr,

                        "normalize_channels": normalize_channels,

                        }, '{}/models/checkpoint_{}_epoch{}.pth'.format(out_dir, tag, curr_epoch))

            if validation_candidates_tsv:

                test(net, curr_epoch, validation_loader, use_cuda)

        if curr_epoch % lr_drop_epochs == 0:

            learning_rate *= lr_drop_ratio

            optimizer = optim.SGD(

                net.parameters(), lr=learning_rate, momentum=momentum)

    logger.info('Finished Training')

    if len(train_sets) == 1:

        train_sets[0].close_candidate_tsvs()

    curr_epoch = n_epoch + prev_epochs

    torch.save({"state_dict": net.state_dict(),

                "tag": tag,

                "epoch": curr_epoch,

                "coverage_thr": coverage_thr,

                "normalize_channels": normalize_channels,

                }, '{}/models/checkpoint_{}_epoch{}.pth'.format(

        out_dir, tag, curr_epoch))

    if validation_candidates_tsv:

        test(net, curr_epoch, validation_loader, use_cuda)

    logger.info("Total Epochs: {}".format(curr_epoch))

    logger.info("Total Epochs: {}".format(curr_epoch))

    logger.info("Training is Done.")

    return '{}/models/checkpoint_{}_epoch{}.pth'.format(out_dir, tag, curr_epoch)

if __name__ == '__main__':

    FORMAT = '%(levelname)s %(asctime)-15s %(name)-20s %(message)s'

    logging.basicConfig(level=logging.INFO, format=FORMAT)

    logger = logging.getLogger(__name__)

    parser = argparse.ArgumentParser(

        description='simple call variants from bam')

    parser.add_argument('--candidates_tsv', nargs="*",

                        help=' train candidate tsv files', required=True)

    parser.add_argument('--out', type=str,

                        help='output directory', required=True)

    parser.add_argument('--checkpoint', type=str,

                        help='pretrained network model checkpoint path', default=None)

    parser.add_argument('--validation_candidates_tsv', nargs="*",

                        help=' validation candidate tsv files', default=[])

    parser.add_argument('--ensemble',

                        help='Enable training for ensemble mode',

                        action="store_true")

    parser.add_argument('--num_threads', type=int,

                        help='number of threads', default=1)

    parser.add_argument('--batch_size', type=int,

                        help='batch size', default=1000)

    parser.add_argument('--max_epochs', type=int,

                        help='maximum number of training epochs', default=1000)

    parser.add_argument('--lr', type=float, help='learning_rate', default=0.01)

    parser.add_argument('--lr_drop_epochs', type=int,

                        help='number of epochs to drop learning rate', default=400)

    parser.add_argument('--lr_drop_ratio', type=float,

                        help='learning rate drop scale', default=0.1)

    parser.add_argument('--momentum', type=float,

                        help='SGD momentum', default=0.9)

    parser.add_argument('--boost_none', type=float,

                        help='the amount to boost none-somatic classification weight', default=100)

    parser.add_argument('--none_count_scale', type=float,

                        help='ratio of the none/somatic canidates to use in each training epoch \

                        (the none candidate will be subsampled in each epoch based on this ratio',

                        default=2)

    parser.add_argument('--max_load_candidates', type=int,

                        help='maximum candidates to load in memory', default=1000000)

    parser.add_argument('--save_freq', type=int,

                        help='the frequency of saving checkpoints in terms of # epochs', default=50)

    parser.add_argument('--merged_candidates_per_tsv', type=int,

                        help='Maximum number of candidates in each merged tsv file ', default=10000000)

    parser.add_argument('--merged_max_num_tsvs', type=int,

                        help='Maximum number of merged tsv files \

                        (higher priority than merged_candidates_per_tsv)', default=10)

    parser.add_argument('--overwrite_merged_tsvs',

                        help='if OUT/merged_tsvs/ folder exists overwrite the merged tsvs',

                        action="store_true")

    parser.add_argument('--train_split_len', type=int,

                        help='Maximum number of candidates used in each split of training (>=merged_candidates_per_tsv)',

                        default=10000000)

    parser.add_argument('--coverage_thr', type=int,

                        help='maximum coverage threshold to be used for network input \

                              normalization. \

                              Will be overridden if pretrained model is provided\

                              For ~50x WGS, coverage_thr=100 should work. \

                              For higher coverage WES, coverage_thr=300 should work.', default=100)

    parser.add_argument('--normalize_channels',

                        help='normalize BQ, MQ, and other bam-info channels by frequency of observed alleles. \

                              Will be overridden if pretrained model is provided',

                        action="store_true")

    args = parser.parse_args()

    logger.info(args)

    use_cuda = torch.cuda.is_available()

    logger.info("use_cuda: {}".format(use_cuda))

    try:

        checkpoint = train_neusomatic(args.candidates_tsv, args.validation_candidates_tsv,

                                      args.out, args.checkpoint, args.num_threads, args.batch_size,

                                      args.max_epochs,

                                      args.lr, args.lr_drop_epochs, args.lr_drop_ratio, args.momentum,

                                      args.boost_none, args.none_count_scale,

                                      args.max_load_candidates, args.coverage_thr, args.save_freq,

                                      args.ensemble,

                                      args.merged_candidates_per_tsv, args.merged_max_num_tsvs,

                                      args.overwrite_merged_tsvs, args.train_split_len,

                                      args.normalize_channels,

                                      use_cuda)

    except Exception as e:

        logger.error(traceback.format_exc())

        logger.error("Aborting!")

        logger.error(

            "train.py failure on arguments: {}".format(args))

        raise e