#!/usr/bin/env python

'''
  Needs Pillow, xrit and numpy
  pip install xrit Pillow

  Numpy is better to be installed using distro packages.
'''

import os, math
import numpy as np
from PIL import Image
from xrit import packetmanager


folder = "."

lritfiles = [f for f in os.listdir(folder) if os.path.isfile(os.path.join(folder, f)) and ".lrit" in f]

grouped = {}

def headersFromFile(filename):
  f = open(filename, "rb")
  try:
    k = packetmanager.readHeader(f)
    type, filetypecode, headerlength, datalength = k
  except Exception,e:
    print("   Header 0 is corrupted for file %s" %filename)
    return None
  f.seek(0, 0)
  data = f.read(headerlength)
  headers = packetmanager.getHeaderData(data)
  f.close()
  return headers


for i in lritfiles:
  headers = headersFromFile(i)
  curLines = 0
  curColumns = 0
  curImageId = 0
  pixelAspect = 1
  coff = 0
  for header in headers:
    if header["type"] == 1:
      curLines += header["lines"]
      curColumns = header["columns"]
    elif header["type"] == 2:
      pixelAspect = float(header["cfac"]) / header["lfac"]
      coff = header["coff"]
    elif header["type"] == 128: # Segment Identification Header
      imageId = header["imageid"]
      segmentId = header["sequence"]
      curImageId = imageId
      if not imageId in grouped:
        grouped[imageId] = { "segments": {}, "lines": 0 }

      grouped[imageId]["segments"][segmentId] = i

  grouped[imageId]["lines"] += curLines
  grouped[imageId]["columns"] = curColumns
  grouped[imageId]["pixelAspect"] = pixelAspect
  grouped[imageId]["startColumn"] = coff

for i in grouped:
  print "Processing image %s" %i
  d = grouped[i]

  # Crop Left
  sc = d["startColumn"]
  hw = min(d["columns"] - sc, sc)
  cl = d["startColumn"] - hw

  # Crop Right
  cf = cl + 2 * hw

  print " Crop Marks: %s %s" % (cl, cf)
  img = Image.new("L", (d["columns"], d["lines"]))
  pixels = np.array([], dtype=np.uint8)
  for z in d["segments"]:
    data = packetmanager.loadData(d["segments"][z])
    data = np.fromstring(data, dtype=np.uint8)
    pixels = np.concatenate((pixels, data))
  img.putdata(pixels)
  newheight = img.size[1] * d["pixelAspect"]
  newheight = int(math.ceil(newheight))
  img = img.crop((cl, 0, cf, img.size[1]))
  img = img.resize((img.size[0], newheight))
  img.save("%s-crop.jpg" % i, "JPEG")