"""
Dynamic Routing Between Capsules
https://arxiv.org/abs/1710.09829

PyTorch implementation by Kenta Iwasaki @ Gram.AI.
CapsRT implemented by horsepurve.
"""

import torch
import torch.nn.functional as F
from torch import nn
import numpy as np
from torch.autograd import Variable
import pickle
from scipy import sparse
from config import *

CNN_EMB = True # False #

BATCH_SIZE = 16 # 16
NUM_CLASSES = 10
NUM_EPOCHS = 20
NUM_ROUTING_ITERATIONS = 1
CUDA = False # True # False #
LR = 0.01

# train_path = 'data/dia_train_269.txt' # 'data/mod_train_2.txt' # 'data/unmod_train_2.txt' # 'data/SCX_train_42.txt' #
# test_path = 'data/dia_test_269.txt' # 'data/mod_test_2.txt' # 'data/unmod_test_2.txt' # 'data/SCX_test_42.txt' #
# result_path = 'dia_pred_269.txt' # 'mod_pred_2.txt' # 'unmod_pred_2.txt' # 'SCX_pred_42.txt' #
RTdata_path = 'dia.pkl' # 'mod.pkl' # 'unmod.pkl' # 'SCX.pkl' #
LOAD_DATA = True # False #
# TODO: add max_length to config.py
## max_length = 66 # 66 # 50 # 38 # 50 #

# log_path = ''
if '' == dict_path:
    dict_path = train_path

from RTdata_emb import Dictionary, RTdata, Pearson, Spearman, Delta_t95, DATA_AUGMENTATION, Corpus
dictionary = Dictionary(dict_path)
'''
if True == LOAD_DATA:
    dictionary = Dictionary(dict_path)
    RTtrain = RTdata(dictionary, max_length, train_path)
    RTtest = RTdata(dictionary, max_length, test_path)
    with open(RTdata_path, 'wb') as output:
        pickle.dump(dictionary, output)
        pickle.dump(RTtrain, output)
        pickle.dump(RTtest, output)
if False == LOAD_DATA:
    with open(RTdata_path, 'rb') as input:
        dictionary = pickle.load(input)
        RTtrain = pickle.load(input)
        RTtest = pickle.load(input)
        print('>> note: load pre-read RTdata from:', RTdata_path)

# DATA_AUGMENTATION = True
SPARSE = True
def desparse(RTtt):
    X = np.zeros((RTtt.number_seq, RTtt.N_aa, RTtt.N_time_step)) # DATA_AUGMENTATION->*2
    for i in range(RTtt.number_seq): # DATA_AUGMENTATION->*2
        # sparse to dense
        X[i,::] = RTtt.X[i].todense()
    RTtt.X = X
if True == SPARSE:
    print('>> note: de-sparse for both train & test data.')
    desparse(RTtrain)
    desparse(RTtest)
'''

def softmax(input, dim=1):
    transposed_input = input.transpose(dim, len(input.size()) - 1)
    # print(transposed_input.contiguous().view(-1, transposed_input.size(-1)).shape)
    '''
    PyTorch 0.3.0:
    UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
    '''
    softmaxed_output = F.softmax(transposed_input.contiguous().view(-1, transposed_input.size(-1)),dim=1)
    return softmaxed_output.view(*transposed_input.size()).transpose(dim, len(input.size()) - 1)

def augmentation(x, max_shift=2):
    _, _, height, width = x.size()

    h_shift, w_shift = np.random.randint(-max_shift, max_shift + 1, size=2)
    source_height_slice = slice(max(0, h_shift), h_shift + height)
    source_width_slice = slice(max(0, w_shift), w_shift + width)
    target_height_slice = slice(max(0, -h_shift), -h_shift + height)
    target_width_slice = slice(max(0, -w_shift), -w_shift + width)

    shifted_image = torch.zeros(*x.size())
    shifted_image[:, :, source_height_slice, source_width_slice] = x[:, :, target_height_slice, target_width_slice]
    return shifted_image.float() # Note float here!

class CapsuleLayer(nn.Module):
    def __init__(self,
                 num_capsules,
                 num_route_nodes,
                 in_channels,
                 out_channels,
                 kernel_size=None,
                 stride=None,
                 num_iterations=NUM_ROUTING_ITERATIONS):

        super(CapsuleLayer, self).__init__()

        self.num_route_nodes = num_route_nodes
        self.num_iterations = num_iterations

        self.num_capsules = num_capsules

        if num_route_nodes != -1:
            self.route_weights = nn.Parameter(torch.randn(num_capsules,
                                                          num_route_nodes,
                                                          in_channels,
                                                          out_channels))
        else:
            self.capsules = nn.ModuleList(
                [nn.Conv2d(in_channels,
                           out_channels,
                           kernel_size=kernel_size,
                           stride=stride,
                           padding=0) for _ in
                 range(num_capsules)])

    def squash(self, tensor, dim=-1):
        squared_norm = (tensor ** 2).sum(dim=dim, keepdim=True)
        scale = squared_norm / (1 + squared_norm)
        return scale * tensor / torch.sqrt(squared_norm)

    def forward(self, x):
        if self.num_route_nodes != -1:
            priors = x[None, :, :, None, :] @ self.route_weights[:, None, :, :, :]

            if True == CUDA:
                logits = Variable(torch.zeros(*priors.size())).cuda()
            if False == CUDA:
                logits = Variable(torch.zeros(*priors.size()))

            for i in range(self.num_iterations):
                probs = softmax(logits, dim=2)
                outputs = self.squash((probs * priors).sum(dim=2, keepdim=True))

                if i != self.num_iterations - 1:
                    delta_logits = (priors * outputs).sum(dim=-1, keepdim=True)
                    logits = logits + delta_logits
        else:
            outputs = [capsule(x).view(x.size(0), -1, 1) for capsule in self.capsules]
            outputs = torch.cat(outputs, dim=-1)
            outputs = self.squash(outputs)

        return outputs

param_2D = {'data' : 'mnist',
            'dim' : 2,
            'conv1_kernel' : 9,
            'pri_caps_kernel' : 9,
            'stride' : 2,
            'digit_caps_nodes' : 32 * 6 * 6,
            'NUM_CLASSES' : NUM_CLASSES}

param_1D = {'data' : 'mnist',
            'dim' : 1,
            'conv1_kernel' : (28, 9),
            'pri_caps_kernel' : (1, 9),
            'stride' : 1,
            'digit_caps_nodes' : 32 * 1 * 12,
            'NUM_CLASSES' : 1}

# conv1_kernel = 15
# conv2_kernel = 15
param_1D_rt = {'data' : 'rt',
               'dim' : 1,
               'conv1_kernel' : (len(dictionary), conv1_kernel),
               'pri_caps_kernel' : (1, conv2_kernel),
               'stride' : 1,
               'digit_caps_nodes' : 32 * 1 * (max_length - conv1_kernel*2 + 2 - conv2_kernel + 1), # 32 # Note: number of conv!
               'NUM_CLASSES' : 1}

param = param_1D_rt

if 2 == param['dim']:
    print('>> note: using image mode.')
if 1 == param['dim']:
    print('>> note: using seq mode.')

class CapsuleNet(nn.Module):
    def __init__(self,conv1_kernel,conv2_kernel):
        super(CapsuleNet, self).__init__()
        EMB_SIZE = 0
        if True == CNN_EMB:
            # self.emb = nn.Embedding(len(dictionary), len(dictionary))
            # Note: if using embedding, EMB_SIZE can be any value, and we choose 20 here
            EMB_SIZE = 20
            self.emb = nn.Embedding(len(dictionary), EMB_SIZE) # we use 20 for all data
        else:
            # Note: if using one-hot encoding, EMB_SIZE must be the same as len(dictionary)
            EMB_SIZE = len(dictionary)
        self.conv1 = nn.Conv2d(in_channels=1,
                               out_channels=256, # 256
                               kernel_size=(EMB_SIZE, conv1_kernel), # param['conv1_kernel'], # (28, 9), # 9,
                               stride=1)
        ''''''
        self.bn1 = nn.BatchNorm2d(256) # Note: do we need this or not?
        self.conv2 = nn.Conv2d(in_channels=256,
                               out_channels=256, # 256
                               kernel_size=(1, conv1_kernel), # (28, 9), # 9,
                               stride=1)
        self.bn2 = nn.BatchNorm2d(256)
        '''
        self.conv3 = nn.Conv2d(in_channels=128,
                               out_channels=256, # 256
                               kernel_size=(1, conv1_kernel), # (28, 9), # 9,
                               stride=1)
        self.bn3 = nn.BatchNorm2d(256)
        '''

        self.primary_capsules = CapsuleLayer(num_capsules=8, # 8
                                             num_route_nodes=-1,
                                             in_channels=256, # 256
                                             out_channels=32, # 32
                                             kernel_size=(1, conv2_kernel), # param['pri_caps_kernel'], # (1, 9), # 9,
                                             stride=param['stride']) # 1) # 2)

        self.digit_capsules = CapsuleLayer(num_capsules=param['NUM_CLASSES'], # 1, #NUM_CLASSES, # DeepRT
                                           num_route_nodes=32 * 1 * (max_length - conv1_kernel*2 + 2 - conv2_kernel + 1), # param['digit_caps_nodes'], # 32 * 1 * 12, # 32 * 6 * 6,
                                           in_channels=8, # 8
                                           out_channels=16) # max_length-conv1_kernel + 1) # 16

        # add dropout:
        # self.dropout = nn.Dropout(0.1) # not good!
        # self.linear = nn.Linear((max_length-conv1_kernel+1)*256,16) # try residue: not good!
        ''' residue is not very good!
        pad = 0
        kernel_h = pad*2+1 # len + pad*2 - (kernel_h - 1) = len
        self.conv_res = nn.Conv2d(in_channels=256,
                                  out_channels=1, # 256
                                  kernel_size=(1, kernel_h), # (28, 9), # 9,
                                  stride=1)
                                  #padding =(0,pad))
        '''
        self.decoder = nn.Sequential(
            nn.Linear(16 * NUM_CLASSES, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 784),
            nn.Sigmoid()
        )

    def forward(self, x, y=None):
        # print('>>dim: input', x.shape) # [batch, 1, 28, 28]
        # print('>>dim: y', y) # [batch, 10] ~ [batch, NUM_CLASSES]
        if True == CNN_EMB:
            x = self.emb(x) # [batch, len] -> [batch, len, dict]
            x = x.transpose(dim0=1, dim1=2) # -> [batch, dict, len]
            x = x[:,None,:,:] # -> [batch, 1, dict, len]

        # ^^^^^ pre-process x ^^^^^
        x = F.relu(self.bn1(self.conv1(x)), inplace=True)

        ''' try residue: not good!
        residue = x.view(x.shape[0],-1)
        residue = self.linear(residue).view(residue.shape[0],1,16)
        # another residue method
        residue = F.relu(self.conv_res(x), inplace=True)
        residue = residue.view(residue.shape[0],1,residue.shape[-1])
        '''

        # x = self.dropout(x)
        x = F.relu(self.bn2(self.conv2(x)), inplace=True) # improvement
        # x = F.relu(self.bn3(self.conv3(x)), inplace=True)
        # print('>>dim: conv1', x.shape) # [batch, 256, 20, 20]
        x = self.primary_capsules(x)
        # print('>>dim: primary_capsules', x.shape) # [batch, 1152, 8] = [batch, 6*6*32, 8]
        # print('>>dim: unsqueezeed', self.digit_capsules(x).shape) # [10, batch, 1, 1, 16] ~ [num_caps, batch, ...]
        if 2 == param['dim']:
            x = self.digit_capsules(x).squeeze().transpose(0, 1) # DeepRT
            # [10, batch, 1, 1, 16] -> squeeze: [10, batch, 16] -> transpose: [batch, 10, 16]
        if 1 == param['dim']:
            x = self.digit_capsules(x).squeeze()[:, None, :]
            # [1, batch, 1, 1, 16] -> squeeze: [batch, 16]
        # print('>>dim: digit_capsules', x.shape) # [batch, 10, 16]

        # add dropout:
        # x = self.dropout(x)
        # x = self.linear(x)
        # x = F.sigmoid(x)

        # x = x + residue # try residue: not good!
        classes = (x ** 2).sum(dim=-1) ** 0.5
        # print('>>dim: classes', classes) # [batch, 10]
        if 2 == param['dim']:
            classes = F.softmax(classes) # DeepRT
        # print('>>dim: softmax', classes)

        if y is None: # Note: not do this during training. Here y is only used for reconstruction
            if 2 == param['dim']:
                # In all batches, get the most active capsule.
                # print('>>dim: reconstruction', classes) # [batch, 10]
                _, max_length_indices = classes.max(dim=1)
                # give: [torch.FloatTensor of size batch] and [torch.FloatTensor of size batch]
                if True == CUDA:
                    y = Variable(torch.sparse.torch.eye(NUM_CLASSES)).cuda().index_select(dim=0, index=max_length_indices.data)
                if False == CUDA:
                    y = Variable(torch.sparse.torch.eye(NUM_CLASSES)).index_select(dim=0, index=max_length_indices.data)
                # generate a new y: [batch, 10] with each column having 1 in batch 0

        if 2 == param['dim']:
            # print('>>dim: x*y', x.shape, y.shape)
            reconstructions = self.decoder((x * y[:, :, None]).view(x.size(0), -1))
           # x: [batch, 10, 16], y: [batch, 10] -> [batch, 10, 1]
            return classes, reconstructions
        if 1 == param['dim']:
            return classes, x # Note here

class CapsuleLoss(nn.Module):
    def __init__(self):
        super(CapsuleLoss, self).__init__()
        self.reconstruction_loss = nn.MSELoss(size_average=False)

    def forward(self, images, labels, classes, reconstructions):
        if 2 == param['dim']:
            # print('>>dim: labels', labels) # [batch, 10]
            # print('>>dim: classes', classes) # [batch, 10]
            left = F.relu(0.9 - classes, inplace=True) ** 2
            right = F.relu(classes - 0.1, inplace=True) ** 2

            margin_loss = labels * left + 0.5 * (1. - labels) * right
            margin_loss = margin_loss.sum()

            reconstruction_loss = self.reconstruction_loss(reconstructions, images)

            loss = (margin_loss + 0.0005 * reconstruction_loss) / images.size(0)
            # print('>>dim: loss', loss) # it's a single value
            return loss
        if 1 == param['dim']:
            # print('>>dim: labels', labels) # torch.cuda.FloatTensor of size batch x 1
            # print('>>dim: classes', classes) # [batch, 1]

            '''
            square = (labels - classes) ** 2
            square = square.sort(dim=0,descending=False)[0]
            cut = int(labels.shape[0]-1)
            loss = (square[:cut]).sum()/cut
            loss = loss ** 0.5 + square[cut] ** 0.5
            '''
            loss = ((labels - classes) ** 2).sum()/labels.shape[0] # MSE # Note: here it must be sum()
            loss = loss ** 0.5 # RMSE

            # print('>>dim: loss', loss)
            return loss

def desparse(RTtt):
    if False == DATA_AUGMENTATION:
        X = np.zeros((RTtt.number_seq, RTtt.N_aa, RTtt.N_time_step)) # DATA_AUGMENTATION->*2
        for i in range(RTtt.number_seq): # DATA_AUGMENTATION->*2
            # sparse to dense
            X[i,::] = RTtt.X[i].todense()
        RTtt.X = X
    else:
        print('>> note: usnig data_augmentation')
        X = np.zeros((RTtt.number_seq*2, RTtt.N_aa, RTtt.N_time_step)) # DATA_AUGMENTATION->*2
        for i in range(RTtt.number_seq*2): # DATA_AUGMENTATION->*2
            # sparse to dense
            X[i,::] = RTtt.X[i].todense()
        RTtt.X = X

if __name__ == "__main__":
    # from torch.autograd import Variable
    from torch.optim import Adam # Adam
    from torchnet.engine import Engine
    # from torchnet.logger import VisdomPlotLogger, VisdomLogger
    # from torchvision.utils import make_grid
    # from torchvision.datasets.mnist import MNIST
    from tqdm import tqdm
    import torchnet as tnt
    import gc
    from time import sleep, time
    import timeit
    T1 = timeit.default_timer()

    # read data ========== ========== ========== ========== ========== ==========
    # CNN_EMB = True
    if False == CNN_EMB:
        print('>> note: using one-hot encoding.')
        if True == LOAD_DATA:
            # dictionary = Dictionary(dict_path)
            RTtrain = RTdata(dictionary, max_length, train_path)
            RTtest = RTdata(dictionary, max_length, test_path)
            with open(RTdata_path, 'wb') as output:
                # pickle.dump(dictionary, output)
                pickle.dump(RTtrain, output)
                pickle.dump(RTtest, output)
        if False == LOAD_DATA:
            with open(RTdata_path, 'rb') as input:
                # dictionary = pickle.load(input)
                RTtrain = pickle.load(input)
                RTtest = pickle.load(input)
                print('>> note: load pre-read RTdata from:', RTdata_path)

        # DATA_AUGMENTATION = True
        SPARSE = True
        # def desparse(RTtt):
        #     X = np.zeros((RTtt.number_seq, RTtt.N_aa, RTtt.N_time_step)) # DATA_AUGMENTATION->*2
        #     for i in range(RTtt.number_seq): # DATA_AUGMENTATION->*2
        #         # sparse to dense
        #         X[i,::] = RTtt.X[i].todense()
        #     RTtt.X = X
        if True == SPARSE:
            print('>> note: de-sparse for both train & test data.')
            desparse(RTtrain)
            desparse(RTtest)
    if True == CNN_EMB:
        print('>> note: using >>>embedding<<< method.')
        corpus = Corpus(dictionary, # format: Corpus(dictionary, train_path, val_path='', test_path='', pad_length=0)
                        train_path,
                        test_path=test_path,
                        pad_length=max_length)

    # read data ========== ========== ========== ========== ========== ==========

    LOG = False
    flog = open(log_path, 'w')

    model = CapsuleNet(conv1_kernel,conv2_kernel)
    if '' == pretrain_path:
        pass
    else:
        model.load_state_dict(torch.load(pretrain_path)) # epoch.pt
        print('>> note: load pre-trained model from:',pretrain_path)

    if True == CUDA:
        model.cuda()

    print("# parameters:", sum(param.numel() for param in model.parameters()))
    flog.write("# parameters:"+str(sum(param.numel() for param in model.parameters()))+'\n')

    optimizer = Adam(model.parameters(), lr = LR)
    # optimizer = SGD(model.parameters(), lr = LR/10., momentum = 0.5)

    engine = Engine()
    meter_loss = tnt.meter.AverageValueMeter()
    if 2 == param['dim']:
        meter_accuracy = tnt.meter.ClassErrorMeter(accuracy=True)
        confusion_meter = tnt.meter.ConfusionMeter(NUM_CLASSES, normalized=True)
    if 1 == param['dim']:
        pass
        # meter_mse = tnt.meter.MSEMeter()

    if True == LOG:
        train_loss_logger = VisdomPlotLogger('line', opts={'title': 'Train Loss'})
        train_error_logger = VisdomPlotLogger('line', opts={'title': 'Train Accuracy'})
        test_loss_logger = VisdomPlotLogger('line', opts={'title': 'Test Loss'})
        test_accuracy_logger = VisdomPlotLogger('line', opts={'title': 'Test Accuracy'})
        confusion_logger = VisdomLogger('heatmap', opts={'title': 'Confusion matrix',
                                                     'columnnames': list(range(NUM_CLASSES)),
                                                     'rownames': list(range(NUM_CLASSES))})
        if 2 == param['dim']:
            ground_truth_logger = VisdomLogger('image', opts={'title': 'Ground Truth'})
            reconstruction_logger = VisdomLogger('image', opts={'title': 'Reconstruction'})

    capsule_loss = CapsuleLoss()

    def get_iterator(mode):
        dataset = MNIST(root='./data', download=True, train=mode)
        data = getattr(dataset, 'train_data' if mode else 'test_data')[:47]
        # [torch.ByteTensor of size number x 28 x 28]
        labels = getattr(dataset, 'train_labels' if mode else 'test_labels')[:47]
        # [torch.LongTensor of size number]
        tensor_dataset = tnt.dataset.TensorDataset([data, labels])

        return tensor_dataset.parallel(batch_size=BATCH_SIZE, num_workers=4, shuffle=mode)

    if False == CNN_EMB:
        data_train = torch.FloatTensor(RTtrain.X)
        label_train = torch.FloatTensor(RTtrain.y)
        print('>> note: delete RTtrain.')
        del RTtrain
        gc.collect()
        print('>> sleeping...')
        for i in range(5):
            print('~.~')
        print('>> wake up!')
    if True == CNN_EMB:
        data_train = corpus.train
        label_train = corpus.train_label

    def get_rt_iterator(mode):
        if mode:
            data = data_train # Note: here must be FloatTensor not ByteTensor!
            labels = label_train
        else:
            if False == CNN_EMB:
                data = torch.FloatTensor(RTtest.X)
                labels = torch.FloatTensor(RTtest.y)
            if True == CNN_EMB:
                data = corpus.test
                labels = corpus.test_label
            # print('>>dim: test data:', data.shape, labels.shape)
        tensor_dataset = tnt.dataset.TensorDataset([data, labels])
        return tensor_dataset.parallel(batch_size=BATCH_SIZE, num_workers=1, shuffle=mode) # 1 for heatmap

    def processor(sample):
        data, labels, training = sample
        # print('>>dim: data, labels, training', data.shape, labels.shape, training)
        # torch.Size([batch, 28, 28]) torch.Size([batch]) True

        if 'mnist' == param['data']:
            data = augmentation(data.unsqueeze(1).float() / 255.0)
        # print('>>dim: data augmentation', data.shape) # torch.Size([batch, 1, 28, 28])
        # print('>>dim: labels', labels) # Note: labels is already LongTensor?
        if 'rt' == param['data']:
            if False == CNN_EMB:
                data = data[:, None, :, :] # Note: add dimension
            if True == CNN_EMB:
                pass

        if 2 == param['dim']:
            # for classification, we use LongTensor
            labels = torch.LongTensor(labels)
            labels = torch.sparse.torch.eye(NUM_CLASSES).index_select(dim=0, index=labels)
        if 1 == param['dim']:
            # for regression, we use FloatTensor
            labels = torch.FloatTensor(labels.numpy())
            labels = labels.view(-1, 1) # from [batch] to [batch, 1]

        if True == CUDA:
            data = Variable(data).cuda()
            labels = Variable(labels).cuda()
        if False == CUDA:
            data = Variable(data)
            labels = Variable(labels)

        if training:
            classes, reconstructions = model(data, labels)
        else:
            classes, reconstructions = model(data)

        loss = capsule_loss(data, labels, classes, reconstructions)

        return loss, classes

    def reset_meters():
        meter_loss.reset()
        if 2 == param['dim']:
            meter_accuracy.reset()
            confusion_meter.reset()
        if 1 == param['dim']:
            pass
            # meter_mse.reset()

    def on_sample(state):
        state['sample'].append(state['train'])

    def on_forward(state):
        '''
        So it is just used for recording?
        '''
        if 1 == param['dim']:
            # print('>>dim: state output', state['output'].data.view(-1))
            # torch.FloatTensor of size [batch x 10]
            # print('>>dim: state sample', state['sample'][1])
            # torch.LongTensor of size [batch]
            # (1): [batch, 1] (2): [batch], so we view (1) as [batch], but no view is fine
            pass
            # meter_mse.add(state['output'].data, torch.FloatTensor(state['sample'][1].numpy()))
        if 2 == param['dim']:
            meter_accuracy.add(state['output'].data, torch.LongTensor(state['sample'][1]))
            confusion_meter.add(state['output'].data, torch.LongTensor(state['sample'][1]))
        meter_loss.add(state['loss'].data[0])

    def on_start_epoch(state):
        reset_meters()
        state['iterator'] = tqdm(state['iterator'])

    def on_end_epoch(state):
        if 2 == param['dim']:
            print('[Epoch %d] Training Loss: %.4f (Accuracy: %.2f%%)' % (
                state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0]))
            flog.write('[Epoch %d] Training Loss: %.4f (Accuracy: %.2f%%)\n' % (
                state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0]))
            if True == LOG:
                train_loss_logger.log(state['epoch'], meter_loss.value()[0])
                train_error_logger.log(state['epoch'], meter_accuracy.value()[0])
        if 1 == param['dim']:
            print('[Epoch %d] Training Loss: %.4f (MSE: %.4f)' % (
                state['epoch'], meter_loss.value()[0], 7)) # meter_mse.value()
            flog.write('[Epoch %d] Training Loss: %.4f (MSE: %.4f)\n' % (
                state['epoch'], meter_loss.value()[0], 7)) # meter_mse.value()

        reset_meters()

        # iterator
        if 'mnist' == param['data']:
            engine.test(processor, get_iterator(False))
        if 'rt' == param['data']:
            engine.test(processor, get_rt_iterator(False))

        if True == LOG:
            test_loss_logger.log(state['epoch'], meter_loss.value()[0])
            if 2 == param['dim']:
                test_accuracy_logger.log(state['epoch'], meter_accuracy.value()[0])
                confusion_logger.log(confusion_meter.value())
            if 1 == param['dim']:
                test_accuracy_logger.log(state['epoch'], 7) # meter_mse.value()

        if 2 == param['dim']:
            print('[Epoch %d] Testing Loss: %.4f (Accuracy: %.2f%%)' % (
                state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0]))
            flog.write('[Epoch %d] Testing Loss: %.4f (Accuracy: %.2f%%)\n' % (
                state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0]))
        if 1 == param['dim']:
            print('[Epoch %d] Testing Loss: %.4f (MSE: %.4f)' % (
                state['epoch'], meter_loss.value()[0], 7)) # meter_mse.value()
            flog.write('[Epoch %d] Testing Loss: %.4f (MSE: %.4f)\n' % (
                state['epoch'], meter_loss.value()[0], 7)) # meter_mse.value()

        if 10 <= state['epoch']: # for heatmap
            torch.save(model.state_dict(), save_prefix+'/epoch_%d.pt' % state['epoch'])
            print('>> model: saved.')

        # prediction:
        # model.load_state_dict(torch.load(PATH))
        # pred_data = Variable(torch.FloatTensor(RTtest.X)[:,None,:,:])
        PRED_BATCH = 16 # 1000 # 16 for heatmap

        if PRED_BATCH > 0:
            '''
            solve memory problem using batch
            '''
            if False == CNN_EMB:
                pred = np.array([])
                # TODO: handle int
                pred_batch_number = int(RTtest.X.shape[0] / PRED_BATCH)+1
                for bi in range(pred_batch_number):
                    test_batch = Variable(torch.FloatTensor(RTtest.X[bi*PRED_BATCH:(bi+1)*PRED_BATCH,:,:])[:,None,:,:])
                    test_batch = test_batch.cuda() # Note: we don't use this block anymore
                    pred_batch = model(test_batch)
                    pred = np.append(pred, pred_batch[0].data.cpu().numpy().flatten())
                # print('>>dim: pred', pred.shape)

                if True == DATA_AUGMENTATION:
                    ''' data augmentation:'''
                    pep_num = int(len(pred) / 2)
                    pred = pred[:pep_num]*0.5 + pred[pep_num:]*0.5
                    obse = RTtest.y[:pep_num]
                    pearson = Pearson(pred,obse)
                    spearman = Spearman(pred,obse)
                else:
                    pearson = Pearson(pred,RTtest.y)
                    spearman = Spearman(pred,RTtest.y)
            if True == CNN_EMB:
                pred = np.array([])
                # TODO: handle int
                pred_batch_number = int(corpus.test.shape[0] / PRED_BATCH)+1
                for bi in range(pred_batch_number):
                    test_batch = Variable(corpus.test[bi*PRED_BATCH:(bi+1)*PRED_BATCH,:])
                    if True == CUDA:
                        test_batch = test_batch.cuda()
                        pred_batch = model(test_batch)
                        pred = np.append(pred, pred_batch[0].data.cpu().numpy().flatten())
                    if False == CUDA:
                        # test_batch = test_batch.cuda()
                        pred_batch = model(test_batch)
                        pred = np.append(pred, pred_batch[0].data.numpy().flatten())
                # print('>>dim: pred', pred.shape)
                obse = corpus.test_label.numpy().flatten()
                pearson = Pearson(pred,obse)
                spearman = Spearman(pred,obse)

        else:
            pred_data = Variable(torch.FloatTensor(RTtest.X)[:,None,:,:])
            if True == CUDA:
                pred_data = pred_data.cuda()
            pred = model(pred_data)
            if True == CUDA:
                # print('>>dim: pred', pred[0].data.cpu().numpy().flatten().shape)
                pearson = Pearson(pred[0].data.cpu().numpy().flatten(),RTtest.y)
                spearman = Spearman(pred[0].data.cpu().numpy().flatten(),RTtest.y)
            if False == CUDA:
                pearson = Pearson(pred[0].data.numpy().flatten(),RTtest.y)
                spearman = Spearman(pred[0].data.numpy().flatten(),RTtest.y)
        ''''''
        print('>> Corr on %d testing samples: %.5f | %.5f' % (len(pred), pearson, spearman))
        flog.write('>> Corr on %d testing samples: %.5f | %.5f\n' % (len(pred), pearson, spearman))
        # writing:
        if True == CNN_EMB:
            obse = corpus.test_label.numpy().flatten()
        if False == CNN_EMB:
            obse = RTtest.y
        with open(result_path, 'w') as fo:
                fo.write('observed\tpredicted\n')
                for i in range(len(pred)):
                    fo.write('%.5f\t%.5f\n' % (obse[i],pred[i]))
        # writing done

        # Reconstruction visualization.
        if 2 == param['dim']:

            # iterator
            if 'mnist' == param['data']:
                test_sample = next(iter(get_iterator(False)))
            if 'rt' == param['data']:
                test_sample = next(iter(get_rt_iterator(False)))
            # print('>>dim: test_sample', test_sample) # [batch, 28, 28]

            ground_truth = (test_sample[0].unsqueeze(1).float() / 255.0)
            # print('>>dim: ground_truth', ground_truth.shape) # torch.FloatTensor of size batch x 1 x 28 x 28

            if True == CUDA:
                pred, reconstructions = model(Variable(ground_truth).cuda())
            if False == CUDA:
                pred, reconstructions = model(Variable(ground_truth))
            # print('>>dim: pred', pred)
            reconstruction = reconstructions.cpu().view_as(ground_truth).data

            if True == LOG:
                ground_truth_logger.log(
                    make_grid(ground_truth, nrow=int(BATCH_SIZE ** 0.5), normalize=True, range=(0, 1)).numpy())
                reconstruction_logger.log(
                    make_grid(reconstruction, nrow=int(BATCH_SIZE ** 0.5), normalize=True, range=(0, 1)).numpy())

    # def on_start(state):
    #     state['epoch'] = 327
    #
    # engine.hooks['on_start'] = on_start

    engine.hooks['on_sample'] = on_sample
    engine.hooks['on_forward'] = on_forward
    engine.hooks['on_start_epoch'] = on_start_epoch
    engine.hooks['on_end_epoch'] = on_end_epoch

    if 'mnist' == param['data']:
        engine.train(processor, get_iterator(True), maxepoch=NUM_EPOCHS, optimizer=optimizer)
    if 'rt' == param['data']:
        engine.train(processor, get_rt_iterator(True), maxepoch=NUM_EPOCHS, optimizer=optimizer)

    T2 = timeit.default_timer()
    print('>> time: %.5f min\n' %((T2-T1)/60.))
    flog.write('>> time: %.5f min\n' %((T2-T1)/60.))
    flog.close()


# -*- coding: utf-8 -*-
"""
.. module: pyAPES
    :synopsis: APES-model component
.. moduleauthor:: Kersti Haahti

Model framework for Atmosphere-Plant Exchange Simulations

Created on Tue Oct 02 09:04:05 2018

Note:
    migrated to python3
    - print on same line
    - dict.keys(), but these are iterated after in for-each-loop

References:
Launiainen, S., Katul, G.G., Lauren, A. and Kolari, P., 2015. Coupling boreal
forest CO2, H2O and energy flows by a vertically structured forest canopy –
Soil model with separate bryophyte layer. Ecological modelling, 312, pp.385-405.

To call model and run single simulation and read results: see example in sandbox.py
    from tools.iotools import read_results
    from pyAPES import driver
    # for NetCDF-outputs
    outputfile = driver(create_ncf=True, result_file='test.nc')
    results = read_results(outputfile) # opens NetCDF-file using xarray

    # for returning results directly
    results = driver(create_ncf=False) # returns dict with integer keys
    results = results[0] # first simulation

LAST EDIT: 15.1.2020 Samuli Launiainen
    * new forestfloor and altered outputs
Todo:

    - make minimal example of handling and plotting outputs using xarray -tools;
      now see tools.iotools.read_forcing for documentation!

"""
import time
import logging

import numpy as np
from pandas import date_range

from tools.iotools import initialize_netcdf,  write_ncf
from canopy.canopy import CanopyModel
from soil.soil import Soil
from canopy.constants import WATER_DENSITY

def driver(parameters,
           create_ncf=False,
           result_file=None):
    """
    Reads parameters as argument, prepares output files, runs model.
    Args:
        parameters (dict/list): either single parameter dictionary or list of parameters
        create_ncf (bool): results saved to netCDF4 file
        result_file (str): name of result file
    """

    # --- CONFIGURATION PARAMETERS of LOGGING and NetCDF -outputs read
    from parameters.outputs import output_variables, logging_configuration
    from logging.config import dictConfig

    # --- LOGGING ---
    dictConfig(logging_configuration)
    logger = logging.getLogger(__name__)

    # --- CHECK PARAMETERS ---

    if isinstance(parameters, dict):
        Nsim = 1
        parameters = [parameters]
    elif isinstance(parameters, list):
        Nsim = len(parameters)
    else:
        raise TypeError('Parameters should be either dict or list.')

    logger.info('Simulation started. Number of simulations: {}'.format(Nsim))

    # --- SIMULATIOS AND OUTPUTS ---

    tasks = []

    for k in range(Nsim):
        tasks.append(
            Model(
                parameters[k]['general']['dt'],
                parameters[k]['canopy'],
                parameters[k]['soil'],
                parameters[k]['forcing'],
                output_variables['variables'],
                nsim=k
            )
        )

    if create_ncf: # outputs to NetCDF-file, returns filename
        gpara = parameters[0]['general'] # same for all tasks
        timestr = time.strftime('%Y%m%d%H%M')
        if result_file:
            filename = result_file
        else:
            filename = timestr + '_pyAPES_results.nc'

        #freq = '{}S'.format(gpara['dt'])
        #time_index = date_range(gpara['start_time'], gpara['end_time'], freq=freq, closed='left')

        time_index = parameters[0]['forcing'].index

        ncf, _ = initialize_netcdf(
                output_variables['variables'],
                Nsim,
                tasks[k].Nsoil_nodes,
                tasks[k].Ncanopy_nodes,
                tasks[k].Nplant_types,
                tasks[k].Nground_types,
                time_index=time_index,
                filepath=gpara['results_directory'],
                filename=filename)

        for task in tasks:
            logger.info('Running simulation number (start time %s): %s' % (
                        time.strftime('%Y-%m-%d %H:%M'), task.Nsim))
            running_time = time.time()
            results = task.run()
            logger.info('Running time %.2f seconds' % (time.time() - running_time))
            write_ncf(nsim=task.Nsim, results=results, ncf=ncf)

            del results

        output_file = gpara['results_directory'] + filename
        logger.info('Ready! Results are in: ' + output_file)

        ncf.close()

        return output_file, tasks[0]

    else: # returns dictionary of outputs
        running_time = time.time()
        results = {task.Nsim: task.run() for task in tasks}

        logger.info('Running time %.2f seconds' % (time.time() - running_time))

        return results, tasks[0] # this would return also 1st Model instance

class Model(object):
    """
    pyAPES - main model class.
    Combines submodels 'CanopyModel' and 'Soil' and handles data-transfer
    between these model components and writing results.

    Last edit: SL 13.01.2020
    """
    def __init__(self,
                 dt,
                 canopy_para,
                 soil_para,
                 forcing,
                 outputs,
                 nsim=0):

        logger = logging.getLogger(__name__)

        self.dt = dt

        self.Nsteps = len(forcing)
        self.forcing = forcing
        self.Nsim = nsim

        self.Nsoil_nodes = len(soil_para['grid']['dz'])
        self.Ncanopy_nodes = canopy_para['grid']['Nlayers']

        # create soil model instance
        self.soil = Soil(soil_para)

        if 'Wa' in forcing and soil_para['water_model']['solve'] is False:
            logger.info("Soil moisture from forcing file")
            soil_para['water_model']['initial_condition']['volumetric_water_content'] = (
                forcing['Wa'].iloc[0])
        if 'Tsa' in forcing and soil_para['heat_model']['solve'] is False:
            logger.info("Soil temperature from forcing file")
            soil_para['heat_model']['initial_condition']['temperature'] = (
                forcing['Tsa'].iloc[0])

        # create canopy model instance
        # initial delayed temperature and degreedaysum for pheno & LAI-models
        if canopy_para['ctr']['pheno_cycle'] and 'X' in forcing:
            for pt in list(canopy_para['planttypes'].keys()):
                canopy_para['planttypes'][pt]['phenop'].update({'Xo': forcing['X'].iloc[0]})
        if canopy_para['ctr']['seasonal_LAI'] and 'DDsum' in forcing:
            for pt in list(canopy_para['planttypes'].keys()):
                canopy_para['planttypes'][pt]['laip'].update({'DDsum0': forcing['DDsum'].iloc[0]})

        self.canopy_model = CanopyModel(canopy_para, self.soil.grid['dz'])

        self.Nplant_types = len(self.canopy_model.planttypes)
        self.Nground_types = len(self.canopy_model.forestfloor.bottomlayer_types)

        # initialize structure to save results
        self.results = _initialize_results(outputs,
                                       self.Nsteps,
                                       self.Nsoil_nodes,
                                       self.Ncanopy_nodes,
                                       self.Nplant_types,
                                       self.Nground_types)

    def run(self):
        """
        Loops through self.forcing and appends to self.results.

        self.forcing variables and units; correspond to uppermost gridpoint:
            precipitation [kg m-2 s-1]
            air_pressure [Pa]
            air_temperature [degC]
            wind_speed [m/s]
            friction_velocity [m/s]
            h2o[mol/mol]
            co2 [ppm]
            zenith_angle [rad]
            lw_in: Downwelling long wave radiation [W/m2]
            diffPar: Diffuse PAR [W/m2]
            dirPar: Direct PAR [W/m2]
            diffNir: Diffuse NIR [W/m2]
            dirNir: Direct NIR [W/m2]
        """

        logger = logging.getLogger(__name__)
        logger.info('Running simulation {}'.format(self.Nsim))
        time0 = time.time()

        #print('RUNNING')
        k_steps=np.arange(0, self.Nsteps, int(self.Nsteps/10))

        for k in range(0, self.Nsteps):
            # --- print progress on screen
            if k in k_steps[:-1]:
                s = str(np.where(k_steps==k)[0][0]*10) + '%'
                print('{0}..'.format(s), end=' ')

            # --- CanopyModel ---
            # run daily loop: updates LAI, phenology and moisture stress ---
            if self.forcing['doy'].iloc[k] != self.forcing['doy'].iloc[k-1] or k == 0:
                self.canopy_model.run_daily(
                        self.forcing['doy'].iloc[k],
                        self.forcing['Tdaily'].iloc[k])

            # compile forcing dict for canopy model: soil_ refers to state of soil model
            canopy_forcing = {
                'wind_speed': self.forcing['U'].iloc[k],            # [m s-1]
                'friction_velocity': self.forcing['Ustar'].iloc[k], # [m s-1]
                'air_temperature': self.forcing['Tair'].iloc[k],    # [deg C]
                'precipitation': self.forcing['Prec'].iloc[k],      # [kg m-2 s-1]
                'h2o': self.forcing['H2O'].iloc[k],                 # [mol mol-1]
                'co2': self.forcing['CO2'].iloc[k],                 # [ppm]
                'PAR': {'direct': self.forcing['dirPar'].iloc[k],   # [W m-2]
                        'diffuse': self.forcing['diffPar'].iloc[k]},
                'NIR': {'direct': self.forcing['dirNir'].iloc[k],   # [W m-2]
                        'diffuse': self.forcing['diffNir'].iloc[k]},
                'lw_in': self.forcing['LWin'].iloc[k],              # [W m-2]
                'air_pressure': self.forcing['P'].iloc[k],          # [Pa]
                'zenith_angle': self.forcing['Zen'].iloc[k],        # [rad]

                # from soil model
                'soil_temperature': self.soil.heat.T[self.canopy_model.ix_roots],       # [deg C]
                'soil_water_potential': self.soil.water.h[self.canopy_model.ix_roots],  # [m] ?
                'soil_volumetric_water': self.soil.heat.Wliq[self.canopy_model.ix_roots], # [m3 m-3]
                'soil_volumetric_air': self.soil.heat.Wair[self.canopy_model.ix_roots],   # [m3 m-3]
                'soil_pond_storage': self.soil.water.h_pond * WATER_DENSITY,    # [kg m-2]
            }

            canopy_parameters = {
                'soil_depth': self.soil.grid['z'][0],   # [m]
                'soil_hydraulic_conductivity': self.soil.water.Kv[self.canopy_model.ix_roots], # [m s-1]
                'soil_thermal_conductivity': self.soil.heat.thermal_conductivity[0],        # [W m-1 K-1]?
                'date': self.forcing.index[k]   # pd.datetime
            }

            # call self.canopy_model.run to solve above-ground part
            out_canopy, out_planttype, out_ffloor, out_groundtype = self.canopy_model.run(
                dt=self.dt,
                forcing=canopy_forcing,
                parameters=canopy_parameters
            )

            # --- Soil model  ---
            # compile forcing for Soil: potential infiltration and evaporation are at from ground surface
            # water fluxes must be in [m s-1]
            soil_forcing = {
                'potential_infiltration': out_ffloor['throughfall'] / WATER_DENSITY,
                'potential_evaporation': ((out_ffloor['soil_evaporation'] +
                                          out_ffloor['capillary_rise']) / WATER_DENSITY),
                'pond_recharge': out_ffloor['pond_recharge'] / WATER_DENSITY,
                'atmospheric_pressure_head': -1.0E6,  # set to large value, because potential_evaporation already account for h_soil
                'ground_heat_flux': -out_ffloor['ground_heat'],
                'date': self.forcing.index[k]}

            if 'Ws' in self.forcing and self.soil.solve_water is False:
                soil_forcing.update({
                    'state_water':{'volumetric_water_content': self.forcing['Ws'].iloc[k]}})
            if 'Tsa' in self.forcing and self.soil.solve_heat is False:
                soil_forcing.update({
                    'state_heat':{'temperature': self.forcing['Tsa'].iloc[k]}})

            # call self.soil to solve below-ground water and heat flow
            soil_flux, soil_state = self.soil.run(
                    dt=self.dt,
                    forcing=soil_forcing,
                    water_sink=out_canopy['root_sink'])

            # --- append results and copy of forcing to self.results
            forcing_output = {
                    'wind_speed': self.forcing['U'].iloc[k],
                    'friction_velocity': self.forcing['Ustar'].iloc[k],
                    'air_temperature': self.forcing['Tair'].iloc[k],
                    'precipitation': self.forcing['Prec'].iloc[k],
                    'h2o': self.forcing['H2O'].iloc[k],
                    'co2': self.forcing['CO2'].iloc[k],
                    'pressure': self.forcing['P'].iloc[k],
                    'par':  self.forcing['dirPar'].iloc[k] + self.forcing['diffPar'].iloc[k],
                    'nir':  self.forcing['dirNir'].iloc[k] + self.forcing['diffNir'].iloc[k],
                    'lw_in': self.forcing['LWin'].iloc[k]
                    }

            soil_state.update(soil_flux)

            self.results = _append_results('forcing', k, forcing_output, self.results)
            self.results = _append_results('canopy', k, out_canopy, self.results)
            self.results = _append_results('ffloor', k, out_ffloor, self.results)
            self.results = _append_results('soil', k, soil_state, self.results)
            self.results = _append_results('pt', k, out_planttype, self.results)
            self.results = _append_results('gt', k, out_groundtype, self.results)

        print('100%')

        ptnames = [pt.name for pt in self.canopy_model.planttypes]

        self.results = _append_results('canopy', None, {'z': self.canopy_model.z,
                                                        'planttypes': np.array(ptnames)}, self.results)

        gtnames = [gt.name for gt in self.canopy_model.forestfloor.bottomlayer_types]

        self.results = _append_results('ffloor', None, {'groundtypes': np.array(gtnames)}, self.results)

        self.results = _append_results('soil', None, {'z': self.soil.grid['z']}, self.results)

        logger.info('Finished simulation %.0f, running time %.2f seconds' % (self.Nsim, time.time() - time0))

        return self.results

def _initialize_results(variables, Nstep, Nsoil_nodes, Ncanopy_nodes, Nplant_types, Nground_types):
    """
    Creates temporary results dictionary to accumulate simulation results
    SL 12.11.2019: removed if 'date' in dimensions and added option to save planttype profiles
    """

    results = {}

    for var in variables:

        var_name = var[0]
        dimensions = var[2]

        if 'canopy' in dimensions:
            if 'planttype' in dimensions:
                var_shape = [Nstep, Nplant_types, Ncanopy_nodes]
            else:
                var_shape = [Nstep, Ncanopy_nodes]

        elif 'soil' in dimensions:
            var_shape = [Nstep, Nsoil_nodes]

        elif 'planttype' in dimensions and 'canopy' not in dimensions:
            var_shape = [Nstep, Nplant_types]

        elif 'groundtype' in dimensions:
            if 'date' not in dimensions:
                var_shape = [Nground_types]
            else:
                var_shape = [Nstep, Nground_types]

        else:
            var_shape = [Nstep]

        results[var_name] = np.full(var_shape, np.NAN)
        # print(var_name, var_shape, dimensions)

    return results

def _append_results(group, step, step_results, results):
    """
    Adds results from each simulation steps to temporary results dictionary
    """

    results_keys = results.keys()
    step_results_keys = step_results.keys()

    for key in step_results_keys:
        variable = group + '_' + key
        if variable in results_keys:
            if key == 'z' or key == 'planttypes' or key == 'groundtypes':
                results[variable] = step_results[key]
            else:
                #print(variable, key, np.shape(results[variable][step]), np.shape(step_results[key]))
                results[variable][step] = step_results[key]

    return results

#if __name__ == '__main__':
#
#    from parameters.parametersets import lettosuo_parameters
#    outputfile=driver(create_ncf=True, parametersets=lettosuo_parameters)
#
#    print(outputfile)


#!/usr/bin/python
# -*- coding: utf-8 -*-
#The script builds an evolutionary history of clones by sharing mutations between samples.

import os
import re
import numpy
import random
import sys, getopt
from Bio import Phylo
from io import StringIO
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

#24 chromosome
all_chr=[]
for i in range(1,24):
	all_chr.append(str(i))
all_chr.append('X')
all_chr.append('Y')

#parameter

opts,args = getopt.getopt(sys.argv[1:], "i:o:",["AF=","input=", "output="])
def get_driver_gene_path():
    '''
    return the driver gene path
    '''
    current_path = os.path.dirname(os.path.realpath(__file__))
    return current_path

AF=''
file_path=''
file_path_out=''
for para,value in opts:
	if para in ('--AF'):
		AF=value
	if para in ('-i','--input'):
		file_path=value
	if para in ('-o','--output'):
		file_path_out=value
print('AF: '+str(AF))
if file_path[-len(os.path.sep):]==os.path.sep:
	file_path=file_path[:-len(os.path.sep)]
print('file_path: '+str(file_path))
if file_path_out[-len(os.path.sep):]==os.path.sep:
	file_path_out=file_path_out[:-len(os.path.sep)]
print('file_path_out: '+str(file_path_out))

def fact(n):
    if n == 0:
        return 1
    else:
        return n*fact(n-1)
def Cmn(n,m):
    return fact(n)/(fact(n-m)*fact(m))

#map driver gene:
# @ a is a list including mutation which is used to map driver gene
# @ b =='FALSE', then output driver gene without mutation
# @ b =='TRUE', then output driver gene with mutation
def search_gene(a,b):
	global count_gene
	count_gene=0
	global gene_list
	gene_list=[]
	for item in a:#read a mutation
		if item in all_mutation.keys():
			for gene in all_mutation[item]:
				if gene in driver_gene.keys():
					if b=='FALSE':#do not output driver gene with mutation
						gene=gene
					if b=='TRUE':
						gene=gene+':'+item
					if gene not in gene_list:
						gene_list.append(gene)
	count_gene=len(gene_list)
	gene_list.append(count_gene)
	return gene_list

#.................................................find best branch..........................
def xunhuan(a):
	num=1
	global first_1
	first_1={}
	global first_3
	first_3={}
	global first_4
	first_4={}
	first_6={}
	while num<=len(a)/2:
		struct_pattern=str(num)+'-'+str(len(a)-num)
		DIC={}#including structure info: Cmn information for a specific branch split
		cishu=int(Cmn(len(a),num))
		while len(DIC.keys())<cishu:
			sel=random.sample(a,num)
			sel.sort()
			if num>1:
				da=','
				for Item in sel:
					da=da+','+Item
				if da[2:] not in DIC.keys():
					DIC.update({da[2:]:[]})
			if num==1:
				DIC.update({sel[0]:[]})
		first_2={}#struct info
		for item in DIC.keys():
			a_1=[]#first =a, then remove item from [a].
			jiaoji_2=[]
			jiaoji_1=[]
			for sample in a:
				a_1.append(sample)
			split=item.split(',')
			split.sort()
			for item5 in split:
				a_1.remove(item5)
			if len(split)==1:
				jiaoji_1=dic[item]#sample mutation which is removed
			else:
				NUM=0
				for item6 in split:
					NUM=NUM+1
					if NUM==1:
						jiaoji_1=dic[item6]
					jiaoji_1=[val for val in jiaoji_1 if val in dic[item6]]#sample mutation which are removed

			i=0
			no_item=','
			a_1.sort()
			for item1 in a_1:
				no_item=no_item+','+item1
				i=i+1
				if i==1:
					jiaoji_2=dic[item1]
				jiaoji_2 = [val for val in jiaoji_2 if val in dic[item1]]#sample mutation which are retained; also is big group

			struct_name=no_item[2:]
			struct_share_mut=len(jiaoji_2)
			num_1=str(len(jiaoji_1))+str(search_gene(jiaoji_1,'FALSE'))
			first_3.update({item:num_1})
			first_4.update({item:jiaoji_1})
			num_2=str(len(jiaoji_2))+str(search_gene(jiaoji_2,'FALSE'))
			first_3.update({struct_name:num_2})
			first_4.update({struct_name:jiaoji_2})
			#if struct: n/2,n/2: we should retained group info which mutation is small
			if num==len(a)/2:
				if len(jiaoji_1)>len(jiaoji_2):
					struct_name=no_item[2:]
					struct_share_mut=len(jiaoji_2)
				if len(jiaoji_1)<len(jiaoji_2):
					struct_name=item
					struct_share_mut=len(jiaoji_1)
				if len(jiaoji_1)==len(jiaoji_2):
					struct_name=item
					struct_share_mut=len(jiaoji_1)
			first_2.update({struct_name:struct_share_mut})
		sort_d=sorted(first_2.items(),key = lambda d:d[1],reverse=True)
		count=0
		#print('max value and second value:')
		max_data=[]
		if num!=len(a)/2:
			for key,value in sort_d:
				count+=1
				if count==1:
					max_data_name=key
					max_data_value=value
					struct_pattern1=struct_pattern+':'+key
					first_6.update({struct_pattern1:value})
				if count<=2:
					max_data.append(value)
					#print(key,value)
				if count>1 and key!=max_data_name and value==max_data_value:
					struct_pattern1=struct_pattern+':'+key
					if struct_pattern1 not in first_6.keys():
						first_6.update({struct_pattern1:value})
		else:
			for key,value in sort_d:
				count+=1
				if count==1:
					max_data_name=key
					max_data_value=value
					max_data.append(value)
					struct_pattern1=struct_pattern+':'+max_data_name
					if struct_pattern1 not in first_6.keys():
						first_6.update({struct_pattern1:max_data_value})
					#print(key,value)
					chongfu={}#judge random chongfu
					while len(chongfu.keys())<(fact(len(key.split(',')))-1):
						list1=key.split(',')
						random.shuffle(list1)
						#print(list1)
						fan_da=','
						for Item1 in list1:
							fan_da=fan_da+','+Item1
						if fan_da[2:]!=key:
							chongfu.update({fan_da[2:]:[]})
				if key != max_data_name and key not in chongfu.keys() and len(max_data)==1:
					max_data.append(value)
					#print(key,value)
				if count>1 and key not in chongfu.keys() and key!=max_data_name and value==max_data_value:
					chongfu2={}
					while len(chongfu2.keys())<(fact(len(key.split(',')))-1):
						list1=key.split(',')
						random.shuffle(list1)
						fan_da=','
						for Item1 in list1:
							fan_da=fan_da+','+Item1
						if fan_da[2:]!=key:
							chongfu2.update({fan_da[2:]:[]})
					chongfu_count=0
					for item_name in chongfu2.keys():
						struct_pattern1=struct_pattern+':'+item_name
						if struct_pattern not in first_6.keys():
							chongfu_count=chongfu_count+1
					if chongfu_count==len(chongfu.keys()):
						struct_pattern1=struct_pattern+':'+key
						if struct_pattern1 not in first_6.keys():
							first_6.update({struct_pattern1:value})
		if int(max_data[1])!=0:
			ratio=int(max_data[0])/int(max_data[1])
		else:
			ratio=int(max_data[0])
		struct_pattern2=struct_pattern+':'+max_data_name
		first_1.update({struct_pattern2:ratio})
		num=num+1
	sort_data=sorted(first_1.items(),key = lambda d:d[1],reverse=True)
	#print(sort_data)
	global return_data
	return_data=[]
	count_data=0
	for key,value in sort_data:
		count_data=count_data+1
		if count_data==1:
			max_data_name=key
			max_data_value=value
		if value==max_data_value:
			return_data.append(key)
	for item in return_data:
		identify=item.split(':')[0]
		for item1 in first_6.keys():
			if identify in item1 :
				if item1 not in return_data:
					return_data.append(item1)
	if max_data_value==0:#ratio==0
		value1=[]
		return_data=[]
		num=divmod(len(a),2)[0]
		while num>0:
			for key in first_4.keys():
				if len(key.split(','))==num and len(first_4[key])!=0:
					value1.append(len(first_4[key]))
			if len(value1)!=0:
				break
			else:
				num=num-1
		max_data=value1[0]
		struct_pattern=str(num)+'-'+str(len(a)-num)
		for data in value1[1:]:
			if int(data)>max_data:
				max_data=int(data)
		for key in first_4.keys():
			if len(key.split(','))==num and len(first_4[key])==max_data:
				struct_pattern1=struct_pattern+':'+key
				return_data.append(struct_pattern1)
	return_data.sort()
	return(return_data)
#split tree
def split_tree(tree_label_split):
	new_tree_label=[]
	num=0
	count1=0
	count2=0
	index1=[]
	new_tree_label=[]
	while num <len(tree_label_split):
		if tree_label_split[num]=='(':
			count1=count1+1
		if tree_label_split[num]==')':
			count2=count2+1
		if count1==count2:
			if tree_label_split[num]==']':
				index1.append(num+1)
				count1=0
				count2=0
		num=num+1
	index_1=0
	for item in index1:
		new_tree_label.append(tree_label_split[index_1:item])
		index_1=item+1
	#print(new_tree_label)
	return(new_tree_label)

#...........................................299 driver gene........................
driver_gene_path=get_driver_gene_path()
print('driver_gene_path: '+driver_gene_path)
if driver_gene_path!='.':
	file=open(driver_gene_path+os.path.sep+'299_driverMutationList_Cell_2018.txt','r')
	global driver_gene
	driver_gene={}
	lines=file.readlines()
	for line in lines[1:]:
		driver_gene.update({line.rstrip():[]})  #uniq driver gene
else:
	print('can not read driver gene file')

#...................................................read file...and mutation................................

filelist=os.listdir(file_path)
print(filelist)
#find all patient
all_patient_ID=[]
for filename in filelist:
	patient_ID=filename.split('.')[0].split('_')[0]
	if patient_ID not in all_patient_ID:
		all_patient_ID.append(patient_ID)
global all_mutation
for patient in all_patient_ID:
	if AF !='':#one sample per file
		print('one sample per file')
		all_mutation={}
		print('patient: '+patient)
		dic={}#[sample_id: all mutation_ID]
		all_data=[]#all_sample_ID
		sample_shu=0#the number of all samples of this patient
		for filename in filelist:
			if filename.split('.')[0].split('_')[0]==patient:
				sample_shu=sample_shu+1
				sample_id=filename.split('.')[0].split('_')[1]
				if sample_id not in all_data:
					all_data.append(sample_id)
				else:
					print('sample_id is not uniq!')
				if sample_id not in dic.keys():
					dic.update({sample_id:[]})
				else:
					print('sample_ID is not uniq!')
				file=open(file_path+os.path.sep+filename,'r')
				lines=file.readlines()
				for line in lines[1:]:
					data=line.rstrip().split('\t')
					if int(data[1])+int(data[2])!=0:
						AF_data=int(data[1])/(int(data[1])+int(data[2]))
						if AF_data>=float(AF) and data[0].split('_')[0] in all_chr:
							if data[0] not in dic[sample_id] :
								dic[sample_id].append(data[0])
								if data[0] not in all_mutation.keys():
									all_mutation.update({data[0]:[]})
								for refer_gene in data[3].split(';'):
									if refer_gene not in all_mutation[data[0]]:
										all_mutation[data[0]].append(refer_gene)
							else:
								print(filename+' contains two or more identical mutation!')
					else:
						print(data[0]+': this sample have mutations that are uncovered!')
	else:
		# 0-1 matrix
		print('0-1 matrix')
		all_mutation={}
		print('patient: '+patient)
		dic={}#[sample_id: all mutation_ID]
		all_data=[]#all_sample_ID
		sample_shu=0#the number of all samples of this patient
		sample_index={}#[index:smaple_name]
		file=open(file_path+os.path.sep+patient+'.txt','r')
		lines=file.readlines()
		data_0=lines[0].rstrip().split('\t')
		for item in data_0[1:-1]:#sample_id
			sample_shu=sample_shu+1
			if item not in dic.keys() :
				dic.update({item:[]})
				all_data.append(item)
				sample_index.update({sample_shu:item})
			else:
				print('sample_id is not uniq!')
		#print(sample_index)
		#print(dic)
		for line in lines[1:]:
			data=line.rstrip().split('\t')
			for i in range(1,len(data)-1):
				if int(data[i])==1 and data[0] not in dic[sample_index[i]]:
					dic[sample_index[i]].append(data[0])
			if data[0] not in all_mutation.keys():
				all_mutation.update({data[0]:[]})
				for refer_gene in data[-1].split(';'):
					if refer_gene not in all_mutation[data[0]]:
						all_mutation[data[0]].append(refer_gene)

	#calculate the length of root trunk
	#remove the sample which have 0 muatation
	print(all_data)
	for item in dic.keys():
		if len(dic[item])==0:
			sample_shu=sample_shu-1
			print(item+' have '+str(len(dic[item]))+' mutation,so remove this sample')
		else:
			print(item+' have mutations: '+str(len(dic[item])))
	for key in list(dic.keys()):
		if not dic.get(key):
			del dic[key]
			all_data.remove(key)
	if len(all_data)>=3:
		all_driver_mutation={}
		for item in dic.keys():
			if len(search_gene(dic[item],'FALSE'))>1:
				for info in search_gene(dic[item],'TRUE')[:-1]:#remove count_map_driver_gene
					split=info.split(':')
					if split[0] not in all_driver_mutation.keys():
						all_driver_mutation.update({split[0]:[split[1]]})
					else:
						if split[1] not in all_driver_mutation[split[0]]:
							all_driver_mutation[split[0]].append(split[1])
		un_uniq_gene=0
		for key in all_driver_mutation.keys():
			if len(all_driver_mutation[key])>1:
				print(key+' have '+str(len(all_driver_mutation[key]))+' uniq mutations: '+str(all_driver_mutation[key]))
			else:
				un_uniq_gene=un_uniq_gene+1
		if un_uniq_gene==len(all_driver_mutation.keys()):
			print('every gene refer to a uniq gene')
		#........................................................................
		i=0
		for item in dic.keys():
			i=i+1
			if i==1:
				all_jiaoji=dic[item]
			all_jiaoji = [val for val in all_jiaoji if val in dic[item]]
		root_node=str(len(all_jiaoji))+'['+str(search_gene(all_jiaoji,'False')[-1])+'-'
		for item in search_gene(all_jiaoji,'False')[:-1]:
			root_node=root_node+item+';'
		root_node=root_node[:-1]+']'
		for item in dic.keys():
			for item1 in all_jiaoji:
				dic[item].remove(item1)

		#...............................................body......................
		dic_baocun={}
		for item in all_data:
			dic_baocun.update({item:[]})
		for item in dic_baocun.keys():
			for item_data in dic[item]:
				dic_baocun[item].append(item_data)
		global all_DATA
		all_result_1=[]
		global branch_path
		branch_path=[]
		global path_string
		path_string=[]
		global final_path_count
		final_path_count=-1
		global DIC_result
		DIC_result={}
		DIC_count=0
		duli_count=0
		duli_state='false'
		chongfu_bianli=0
		while final_path_count!=len(branch_path) or len(path_string)!=len(branch_path):
			chongfu_bianli=chongfu_bianli+1
			dic={}
			for item in all_data:
				dic.update({item:[]})
			for item in dic.keys():
				for item_data in dic_baocun[item]:
					dic[item].append(item_data)
			all_DATA=[]
			all_DATA.append(all_data)
			result_1=[]
			path=','
			dic_result={}
			for all_item in all_DATA:
				dic1={}
				for item in all_item:
					dic1.update({item:[]})
				for item in dic1.keys():
					for item_data in dic[item]:
						dic1[item].append(item_data)
				#check if there are any-two sample that do not share any mutation?
				while len(all_item)>=2:
					sample_2={}
					count2=0
					cishu2=int(Cmn(len(all_item),2))
					while len(sample_2.keys())<cishu2:
						sel2=random.sample(all_item,2)
						original_data=sel2[0]+','+sel2[1]
						fan_original_data=sel2[1]+','+sel2[0]
						if fan_original_data not in sample_2.keys():
							sample_2.update({original_data:[]})
					for item_2 in sample_2.keys():
						item_2_split=item_2.split(',')
						sample_2_jiaoji=[val for val in dic1[item_2_split[0]] if val in dic1[item_2_split[1]]]
						if len(sample_2_jiaoji)!=0:
							count2=count2+1
					if count2>0:
						duli_count=duli_count+1
						result_all=xunhuan(all_item)
						path_state='false'
						for result_path in result_all:
							same_path_count=0
							identify1=result_path.split(':')[0]
							for result_path_1 in result_all:
								if identify1 == result_path_1.split(':')[0]:
									same_path_count=same_path_count+1
							if same_path_count>=2:
								path_state='true'
						branch_path_iter=[]
						for path_item in branch_path:
							if path_item not in branch_path_iter:
								branch_path_iter.append(path_item)
							else:
								print('branch_path error')
						#print('path_state: '+path_state)
						for result_all_1 in result_all:
							if path_state=='false':
								path_merge=path+';'+result_all_1.split(':')[0]
							else:
								path_merge=path+';'+result_all_1
							#print('path_merge: '+path_merge)
							if path_merge[0]==',' and chongfu_bianli==1:
								if path_merge[2:] not in branch_path:
									branch_path.append(path_merge[2:])
							else:
								if path_merge[0]==',':
									path_merge=path_merge[2:]
								path_all=len(branch_path_iter)
								path_count=0
								while path_count<path_all:
									if len(branch_path_iter[path_count].split(';'))==len(path_merge.split(';'))-1 and branch_path_iter[path_count] ==path_merge[:len(branch_path_iter[path_count])] :
										if  branch_path_iter[path_count] == path_merge[:len(branch_path_iter[path_count])]:
											if branch_path_iter[path_count] in branch_path:
												branch_path.remove(branch_path_iter[path_count])
										if path_merge not in branch_path:
											branch_path.append(path_merge)
									path_count=path_count+1

						for result_all_1 in result_all:
							if path_state=='false':
								if path[0]==',':
									#print('ok1, in branch_path')
									path_merge=path+';'+result_all_1.split(':')[0]
									path_merge=path_merge[2:]
									path_copy_state='FALSE'
									for path_copy1 in branch_path:
										if path_copy1[-3:]!='end' and path_merge==path_copy1[:len(path_merge)] :
											path=path_merge
											path_copy_state='TRUE'
											break
									if path_copy_state=='TRUE':
										break
								else:
									path_merge=path+';'+result_all_1.split(':')[0]
									path_copy_state='FALSE'
									for path_copy1 in branch_path:
										if path_copy1[-3:]!='end' and path_merge==path_copy1[:len(path_merge)] :
											path=path_merge
											path_copy_state='TRUE'
											break
									if path_copy_state=='TRUE':
										break
							else:
								if path[0]==',':
									#print('ok3, in branch_path')
									path_merge=path+';'+result_all_1
									path_merge=path_merge[2:]
									path_copy_state='FALSE'
									for path_copy1 in branch_path:
										if path_copy1[-3:]!='end' and path_merge==path_copy1[:len(path_merge)] :
											path=path_merge
											path_copy_state='TRUE'
											break
									if path_copy_state=='TRUE':
										break
								else:
									path_merge=path+';'+result_all_1
									path_copy_state='FALSE'
									for path_copy1 in branch_path:
										if path_copy1[-3:]!='end' and path_merge==path_copy1[:len(path_merge)] :
											path=path_merge
											path_copy_state='TRUE'
											break
									if path_copy_state=='TRUE':
										break
						result=result_all_1
						result_split=result.split(':')
						result_split=result_split[1].split(',')
						first='.'
						#print('max_pair:')
						for item_split in result_split:
							first=first+','+item_split
						first=first[2:]
						return_2=[]
						for item_split in all_item:
							if item_split not in result_split:
								return_2.append(item_split)
						return_2.sort()
						#print(return_2)
						return_2_zuhe=[]
						while len(return_2_zuhe) <fact(len(return_2)):
							sel=random.sample(return_2,len(return_2))
							da=','
							for Item in sel:
								da=da+','+Item
							if da[2:] not in return_2_zuhe:
								return_2_zuhe.append(da[2:])
						#print(return_2_zuhe)
						second='.'
						for item3 in return_2_zuhe:
							if item3 in first_3.keys():
								second=item3
								break

						if int(first_3[first].split('[')[0])!=len(first_4[first]):
							print('false')
						result_item_1=first+':'+str(len(first_4[first]))+'['+str(search_gene(first_4[first],'False')[-1])+'-'
						for item in search_gene(first_4[first],'False')[:-1]:
							result_item_1=result_item_1+item+';'
						result_item_1=result_item_1[:-1]+']'
						result_1.append(result_item_1)
						dic_result.update({first:first_4[first]})
						result_item_2=second+':'+str(len(first_4[second]))+'['+str(search_gene(first_4[second],'False')[-1])+'-'
						for item in search_gene(first_4[second],'False')[:-1]:
							result_item_2=result_item_2+item+';'
						result_item_2=result_item_2[:-1]+']'
						result_1.append(result_item_2)
						dic_result.update({second:first_4[second]})
						if int(result[0:1])<2:
							result_split=result.split(':')
							result_split_split=result_split[1].split(',')
							all_item=[]
							for item in result_split_split:
								all_item.append(item)
							for key in list(dic1.keys()):
								if key not in result_split_split:
									del dic1[key]
							#filter key which value=null
							for key in list(dic1.keys()):
								if not dic1.get(key):
									del dic1[key]
									all_item.remove(key)
							#print(all_item)
							i=0
							for item_1 in dic1.keys():
								i=i+1
								if i==1:
									jiaoji=dic1[item_1]
								jiaoji=[val for val in jiaoji if val in dic1[item_1]]

							for item in dic1.keys():
								for item1 in jiaoji:
									dic1[item].remove(item1)

							for item in dic1.keys():
								for item1 in jiaoji:
									dic[item].remove(item1)
						else:
							group1=[]
							group2=[]
							result_split=result.split(':')
							result_split_split=result_split[1].split(',')
							for item in result_split_split:
								group1.append(item)
							group1.sort()
							i=0
							for item_1 in group1:
								i=i+1
								if i==1:
									jiaoji=dic[item_1]
								jiaoji=[val for val in jiaoji if val in dic[item_1]]

							for item in group1:
								for item1 in jiaoji:
									dic[item].remove(item1)

							for item in all_item:
								if item not in result_split_split:
									group2.append(item)
							group2.sort()
							i=0
							for item_1 in group2:
								i=i+1
								if i==1:
									jiaoji=dic[item_1]
								jiaoji=[val for val in jiaoji if val in dic[item_1]]

							for item in group2:
								for item1 in jiaoji:
									dic[item].remove(item1)

							all_DATA.append(group1)
							all_DATA.append(group2)

							break
					else:
						for item in all_item:
							result_item =item+':'+str(len(dic1[item]))+'['+str(search_gene(dic1[item],'False')[-1])+'-'
							for item0 in search_gene(dic1[item],'False')[:-1]:
								result_item=result_item+item0+';'
							result_item=result_item[:-1]+']'
							#print(result_item)
							result_1.append(result_item)
							dic_result.update({item:dic1[item]})
							if duli_count==0:
								result_item=result_item+';end'
								branch_path.append(result_item)
						if duli_count==0:
							all_result_1.append(result_1)
							duli_state='true'
						#print(sample_name)
						break
			if duli_count>0:
				path=path+';end'
				if path==',;end':
					print('error: '+path)
				for path_2 in branch_path:
					if path[:-4]==path_2:
						branch_path.remove(path_2)
				if path not in branch_path:
					branch_path.append(path)
			for i in branch_path:
				if branch_path.count(i)>1:
					print('branch_path have chongfu item: '+i)
			final_path_count=0
			for path_data in branch_path:
				if path_data.split(';')[-1]=='end' or path_data[0].isalpha():
					final_path_count=final_path_count+1

			if duli_state=='false':
				if path not in path_string and path[-3:]=='end':
					path_string.append(path)

					DIC_result.update({DIC_count:dic_result})
					DIC_count=DIC_count+1
					all_result_1.append(result_1)
			else:
				for path_1 in branch_path:
					if path_1 not in path_string and path_1[-3:]=='end':
						path_string.append(path_1)
					else:
						print('The tree structure that completely branches from the root node is wrong!')
				DIC_result.update({DIC_count:dic_result})
				DIC_count=DIC_count+1
		print('branch_path count: '+str(len(branch_path)))
		print('path_string: '+str(len(path_string)))
		print('all_result: '+str(len(all_result_1)))
		print('DIC_result: '+str(len(DIC_result.keys())))
		for item in all_result_1:
			item_copy = item[:]

			for data in item_copy:
				if int(data.split(':')[1].split('[')[0])==0 and len(data.split(':')[0].split(','))>1:
					item.remove(data)
		all_result_weight=[]
		for item in all_result_1:
			weight=0

			for data in item:
				length=len(data.split(':')[0].split(','))
				data=data.split(':')[1].split('[')[0]
				if length>=2:
					weight=weight+length*int(data)
			all_result_weight.append(weight)

		i=0
		max_weight_data=0
		while i <len(all_result_weight):
			if i==0:
				max_weight_data=all_result_weight[i]
			if all_result_weight[i]>max_weight_data:
				max_weight_data=all_result_weight[i]
			i=i+1
		result_max_weight=[]
		if max_weight_data==0:
			DIC_result_max_weight={}
			for item in all_result_1:
				result_max_weight.append(item)
			DIC_result_max_weight.update({0:DIC_result[0]})
		else:
			DIC_result_max_weight={}
			i=0
			i1=0
			while i <len(all_result_weight):
				if all_result_weight[i]==max_weight_data:
					all_result_1[i].sort()
					if all_result_1[i] not in result_max_weight:
						result_max_weight.append(all_result_1[i])
						DIC_result_max_weight.update({i1:DIC_result[i]})

						i1=i1+1
				i=i+1

		tree_count=0
		for result_1 in result_max_weight:
			dic_result=DIC_result_max_weight[tree_count]
			dic_final={}
			ll=[]#header
			ll.append('mutation_id')
			for item in dic_result.keys():
				if len(dic_result[item])!=0 or len(item.split(','))==1:
					ll.append(item)#all branch
			for item in ll[1:] :
				for mut in dic_result[item]:
					if mut not in dic_final.keys():
						LL=[]
						LL.append(mut)
						for item1 in ll[1:]:
							if mut in dic_result[item1]:
								LL.append(1)
							else:
								LL.append(0)
						driver_string='.'
						for driver_item in search_gene([mut],'FALSE')[:-1]:
							driver_string=driver_string+driver_item
						if len(driver_string)==1:
							LL.append('.')
						else:
							LL.append(driver_string[1:])
						dic_final.update({mut:LL})
			#print(dic_final)
			ll.append('driver_gene')
			result_2=[]
			number=1
			#print(str(sample_shu))
			while number<=sample_shu:
				for item in result_1:
					if number==1:
						if len(item.split(':')[0].split(','))==number:
							#print(item.split(':')[0].split(','))
							result_2.append(item)
					else:
						if len(item.split(':')[0].split(','))==number:
							new_item=','
							if int(item.split(':')[1].split('[')[0])!=0:
								for item1 in item.split(':')[0].split(','):
									for item2 in result_2:
										num=0
										while num<len(item2):
											if item2[num]=='(':
												num=num+1
											else:
												break
										if item2[len(str(item1))+num]==':':
											if item1 ==item2[num:len(str(item1))+num]:
												new_item=new_item+','+item2
												result_2.remove(item2)
								new_item='('+new_item[2:]+'):'
								for item3 in item.split(':')[1:]:
									new_item=new_item+item3+':'
								#print(new_item[:-1])
								result_2.append(new_item[:-1])
					#print(result_2)
				number=number+1
			#print(result_2)
			final_result=','
			for item in result_2:
				final_result=final_result+','+item
			final_result='('+final_result[2:]+'):'+root_node
			print('tree_structure:')
			print(final_result)
			if tree_count==0:
				first_final=final_result
			tree_structure=re.sub(u"\\[.*?]", "", final_result)
			tree_structure=tree_structure+';'
			#print(tree_structure)
			#.......................................plot.........................................
			if tree_count==0 or (tree_count>0 and final_result!=first_final):
				out=open(file_path_out+os.path.sep+patient+'_info_'+str(tree_count)+'.txt','w')
				for item in ll:
					out.write(str(item))
					out.write('\t')
				out.write('\n')
				for item in dic_final.keys():
					count=0
					for data in dic_final[item][1:-1]:
						count=count+int(data)
					if count>=2:
						out.write(item)
						out.write('\t')
						for data in dic_final[item][1:]:
							out.write(str(data))
							out.write('\t')
						out.write('\n')
				out.close()
				tree = Phylo.read(StringIO(tree_structure), "newick")
				#tree.ladderize()# Flip branches so deeper clades are displayed at top
				tree.rooted = True
				tree.name=patient
				tree = tree.as_phyloxml()
				tree.root.branch_labels=root_node
				#global tree_label
				if ('[' in final_result[1:].split(':')[-1]) and (']' in final_result[1:].split(':')[-1]):
					tree_label=final_result[1:].split(':')[:-1]
				else:
					tree_label=final_result[1:].split(':')[:-2]
				tree_label_1=[str(i) for i in tree_label]
				tree_labels = ':'.join(tree_label_1)
				tree_label_split=tree_labels[:-1]
				global new_tree_labels
				new_tree_labels=split_tree(tree_label_split)
				original_count=len(new_tree_labels)
				original_count1=original_count
				num=0
				for item in new_tree_labels:
					if num <=original_count-1:
						tree.clade[num].branch_labels=item.split(':')[-1]
						if ('(' in item) or (')' in item):
							new_item_1=[str(i) for i in item.split(':')[:-1]]
							new_item_split=':'.join(new_item_1)
							new_item_split=new_item_split[1:-1]
							new_items=split_tree(new_item_split)
							index=0
							for branch in new_items:
								branch=str(num)+'*'+str(index)+'*'+branch
								new_tree_labels.append(branch)
								index=index+1
						original_count1=original_count1-1
					else:
						branch=item.split('*')[-1]
						branch_count=item.split('*')[:-1]
						#print(branch_count)
						branch_count_merge=','
						for index in branch_count:
							branch_count_merge=branch_count_merge+'*'+index
						branch_count_merge=branch_count_merge[2:]
						str_count='tree.clade['
						for ITEM in branch_count:
							str_count=str_count+ITEM+','
						str_count=str_count[:-1]+'].branch_labels = '+'"'+item.split(':')[-1]+'"'
						#print(str_count)
						exec(str_count)
						#print(tree)
						if ('(' in item) or (')' in item):
							new_item_1=[str(i) for i in item.split(':')[:-1]]
							new_item_split=':'.join(new_item_1)
							new_item_split=new_item_split.split('*')[-1]
							new_item_split=new_item_split[1:-1]
							new_items=split_tree(new_item_split)
							index=0
							for branch1 in new_items:
								branch1=str(branch_count_merge)+'*'+str(index)+'*'+branch1
								new_tree_labels.append(branch1)
								index=index+1
					num=num+1
				matplotlib.rc('font', size=6)
				plt.rcParams['lines.linewidth'] = 0.7
				Phylo.draw(tree,do_show=False,branch_labels=lambda c: c.branch_labels)
				#Phylo.draw(tree,do_show=False,branch_labels=lambda c: int(c.branch_length))
				plt.savefig(file_path_out+os.path.sep+patient+'_tree_'+str(tree_count)+'.pdf',dpi=500)
				plt.close()
				tree_count=tree_count+1
		print(patient+' ok!')
	else:
		print(patient+' have less patients!')