code: # make deterministic
  from mingpt.utils import set_seed
  set_seed(42)
  #frompc
  
  import numpy as np
  import torch
  import torch.nn as nn
  from torch.nn import functional as F
  import math
  from torch.utils.data import Dataset
  
  from mingpt.model import *
  
  from mingpt.trainer import Trainer, TrainerConfig
  from mingpt.utils import sample
  
  import logging
  logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO,
  )
  import pdb
  
  from dataset.tcga_dataset import TCGA
  from dataset.tcga_conv_dataset import TCGAConv
  from dataset.pcawg_conv_dataset import *
  from dataset.pcawg_dataset import PCAWG
  from dataset.pcawg_emb_dataset import PCAWGEmb
  from dataset.pcawg_sepdataset import PCAWGSep
  from dataset.pcawg_2stream import PCAWG2Stream
  from dataset.tcgadisttabletoemb_dataset import TCGADist
  from dataset.tcgamutdist_dataset import TCGAMutDist
  from dataset.tcgamutdistasone_dataset import TCGAMutDistasOne
  from dataset.tcgapcawg_dataset import TCGAPCAWG
  from dataset.newtcgapcawg_dataset import NewTCGAPCAWG
  from dataset.finaltcgapcawg_dataset import FinalTCGAPCAWG
  
  from mingpt.bert import *
  from preprocessing.dmm.dmm import *
  from preprocessing.fromvcffiles import *
  
  import argparse
  import os
  import pandas as pd
  
  def translate_args(args):
  
      cwd = os.getcwd()
      args.cwd = cwd
  
      args.mutation_coding = cwd + '/preprocessing/dmm/data/mutation_codes_sv.tsv'
      args.input = args.data_dir
  
      args.output = cwd + '/data/raw/out/00b9d0e6-69dc-4345-bffd-ce32880c8eef.consensus.20160830.somatic.snv_mnv.tsv.gz' 
  
      args.reference = '/csc/epitkane/data/ref_genomes/hs37d5_1000GP/hs37d5_1000GP.fa'
      args.context = 8
  
      args.sample_id = 'submitted_sample_id'
  
      args.tmp = cwd + '/data/raw/tmp/'
      args.verbose = 1
      args.generate_negatives = 1
      args.report_interval = 100000
  
      return args
  
  
  def get_args():
          parser = argparse.ArgumentParser(description='TCGA / PEACOCK experiment')
  
          # DATASET
          parser.add_argument('--cwd', type=str,help='project dir')
  
          parser.add_argument('--dataset', type=str, default='pcawg',
                          help='dataset')
          # MODEL
          parser.add_argument('--arch', type=str, default=None,
                          help='architecture')
          # DIRECTORY
          parser.add_argument('--data-dir', type=str, default=None,
                          help='data directory')
          parser.add_argument('--crossdata-dir', type=str, default=None,
                          help='data directory')
          parser.add_argument('--adddata-dir', type=str, default=None,
                          help='data directory')
  
          parser.add_argument('--n-class', type=int, default=None,
                          help='number of class')
  
          parser.add_argument('--batch-size', type=int, default=1,
                          help='batch size')
  
          parser.add_argument('--block-size', type=int, default=1000,
                          help='block of sequence')
  
          parser.add_argument('--context-length', type=int, default=256,
                          help='length of sequence')
          parser.add_argument('--n-layer', type=int, default=1,
                          help='attention layer')
          parser.add_argument('--n-head', type=int, default=8,
                          help='attention head')
          parser.add_argument('--n-emb', type=int, default=128,
                          help='embedding dimension')
          parser.add_argument('--n-vocab-type', type=int, default=1,
                          help='embedding dimension')
  
          parser.add_argument('--tag', type=str, default='myexperiment',
                          help='dataset')
          
          parser.add_argument('--train', action='store_true', default=False)
          parser.add_argument('--predict', action='store_true', default=False)
          parser.add_argument('--trainbp', action='store_true', default=False)
          parser.add_argument('--vis-weight', action='store_true', default=False)
          parser.add_argument('--top-weight', action='store_true', default=False)
  
          parser.add_argument('--visval', action='store_true', default=False)
  
  
          parser.add_argument('--single-predict', action='store_true', default=False)
          parser.add_argument('--create-dataset', action='store_true', default=False)
          parser.add_argument('--two-streams', action='store_true', default=False)
          parser.add_argument('--three-streams', action='store_true', default=False)
  
          parser.add_argument('--filter', action='store_true', default=False)
  
          parser.add_argument('--bert', action='store_true', default=False)
          parser.add_argument('--withclass', action='store_true', default=False)
          parser.add_argument('--default', action='store_true', default=False)
          parser.add_argument('--addposition', action='store_true', default=False)
          parser.add_argument('--oneDhot', action='store_true', default=False)
          parser.add_argument('--addorder', action='store_true', default=False)
          parser.add_argument('--addtoken', action='store_true', default=False)
          parser.add_argument('--addtriplet', action='store_true', default=False)
          parser.add_argument('--addtriplettoken', action='store_true', default=False)
          parser.add_argument('--addgestoken', action='store_true', default=False)
          parser.add_argument('--addrt', action='store_true', default=False)
          parser.add_argument('--addlongcontext', action='store_true', default=False)
          parser.add_argument('--tokenizedlongcontext', action='store_true', default=False)
          parser.add_argument('--ohlongcontext', action='store_true', default=False)
          parser.add_argument('--flattenohlongcontext', action='store_true', default=False)
          parser.add_argument('--addpostoken', action='store_true', default=False)
          parser.add_argument('--addrttoken', action='store_true', default=False)
          parser.add_argument('--balance', action='store_true', default=False)
  
          parser.add_argument('--l1', action='store_true', default=False)
  
          parser.add_argument('--fold', type=int, default=1, 
                              help='number of mutation')
  
          parser.add_argument('--output-mode', type=str, default='token',help='dataset')
  
          parser.add_argument('--rbm', action='store_true', default=False)
  
          parser.add_argument('--newtraining', action='store_true', default=False)
          parser.add_argument('--newpredict', action='store_true', default=False)
          parser.add_argument('--newpredict2', action='store_true', default=False)
          parser.add_argument('--normal', action='store_true', default=False)
  
          parser.add_argument('--freezeemb', action='store_true', default=False)
  
          parser.add_argument('--predictvis', action='store_true', default=False)
  
          parser.add_argument('--crossdata', action='store_true', default=False)
  
          parser.add_argument('--nummut', type=int, default=0,
                          help='number of mutation')
          parser.add_argument('--frac', type=float, default=0,
                          help='frac')
  
          parser.add_argument('--mutratio', type=str, default='',
                          help='mutation ratio')
  
          parser.add_argument('--spectral', action='store_true', default=False)
          parser.add_argument('--finalpredict', action='store_true', default=False)
  
          parser.add_argument('--finalpredictnewdata', action='store_true', default=False)
          parser.add_argument('--single-pred-vcf', action='store_true', default=False)
  
  
          parser.add_argument('--vis-attention', action='store_true', default=False)
  
  
          #dmm_parser
          parser.add_argument('-v', '--verbose', type=int, help='Try to be more verbose')
          parser.add_argument('--mutation-coding', help='Mutation coding table ("ref alt code"/line) [{}]'.format(\
                                  defaults['mutation_coding']), metavar='fn')
          parser.add_argument('--config', help='Read parameters from a JSON file')
          parser.add_argument('--data-config',
                                  help='Column specification for --input, --validation and --aux-data  [{}]'.format(\
                                  defaults['data_config']))
          parser.add_argument('--random-seed', default=None, type=int, metavar='seed')
          parser.add_argument('--tmp')
          
          parser.add_argument('-i', '--input', action='append', metavar='dir(s)',
                                  help='Either a directory with vcf/maf[.gz] files or a vcf/maf[.gz] file (-i may be given more than once)')
          parser.add_argument('-o', '--output', metavar='fn', help='Preprocessed mutation data')
          parser.add_argument('-r', '--reference', metavar='ref', help='Reference genome (fasta) [{}]'.format(\
                                  defaults['reference']))
          parser.add_argument('-k', '--context', help='Sequence context length (power of 2) [{}]'.format(\
                                  defaults['context']), metavar='bp', type=int,default=8)
          parser.add_argument('-e', '--errors', metavar='fn',
                                  help='File where to log errors [{}]'.format(defaults['errors']))
          parser.add_argument('--no-ref-preload', help='Use samtools to read reference on demand (slow but fast startup) [false]',
                                  action='store_true')
          parser.add_argument('--no-filter', help='Process all variants [default=only PASS/. variants]',
                                  action='store_true')
          parser.add_argument('--sample-id', help='Sample identifier column name in MAF file')
          parser.add_argument('-n', '--generate_negatives', help='Ratio of negative to positive examples [{}]. Two passes on data are required for n>0.'.format(\
                                  defaults['negative_ratio']), type=float)
          parser.add_argument('--median-variant-type-negatives', action='store_true',
                                  help='Generate median number of each variant type as negative examples for each sample')
          parser.add_argument('--median-variant-type-file', help='Load median variant numbers from a file')
          parser.add_argument('--negative-generation-mode', help='[generate] output in one go (default), [augment] input files or [process] augmented files', default='generate')
          parser.add_argument('--info-column', help='Input column name to write toutputo output (MAF input only). May be specified more than once.', action='append')
          parser.add_argument('--report-interval', help='Interval to report number of variants processed',
                                  type=int)
          parser.add_argument('--array-jobs', help='How many array jobs in total', type=int)
          parser.add_argument('--array-index', help='Index of this job', type=int)
          parser.add_argument('--nope', help='Only one variant per output sequence', action='store_true')
          parser.add_argument('--no-overwrite', help='Do not overwrite if output exists', action='store_true')
  
  
          args = parser.parse_args()
          return args
  
  def get_dataloader(args,train_val,load):
  
      if args.dataset == 'finalpcawg' or args.dataset == 'wgspcawg':
          if train_val=='training':
              dataloader_class = FinalTCGAPCAWG(dataset_name = args.dataset, 
                                  data_dir=args.data_dir, 
                                  mode='training', 
                                  curr_fold=args.fold, 
                                  block_size=args.block_size, 
                                  load=False,
                                  mutratio = args.mutratio,
                                  addtriplettoken=args.addtriplettoken,
                                  addpostoken=args.addpostoken,
                                  addgestoken=args.addgestoken,
                                  addrt=args.addrt,
                                  nummut = args.nummut,
                                  frac = args.frac,
                                  crossdata = args.crossdata,
                                  crossdatadir = args.crossdata_dir,
                                  adddatadir = args.adddata_dir
                                  )
  
          elif train_val=='validation':
              dataloader_class = FinalTCGAPCAWG(dataset_name = args.dataset, 
                                  data_dir=args.data_dir, 
                                  mode='validation', 
                                  curr_fold=args.fold, 
                                  block_size=args.block_size, 
                                  load=False,
                                  mutratio = args.mutratio,
                                  addtriplettoken=args.addtriplettoken,
                                  addpostoken=args.addpostoken,
                                  addgestoken=args.addgestoken,
                                  addrt=args.addrt,
                                  nummut = args.nummut,
                                  frac = args.frac,
                                  crossdata = args.crossdata,
                                  crossdatadir = args.crossdata_dir,
                                  adddatadir = args.adddata_dir)
  
      elif args.dataset == 'finaltcga' or args.dataset == 'westcga':
          if train_val=='training':
              dataloader_class = FinalTCGAPCAWG(dataset_name = args.dataset, 
                                  data_dir=args.data_dir, 
                                  mode='training', 
                                  curr_fold=args.fold, 
                                  block_size=args.block_size, 
                                  load=False,
                                  mutratio = args.mutratio,
                                  addtriplettoken=args.addtriplettoken,
                                  addpostoken=args.addpostoken,
                                  addgestoken=args.addgestoken,
                                  addrt=args.addrt,
                                  nummut = args.nummut,
                                  frac = args.frac,
                                  crossdata = args.crossdata,
                                  crossdatadir = args.crossdata_dir,
                                  adddatadir = args.adddata_dir)
  
          elif train_val=='validation':
              dataloader_class = FinalTCGAPCAWG(dataset_name = args.dataset, 
                                  data_dir=args.data_dir, 
                                  mode='validation', 
                                  curr_fold=args.fold, 
                                  block_size=args.block_size, 
                                  load=False,
                                  mutratio = args.mutratio,
                                  addtriplettoken=args.addtriplettoken,
                                  addpostoken=args.addpostoken,
                                  addgestoken=args.addgestoken,
                                  addrt=args.addrt,
                                  nummut = args.nummut,
                                  frac = args.frac,
                                  crossdata = args.crossdata,
                                  crossdatadir = args.crossdata_dir,
                                  adddatadir = args.adddata_dir)
                      
      return dataloader_class
  
  def get_model(args,mconf):
          if args.arch == 'GPTConv':
                  model = GPTConv(mconf)
          elif args.arch == 'GPTConvDeeper':
                  model = GPTConvDeeper(mconf)
          elif args.arch == 'GPTNonPosition':
                  model = GPTNonPosition(mconf)
          elif args.arch == 'CTransformer':
                  model = CTransformer(mconf)
          elif args.arch == 'ConvTransformer':
                  model = ConvTransformer(mconf)
          elif args.arch == 'Conv2DTransformer':
                  model = Conv2DTransform
          elif args.arch == 'Transformer2Stream':
                  model = Transformer2Stream(mconf)
          elif args.arch == 'CTransformerDNN':
                  model = CTransformerDNN(mconf)
          elif args.arch == 'CTransformerMutDist':
                  model = CTransformerMutDist(mconf)
          elif args.arch == 'SimpleAttention':
                  model = SimpleAttention(mconf)
          elif args.arch == 'BertForSequenceClassification':
                  model = BertForSequenceClassification(mconf)
          elif args.arch == 'BertwithPosition':
                  model = BertwithPosition(mconf)
          elif args.arch == 'CTransformerWithPaddingIDX':
                  model = CTransformerWithPaddingIDX(mconf)
          elif args.arch == 'Conv2DTransformerOnehot':
                  model = Conv2DTransformerOnehot(mconf)
          elif args.arch == 'CTransformerWithPaddingIDXandfirstvec':
                  model = CTransformerWithPaddingIDXandfirstvec(mconf)
          elif args.arch == 'Conv2DTransformerOnehotDeeper':
                  model = Conv2DTransformerOnehotDeeper(mconf)
          elif args.arch == 'DNNTransformerOnehotDeeper':
                  model = DNNTransformerOnehotDeeper(mconf)
          elif args.arch == 'CTransformerWithPosition':
                  model = CTransformerWithPosition(mconf)
          elif args.arch == 'CTransformerWithPositionConcate':
                  model = CTransformerWithPositionConcate(mconf)
          elif args.arch == 'DNNTransformerOnehotDeeperwithPosition':
                  model = DNNTransformerOnehotDeeperwithPosition(mconf)
          elif args.arch == 'DNNTransformerOnehotDeeperwithPositionwithOrder':
              model = DNNTransformerOnehotDeeperwithPositionwithOrder(mconf)
          elif args.arch == 'CTransformerDNNWithPositionConcateToken':
              model = CTransformerDNNWithPositionConcateToken(mconf)
          elif args.arch == 'CTransformerDNNWithPositionConcateTokenSep':
              model = CTransformerDNNWithPositionConcateTokenSep(mconf)
          elif args.arch == 'CTransformerRBMWithPositionConcate':
              model = CTransformerRBMWithPositionConcate(mconf)
          elif args.arch == 'TripletPositionTokenandOnehot':
              model = TripletPositionTokenandOnehot(mconf) 
          elif args.arch == 'PositionToken':
              model = PositionToken(mconf) 
          elif args.arch == 'TripletPositionTokenandOnehotConcAfter':
              model = TripletPositionTokenandOnehotConcAfter(mconf)
          elif args.arch == 'TripletPositionRTToken':
              model = TripletPositionRTToken(mconf)
          elif args.arch == 'FullConvTransformer':
              model = FullConvTransformer(mconf)
          elif args.arch == 'TripletPositionTokenBest':
              model = TripletPositionTokenBest(mconf)
          elif args.arch == 'TripletPositionTokenRT':
              model = TripletPositionTokenRT(mconf)
          elif args.arch == 'EmbFC':
              model = EmbFC(mconf)    
          elif args.arch == 'TripletPositionTokenOldBest':
              model = TripletPositionTokenOldBest(mconf)
          elif args.arch == 'CTransformerPCAWGtoTCGA_TPGES':
              model = CTransformerPCAWGtoTCGA_TPGES(mconf)
          elif args.arch == 'CTransformerPCAWGtoTCGA_T':
              model = CTransformerPCAWGtoTCGA_T(mconf)
          elif args.arch == 'TripletPosition':
              model = TripletPosition(mconf)    
          elif args.arch == 'TripletPositionGES':
              model = TripletPositionGES(mconf)
          elif args.arch == 'TripletPositionGESRT':
              model = TripletPositionGESRT (mconf)   
          elif args.arch == 'TripletPositionF':
              model = TripletPositionF(mconf)    
          elif args.arch == 'TripletPositionGESF':
              model = TripletPositionGESF(mconf)
          elif args.arch == 'CTransformerF':
              model = CTransformerF(mconf)
          elif args.arch == 'EmbFCPos':
              model = EmbFCPos(mconf)
          elif args.arch == 'EmbFCPosGES':
              model = EmbFCPosGES(mconf)
  
          return model
  
  def fold_split(args):
  
          num_class = os.listdir(args.data_dir)
          class_name = [i for i in num_class if len(i.split('.'))==1]
          class_name = sorted(class_name)
  
          num_samples = []
  
          for i in class_name:
                  ns = len(os.listdir(args.data_dir+i))
                  num_samples.append(ns)
  
          d = {'class_name':class_name,'n_samples':num_samples}
          pd_class_info = pd.DataFrame(d)
          
          folds=10
  
          class_used = pd_class_info.loc[pd_class_info['n_samples']>=folds]
          class_used = class_used.rename_axis('class_index').reset_index()
          class_used.to_csv(args.data_dir + 'sample_info_' + args.dataset + '.csv', index=False)
  
          num_class=len(class_used)
  
          tuple_list = []
  
          for nm_class in class_used['class_name']:
                  num_sample = class_used.loc[class_used['class_name']==nm_class]['n_samples'].values[0]
                  class_idx = class_used.loc[class_used['class_name']==nm_class]['class_index'].values[0]
                  samples = os.listdir(args.data_dir+nm_class)
                  count_split = 0
  
                  for i in range(0,num_sample):
                          count_split = count_split+1
                          if count_split > folds:
                                  count_split = 1
  
                          tuple_onerow = tuple([nm_class,class_idx,samples[i],count_split])
                          tuple_list.append(tuple_onerow)
                  
          all_split = pd.DataFrame(tuple_list,columns = ['class_name','class_index','name_samples','split'])
  
          test_split = pd.DataFrame(columns = all_split.columns)
          train_split = pd.DataFrame(columns = all_split.columns)
          validation_split = pd.DataFrame(columns = all_split.columns)
  
          for i in range(1,folds):
              test = all_split.loc[all_split['split']==i]
              train = all_split.loc[all_split['split']!=i]
              split_min = i + 1
              if split_min >= folds:
                  split_min = 1
              validation = train.loc[train['split']==split_min]
              train = train.loc[train['split']!=split_min]
              train['split'] = i
              validation['split'] = i
  
              test_split = test_split.append(test)
              validation_split = validation_split.append(validation)
              train_split = train_split.append(train)
  
          train_split.to_csv(args.data_dir + 'train_split.csv', index=False)
          validation_split.to_csv(args.data_dir + 'validation_split.csv', index=False)
          test_split.to_csv(args.data_dir + 'test_split.csv', index=False)
  
  
  if __name__ == '__main__':
  
          best_accuracy=0
  
          args = get_args()
  
          if args.train:
  
                  #class_info = fold_split(args)
  
                  block_size = args.block_size # spatial extent of the model for its context
                  train_dataset = get_dataloader(args=args,train_val='training',load= not args.create_dataset)
                  validation_dataset = get_dataloader(args=args,train_val='validation',load= not args.create_dataset)
  
                  if args.bert:
                      if args.default:
                          mconf = BertConfig(vocab_size_or_config_json_file = train_dataset.vocab_size,num_class=args.n_class)
                      else:
                          if args.addposition:
                              mconf = BertConfig(vocab_size_or_config_json_file = train_dataset.vocab_size,num_class=args.n_class,num_hidden_layers=args.n_layer,hidden_size=args.n_emb,num_attention_heads=args.n_head,type_vocab_size=args.n_vocab_type,position_size=train_dataset.position_size)
                          else:    
                              mconf = BertConfig(vocab_size_or_config_json_file = train_dataset.vocab_size,num_class=args.n_class,num_hidden_layers=args.n_layer,hidden_size=args.n_emb,num_attention_heads=args.n_head,type_vocab_size=args.n_vocab_type)
  
                      model = get_model(args,mconf)
  
                      string_logs = f"{args.tag}_{args.arch}_bs{args.block_size:.0f}_nl{args.n_layer:.0f}_nh{args.n_head:.0f}_ne{args.n_emb:.0f}_cl{args.context_length:.0f}/"
                      tconf = TrainerConfig(max_epochs=150, batch_size=1, learning_rate=0.001,
                              lr_decay=True, warmup_tokens=1*150, final_tokens=150*len(train_dataset)*block_size,
                              num_workers=1,string_logs=string_logs, args=args)
                      trainer = Trainer(model, train_dataset, validation_dataset, tconf)
                      trainer.bert_train()
  
                  if args.rbm:
                      num_class=args.n_class
                      mconf = GPTConfig(vocab_size=train_dataset.vocab_size, block_size=block_size,num_class=num_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
  
                      if args.addposition:
                          mconf = GPTConfig(vocab_size=train_dataset.vocab_size, block_size=block_size,num_class=num_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=train_dataset.position_size)
  
                      model = get_model(args,mconf)
  
                      string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                      tconf = TrainerConfig(max_epochs=150, batch_size=1, learning_rate=6e-4,
                              lr_decay=True, warmup_tokens=1*150, final_tokens=150*len(train_dataset)*block_size,
                              num_workers=1,string_logs=string_logs, args=args)
                      trainer = Trainer(model, train_dataset, validation_dataset, tconf)
  
                      output_mode = args.output_mode.split("_")
  
                      if len(output_mode)>1:
                          trainer.multi_stream_rbm(len(output_mode))
                      else:
                          trainer.basic_train()
                  else:
                      num_class=args.n_class
                      mconf = GPTConfig(vocab_size=train_dataset.vocab_size, block_size=block_size,num_class=num_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
  
                      if args.addposition:
                          mconf = GPTConfig(vocab_size=train_dataset.vocab_size, block_size=block_size,num_class=num_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=train_dataset.position_size)
  
                      model = get_model(args,mconf)
  
                      string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                      tconf = TrainerConfig(max_epochs=150, batch_size=1, learning_rate=6e-4,
                              lr_decay=True, warmup_tokens=1*150, final_tokens=150*len(train_dataset)*block_size,
                              num_workers=1,string_logs=string_logs, args=args)
                      trainer = Trainer(model, train_dataset, validation_dataset, tconf)
  
                      output_mode = args.output_mode.split("_")
  
                      if len(output_mode)>1:
                          trainer.multi_stream(len(output_mode))
                      else:
                          trainer.basic_train()
  
          if args.newtraining:
              block_size = args.block_size # spatial extent of the model for its context
              train_dataset = get_dataloader(args=args,train_val='training',load= not args.create_dataset)
              validation_dataset = get_dataloader(args=args,train_val='validation',load= not args.create_dataset)
  
              mconf = GPTConfig(vocab_size=train_dataset.vocab_size, block_size=block_size,num_class=args.n_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
                      
              if args.addpostoken:
                  mconf = GPTConfig(vocab_size=train_dataset.vocab_size, block_size=block_size,num_class=args.n_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=train_dataset.position_size,rt_size = train_dataset.rt_size)
  
              if args.addgestoken:
                  mconf = GPTConfig(vocab_size=train_dataset.vocab_size, block_size=block_size,num_class=args.n_class,
                  n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=train_dataset.position_size, ges_size =  train_dataset.ges_size,rt_size = train_dataset.rt_size)
  
              model = get_model(args,mconf)
  
              string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
              tconf = TrainerConfig(max_epochs=150, batch_size=1, learning_rate=6e-4,
                      lr_decay=True, warmup_tokens=1*150, final_tokens=150*len(train_dataset)*block_size,
                      num_workers=1,string_logs=string_logs, args=args)
              trainer = Trainer(model, train_dataset, validation_dataset, tconf)
  
              output_mode = args.output_mode.split("_")
  
              trainer.dynamic_stream()
  
          
          if args.predict:
  
                  class_info = fold_split(args)
                  block_size = args.block_size # spatial extent of the model for its context
  
                  training_dataset = get_dataloader(args=args,train_val='training',load=True)
                  validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
                  test_dataset = get_dataloader(args=args,train_val='testing',load=True)
  
                  num_class=args.n_class
  
                  mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                  n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
  
                  if args.addposition:
                          mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                          n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=validation_dataset.position_size)
  
                  model = get_model(args,mconf)
  
                  string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                  tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                          lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                          num_workers=20,string_logs=string_logs, args=args)
  
                  trainer = Trainer(model, None,[validation_dataset], tconf)
  
                  output_mode = args.output_mode.split("_")
  
                  if len(output_mode)>1:
                      trainer.predict_multi_stream(len(output_mode))
                  else:
                      trainer.predict()
  
          if args.newpredict:
  
                  class_info = fold_split(args)
                  block_size = args.block_size # spatial extent of the model for its context
  
                  training_dataset = get_dataloader(args=args,train_val='training',load=True)
                  validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
                  test_dataset = get_dataloader(args=args,train_val='testing',load=True)
  
                  num_class=args.n_class
  
                  mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                  n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=validation_dataset.position_size, rt_size = validation_dataset.rt_size)
  
                  model = get_model(args,mconf)
  
                  string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                  tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                          lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                          num_workers=20,string_logs=string_logs, args=args)
  
                  trainer = Trainer(model, None,[validation_dataset], tconf)
  
                  if args.visval:
                      trainer.vis_embed()
  
                  if args.crossdata:
                      trainer.newpredict_dynamic_streamc(args.predictvis)
                  else:
                      trainer.newpredict_dynamic_stream(args.predictvis)
  
          if args.finalpredict:
  
                  class_info = fold_split(args)
                  block_size = args.block_size # spatial extent of the model for its context
  
                  validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
                  train_dataset = get_dataloader(args=args,train_val='training',load=True)
  
                  num_class=args.n_class
             
                  mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=args.n_class, n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,
                  position_size=validation_dataset.position_size, ges_size =  validation_dataset.ges_size,rt_size = validation_dataset.rt_size)
  
                  model = get_model(args,mconf)
  
                  string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                  tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                          lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                          num_workers=20,string_logs=string_logs, args=args)
  
                  trainer = Trainer(model, None,[validation_dataset], tconf)
  
                  if args.vis_attention:
                      trainer = Trainer(model, None,[train_dataset, validation_dataset], tconf) 
  
                      trainer.visualize_attention(args.vis_attention)
  
                  else:                    
                      if args.visval:                   
                          trainer.vis_embed()
                      
                      if args.predictvis:
                          trainer = Trainer(model, None,[train_dataset,validation_dataset], tconf) 
  
                      trainer.finalpredict_dynamic_stream(args.predictvis,args.adddata_dir)
  
          if args.finalpredictnewdata:
  
                  class_info = fold_split(args)
                  block_size = args.block_size # spatial extent of the model for its context
  
                  validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
             
                  mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=args.n_class, n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,
                  position_size=validation_dataset.position_size, ges_size =  validation_dataset.ges_size,rt_size = validation_dataset.rt_size)
  
                  model = get_model(args,mconf)
  
                  string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                  tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                          lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                          num_workers=20,string_logs=string_logs, args=args)
  
                  trainer = Trainer(model, None,[validation_dataset], tconf)
  
                  if args.vis_attention:
                      trainer = Trainer(model, None,[validation_dataset], tconf) 
                      trainer.visualize_attention(args.vis_attention)
  
                  else:                    
                      if args.visval:                   
                          trainer.vis_embed()
                      
                      if args.predictvis:
                          trainer = Trainer(model, None,[validation_dataset], tconf) 
  
                      trainer.finalpredict_newdata(args.predictvis,args.adddata_dir)
  
          if args.newpredict2:
  
                  class_info = fold_split(args)
                  block_size = args.block_size # spatial extent of the model for its context
  
                  training_dataset = get_dataloader(args=args,train_val='training',load=True)
                  validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
  
                  num_class=args.n_class
  
                  mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                  n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=validation_dataset.position_size, ges_size =  validation_dataset.ges_size, rt_size = validation_dataset.rt_size)
  
                  model = get_model(args,mconf)
  
                  string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                  tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                          lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                          num_workers=20,string_logs=string_logs, args=args)
  
                  trainer = Trainer(model, None,[validation_dataset], tconf)
  
                  if args.visval:
                      trainer.vis_embed2()
  
                  trainer.newpredict_dynamic_stream(args.predictvis)
  
          if args.single_predict:
  
                  block_size = args.block_size
                  num_class=args.n_class
  
                  validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
  
                  test_dataset = SinglePrediction(data_dir = args.data_dir)
  
                  mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                  n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
  
                  model = get_model(args,mconf)
                  string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
                  tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-4,
                          lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                          num_workers=20,string_logs=string_logs, args=args)
  
                  trainer = Trainer(model, None,[test_dataset], tconf)
  
                  trainer.single_predict()
  
          if args.vis_weight:
              class_info = fold_split(args)
              block_size = args.block_size # spatial extent of the model for its context
  
              training_dataset = get_dataloader(args=args,train_val='training',load=True)
              validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
              test_dataset = get_dataloader(args=args,train_val='testing',load=True)
  
              num_class=args.n_class
  
              mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
              n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
  
              if args.addposition:
                      mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=validation_dataset.position_size)
  
              model = get_model(args,mconf)
  
              string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
              tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                      lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                      num_workers=20,string_logs=string_logs, args=args)
  
              trainer = Trainer(model, None,[validation_dataset,test_dataset], tconf)
  
              output_mode = args.output_mode.split("_")
  
              if len(output_mode)>1:
                  trainer.predict_vis(len(output_mode))
              else:
                  trainer.predict()
  
          if args.top_weight:
              class_info = fold_split(args)
              block_size = args.block_size # spatial extent of the model for its context
  
              training_dataset = get_dataloader(args=args,train_val='training',load=True)
              validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
              test_dataset = get_dataloader(args=args,train_val='testing',load=True)
  
              num_class=args.n_class
  
              mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
              n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
  
              if args.addposition:
                      mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=validation_dataset.position_size)
  
              model = get_model(args,mconf)
  
              string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
              tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                      lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                      num_workers=20,string_logs=string_logs, args=args)
  
              trainer = Trainer(model, None,[training_dataset,validation_dataset,test_dataset], tconf)
  
              output_mode = args.output_mode.split("_")
  
              if len(output_mode)>1:
                  trainer.topweight_vis(len(output_mode))
              else:
                  trainer.predict()
  
          if args.single_pred_vcf:
  
              args = translate_args(args)
  
              #cmd_preprocess(args)
              preprocessing_fromdmm(args)
  
              pdb.set_trace()
  
              class_info = fold_split(args)
              block_size = args.block_size # spatial extent of the model for its context
  
              training_dataset = get_dataloader(args=args,train_val='training',load=True)
              validation_dataset = get_dataloader(args=args,train_val='validation',load=True)
              test_dataset = get_dataloader(args=args,train_val='testing',load=True)
  
              num_class=args.n_class
  
              mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
              n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256)
  
              if args.addposition:
                      mconf = GPTConfig(vocab_size=validation_dataset.vocab_size, block_size=block_size,num_class=num_class,
                      n_layer=args.n_layer,n_head=args.n_head, n_embd=args.n_emb,context_length=args.context_length,conv_filter=256,position_size=validation_dataset.position_size)
  
              model = get_model(args,mconf)
  
              string_logs = f"{args.tag}_{args.arch}_bs{mconf.block_size:.0f}_nl{mconf.n_layer:.0f}_nh{mconf.n_head:.0f}_ne{mconf.n_embd:.0f}_cl{mconf.context_length:.0f}/"
  
              tconf = TrainerConfig(max_epochs=200, batch_size=1, learning_rate=6e-3,
                      lr_decay=True, warmup_tokens=1*200, final_tokens=200*len(validation_dataset)*block_size,
                      num_workers=20,string_logs=string_logs, args=args)
  
              trainer = Trainer(model, None,[training_dataset,validation_dataset,test_dataset], tconf)
  
              output_mode = args.output_mode.split("_")
  
              if len(output_mode)>1:
                  trainer.topweight_vis(len(output_mode))
              else:
                  trainer.predict()
  
  
  
  
  
           
  import torch
  import torch.nn as nn
  from torch.nn import functional as F
  import math
  from torch.utils.data import Dataset
  import os
  import pandas as pd
  import pdb
  import numpy as np
  import math
  import pickle
  import random
  from sklearn.utils import shuffle
  
  class TCGAPCAWG_Dataloader(Dataset):
  
      def __init__(self, dataset_name = None, 
                  data_dir=None, 
                  mode='training', 
                  curr_fold=1, 
                  block_size=5000, 
                  load=False,
                  addtriplettoken=False,
                  addpostoken=False,
                  addgestoken=False,
                  addrt=False,
                  nummut = 0,
                  frac = 0,
                  crossdata=False,
                  crossdatadir=None,
                  pcawg2tgca_class=False,
                  tcga2pcawg_class=False,
                  mutratio = '1-0-0-0-0',
                  adddatadir = None,
                  input_filename=None,
                  args = None,
                  gx_dir=None,
                  addepigen=False):
  
          self.dataset_name = dataset_name
          self.data_dir=data_dir
          self.mode=mode
          self.curr_fold=int(curr_fold)
          self.block_size=block_size
          self.load=load
          self.addtriplettoken=addtriplettoken
          self.addpostoken=addpostoken
          self.addrt=addrt
          self.nummut = nummut
          self.frac = frac
          self.addgestoken = addgestoken
          self.crossdata= crossdata
          self.crossdatadir = crossdatadir
          self.adddatadir = adddatadir
          self.args = args
          self.gx_dir = gx_dir
          self.mutratio = mutratio
  
          self.pcawg2tgca_class=pcawg2tgca_class
          self.tcga2pcawg_class=tcga2pcawg_class
  
          self.newformat = True
          self.newformat = False
  
          self.NiSi = False
          self.SNV = False
          self.indel = False
          self.SVMEI = False
          self.Normal = False
  
          self.dnn_input = 1
  
          if self.args == None:
              self.single_pred_vcf = False
              self.cwd = str(os.path.abspath('..')) + '/'
          else:
              self.single_pred_vcf = self.args.single_pred_vcf
              self.cwd = self.args.cwd
  
  
          self.input_filename = input_filename
  
          if self.nummut > 0 :
              self.block_size = self.nummut
  
          if self.dataset_name == 'pcawg':
              if self.args.multi_pred_vcf:
                  fulltuple = []
                  for idx in range(len(input_filename)):
                      va = input_filename[idx]
                      onetup = (va[:-4],'',1,1)
                      #print(onetup)
                      fulltuple.append(onetup)
                  self.validation_fold = pd.DataFrame(fulltuple,columns =['samples', 'nm_class', 'slices','fold'])
                  self.test_fold = self.validation_fold
                  self.newformat = True
              
              if self.single_pred_vcf:
                  self.onlyfilename = self.args.input_filename[:-4]
                  onetup = [(self.onlyfilename,'',1,1)]
                  self.validation_fold = pd.DataFrame(onetup,columns =['samples', 'nm_class', 'slices','fold'])
                  self.test_fold = self.validation_fold
                  
                  self.newformat = True
              '''
              else:
                  if self.newformat:
                      self.training_fold = pd.read_csv('./dataset_utils/pcawg_train.csv',index_col=0)
                      self.training_fold = self.training_fold.loc[self.training_fold['fold']==self.curr_fold]
                      self.validation_fold = pd.read_csv('./dataset_utils/pcawg_val.csv',index_col=0)
                      self.validation_fold = self.validation_fold.loc[self.validation_fold['fold']==self.curr_fold]
                  else:
                      self.training_fold = pd.read_csv('./oldformat/pcawg_trainfold' + str(self.curr_fold) + '.csv',index_col=0)
                      self.validation_fold = pd.read_csv('./oldformat/pcawg_valfold' + str(self.curr_fold) + '.csv',index_col=0)
              '''
          elif self.dataset_name == 'tcga':
              self.training_fold = pd.read_csv('./dataset_utils/tcga_trainfold' + str(self.curr_fold) + '.csv',index_col=0)
              self.validation_fold = pd.read_csv('./dataset_utils/tcga_valfold' + str(self.curr_fold) + '.csv',index_col=0)
          elif self.dataset_name == 'westcga':
              self.training_fold = pd.read_csv('./dataset_utils/tcgawes_trainfold' + str(self.curr_fold) + '.csv',index_col=0)
              self.validation_fold = pd.read_csv('./dataset_utils/tcgawes_valfold' + str(self.curr_fold) + '.csv',index_col=0)
          elif self.dataset_name == 'wgspcawg':
              self.training_fold = pd.read_csv('./dataset_utils/pcawgwgs_trainfold' + str(self.curr_fold) + '.csv',index_col=0)
              self.validation_fold = pd.read_csv('./dataset_utils/pcawgwgs_valfold' + str(self.curr_fold) + '.csv',index_col=0)
  
          if self.dataset_name == 'wgsgx':
              self.gx = pd.read_csv(self.gx_dir + 'PCAWG_gene_expression.tsv',sep='\t',index_col=0)
              #self.gx = self.gx.iloc[:,-100:]
  
              self.training_fold = pd.read_csv(self.cwd + 'dataset_utils/wgsgx_train.csv',index_col=0)
              self.training_fold = self.training_fold.loc[self.training_fold['fold'] == self.curr_fold]
  
              self.validation_fold = pd.read_csv(self.cwd + 'dataset_utils/wgsgx_val.csv',index_col=0)
              self.validation_fold = self.validation_fold.loc[self.validation_fold['fold'] == self.curr_fold]  
              self.newformat = True
  
              self.dnn_input = len(self.gx.iloc[0,:-2])
              #pdb.set_trace()
  
          if self.adddatadir is not None:
              adddata = pd.DataFrame(columns=self.validation_fold.columns)
              adddata.columns = self.validation_fold.columns
  
              folder = os.listdir(self.adddatadir)
  
              for i in folder:
  
                  samples = os.listdir(self.adddatadir + i )
                  for j in samples:
                      if j[0:3] == 'new':
                          counter = pd.read_csv(self.adddatadir + i + '/count_new_' + j[4:],index_col=0)
  
                          listall = [i,j[4:]] + counter['0'].values.tolist() + [1]
  
                          pds = pd.DataFrame(listall)
                          pds = pds.T
                          pds.columns=self.validation_fold.columns
  
                          adddata = adddata.append(pds)
  
              adddata = adddata.reset_index(drop=True)
  
              self.adddata = adddata
  
              #self.validation_fold = self.validation_fold.append(self.adddata)
              self.validation_fold = self.adddata
              self.data_dir = self.adddatadir
  
          if self.single_pred_vcf:
              samples_names = input_filename[:-4]
              pd_count = pd.read_csv(args.tmp_dir + 'count_' + input_filename[:-4] + '.csv', index_col=0)['0'].to_list()
              onerow = ['',samples_names] + pd_count + [1]
              pd_data = pd.DataFrame(onerow).T
              pd_data.columns = ['nm_class','samples','NiSi','SNV','indel','SVMEI','Normal','fold']
              self.validation_fold = pd_data
              self.test_fold = self.validation_fold
  
          self.load_classinfo()
  
          self.vocab_mutation = pd.read_csv(self.cwd + 'extfile/dictMutation.csv',index_col=0)
          self.allSNV_index = 0
  
          if self.mutratio is not None:
              self.mutratio = mutratio.split('-')
              self.mutratio = [float(i) for i in self.mutratio]
          
              if self.mutratio[0]>0:
                  self.NiSi = True 
              if self.mutratio[1]>0:
                  self.SNV = True
              if self.mutratio[2]>0:
                  self.indel = True
              if self.mutratio[3]>0:
                  self.SVMEI = True
              if self.mutratio[4]>0:
                  self.Normal = True
  
          vocabsize = 0
          if self.NiSi:
              vocabsize = len(self.vocab_mutation.loc[self.vocab_mutation['typ']=='NiSi'])
              if self.SNV:
                  vocabsize = vocabsize + len(self.vocab_mutation.loc[self.vocab_mutation['typ']=='SNV'])
                  if self.indel:
                      vocabsize = vocabsize + len(self.vocab_mutation.loc[self.vocab_mutation['typ']=='indel'])                   
                      if self.SVMEI:
                          vocabsize = vocabsize + len(self.vocab_mutation.loc[self.vocab_mutation['typ'].isin(['MEI','SV'])])
                          if self.Normal:
                              vocabsize = vocabsize + len(self.vocab_mutation.loc[self.vocab_mutation['typ']=='Normal'])
  
          self.vocab_size = vocabsize + 1
          #print(self.vocab_size)
  
          #pdb.set_trace()
  
          self.pd_position_vocab = pd.read_csv(self.cwd + 'extfile/dictChpos.csv',index_col=0)
          self.pd_ges_vocab = pd.read_csv(self.cwd + 'extfile/dictGES.csv',index_col=0)
  
          self.position_size = len(self.pd_position_vocab) + 1
          self.ges_size = len(self.pd_ges_vocab) + 1
          
          self.rt_size =  1
  
          self.midstring = '.' + self.dataset_name + str(mutratio) + str(int(self.addtriplettoken)) + str(int(self.addpostoken)) +  str(int(self.addgestoken)) +  str(int(self.addrt)) + '/' 
          
          if self.mode == 'validation' or self.mode == 'testing':
              if self.crossdata:
                  os.makedirs(self.crossdatadir + self.midstring, exist_ok=True)
                  self.data_dir = self.crossdatadir
                  #pdb.set_trace()
                  
              else:
                  os.makedirs(self.data_dir + self.midstring, exist_ok=True)
  
      def load_classinfo(self):
          if self.dataset_name == 'pcawg':
              pd_data = pd.read_csv(self.cwd + 'dataset_utils/classinfo_pcawg.csv',index_col = 0)
              self.pd_class_info = pd.DataFrame(pd_data)
          elif self.dataset_name == 'wgsgx':
              pd_data = pd.read_csv(self.cwd + 'dataset_utils/classinfo_wgsgx.csv',index_col = 0)
              self.pd_class_info = pd.DataFrame(pd_data)
          else:
              num_class = os.listdir(self.data_dir)
              name_class = [i for i in num_class if len(i.split('.'))==1]
              name_class = sorted(name_class)
              n_samples = []
              for idx,nm_class in enumerate(name_class):
                  samples = os.listdir(self.data_dir+nm_class)
                  samples = [x for x in samples if x[:10]=='count_new_']
                  n_samples.append(len(samples))
              data = list(zip(name_class, np.arange(len(name_class)),n_samples))  
              self.pd_class_info = pd.DataFrame(data,columns=['class_name','class_index','n_samples'])
  
      def get_data(self,idx):
              
          if self.mode=='training':
              instances=self.training_fold.iloc[idx]     
          elif self.mode=='validation':
              instances=self.validation_fold.iloc[idx]
          elif self.mode == 'testing':
              instances=self.test_fold.iloc[idx]
  
          if self.newformat:
              samples = instances['samples'] + '.csv'
              target_name = instances['nm_class']
  
              if self.single_pred_vcf:
                  pd_row = pd.read_csv(self.data_dir +'/count_' + samples,index_col=0).T
                  row_count = pd_row.values[0]
              else:
                  pd_row = pd.read_csv(self.data_dir + target_name +'/count_' + samples,index_col=0).T
                  row_count = pd_row.values[0]
              
          else:
              target_name = instances['nm_class']
              samples = instances[1]
              row_count = instances[['NiSi','SNV','indel','SVMEI','Normal']].to_numpy()
  
          if self.mutratio is not None:
              avail_count = np.asarray(self.mutratio) * self.block_size      
              
              diff = avail_count - row_count
              pos = diff>0
              avail_count1 = row_count * pos
              diff = row_count > avail_count
  
              avail_count2 = avail_count * diff
              avail_count3 = avail_count1 + avail_count2
              shadowavail_count3 = avail_count3
              shadowavail_count3[0] = row_count[0]
  
              if sum(shadowavail_count3) > self.block_size:
                  diff = self.block_size - sum(avail_count3) 
                  shadowavail_count3[0] = diff + avail_count3[0]
                  
              avail_count2 = shadowavail_count3.astype(int)
  
              if avail_count2[0]<0:
      
                  secondmax = avail_count2[np.argmax(avail_count2)]
                  avail_count2 = avail_count2 * 0.7
  
                  avail_count = avail_count2
  
                  diff = avail_count - row_count
                  pos = diff>0
                  avail_count1 = row_count * pos
                  diff = row_count > avail_count
  
                  avail_count2 = avail_count * diff
                  avail_count3 = avail_count1 + avail_count2
                  shadowavail_count3 = avail_count3
                  shadowavail_count3[0] = row_count[0]
  
                  if sum(shadowavail_count3) > self.block_size:
                      diff = self.block_size - sum(avail_count3) 
                      shadowavail_count3[0] = diff + avail_count3[0]
                      
                  avail_count2 = shadowavail_count3.astype(int)
  
              avail_count = avail_count2
  
      
          def grab(pd_input,grabcol):
              return pd_input[grabcol]
  
          def allgrab(grabcol):
              
  
              if self.NiSi:
                  #pdb.set_trace()
                  if self.newformat:
                      pd_nisi = pd.read_csv(self.data_dir + target_name + '/' + 'SNV_' + samples,index_col=0)
                  else:
                      pd_nisi = pd.read_csv(self.data_dir + target_name + '/' + 'NiSi_new_' + samples,index_col=0)
                  pd_nisi = pd_nisi.sample(n = avail_count[0], replace = False)
                  pd_nisi = grab(pd_nisi,grabcol)
  
                  if self.SNV:
                      if self.newformat:
                          pd_SNV = pd.read_csv(self.data_dir + target_name + '/' + 'MNV_' + samples,index_col=0)
                      else:
                          pd_SNV = pd.read_csv(self.data_dir + target_name + '/' + 'SNV_new_' + samples,index_col=0)
                      pd_SNV = pd_SNV.sample(n = avail_count[1], replace = False)
                      pd_SNV = grab(pd_SNV,grabcol)
                      pd_nisi = pd_nisi.append(pd_SNV)
  
                      if self.indel:
                          pd_indel = pd.read_csv(self.data_dir + target_name + '/' + 'indel_' + samples,index_col=0)
                          pd_indel = pd_indel.sample(n = avail_count[2], replace = False)
                          pd_indel = grab(pd_indel,grabcol)
                          pd_nisi = pd_nisi.append(pd_indel)
                          
                          if self.SVMEI:
                              if self.newformat:
                                  pd_meisv = pd.read_csv(self.data_dir + target_name + '/' + 'MEISV_' + samples,index_col=0)
                              else:
                                  pd_meisv = pd.read_csv(self.data_dir + target_name + '/' + 'MEISV_new_' + samples,index_col=0)
                              pd_meisv = pd_meisv.sample(n = avail_count[3], replace = False)
                              pd_meisv = grab(pd_meisv,grabcol)
                              pd_nisi = pd_nisi.append(pd_meisv)
  
                              if self.Normal:
                                  if self.newformat:
                                      pd_normal = pd.read_csv(self.data_dir + target_name + '/' + 'Neg_' + samples,index_col=0)
                                  else:
                                      pd_normal = pd.read_csv(self.data_dir + target_name + '/' + 'Normal_new_' + samples,index_col=0)
                                  pd_normal = pd_normal.sample(n = avail_count[4], replace = False)
                                  pd_normal = grab(pd_normal,grabcol)
                                  pd_nisi = pd_nisi.append(pd_normal) 
  
                  pd_nisi = pd_nisi.fillna(0)
              return pd_nisi
  
          pd_nisi = pd.DataFrame()
          if self.addtriplettoken:
              if self.mode=='training' :
                  pd_nisi = allgrab(['triplettoken'])
              else:
                  filename = self.data_dir + self.midstring + 'val_' + samples
                  if os.path.isfile(filename):
                      try:
                          pd_nisi = pd.read_csv(filename,index_col=0)
                      except:
                          pd_nisi = allgrab(['triplettoken'])
                          pd_nisi = pd_nisi.dropna()
                          pd_nisi.to_csv(filename)                        
                      
                  else:
                      pd_nisi = allgrab(['triplettoken'])
                      pd_nisi.to_csv(filename)
  
          if self.addpostoken:
              if self.mode=='training' :
                  pd_nisi = allgrab(['triplettoken','postoken'])
              else:
                  #pdb.set_trace()
                  filename = self.data_dir + self.midstring + 'val_' + samples
                  if os.path.isfile(filename):
                      try:
                          pd_nisi = pd.read_csv(filename,index_col=0)
                      except:
                          pd_nisi = allgrab(['triplettoken','postoken'])
                          pdb.set_trace()
                          pd_nisi.to_csv(filename)
                  else:
                      pd_nisi = allgrab(['triplettoken','postoken'])
                      pd_nisi.to_csv(filename)
        
          if self.addgestoken:
              if self.mode=='training' :
                  pd_nisi = allgrab(['triplettoken','postoken','gestoken'])
              else:
                  filename = self.data_dir + self.midstring + 'val_' + samples
                  if os.path.isfile(filename):
                      try:
                          pd_nisi = pd.read_csv(filename,index_col=0)
                      except:
                          pd_nisi = allgrab(['triplettoken','postoken','gestoken'])
                          pd_nisi.to_csv(filename)
  
                  else:
                      pd_nisi = allgrab(['triplettoken','postoken','gestoken'])
                      pd_nisi.to_csv(filename)
  
          if self.addrt:
              if self.mode=='training' :
                   pd_nisi = allgrab(['triplettoken','postoken','gestoken','rt'])
              else:
                  filename = self.data_dir + self.midstring + 'val_' + samples
                  if os.path.isfile(filename):
                      try:
                          pd_nisi = pd.read_csv(filename,index_col=0)
                      except:
                          pd_nisi = allgrab(['triplettoken','postoken','gestoken','rt'])
                          pd_nisi.to_csv(filename)
  
                  else:
                      pd_nisi = allgrab(['triplettoken','postoken','gestoken','rt'])
                      pd_nisi.to_csv(filename)
  
          #pdb.set_trace()
          pd_nisi = pd_nisi.dropna()
          
          if self.nummut > 0:
              if self.nummut < len(pd_nisi):
                  pd_nisi = pd_nisi.sample(n = self.nummut, replace = False)
              else:
                  pd_nisi = pd_nisi.sample(n = len(pd_nisi), replace = False)
                 
          #pdb.set_trace()
  
          if self.frac > 0:
              pd_nisi = pd_nisi.sample(frac = self.frac)
  
          if self.mode =='training':
              pd_nisi = pd_nisi.sample(frac = 1)
  
          #pdb.set_trace()
  
          np_triplettoken = pd_nisi.to_numpy()    
  
          is_padding = False
          if len(pd_nisi) < self.block_size:
              mins = self.block_size - len(np_triplettoken)
              is_padding = True
              
          datanumeric = []
          #pdb.set_trace()
          for i in pd_nisi.columns:
              np_data = pd_nisi[i].to_numpy() 
              if is_padding:
                  np_data = np.copy(np.pad(np_data, ((0, mins)), mode='constant', constant_values=0))
              
              if i == 'rt':
                  tensordata = torch.tensor(np.round(np_data, 1), dtype=torch.half)
                  #tensordata = np.round(np_data, 3)
  
              if len(np_data) > self.block_size:
                  np_data = np.asarray(np_data[:self.block_size],dtype=int)
                  tensordata = torch.tensor(np_data, dtype=torch.long)
              else:
                  np_data = np.asarray(np_data,dtype=int)
                  tensordata = torch.tensor(np_data, dtype=torch.long)
              datanumeric.append(tensordata)
          
          datastring = samples
  
          if self.dataset_name=='wgsgx':
              #pdb.set_trace()
              gx_data = self.gx.loc[self.gx['samples']==samples[:-4]]
              gx_data = gx_data.iloc[:,:-2].values
              tensorgx_data = torch.tensor(gx_data, dtype=torch.float)
  
              datanumeric.append(tensorgx_data)
  
          #print(datanumeric)
          data=[datastring,datanumeric]
          #pdb.set_trace()
  
          if target_name != '':
              if self.crossdata:
                  #pdb.set_trace()
                  target = self.pd_class_infoto.loc[self.pd_class_infoto['class_name']==target_name].class_index.values[0]
              else:
                  target = self.pd_class_info.loc[self.pd_class_info['class_name']==target_name].class_index.values[0]
              target = target.astype(np.int16)
              target = torch.tensor(target, dtype=torch.long)
          else:
              target = ''
  
  
          if self.adddatadir is not None:
              return data,[target,target_name]
          else:    
              return data,target
  
      def __len__(self):
  
          if self.mode=='training':
              return len(self.training_fold)
          elif self.mode=='validation':
              return len(self.validation_fold)
          elif self.mode=='testing':
              return len(self.test_fold)
  
      def __getitem__(self, idx):       
  
          data,target = self.get_data(idx)
  
          return data, target
  
  
  if __name__ == '__main__':
  
      #dataloader = PCAWG(dataset_name = 'PCAWG', data_dir='/csc/epitkane/projects/PCAWG/shuffled_samples/', mode='training',portion = [8,1,1], folds=10, curr_fold=1,load=True,load_token=True)
  
      #dataloader = PCAWG(dataset_name = 'pcawg_mut3_comb0', data_dir='/csc/epitkane/projects/PCAWG20191001/data/modified_data/train/all24classes/', mode='training',portion = [8,1,1], folds=10, curr_fold=1,load=True,load_token=True,ncontext=3,addposition=False,filter=False,topk=5000)
      #dataloaderVal = PCAWG(dataset_name = 'pcawg_mut3_comb0', data_dir='/csc/epitkane/projects/PCAWG20191001/data/modified_data/train/all24classes/', mode='validation',portion = [8,1,1], folds=10, curr_fold=1,load=True,load_token=True,ncontext=3,addposition=False,filter=False,topk=5000)
      #/csc/epitkane/projects/tcga/new23classes/
      #/csc/epitkane/projects/PCAWG20191001/data/modified_data/train/new24classes/
  
      #G:/experiment/data/new24classes/
      '''
      dataloaderVal = FinalTCGAPCAWG(dataset_name = 'finalpcawg', 
                                  data_dir='G:/experiment/data/new24classes/', 
                                  mode='validation', 
                                  curr_fold=1, 
                                  block_size=5000, 
                                  load=False,
                                  mutratio = '0.3-0.3-0.3-0-0',
                                  addtriplettoken=False,
                                  addpostoken=False,
                                  addgestoken=True,
                                  addrt=False,
                                  nummut = 0,
                                  frac = 0,
                                  adddatadir='G:/experiment/data/icgc/')
  
      #pdb.set_trace()
      data,target = dataloaderVal.__getitem__(0)
      pdb.set_trace()
  
      for k in range(0,len(dataloaderVal)):
          print(k)
          data,target = dataloaderVal.__getitem__(k)
      '''
  
  
  
      '''
      WGS GX
      '''
  
      #/scratch/project_2001668/data/pcawg
  
      dataloaderVal = TCGAPCAWG_Dataloader(dataset_name = 'wgsgx', 
                                          data_dir='/scratch/project_2001668/data/pcawg/allclasses/newformat/', 
                                          mode='training', 
                                          curr_fold=1, 
                                          block_size=5000, 
                                          load=False,
                                          addtriplettoken=True,
                                          addpostoken=False,
                                          addgestoken=False,
                                          addrt=False,
                                          nummut = 0,
                                          frac = 0,
                                          mutratio = '1-0-0-0-0',
                                          adddatadir = None,
                                          input_filename=None,
                                          args = None,
                                          gx_dir = '/scratch/project_2001668/data/pcawg/PCAWG_geneexp/')
      
      data,target = dataloaderVal.__getitem__(0)
      pdb.set_trace()
  
      '''
      fold = [1,2,3,4,5,6,7,8,9,10]
      mutratios = ['1-0-0-0-0','0.5-0.5-0-0-0','0.4-0.3-0.3-0-0','0.3-0.3-0.20-0.20-0','0.25-0.25-0.25-0.15-0.1']
  
      retrieve = ['addtriplettoken','addpostoken','addgestoken','addrt']
  
      for fo in fold:
          for i in retrieve:
              if i == 'addtriplettoken':
                  addtriplettoken = True
              else:
                  addtriplettoken = False
              
              if i == 'addpostoken':
                  addpostoken = True
              else:
                  addpostoken = False
  
              if i == 'addgestoken':
                  addgestoken = True
              else:
                  addgestoken = False
  
              if i == 'addrt':
                  addrt = True
              else:
                  addrt = False
  
              for j in mutratios:
                  dataloaderVal = FinalTCGAPCAWG(dataset_name = 'finalpcawg', 
                                      data_dir='G:/experiment/data/new24classes/', 
                                      mode='validation', 
                                      curr_fold=1, 
                                      block_size=5000, 
                                      load=False,
                                      mutratio = j,
                                      addtriplettoken=addtriplettoken,
                                      addpostoken=addpostoken,
                                      addgestoken=addgestoken,
                                      addrt=addrt,
                                      nummut = 0,
                                      frac = 0)
                  for k in range(0,len(dataloaderVal)):
                      print(str(fo) + ' ' + str(k) + ' ' + i + ' ' + j + ' ' + str(addtriplettoken) + str(addpostoken) + str(addgestoken) + str(addrt))
                      data,target = dataloaderVal.__getitem__(k)
      pdb.set_trace()
  
      dataloaderVal = TCGA(dataset_name = 'tcga_emb', data_dir='/csc/epitkane/projects/tcga/all23classes/', mode='validation',portion = [8,1,1], folds=10, curr_fold=1,load=True,load_token=True,ncontext=64,addposition=True,filter=True,block_size=300,withclass=True,twostream=False)
  
      for i in range(len(dataloaderVal)):
          data,target = dataloaderVal.__getitem__(i)
  
      dataloaderVal = TCGA(dataset_name = 'tcga_emb', data_dir='/csc/epitkane/projects/tcga/all23classes/', mode='testing',portion = [8,1,1], folds=10, curr_fold=1,load=True,load_token=True,ncontext=64,addposition=True,filter=True,block_size=300,loaddist=False,withclass=True,twostream=False)
  
      for i in range(len(dataloaderVal)):
          data,target = dataloaderVal.__getitem__(i)
      
      pdb.set_trace()
      '''
  
      
apis: pandas.DataFrame
apis: pandas.read_csv pandas.DataFrame torch.tensor numpy.asarray numpy.argmax numpy.round numpy.pad