#!/usr/bin/env python3
from tabulate import tabulate
import matplotlib.pyplot as plt

with open('input.txt') as f:
  input_string = f.read().split('\n')

def twos_comp(num, length):
  if num == 0:
    return 0
  return abs((num ^ ((1 << length) - 1)) + 1)

def arithmatic_shiftr(num, length, times):
  for t in range(times):
    num = (num >> 1) | ((1 << length - 1) & num)
  return num

def arithmatic_shiftl(num, length):
  if num & (1 << length - 1):
    return (num << 1) | (1 << length - 1)
  else:
    return (num << 1) & ~(1 << length - 1)


def twoscomp_to_int(num, length):
  if num & (1 << length - 1):
    return (-1 * twos_comp(num, length))
  return num & (1 << length) - 1

def debug(results):
  headers = ['multiplicand bin', 'multiplier bin', 'multiplicand dec', 'multiplier dec', 'expected bin', 'expected dec', 'booth if correct', 'booth mod if correct']
  table = []
  for [multiplicand_bin, multiplier_bin, result_booth, result_booth_mod, length] in results:
    multiplicand = twoscomp_to_int(multiplicand_bin, length)
    multiplier = twoscomp_to_int(multiplier_bin, length)
    expected = multiplicand * multiplier
    expected_bin = (twos_comp(expected, length * 2), expected) [expected > 0]
    success_b = [bin(result_booth), "PASS"] [result_booth == expected_bin]
    success_bm = [bin(result_booth_mod), "PASS"] [result_booth_mod == expected_bin]
    
    table.append([bin(multiplicand_bin), bin(multiplier_bin), multiplicand, multiplier, bin(expected_bin), expected, success_b, success_bm])
  print("\nCHECKS: \n", tabulate(table, headers), "\n")


  
def booth(multiplier, multiplicand, length):
  operations = 0
  multiplicand_twos_comp = twos_comp(multiplicand, length)
  result = multiplier << 1 # extended bit
  for i in range(length):
    op = result & 0b11
    if op == 0b01:
      operations += 1
      result += multiplicand << (length + 1)
    if op == 0b10:
      operations += 1
      result += multiplicand_twos_comp << (length + 1)
    result &= (1 << (length * 2) + 1) - 1 # get rid of any overflows
    result = arithmatic_shiftr(result, (length * 2) + 1, 1)
  result = result >> 1
  return (result, operations)

def booth_mod(multiplier, multiplicand, length):
  operations = 0
  multiplicand |= ((1 << length - 1) & multiplicand) << 1 # extend multiplicand sign to prevent overflow when mult/sub by 2
  multiplicand_twos_comp = twos_comp(multiplicand, length + 1)
  result = multiplier << 1 # extended bit
  for i in range(int((length) / 2)):
    op = result & 0b111
    match op:
      case 0b010 | 0b001: # add
        print("add")
        result += multiplicand << (length + 1)
        operations += 1
      case 0b011:         # add * 2
        print("add * 2")
        result += arithmatic_shiftl(multiplicand, length + 1) << (length + 1) 
        operations += 1
      case 0b100:         # sub * 2
        print("sub * 2")
        result += arithmatic_shiftl(multiplicand_twos_comp, length + 1) << (length + 1)
        operations += 1
      case 0b101 | 0b110: # sub
        print("sub ")
        result += multiplicand_twos_comp << (length + 1)
        operations += 1
    result &= (1 << ((length * 2) + 2)) - 1 # get rid of any overflows
    result = arithmatic_shiftr(result, (length * 2) + 2, 2)
  # *barfs on your dog*
  result = ((result | ((1 << ((length * 2) + 2)) >> 1)) & ((1 << ((length * 2) + 1)) - 1)) >> 1
  return (result, operations)

if __name__ == "__main__":
  result_headers = ['multiplicand', 'multiplier', 'result (bin)', 'result (hex)']
  result_table = []

  opcount_headers = ['multiplicand', 'multiplier', 'length', 'booth', 'modified booth']
  opcount_table = []

  lengths = [] # for matplotlib plot
  ops_booth = []
  ops_mod_booth = []

  debug_results = []

  for operation in input_string:
    if operation == '' or operation[0] == '#':
      continue
    length = len(operation.split(" ")[0])
    multiplicand = int(operation.split(" ")[0], 2)
    multiplier = int(operation.split(" ")[1], 2)

    # get result and operation count of both algorithims
    result_booth = booth(multiplier, multiplicand, length)
    result_mod_booth = booth_mod(multiplier, multiplicand, length)

    # gather data for matplotlib
    ops_booth.append(result_booth[1])
    ops_mod_booth.append(result_mod_booth[1])
    lengths.append(length)
    
    #gather data for report results table
    result_table.append([bin(multiplicand), bin(multiplier), bin(result_booth[0]), hex(result_booth[0])])

    #gather data for test function to check if simulator is working
    debug_results.append([multiplicand, multiplier, result_booth[0], result_mod_booth[0], length])

    #gather data for operation count table
    opcount_table.append([bin(multiplicand), bin(multiplier), length, result_booth[1], result_mod_booth[1]])

  debug(debug_results)
  print(tabulate(result_table, result_headers))
  print(tabulate(opcount_table, opcount_headers))

  
  # generate graph
  plt.plot(lengths, ops_booth, '^--m', label='booths algorithim')
  plt.plot(lengths, ops_mod_booth, 'v--c', label='modified booths algorithim')
  plt.gca().set_xlabel("Length of Operands")
  plt.gca().set_ylabel("Number of Additions and Subtractions")
  plt.legend(loc='upper left')
  plt.show()