# -*- coding: utf-8 -*-
"""deepalu.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1sTG4i6kvVYFm3keyIIkb58BdeXwtIKrS
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow as tf
import pysam
import sys
import random
import os.path
import numpy as np
import datetime, os
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import callbacks
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.io import TFRecordWriter


_TRAIN = '/DeepAlu/DeepAlu_model/train_dataset/train.tfrecord'
_EVAL = '/DeepAlu/DeepAlu_model/train_dataset/eval.tfrecord'
_TEST = '/DeepAlu/DeepAlu_model/train_dataset/test.tfrecord'        
class PileupImageOptions:
  """Creates a PileupImageOptions populated with good default values."""
  out_dir="./"
  model_base_dir="/DeepAlu/DeepAlu_model"
  learning_rate=0.004
  l2=0.001
  #batch_size=256,
  #window_size=21,
  #ref_path='hs37d5.fa.gz',
  #vcf_path='NA12878_calls.vcf.gz',
  #bam_path='NA12878_sliced.bam',
  #out_dir='examples',
  model_dir='ngs_model'
  log_dir='logs'
  reference='reference/Homo_sapiens_assembly37.fasta'
  reference=model_base_dir+'/reference/Homo_sapiens_assembly38.fasta'
  reference=model_base_dir+'/reference/hs37d5.fa'
  reference_hg19=model_base_dir+'/reference/ucsc.hg19.fasta'
  reference_hs37d5=model_base_dir+'/reference/hs37d5.fa'
  alu_reference=model_base_dir+'/reference/ME_all.fa'
  #reference='drive/MyDrive/reference/Homo_sapiens_assembly38.fasta'
  #reference_hg19='drive/MyDrive/reference/ucsc.hg19.fasta'
  #alu_reference='drive/MyDrive/reference/ME_all.fa'
  support_alt_reads='support_alt_reads'
  regions='regions'
  split_softclipped='/DeepAlu/data_cluster/split_softclipped_sort'
  discord_read_bwa='discord_read_bwa'
#  support_alt_reads='support_alt_reads_NA12878'
#  regions='NA12878_valiate'
  reference_band_height=5
  base_color_offset_a_and_g=40
  base_color_offset_t_and_c=30
  base_color_stride=70
  allele_supporting_read_alpha=1.0
  allele_unsupporting_read_alpha=0.6
  other_allele_supporting_read_alpha=0.6
  reference_matching_read_alpha=0.2
  reference_mismatching_read_alpha=1.0
  indel_anchoring_base_char='*'
  reference_alpha=0.4
  total_epochs=1000
  height=100
  batch_size=300
  reference_base_quality=60
  positive_strand_color=70
  negative_strand_color=240
  base_quality_cap=40
  mapping_quality_cap=60
  min_base_quality=8
  min_mapping_quality=7
  #height=dv_constants.PILEUP_DEFAULT_HEIGHT,
  #width=dv_constants.PILEUP_DEFAULT_WIDTH,
  num_channels=6
  read_overlap_buffer_bp=5
  #read_requirements=read_requirements
  #multi_allelic_mode=deepvariant_pb2.PileupImageOptions.ADD_HET_ALT_IMAGES,
  # Fixed random seed produced with 'od -vAn -N4 -tu4 < /dev/urandom'.
  random_seed=2101079370
  #sequencing_type=deepvariant_pb2.PileupImageOptions.UNSPECIFIED_SEQ_TYPE,
  alt_aligned_pileup='none'
  types_to_alt_align='indels'
  min_non_zero_allele_frequency=0.00001
  use_allele_frequency=False
  kMaxPixelValueAsFloat = 254.0
  window_size=350
  def __init__(self,a):
      self.name=a
options_=PileupImageOptions(3)
print(options_.base_color_stride)


class ImageRow:
  base=list()
  base_quality=list()
  mapping_quality=list()
  on_positive_strand=list()
  supports_alt=list()
  matches_ref=list()
  op_length=list()
  sequencing_type=list()
  num_channels=list()
  null_size=0
  height=0
  def __init__(self):
    self.base=list()        
    self.base_quality=list()        
    self.mapping_quality=list()        
    self.supports_alt=list()       
    self.matches_ref=list()    
    self.on_positive_strand=list()   
    self.height=0 
  def place_hold(self,ishead,size_t):
    if ishead:
      self.base=[0]*size_t        
      self.base_quality=[0]*size_t        
      self.mapping_quality=[0]*size_t        
      self.supports_alt=[0]*size_t       
      self.matches_ref=[0]*size_t
      self.on_positive_strand=[0]*size_t
    else:
      if size_t>0 :
        self.base.extend([0]*size_t)        
        self.base_quality.extend([0]*size_t)        
        self.mapping_quality.extend([0]*size_t)        
        self.supports_alt.extend([0]*size_t)        
        self.matches_ref.extend([0]*size_t)
        self.on_positive_strand.extend([0]*size_t)     
      else  :
        self.base=self.base[:351]        
        self.base_quality=self.base_quality[:351]        
        self.mapping_quality=self.mapping_quality[:351]        
        self.supports_alt=self.supports_alt[:351]        
        self.matches_ref=self.matches_ref[:351]        
        self.on_positive_strand=self.on_positive_strand[:351]        

class DV_call:
  break_pos=''
  image_start_pos=''
  region_reads=''
  reads_mapRef=''
  reads_mapClipL=''
  reads_mapClipR=''
  reads_mapDiscord=''
  reference_seq=''
  alu_reference_seq=''
  pos=''
  def __init__(self,options_,chr,pos,sample):
  #  print(chr+':'+pos)
    self.pos=pos
    self.break_pos=str(chr)+":"+str(pos)
    self.image_start_pos=int(pos)-176
#    print(str(sample)+'_'+str(chr)+":"+str(int(pos)-50)+"-"+str(int(pos)+50)+"_mapClipR.sam")
    self.reads_mapRef=pysam.AlignmentFile(options_.split_softclipped+'/'+str(sample)+'_'+str(chr)+":"+str(int(pos)-50)+"-"+str(int(pos)+50)+"_mapRef.sam", "rb",check_sq=False)
    self.reads_mapClipL=pysam.AlignmentFile(options_.split_softclipped+'/'+str(sample)+'_'+str(chr)+":"+str(int(pos)-50)+"-"+str(int(pos)+50)+"_mapClipL.sam", "rb",check_sq=False)
    self.reads_mapClipR=pysam.AlignmentFile(options_.split_softclipped+'/'+str(sample)+'_'+str(chr)+":"+str(int(pos)-50)+"-"+str(int(pos)+50)+"_mapClipR.sam", "rb",check_sq=False)
  #  self.reads_mapDiscord=pysam.AlignmentFile(options_.discord_read_bwa+'/'+str(sample)+'_'+str(chr)+":"+str(int(pos)-50)+"-"+str(int(pos)+50)+"_discord.sam", "rb",check_sq=False)
    self.reference_seq=pysam.FastaFile(options_.reference)
    self.alu_reference_seq=pysam.FastaFile(options_.alu_reference)
    if sample=='HG002':
      self.reference_seq=pysam.FastaFile(options_.reference_hs37d5)

def BaseColor(base):
  if len(base)>1:
    case=base[0]
  else:
    case=base
  if case=='I':
	  return (10 + options_.base_color_stride * 2);
  elif case== 'D': 
	  return (60 + options_.base_color_stride * 1);
  elif case== 'A' or case== 'a': 
    return (options_.base_color_offset_a_and_g +
                      options_.base_color_stride * 3);
  elif case== 'G' or case== 'g': 
    return (options_.base_color_offset_a_and_g +
                      options_.base_color_stride * 2);
  elif case== 'T' or case== 't': 
    return (options_.base_color_offset_t_and_c +
                      options_.base_color_stride * 1)
  elif case== 'C' or case== 'c': 
    return (options_.base_color_offset_t_and_c +
                      options_.base_color_stride * 0);
  else:	
    return 0;
def MatchesRefColor(base_matches_ref):
	alpha = options_.reference_matching_read_alpha if base_matches_ref else options_.reference_mismatching_read_alpha
	return int(options_.kMaxPixelValueAsFloat * alpha)

def SupportsAltColor(read_supports_alt):
  alpha = options_.allele_supporting_read_alpha if read_supports_alt else options_.allele_unsupporting_read_alpha
  return int(options_.kMaxPixelValueAsFloat * alpha);

def BaseQualityColor(base_qual):
	capped = float(min(options_.base_quality_cap, base_qual))
	return int(options_.kMaxPixelValueAsFloat *(capped / options_.base_quality_cap))

def MappingQualityColor(mapping_qual):
	capped = float(min(options_.mapping_quality_cap, mapping_qual))
	return int(options_.kMaxPixelValueAsFloat *(capped / options_.mapping_quality_cap))

def StrandColor(on_positive_strand):
  return options_.positive_strand_color if on_positive_strand else options_.negative_strand_color

BaseColor_value=MatchesRefColor(True)
print(BaseColor_value)



def action_per_cigar_unit(img_row,read,ref_i,read_i,cigar_op,image_start_pos,alt_color,mapping_color,strand_color,ref_bases):
  read_base = 0;
  if cigar_op == 1:
  #  insert
      read_base = 'I'
  elif cigar_op == 2:
  #del 
      ref_i -= 1;  # Adjust anchor base on reference
      read_base = 'D'
  elif cigar_op == 0 or cigar_op == 7 or cigar_op == 8: 
      read_base = read.query_sequence[read_i]
  elif cigar_op == 4 : 
      read_base = read.query_sequence[read_i]
  if read_base !=0 :
      base_quality = read.query_qualities[read_i]
 #     if ref_i == dv_call.variant().start() and base_quality < min_base_quality :
 #       return false;
      matches_ref = False
#      print('base_len:'+str(len(ref_bases))+",ref_i:"+str(ref_i))
      if read_base == ref_bases[ref_i] or read_base==ref_bases[ref_i].upper():
        matches_ref=True

      img_row.base.append(BaseColor(read_base))
      img_row.base_quality.append(BaseQualityColor(base_quality))
      img_row.mapping_quality.append(mapping_color)
      img_row.on_positive_strand.append(strand_color)
      img_row.supports_alt.append(alt_color)
      img_row.matches_ref.append(MatchesRefColor(matches_ref))
 #     if img_row.use_allele_frequency:
 #       img_row.allele_frequency[col] = allele_frequency_color
  return True

def encode_ref(refbase):
  base_quality_color= BaseQualityColor(options_.reference_base_quality) 
  mapping_quality_color = MappingQualityColor(options_.reference_base_quality)
  strand_color = StrandColor(True)
  alt_color = SupportsAltColor(False)
  match_ref_color=MatchesRefColor(True)
  img_row=ImageRow()
  #print(strand_color)
  min_base_quality = options_.min_base_quality
  min_mapping_quality =options_.min_mapping_quality
  for ref_i in range(0,len(refbase)):
    img_row.base.append(BaseColor(refbase[ref_i]))
    img_row.base_quality.append(base_quality_color)
    img_row.mapping_quality.append(mapping_quality_color)
    img_row.on_positive_strand.append(strand_color)
    img_row.supports_alt.append(alt_color)
    img_row.matches_ref.append(match_ref_color)
  return img_row

def EncodeRead(map_type,clip_pad,alu_call,read,image_start_pos,debug_mode):
  mapping_quality = read.mapping_quality
  is_forward_strand = False if read.is_reverse else True
  #print(is_forward_strand)
  if map_type == -1:
    supports_alt=True
#  alt_color = SupportsAltColor(supports_alt)
  mapping_color = MappingQualityColor(mapping_quality)
  strand_color = StrandColor(is_forward_strand)
  #print(strand_color)
  min_base_quality = options_.min_base_quality
  min_mapping_quality =options_.min_mapping_quality
  clip_start=0
  clip_end=0
#  print(alu_call.pos+";"+read.query_name+";"+str(read.cigartuples))
  if mapping_quality < min_mapping_quality:
    return False
  if read.cigartuples[0][0]==4 :
    clip_start=read.cigartuples[0][1]
  if read.cigartuples[-1][0]==4 :
    clip_end=read.cigartuples[-1][1]
  #print(ref_bases)
  left_cut_num=0
  if (int(alu_call.pos)-176)>(int(read.reference_start)-clip_start):
    left_cut_num=(int(alu_call.pos)-176)-(read.reference_start-clip_start)
  if map_type == 3 :
    if (int(read.reference_start)-clip_start)>=0:
      left_cut_num=0
    else:
      left_cut_num=clip_start-int(read.reference_start)
  if map_type != 3:
#    print(read.reference_end)
#    print(read.reference_name)
    ref_bases=alu_call.reference_seq.fetch(reference=read.reference_name,
                                           start=read.reference_start-(clip_start-left_cut_num),
                                           end=read.reference_end+clip_end)
  else:
    start=read.reference_start-clip_start
    if start <0:
      start=0
    ref_bases=alu_call.alu_reference_seq.fetch(reference=read.reference_name,
                                               start=start,
                                               end=read.reference_end+clip_end)   
#    print(read.cigartuples)
  for tmp_i in range(0,len(ref_bases)):
    if ref_bases[tmp_i]=='a':
      ref_bases[tmp_i]=='A' 
    if ref_bases[tmp_i]=='t':
      ref_bases[tmp_i]=='T' 
    if ref_bases[tmp_i]=='c':
      ref_bases[tmp_i]=='C' 
    if ref_bases[tmp_i]=='g':
      ref_bases[tmp_i]=='G' 

  img_row=ImageRow()
  place_hold_left = read.reference_start-clip_start-image_start_pos;
  if map_type == 3:
    place_hold_left = read.reference_start-clip_start;
  if place_hold_left>0:
    img_row.place_hold(True,place_hold_left)

  read_i = 0;
  ref_i = 0;
  ok = True;
  op_ig=0

  for op,op_len in read.cigartuples:
    if read_i > int(options_.window_size)/2 or read_i >=len(ref_bases) or ref_i >=len(ref_bases):
        break 
    if op==0 or op==8 or op==9 or op==4 :
        # Alignment and softclipped op.
      for i in range(0,op_len) :
        if read_i > int(options_.window_size)/2 or read_i >=len(ref_bases) or ref_i >=len(ref_bases):
            break 
        if read_i < left_cut_num:
          if map_type == 3:
              place_hold_left = read.reference_start-clip_start;
          read_i=read_i+1
          continue 
        if map_type == 1 and read_i <clip_pad :
          supports_alt=True
        elif map_type == 2 and read_i >=len(read.query_sequence)-clip_pad :
          supports_alt=True
        elif map_type == 3:
          supports_alt=True
        else:
          supports_alt=False
        alt_color=SupportsAltColor(supports_alt)
        ok = ok and action_per_cigar_unit(img_row,read,ref_i,read_i,op,image_start_pos,alt_color,mapping_color,strand_color,ref_bases)
        ref_i=ref_i+1
        read_i=read_i+1
    elif op==1:
        # Insert op.
      if map_type == 1 and read_i <clip_pad :
        supports_alt=True
      elif map_type == 2 and read_i >=clip_pad :
        supports_alt=True
      elif map_type==3:
        supports_alt=True
      else:
        supports_alt=False
      alt_color=SupportsAltColor(supports_alt)
      ok = action_per_cigar_unit(img_row,read,ref_i-1,read_i,op,image_start_pos,alt_color,mapping_color,strand_color,ref_bases)
      read_i += op_len;
    elif op==2 or op==3 :
        # Delete op.
      for i in range(0,op_len) :
        if ref_i >=len(ref_bases):
            break 
        if map_type == 1 and read_i <clip_pad :
          supports_alt=True
        elif map_type == 2 and read_i >=clip_pad :
          supports_alt=True
        elif map_type==3:
          supports_alt=True
        else:
          supports_alt=False
        alt_color=SupportsAltColor(supports_alt)
        ok = action_per_cigar_unit(img_row,read,ref_i,read_i - 1,op,image_start_pos,alt_color,mapping_color,strand_color,ref_bases)
        ref_i= ref_i+1;
    elif op==5 or op==6 :
      op_ig=op_ig+1
        # Hardclipped and pad ops.  Do nothing.
      #  print("Ignored ops")       
    else :
      print(str(op)+":Unrecognized CIGAR op")

    # Bail out if we found this read had a low-quality base at the
    # call site.
    if not ok:
      return False
  #print("base_len:"+str(len(img_row.base)))
  #print(img_row.base)
  if len(img_row.base) == 0:
    return False
  place_hold_right=351-len(img_row.base)
  if place_hold_right>0:
#    if map_type == 3:
#       print('place_hold_right:'+str(place_hold_right))
    img_row.place_hold(False,place_hold_right)
  if place_hold_right<0:
    img_row.place_hold(False,place_hold_right)
 #img_row.place_hold(False,col)
 # print("base:"+str(len(img_row.base) ))
  return img_row


def make_alu_examples(hparams,genotype_file,debug_mode=0,shuffle_depth=95):
  record=[]
  print('make Alu trfrecord ....')
  record_num=0
  with open(genotype_file, 'r', encoding='utf-8') as f:
    for record_each in f:
      record_num=record_num+1
#      if record_num >100:
#        break
      chr,pos,gt,data_type=record_each.strip().split('\t')
      print(record_each)
      if not os.path.isfile(hparams.split_softclipped+'/'+str(data_type)+'_'+str(chr)+":"+str(int(pos)-50)+"-"+str(int(pos)+50)+"_mapRef.sam"):
        print("Do not find sam file:"+hparams.split_softclipped+'/'+str(data_type)+'_'+str(chr)+":"+str(int(pos)-50)+"-"+str(int(pos)+50)+"_mapRef.sam")
        sys.exit()
#        continue
#      if data_type=='NA12878' and gt==0:
#        continue
      #record_num=record_num+1
      if record_num%100==0:
        print("Generate tfrecord :{0}%,{1} /26110".format(round((record_num + 1) * 100 / 26110),record_num), end="\n")
#      if record_num >500:
#        break
      #print(data_type)
      img_total=ImageRow()
      alu_call=DV_call(hparams,chr,pos,data_type)
      img_row_num=0
      ref_bases=alu_call.reference_seq.fetch(reference=chr,start=int(pos)-176,end=int(pos)+175)
      ref_rows = ([encode_ref(ref_bases)] *options_.reference_band_height)
      for img_row_each in ref_rows:
        img_row_num=img_row_num+1
        img_total.base.extend(img_row_each.base)
        img_total.base_quality.extend(img_row_each.base_quality)
        img_total.mapping_quality.extend(img_row_each.mapping_quality)
        img_total.on_positive_strand.extend(img_row_each.on_positive_strand)
        img_total.supports_alt.extend(img_row_each.supports_alt)
        img_total.matches_ref.extend(img_row_each.matches_ref)
        img_total.height=img_total.height+1

       
      example = tf.train.Example()
      features = example.features
      for read in alu_call.reads_mapRef:
        image_start_pos=alu_call.image_start_pos
        img_row_each=EncodeRead(-1,0,alu_call,read,image_start_pos,debug_mode)
        if img_row_each:
          img_row_num=img_row_num+1
          img_total.base.extend(img_row_each.base)
          img_total.base_quality.extend(img_row_each.base_quality)
          img_total.mapping_quality.extend(img_row_each.mapping_quality)
          img_total.on_positive_strand.extend(img_row_each.on_positive_strand)
          img_total.supports_alt.extend(img_row_each.supports_alt)
          img_total.matches_ref.extend(img_row_each.matches_ref)
          img_total.height=img_total.height+1
 
      for read in alu_call.reads_mapClipR:
        clip_pad=read.next_reference_start
        #print(support_alt)
        image_start_pos=alu_call.image_start_pos
        img_row_each=EncodeRead(2,clip_pad,alu_call,read,image_start_pos,debug_mode)
       #print(img_row_each.base)
        #sys.exit()
        if img_row_each:
          img_row_num=img_row_num+1
          img_total.base.extend(img_row_each.base)
          img_total.base_quality.extend(img_row_each.base_quality)
          img_total.mapping_quality.extend(img_row_each.mapping_quality)
          img_total.on_positive_strand.extend(img_row_each.on_positive_strand)
          img_total.supports_alt.extend(img_row_each.supports_alt)
          img_total.matches_ref.extend(img_row_each.matches_ref)
          img_total.height=img_total.height+1
      #    print(img_row_each.supports_alt)

      for read in alu_call.reads_mapClipL:
        clip_pad=read.next_reference_start
        #print(support_alt)
        image_start_pos=alu_call.image_start_pos
        img_row_each=EncodeRead(1,clip_pad,alu_call,read,image_start_pos,debug_mode)
        if img_row_each:
          img_row_num=img_row_num+1
          img_total.base.extend(img_row_each.base)
          img_total.base_quality.extend(img_row_each.base_quality)
          img_total.mapping_quality.extend(img_row_each.mapping_quality)
          img_total.on_positive_strand.extend(img_row_each.on_positive_strand)
          img_total.supports_alt.extend(img_row_each.supports_alt)
          img_total.matches_ref.extend(img_row_each.matches_ref)
          img_total.height=img_total.height+1
      #    print(img_row_each.supports_alt)


      for read in alu_call.reads_mapDiscord:
        clip_pad=0
        #print(support_alt)
        image_start_pos=alu_call.image_start_pos
#        print('img_row_num:'+str(img_row_num))
        img_row_each=EncodeRead(3,clip_pad,alu_call,read,image_start_pos,debug_mode)
        #print(img_row_each.base)
        #sys.exit()
        if img_row_each:
          img_row_num=img_row_num+1
          img_total.base.extend(img_row_each.base)
          img_total.base_quality.extend(img_row_each.base_quality)
          img_total.mapping_quality.extend(img_row_each.mapping_quality)
          img_total.on_positive_strand.extend(img_row_each.on_positive_strand)
          img_total.supports_alt.extend(img_row_each.supports_alt)
          img_total.matches_ref.extend(img_row_each.matches_ref)
          img_total.height=img_total.height+1

      feature_total=[]
      img_total.base=np.array(img_total.base).reshape((img_total.height,-1))
      img_total.base_quality=np.array(img_total.base_quality).reshape((img_total.height,-1))
      img_total.mapping_quality=np.array(img_total.mapping_quality).reshape((img_total.height,-1))
      img_total.on_positive_strand=np.array(img_total.on_positive_strand).reshape((img_total.height,-1))
      img_total.supports_alt=np.array(img_total.supports_alt).reshape((img_total.height,-1))
      img_total.matches_ref=np.array(img_total.matches_ref).reshape((img_total.height,-1))

      if img_total.height-5 >shuffle_depth:
        rs = np.array(random.sample(range(5, img_total.height), shuffle_depth))
        rs=np.append(rs,[0,1,2,3,4])
        rs.sort()
        img_total.base=img_total.base[rs,]
        img_total.base_quality=img_total.base_quality[rs,]
        img_total.mapping_quality=img_total.mapping_quality[rs,]
        img_total.on_positive_strand=img_total.on_positive_strand[rs,]
        img_total.supports_alt=img_total.supports_alt[rs,]
        img_total.matches_ref=img_total.matches_ref[rs,]
        img_total.height=shuffle_depth
      img_raw_len=len(img_total.base)
      for i in range(img_raw_len,hparams.height):
        img_total.height=img_total.height+1
        img_row_num=img_row_num+1
        img_total.base=np.concatenate((img_total.base,np.array([[0]*351])),axis=0)
        img_total.base_quality=np.concatenate((img_total.base_quality,np.array([[0]*351])),axis=0)
        img_total.mapping_quality=np.concatenate((img_total.mapping_quality,np.array([[0]*351])),axis=0)
        img_total.on_positive_strand=np.concatenate((img_total.on_positive_strand,np.array([[0]*351])),axis=0)
        img_total.supports_alt=np.concatenate((img_total.supports_alt,np.array([[0]*351])),axis=0)
        img_total.matches_ref=np.concatenate((img_total.matches_ref,np.array([[0]*351])),axis=0)
      if len(img_total.base) >100:
        print(np.array(img_total.base).shape)
        break
      feature_total.append([img_total.base])
      feature_total.append([img_total.base_quality])
      feature_total.append([img_total.mapping_quality])
      feature_total.append([img_total.on_positive_strand])
      feature_total.append([img_total.supports_alt])
      feature_total.append([img_total.matches_ref])
      feature_total=np.stack(feature_total,axis=-1).flatten()

      example = tf.train.Example(
        features=tf.train.Features(
            feature={
                'img_data' : tf.train.Feature(float_list=tf.train.FloatList(value=feature_total)),
                'label' : tf.train.Feature(int64_list=tf.train.Int64List(value=[int(gt)])),
                'chr' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[bytes(chr,'utf-8')])),
                'pos' : tf.train.Feature(int64_list=tf.train.Int64List(value=[int(pos)])),
                'height' : tf.train.Feature(int64_list=tf.train.Int64List(value=[img_total.height]))
            })).SerializeToString()

      yield data_type,example


def generate_tfrecord_datasets(hparams):
  """Writes out TFRecords files for training, evaluation, and test datasets."""
  if not os.path.exists(hparams.out_dir):
    os.makedirs(hparams.out_dir)

  # Fraction of examples in each dataset.
  train_eval_test_split = [0.7, 0.2, 0.1]
  num_train_examples = 0
  num_eval_examples = 0
  num_test_examples = 0
  # Generate training, test, and evaluation examples.
  with TFRecordWriter(os.path.join(hparams.out_dir, _TRAIN)) as train_out, \
       TFRecordWriter(os.path.join(hparams.out_dir, _EVAL)) as eval_out, \
       TFRecordWriter(os.path.join(hparams.out_dir, _TEST)) as test_out:
#    hparams.reference='reference/Homo_sapiens_assembly38.fasta'
    all_examples = make_alu_examples(hparams,'input_gt.txt')
  
    #print(all_examples.on_positive_strand[0])
    for data_type,example in all_examples:
      #print(example.on_positive_strand[0])
      #break
      ran_i=random.random()
      if data_type == 'HG002':
        test_out.write(example)
        num_test_examples += 1
      if ran_i <0.2:
        eval_out.write(example)
        num_eval_examples += 1
      else:
        train_out.write(example)
        num_train_examples += 1

  print('# of training examples: %d' % num_train_examples)
  print('# of evaluation examples: %d' % num_eval_examples)
  print('# of test examples: %d' % num_test_examples)

def get_dataset(hparams, filename, num_epochs):
  """Reads in and processes the TFRecords dataset.

  Builds a pipeline that returns pairs of features, label.
  """

  # Define field names, types, and sizes for TFRecords.
  proto_features = {
      'img_data':
          tf.io.FixedLenSequenceFeature(shape=[], dtype=tf.float32,allow_missing=True),
      'label':
          tf.io.FixedLenFeature(shape=[1], dtype=tf.int64),
      'chr':
          tf.io.FixedLenFeature(shape=(), dtype=tf.string),
      'pos':
          tf.io.FixedLenFeature(shape=[1], dtype=tf.int64),
      'height':
          tf.io.FixedLenFeature(shape=[1], dtype=tf.int64),
  }

  def _process_input(proto_string):
    """Helper function for input function that parses a serialized example."""

    parsed_features = tf.io.parse_single_example(
        serialized=proto_string, features=proto_features)

    feature_columns = []
    features=[]
    feature_key=[
                 'base',
                 'base_quality',
                 'mapping_quality',
                 'on_positive_strand',
                 'supports_alt',
                 'matches_ref'
                 ]
    label = parsed_features['label']
    chr= parsed_features['chr']
    pos= parsed_features['pos']
    features = tf.reshape(parsed_features['img_data'],[hparams.height,351,6])
    return features,label

  ds = tf.data.TFRecordDataset(filenames=filename)
  ds = ds.map(map_func=_process_input)
  return ds

