# afd_udi.py
#
# Arduino host UDI capturing utility, early pre-alpha
#
# Written by Denis Dratov <volutar@gmail.com>
#
# This is free and unencumbered software released into the public domain.
# See the file COPYING for more details, or visit <http://unlicense.org>.

import sys, struct, glob, serial, argparse, time, math
import pickle
from demfm import DeMFM

tr_peakarr=[]
tr_statscol=[]
tr_statscols=[]
mfm_th1=0
mfm_th2=0

tracks=[]
syncs=[]
sectors=[]
stats=[]
rpm=0
cyls=0
sides=0
ser=None
_CRC=0xffffffff


class Tee(object):
    def __init__(self, *files):
        self.files = files
    def write(self, obj):
        for f in self.files:
            f.write(obj)

def udicrc(buf):
  global _CRC
  for i in range(len(buf)):
    _CRC ^= 0xffffffff ^ buf[i]
    for k in range(8):
      temp=_CRC&1
      _CRC=(_CRC&0x80000000)|_CRC>>1
      if temp: _CRC^=0xedb88320
    _CRC ^= 0xffffffff
  return _CRC


def serial_ports():
  """ Lists serial port names

      :raises EnvironmentError:
          On unsupported or unknown platforms
      :returns:
          A list of the serial ports available on the system
  """
  if sys.platform.startswith('win'):
    ports = ['COM%s' % (i + 1) for i in range(256)]
  elif sys.platform.startswith('linux') or sys.platform.startswith('cygwin'):
    # this excludes your current terminal "/dev/tty"
    ports = glob.glob('/dev/tty[A-Za-z]*')
  elif sys.platform.startswith('darwin'):
    ports = glob.glob('/dev/tty.*')
  else:
    raise EnvironmentError('Unsupported platform')

  result = []
  for port in ports:
    try:
      s = serial.Serial(port)
      s.close()
      result.append(port)
    except (OSError, serial.SerialException):
      pass
  return result

def print_ba(ba):
  for a in ba:
    b=a
    if a==0: c=' '
    elif a<4: c='.'
    elif a<8: c=','
    elif a<16: c='-'
    elif a<32: c='+'
    elif a<64: c='x'
    elif a<128: c='#'
    elif a<192: c='X'
    else: c='%'
    print(c,end='')
  print()


def tr_printstats(stats):
  global tr_statscol,mfm_th1,mfm_th2
  statscol2=[]
  mxc=1
  for a in stats:
    if a>mxc: mxc=a
  for a in range(64):
    statscol2.append(math.ceil(((stats[a]-0.25)/mxc)*16))
  for b in range(8):
    for a in range(64):
     if b==0 and a==mfm_th1//2-8: print('V',end='')
     elif b==0 and a==mfm_th2//2-8: print('V',end='')
     else:
       c=7-b
       if c*2+1==statscol2[a]:
         print('.',end='')
       elif c*2<statscol2[a]:
         print(':',end='')
       else:
         print(' ',end='')
    print()
  print('us------2-------3-------4-------5-------6-------7-------8-------9')


def tr_peakextract(st):
  global tr_peakarr,tr_statscols
  peak=0
  peakval=0
  for a in range(st,64):
    if tr_statscols[a]>peakval:
      peak=a
      peakval=tr_statscols[a]
    delta=peakval-tr_statscols[a]
    if delta>3: break
#  print(f"peak={peak+8}, peakval={peakval}")
  lpk=peak
  pkval=peakval
  while lpk>st:
    lpk-=1
    if tr_statscols[lpk]>pkval: break
    pkval=tr_statscols[lpk]
    tr_statscols[lpk]=0
  rpk=peak
  pkval=peakval
  while rpk<63:
    rpk+=1
    if tr_statscols[rpk]>pkval: break
    pkval=tr_statscols[rpk]
    tr_statscols[rpk]=0
  tr_statscols[peak]=0
  rpk-=1
  tr_peakarr.append(lpk+8)
  tr_peakarr.append(peak+8)
  tr_peakarr.append(rpk+8)
#  print(f"lpk={lpk+8}, rpk={rpk+8}")
  return rpk


def probetrack(trno):
  global ser,tr_statscol,mfm_th1,mfm_th2
#  print("Probing track stats, to figure thresholds..")
  ser.write("S".encode())
  while ser.inWaiting()==0:
    pass
  a=ord(ser.read()) #254 as start if stream, and 255 is the last
  if a!=254:
    print("Errorneous probing stream:",a)
    exit
  stats=bytearray()
#  stats4=bytearray()
  a=0
  inf=0
  col=0
#  print('------16|-----24|-----32|-------------48|-------------64|-----72|')
#  print('us-----2|------3|------4|------5|------6|------7|------8|------9|')
#  print('us------2-------3-------4-------5-------6-------7-------8-------9')
  while inf<64*512:
    rb = ord(ser.read())
    inf+=1
    if inf<=64:
      continue #skip first garbage 64
    if rb==255:
      break
    stats.append(rb)
    a+=1
    col+=1
    if col==64:
      col=0
#      print_ba(stats[a-64:a])

#collect overall track timing stats
  ln=len(stats)
  statsall=[]
  for a in range(64):
    row=0
    c=0
    while (row*64+a)<ln:
      c+=stats[row*64+a] 
      row+=1
    statsall.append(c)

#convert to logarithmic values
  tr_statscol.clear()
  mxc=0
  for a in statsall:
    if a>0: c=math.log2(a+1)
    else: c=0
    if c>mxc: mxc=c
    tr_statscol.append(c)
#    print(format(c,'1.4f'),end='; ')

#  with open(f'statscol-{trno}.txt', 'wb') as f:
#    pickle.dump(tr_statscol,f)
#  f.close()

#smoothing stats
  tr_statscols.clear()
  for a in range(64):
    vl=tr_statscol[a]
    if a-2>0: vl+=tr_statscol[a-2]*0.4
    if a-1>0: vl+=tr_statscol[a-1]*0.8
    if a+1<64: vl+=tr_statscol[a+1]*0.8
    if a+2<64: vl+=tr_statscol[a+2]*0.4
    tr_statscols.append(vl/4)

#extract 3 peaks
  tr_peakarr.clear()
  lpk=0
  lpk=tr_peakextract(lpk)
  lpk=tr_peakextract(lpk)
  lpk=tr_peakextract(lpk)

#extract speed info from peaks
  mfm=[]
  mfm.append(tr_peakarr[4]/tr_peakarr[1])#first ratio =1.5
  mfm.append(tr_peakarr[7]/tr_peakarr[1])#second ratio =2.0
  mfm.append(tr_peakarr[1]/8)            #first peak in us
  mfm.append((tr_peakarr[3]+tr_peakarr[2])/2) #first threshold
  mfm.append((tr_peakarr[6]+tr_peakarr[5])/2) #second threshold
  mfm_th1=int(mfm[3]*2)
  mfm_th2=int(mfm[4]*2)

  tr_printstats(tr_statscol)

  '''
  for a in range(int(mfm[3])-8): print(' ',end='')
  print('|',end='')
  for a in range(int(mfm[4]-mfm[3]-1)): print(' ',end='')
  print('|',end='')
  print()
  '''
#  for a in range(int(64-(mfm[4]-8))): print('-',end='')

#  print(tr_peakarr)

#recreate smoothed stats
  tr_statscols.clear()
  for a in range(64):
    vl=tr_statscol[a]
    if a-2>0: vl+=tr_statscol[a-2]*0.4
    if a-1>0: vl+=tr_statscol[a-1]*0.8
    if a+1<64: vl+=tr_statscol[a+1]*0.8
    if a+2<64: vl+=tr_statscol[a+2]*0.4
    tr_statscols.append(vl/4)

#speed analyze
  if mfm[0]>=1.3 and mfm[0]<=1.7:
    isMFM=True
  else: isMFM=False
  if mfm[2]>=1 and mfm[2]<3:
    MFMspd=1
  elif mfm[2]>=3 and mfm[2]<=5:
    MFMspd=2
  else: MFMspd=-1
  MFMqlt = (1- tr_statscols[tr_peakarr[2]-8] / tr_statscols[tr_peakarr[1]-8])*100

  print('Track type: ',end='')
  if isMFM: print('MFM/',end='')
  else: print('FM(GCR)/',end='')
  if MFMspd==1: print('HD   ',end='')
  elif MFMspd==2: print('DD   ',end='')
  else: print('garbage?    ',end='')
  print(f'Period={mfm[2]:2.2f}us    Quality: {MFMqlt:2.2f}%')

#  print(f"Setting Threshold T{mfm_th1:03d}{mfm_th2:03d}".encode())
  ser.write(f"T{mfm_th1:03d}{mfm_th2:03d}".encode()) #set const threshold
  ser.read()


  '''
  row=0
#  print('Len=',len(stats))
  while True:
    for a in range(64):
      c=0
      for b in range(4):
       if (row*4+b)*64 < len(stats):
         c+=stats[(row*4+b)*64+a]
      if c>250: c=250
      stats4.append(c)
    row+=1
    if row*4*64>=len(stats): break
  print()
  
#  print('Len4=',len(stats4))
  for a in range(0,len(stats4)//(64)):
    print_ba(stats4[a*64:(a+1)*64])
  '''



def exportudi(args):
  global tracks, syncs, sectors, cyls, sides, _CRC

  print ("Exporting UDI file v1.0 \"%s\"..." %args.outfile)
  print ("Cyls:", cyls," Sides:", sides)


  data=bytearray()
  length=0
  track_range=len(tracks)
  for tr in range(track_range):
    tlen=len(tracks[tr])
    length+=(3+tlen+len(syncs[tr]))
    data+=struct.pack("<BH",0,tlen) #0=mfm
    data+=tracks[tr]
    data+=syncs[tr]
#    if len(sectors[tr])!=32: print(sectors[tr])
    print('.',end='')
#  print(length)

  header = struct.pack("<4sI4BI",
    b"UDI!",   # Signature
    length+16, # File length - last 4 bytes of CRC (head+tracks)
    0,         # Version
    cyls-1,    # Cylinders-1 (0x4f=80)
    sides-1,   # Sides-1 (0=1side, 1=2sides, 3..ff-reserved
    0,         # unused
    0)         # EXTHDL size


  _CRC=0xffffffff
  print("\nCalculating CRC...")
  udicrc(header)
  udicrc(data)
  print("Writing file...")
  with open(args.outfile, "wb") as f:
    f.write(header)
    f.write(data)
    f.write(struct.pack("<I",_CRC))


def findArduino(s_ports):
  global ser
  for sprt in s_ports:
    print('  Port: ',sprt,end=' - ',flush=True)
    
    ser=serial.Serial()
    ser.port=sprt
    ser.baudrate=2000000
    ser.timeout=3
    ser.open()
    #probing serial device
    s = ser.read()
    ser.write("?".encode())
    ser.flush()
    time.sleep(0.1)

    s = ser.readline().decode().strip()
    if s[0:3]=='AFD':
      print('AFD device, firmware version',s[3:])
      print()
      return True
    else: #wrong answer
      print('unknown serial device:',s)
      ser.close()
      continue
  return False


def _main(argv):
  global cyls, sides, ser, rpm

  cyls=83
  sides=2
  
  print("Arduino floppy dump utility v0.01b (UDI)    (c) 2019-2020 by Denis Dratov")
  print()

  parser = argparse.ArgumentParser()
  parser.add_argument('outfile',nargs='?', default='*', help='output filename')
  args = parser.parse_args(argv[1:])

  s_ports=serial_ports()
  if s_ports==[]: 
    print('No COM ports found')
    ser=None
  elif not findArduino(s_ports):
    print('COM ports are found, but there''s no AFD device')
    ser=None

  if ser is None: #oops, no AFD device
    exit(-1)

  loaded=False
  try:
    f = open('dump.log', 'w')
    backup = sys.stdout
    sys.stdout = Tee(sys.stdout, f)

    ser.write("+".encode())
    ser.read()
    ser.write("&3".encode())
    period=ser.readline().decode().strip()
    if period=='0':
      print('There''s no index impulse (disk is not inserted)')
    else:
      rpm=60*62500/int(period) #current drive rpm
      print('Rotation speed:',format(rpm,'4.4f'),'rpm')

      ser.write(".".encode())
      ser.read()

#      probetrack()
#      raise 'probed'
      for tr in range(cyls*sides):
        track=tr//2
        side=tr%2
        if side==1:
          sids='['
        else:
          sids=']'

        ser.write(("#"+format(track,'02')).encode())
        ser.read()
        ser.write(sids.encode())
        ser.read()

        print()
        probetrack(tr)
        
        ser.write("R".encode())
        while ser.inWaiting()==0:
          pass

        # first pass - receiving packed track data from capturing device
        bytes=bytearray()
        a=0
        inf=0
        while inf<200000:
          rb = ord(ser.read())
          bytes.append(rb)
          if rb==0:
            break
          a+=1
          inf+=1

        # converting data to MFM stream
        revsize=len(bytes)-256
    #    print('Revolution number of domains = ',revsize)
        revsizel=0
        bytesl=bytearray()
        mfm=0
        mfml=8
        curbyt=0
        for byt in bytes:
    #      print(len(bytesl),end=' ')
          for c in range(4): #4 MFM distances in byte (2bits, 01/10/11)
            bytl=(byt>>6)&0x03
            byt<<=2
    #        print(bytl+1,end='')

            if bytl>0:
              for t in range(bytl):
                mfm>>=1
                mfml-=1
                if mfml==0:
                  bytesl.append(mfm)
                  mfm=0
                  mfml=8
            mfm>>=1
            mfm+=0x80
            mfml-=1
            if mfml==0:
              bytesl.append(mfm)
              mfm=0
              mfml=8

          curbyt+=1
          if curbyt==revsize:
            revsizel=len(bytesl) #store revolution size (because we are getting extra 256 bytes after index)

  #      print('Revolution size in MFM = ',revsizel)

    #    for x in range(0,len(bytesl)):
    #      print(format(bytesl[x],"02X"),end=" ")

        demfm=DeMFM()
        bytesl=bytesl[:revsizel+1]
        demfm.decode_mfm(bytesl)
        tracks.append(demfm.bytes)
        syncs.append(demfm.syncs)
        sectors.append(demfm.cat)
        
        print(f"{tr:03}: ",end="") 
        demfm.print_cat_short()
#        demfm.print_cat()
      
      if len(tracks)>160:
        deleted=0
        for tr in range(len(tracks)-1,159,-1):
          ok=False
          for i in range(len(sectors[tr])):
            if sectors[tr][i][1]>=0xf8 or sectors[tr][i][1]==1 or sectors[tr][i][1]==2:
              ok=True
              break
          if not ok:
            print("Removing unformatted track:",tr)
            deleted+=1
            tracks.pop(tr)
            syncs.pop(tr)
            sectors.pop(tr)
        cyls-=deleted//2
        if deleted%2==1:
          print("Returning uneven empty track")
          tracks.append(bytearray([0x4e,0x4e,0x4e,0x4e,0x4e,0x4e,0x4e,0x4e]))
          syncs.append(bytearray(0))
          sectors.append([])
 
      loaded=True
     
  finally:
    ser.write(".".encode())
    ser.read()
    ser.write("-".encode())
    ser.read()

    sys.stdout = backup

  if loaded:
    if args.outfile=='*':
      args.outfile=input("Enter output filename: ")
    exportudi(args)



  '''
  for i in range(0,len(bytes)-1):
    ch=bytes[i]
    print('%02x '%ch,end='')
  '''



if __name__ == "__main__":
  _main(sys.argv)
