from mat import Mat
from vec import Vec
from matutil import rowdict2mat, mat2coldict, coldict2mat
from QR_solve import QR_solve
from image_mat_util import *
from image import image2file
from math import sqrt, floor
from os import stat
from feedback import printtolog, setlog, refresh

def move2board(v):
  '''
    Input:
    - v: a vector with domain {'y1','y2','y3'}, the coordinate
    representation of a point q.
    Output:
    - A {'y1','y2','y3'}-vector z, the coordinate representation
    in whiteboard coordinates of the point p such that the line
    through the origin and q intersects the whiteboard plane at p.
    '''
  return Vec({'y1','y2','y3'},
             {'y1':v['y1']/v['y3'],'y2':v['y2']/v['y3'],'y3':1})

def make_equations(x1, x2, w1, w2):
  '''
    Input:
    - x1 & x2: photo coordinates of a point on the board
    - y1 & y2: whiteboard coordinates of a point on the board
    Output:
    - List [u,v] where u*h = 0 and v*h = 0
    '''
  domain = {(a, b) for a in {'y1', 'y2', 'y3'}
            for b in {'x1', 'x2', 'x3'}}
  u = Vec(domain,
        {('y1','x1'):-x1, ('y1','x2'):-x2, ('y1','x3'):-1,
         ('y3','x1'):w1 * x1, ('y3','x2'):w1 * x2, ('y3','x3'):w1})
  v = Vec(domain,
        {('y2','x1'):-x1, ('y2','x2'):-x2, ('y2','x3'):-1,
         ('y3','x1'):w2 * x1, ('y3','x2'):w2 * x2, ('y3','x3'):w2})
  return [u, v]

def mat_move2board(Y):
  '''
    Input:
    - Y: Mat instance, each column of which is a 'y1', 'y2', 'y3'
    vector giving the whiteboard coordinates of a point q.
    Output:
    - Mat instance, each column of which is the corresponding
    point in the whiteboard plane (the point of intersection
    with the whiteboard plane of the line through the origin
    and q).
    '''
  refresh()

  cd = mat2coldict(Y)

  refresh()

  md = { i: move2board(v) for (i,v) in cd.items()}

  refresh()

  return coldict2mat(md)

def coord_map_matrix(orgx, outx, orgd, outd, tl, bl, tr, br, xscale):
  domain = {(a, b) for a in outd for b in orgd}
  tle = make_equations(tl[0],tl[1],0,0)
  ble = make_equations(bl[0],bl[1],0,1)
  tre = make_equations(tr[0],tr[1],xscale,0)
  bre = make_equations(br[0],br[1],xscale,1)
  f = {(outx,orgx):1}
  w = Vec(domain, f)
  
  m = {0:tle[0],1:tle[1],2:ble[0],3:ble[1],4:tre[0],
                  5:tre[1],6:bre[0],7:bre[1],8:w}
  L = rowdict2mat(m)
  b = Vec(set(range(9)), {8:1})

  refresh()

  h = QR_solve(L, b)
  return Mat((outd,orgd), h.f)

def deskew_image(name, outname, tl, bl, tr, br, xscale):
  '''
    Input:
    - name: name (may include path) of the input image file.
    - outname: name (may include path) of the output image file.
    - tl: x,y coords of top left corner to straighten.
    - bl: x,y coords of bottom left corner to straighten.
    - tr: x,y coords of top right corner to straighten.
    - br: x,y coords of bottom right corner to straighten.
    - xscale: scale of x direction (relative to y==1).
    '''

  # Some constants
  
  orgx,orgy,orgu = 'x1','x2','x3'
  outx,outy,outu = 'y1','y2','y3'
  orgd = {orgx, orgy, orgu}
  outd = {outx, outy, outu}
  orgrows = tuple(sorted(orgd))
  outrows = tuple(sorted(orgd))
  
  printtolog("Setting up the math", True)

  H = coord_map_matrix(orgx, outx, orgd, outd,
                       tl, bl, tr, br, xscale)
  
  printtolog("Loading the image from file", True)
  
  ext=((min(tl[0],bl[0],tr[0],br[0])-2,
        min(tl[1],bl[1],tr[1],br[1])-2),
       (max(tl[0],bl[0],tr[0],br[0])+2,
        max(tl[1],bl[1],tr[1],br[1])+2))

  img = image.file2image(name)
  (orgpts, colors) = image2mat(img, ext, orgrows)

  printtolog("Transforming the points", True)

  outpts = mat_move2board(H * orgpts)

  # Calculating output width and height
  
  orgwid = max(abs(tr[0]-tl[0]), abs(br[0]-bl[0]))
  orghgt = max(abs(br[1]-tr[1]), abs(bl[1]-tl[1]))
  outhgt = floor(max(sqrt(orgwid**2 + orghgt**2), orghgt))

  if not (tl[0],tl[1]) in outpts.D[1]:
    printtolog("Not found: {0}".format((tl[0],tl[1])))
  ylist = [outpts[(outx,(tl[0],tl[1]))],
            outpts[(outx,(bl[0],bl[1]))],
            outpts[(outx,(tr[0],tr[1]))],
            outpts[(outx,(br[0],br[1]))]]
  xlist = [outpts[(outy,(tl[0],tl[1]))],
            outpts[(outy,(bl[0],bl[1]))],
            outpts[(outy,(tr[0],tr[1]))],
            outpts[(outy,(br[0],br[1]))]]
  xmin = min(ylist)
  xmax = max(ylist)
  ymin = min(xlist)
  ymax = max(xlist)

  img = mat2image(outpts, colors, ext, outhgt, xscale, outrows,
                     xmin, ymin, xmax, ymax, tl, bl)

  printtolog("Saving image to disk", True)

  image2file(img, outname)

  return outname, stat(outname).st_size

def deskew(name, outname, x1, y1, x2, y2, x3, y3, x4, y4,
           fac, log=None):
  setlog(log, False)
  return deskew_image(name, outname, (x1,y1), (x2,y2),
                      (x3,y3), (x4,y4), fac)

def deskew1(log=None):
  setlog(log, False)
  return deskew_image('board.png', 'brd-output.png',
                       (362,38),(336,589),(592,158),(577,485), 2.5)

def deskew2(log=None):
  setlog(log, False)
  return deskew_image('cit.png', 'cit2-output.png',
                       (82,73),(81,103),(105,69),(105,102), 1)

def deskew3(log=None):
  setlog(log, False)
  return deskew_image('cit.png', 'cit3-output.png',
                       (136,64),(136,98),(152,71),(152,103), 1)

def deskew4(log=None):
  setlog(log, False)
  return deskew_image('painting.png', 'pntg-output.png',
                       (133,144),(142,512),(335,36),(347,644), 0.8)

def deskew5(log=None):
  setlog(log, False)
  return deskew_image('highres.png', 'hires-output.png',
                       (663,722),(714,2559),(1679,173),(1739,3223),
                       0.8)