#!/usr/bin/env cctbx.python

# code to mend BKGINIT.py under special circumstances

import os
import sys
import binascii
import copy

from scitbx.array_family import flex
from cbflib_adaptbx import uncompress, compress

def recompute_BKGINIT(bkginit_in, init_lp, bkginit_out):

  start_tag = binascii.unhexlify('0c1a04d5')

  data = open(bkginit_in, 'rb').read()
  data_offset = data.find(start_tag) + 4
  cbf_header = data[:data_offset - 4]

  fast = 0
  slow = 0
  length = 0

  for record in cbf_header.split('\n'):
    if 'X-Binary-Size-Fastest-Dimension' in record:
      fast = int(record.split()[-1])
    elif 'X-Binary-Size-Second-Dimension' in record:
      slow = int(record.split()[-1])
    elif 'X-Binary-Number-of-Elements' in record:
      length = int(record.split()[-1])

  assert(length == fast * slow)

  pixel_values = uncompress(packed = data[data_offset:],
                            fast = fast, slow = slow)

  untrusted = []

  for record in open(init_lp):
    if 'UNTRUSTED_RECTANGLE=' in record:
      untrusted.append(map(int, record.replace('.', ' ').split()[1:5]))

  modified_pixel_values = copy.deepcopy(pixel_values)

  for s in range(5, slow - 5):
    y = s + 1
    for f in range(5, fast - 5):
      x = f + 1
      trusted = True
      for x0, x1, y0, y1 in untrusted:
        if (x >= x0) and (x <= x1) and (y >= y0) and (y <= y1):
          trusted = False
          break

      if trusted:
        pixel = pixel_values[s * fast + f]
        if pixel < 0:
          pixels = []
          for j in range(-2, 3):
            for i in range(-2, 3):
              p = pixel_values[(s + j) * fast + f + i]
              if p > 0:
                pixels.append(p)
          modified_pixel_values[s * fast + f] = int(
                                  sum(pixels) / len(pixels))

  open(bkginit_out, 'wb').write(cbf_header + start_tag +
                                compress(modified_pixel_values))

  return

if __name__ == '__main__':

  recompute_BKGINIT('BKGINIT.cbf', 'INIT.LP', sys.argv[1])