#!/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
  for header in headers:
    if header["type"] == 1:
      curLines += header["lines"]
      curColumns = header["columns"]
    elif header["type"] == 2:
      pixelAspect = float(header["cfac"]) / header["lfac"]
    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

for i in grouped:
  print "Processing image %s" %i
  d = grouped[i]
  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.resize((img.size[0], newheight))
  img.save("%s.jpg" % i, "JPEG")