#!/usr/bin/env python3
#
# Copyright (C) 2021 Jonathan A. Poritz
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
import argparse
import numpy as np
import math
from PIL import Image
parser = argparse.ArgumentParser(description='Show a square of pixels from a genetic sequence.')
parser.add_argument("--input_file", help="File that contains the genetic sequence; should have no characters other than the genetic data, including no spaces or newlines.  Default is 'sequence'.", default="sequence")
parser.add_argument("--output_file", help="Filename for the output image.  Default is 'output.jpg'.", default="output.jpg")
parser.add_argument("--color", help="Make a color representation of the genetic sequence.  Default is 'True'.", dest='use_color', action='store_true')
parser.add_argument("--bw", help="Make a B&W representation of the genetic sequence.  Default is 'False'.", dest='use_color', action='store_false')
parser.set_defaults(use_color=True)
args = parser.parse_args()
f=open(args.input_file,"r")
s=f.read()
seq_len=len(s)
if args.use_color:
  pix_len=math.ceil(seq_len/12)
else:
  pix_len=math.ceil(seq_len*2)
full_pix_seq_len=seq_len-(seq_len%12)
sqr=math.ceil(math.sqrt(pix_len))
conv={'a':0,
      'A':0,  # sequence data can be in caps or lowercase
      't':1,
      'T':1,
      'u':1,  # works with DNA or RNA sequences
      'U':1,
      'c':2,
      'C':2,
      'g':3,
      'G':3}
palette = np.zeros((sqr,sqr,3),dtype=np.uint8)
index = 0
print("image will be a square of side "+str(sqr))
if args.use_color:
  for i in range(sqr):
    for j in range(sqr):
      for k in range(3):
        color_val=0
        for l in range(4):
          color_val = 4*color_val + conv[s[index]]
          index += 1
          if index == seq_len:
            break
        palette[i,j,k] = color_val
        if index == seq_len:
          break
      if index == seq_len:
        break
    if index == seq_len:
      break
else:
  toggle = False
  for i in range(sqr):
    for j in range(sqr):
      toggle = not toggle
      if toggle:
        if conv[s[index]]>1:
          palette[i,j,0] = palette[i,j,1] = palette[i,j,2] = 255
      else:
        if conv[s[index]]%2:
          palette[i,j,0] = palette[i,j,1] = palette[i,j,2] = 255
        index += 1
      if index == seq_len:
        break
    if index == seq_len:
      break
pil_img=Image.fromarray(palette)
pil_img.save(args.output_file)