// modified
// 
// verify int->flt in header while writing
// rename burkett stuff to load/saveBFlow.  
// Avoid using burkitt except for ground truth translation - the flipped Y
// is too confusing, already bit me.

import java.io.*;

import GR.*;
import zlib.*;

public class FlowUtil
{
  /**
   * Save to Burkitt/Barron flow format (*.flow by convention)
   * beware, flip *side-effects* the flow fields.
   * Call save with flip only once!
   *

Burkitt format has Y=0 at the top, but has V rightside up!!!!
yos groundtruth 20,20   2,0	# clearly sky, 
yos groundtruth 20,200  -3.2,-2.96 # flow in lower left is to the left and down
so burkitt format has v flipped!  (because with Y inverted, down should be 
positive)

   * Also beware: Barron calc_error and psflow seem to read the header as
   * x,y,x,y,x,y
   * but Barron lucas puts it out as y,x,y,x,y,x ????
   * Looking at printed output of lucas, and at the ground truth flow,
   * it seems like x,y order is correct and that lucas is swapping from
   * an internal convention where x,y are flipped from the usual?
   */
  public static void saveBFlow(String pathprefix,
			       float[][] u, float[][] v,
			       boolean flip,
			       int xwin, int ywin)
    throws IOException
  {
    System.out.println("saving flow field to " + pathprefix);
    if (flip) {
      grfArrUtil.fscale(u, -1.f);
      grfArrUtil.fscale(v, -1.f);
    }

    int yres = u.length;
    int xres = u[0].length;

    DataOutputStream dout
      = new DataOutputStream(new BufferedOutputStream
	(new FileOutputStream(pathprefix + ".flow")));

    // header gives picture size x,y,
    // computed flow size x,y   (may be less than the picture size)
    // offset of the flow in the picture
    int xoff = (xres-xwin) / 2;
    int yoff = (yres-ywin) / 2;
    System.out.println(" resolution = "+xres+","+yres);
    System.out.println(" window =     "+xwin+","+ywin);
    System.out.println(" offset =     "+xoff+","+yoff);

    float val = (float)xres;  zliberror.assert((int)val == xres);
    dout.writeFloat(val);	// totalx

    val = (float)yres;  zliberror.assert((int)val == yres);
    dout.writeFloat(val);	// totaly	

    val = (float)xwin;  zliberror.assert((int)val == xwin);
    dout.writeFloat(val);	// numx

    val = (float)ywin;  zliberror.assert((int)val == ywin);
    dout.writeFloat(val);	// numy

    val = (float)xoff;  zliberror.assert((int)val == xoff);
    dout.writeFloat(val);		// offx

    val = (float)yoff;  zliberror.assert((int)val == yoff);
    dout.writeFloat(val);		// offy

    for( int iy=0; iy < ywin; iy++ ) {
      int iyoff = iy+yoff;
      for( int ix=0; ix < xwin; ix++ ) {
	int ixoff = ix+xoff;
	dout.writeFloat(u[iyoff][ixoff]);
	//dout.writeFloat(v[iyoff][ixoff]);
	dout.writeFloat(-v[iyoff][ixoff]);	// flip v, see above
      }
    }

    dout.close();
  } //saveBFlow

  //----------------------------------------------------------------

  /**
   * read the header from Barron flow format 
   */
  public static void readBFlowHeader(String path, int[] header)
    throws IOException
  {
    DataInputStream fin
      = new DataInputStream(new BufferedInputStream(new FileInputStream(path)));

    int xres = (int)fin.readFloat();
    int yres = (int)fin.readFloat();
    int xwin = (int)fin.readFloat();
    int ywin = (int)fin.readFloat();
    int xoff = (int)fin.readFloat();
    int yoff = (int)fin.readFloat();

    header[0] = xres;
    header[1] = yres;
    header[2] = xwin;
    header[3] = ywin;
    header[4] = xoff;
    header[5] = yoff;

    fin.close();
  } //readBFlowHeader

  //----------------------------------------------------------------

  /**
   * load from Barron flow format to an array of 2 float arrays.
   */
  public static float[][][] loadBFlow(String path)
    throws IOException
  {
    float[][][] uv = new float[2][][];
    DataInputStream fin
      = new DataInputStream(new BufferedInputStream(new FileInputStream(path)));

    int xres = (int)fin.readFloat();
    int yres = (int)fin.readFloat();
    int xwin = (int)fin.readFloat();
    int ywin = (int)fin.readFloat();
    int xoff = (int)fin.readFloat();
    int yoff = (int)fin.readFloat();
    System.out.println("flow "+path+
		       " res="+xres+","+yres+
		       " win="+xwin+","+ywin+
		       " off="+xoff+","+yoff);

    float[][] u = new float[yres][xres];
    float[][] v = new float[yres][xres];
    uv[0] = u;
    uv[1] = v;

    for( int iy=0; iy < ywin; iy++ ) {
      int iyoff = iy+yoff;
      for( int ix=0; ix < xwin; ix++ ) {
	int ixoff = ix+xoff;
	float u1 = fin.readFloat();
	float v1 = fin.readFloat();
	u[iyoff][ixoff] = u1;
	v[iyoff][ixoff] = v1;
      }
    }

    fin.close();
    return uv;
  } //loadBFlow
  
  //----------------------------------------------------------------
  //----------------------------------------------------------------

  final public static int FLOMAGIC = 76567;

  /**
   * save to zilla single-file flow format (*.flo by convention).
   * This is a rationalized version of the burkitt format,
   * has y consistent rather than upside-down,
   * has integers in the header rather than floats,
   * has a magic word.
   *
   * BEWARE calling flip side-effects the arrays.
   */
  public static void saveFlow(String pathprefix,
			      float[][][] uv,
			      boolean flip,
			      int xwin, int ywin)
    throws IOException
  {
    saveFlow(pathprefix, uv[0], uv[1], flip, xwin, ywin);
  }


  /**
   */
  public static void saveFlow(String pathprefix, float[][][] uv)
    throws IOException
  {
    int yres = uv[0].length;
    int xres = uv[0][0].length;
    saveFlow(pathprefix, uv[0], uv[1], false, xres, yres);
  }


  /**
   */
  public static void saveFlow(String path,
			      float[][] u, float[][] v,
			      boolean flip,
			      int xwin, int ywin)
    throws IOException
  {
    path = zlib.addExtensionIfMissing(path, ".flo");

    System.out.println("saving flow field to " + path);
    if (flip) {
      grfArrUtil.fscale(u, -1.f);
      grfArrUtil.fscale(v, -1.f);
    }

    int yres = u.length;
    int xres = u[0].length;

    DataOutputStream dout
      = new DataOutputStream(new BufferedOutputStream
	(new FileOutputStream(path)));

    // header gives picture size x,y,
    // computed flow size x,y   (may be less than the picture size)
    // offset of the flow in the picture
    int xoff = (xres-xwin) / 2;
    int yoff = (yres-ywin) / 2;
    System.out.println(" resolution = "+xres+","+yres);
    System.out.println(" window =     "+xwin+","+ywin);
    System.out.println(" offset =     "+xoff+","+yoff);

    dout.writeInt(FLOMAGIC);	

    dout.writeInt(xres);	// totalx
    dout.writeInt(yres);	// totaly	
    dout.writeInt(xwin);	// numx
    dout.writeInt(ywin);	// numy
    dout.writeInt(xoff);	// offx
    dout.writeInt(yoff);	// offy

    for( int iy=0; iy < ywin; iy++ ) {
      int iyoff = iy+yoff;
      for( int ix=0; ix < xwin; ix++ ) {
	int ixoff = ix+xoff;
	dout.writeFloat(u[iyoff][ixoff]);
	dout.writeFloat(v[iyoff][ixoff]);
      }
    }

    dout.close();
  } //saveFlow

  //----------------------------------------------------------------

  /**
   * read the header from single file .flo format 
   */
  public static void readFlowHeader(String path, int[] header)
    throws IOException
  {
    path = zlib.addExtensionIfMissing(path, ".flo");

    DataInputStream fin
      = new DataInputStream(new BufferedInputStream(new FileInputStream(path)));

    int magic = fin.readInt();
    if (magic != FLOMAGIC) {
      String err = "file is not a .flo: " + path;
      System.err.println(err);
      throw new IOException(err);
    }

    int xres = fin.readInt();
    int yres = fin.readInt();
    int xwin = fin.readInt();
    int ywin = fin.readInt();
    int xoff = fin.readInt();
    int yoff = fin.readInt();

    header[0] = xres;
    header[1] = yres;
    header[2] = xwin;
    header[3] = ywin;
    header[4] = xoff;
    header[5] = yoff;

    fin.close();
  } //readFlowHeader

  //----------------------------------------------------------------

  /**
   * load from single-file .flo format to an array of 2 float arrays.
   */
  public static float[][][] loadFlow(String path)
    throws IOException
  {
    path = zlib.addExtensionIfMissing(path, ".flo");

    float[][][] uv = new float[2][][];
    DataInputStream fin
      = new DataInputStream(new BufferedInputStream(new FileInputStream(path)));

    int magic = fin.readInt();
    if (magic != FLOMAGIC) {
      String err = "file is not a .flo: " + path;
      System.err.println(err);
      throw new IOException(err);
    }

    int xres = fin.readInt();
    int yres = fin.readInt();
    int xwin = fin.readInt();
    int ywin = fin.readInt();
    int xoff = fin.readInt();
    int yoff = fin.readInt();
    System.out.println("loadFlow "+path+
		       " res="+xres+","+yres+
		       " win="+xwin+","+ywin+
		       " off="+xoff+","+yoff);

    float[][] u = new float[yres][xres];
    float[][] v = new float[yres][xres];
    uv[0] = u;
    uv[1] = v;

    for( int iy=0; iy < ywin; iy++ ) {
      int iyoff = iy+yoff;
      for( int ix=0; ix < xwin; ix++ ) {
	int ixoff = ix+xoff;
	float u1 = fin.readFloat();
	float v1 = fin.readFloat();
	u[iyoff][ixoff] = u1;
	v[iyoff][ixoff] = v1;
      }
    }

    fin.close();
    return uv;
  } //loadFlow


  /**
   * load from single-file .flo format to an array of 2 float arrays.
   */
  public static void loadFlow(String path, float[][][] uv)
    throws IOException
  {
    path = zlib.addExtensionIfMissing(path, ".flo");

    DataInputStream fin
      = new DataInputStream(new BufferedInputStream(new FileInputStream(path)));

    int magic = fin.readInt();
    if (magic != FLOMAGIC) {
      String err = "file is not a .flo: " + path;
      System.err.println(err);
      throw new IOException(err);
    }

    int xres = fin.readInt();
    int yres = fin.readInt();
    int xwin = fin.readInt();
    int ywin = fin.readInt();
    int xoff = fin.readInt();
    int yoff = fin.readInt();
    System.out.println("flow "+path+
		       " res="+xres+","+yres+
		       " win="+xwin+","+ywin+
		       " off="+xoff+","+yoff);

    zliberror.assert(uv[0].length >= yres);
    zliberror.assert(uv[0][0].length >= xres);
    float[][] u = uv[0];
    float[][] v = uv[1];

    for( int iy=0; iy < ywin; iy++ ) {
      int iyoff = iy+yoff;
      for( int ix=0; ix < xwin; ix++ ) {
	int ixoff = ix+xoff;
	float u1 = fin.readFloat();
	float v1 = fin.readFloat();
	u[iyoff][ixoff] = u1;
	v[iyoff][ixoff] = v1;
      }
    }

    fin.close();
  } //loadFlow


  //----------------------------------------------------------------

  public static void flipFlow(float[][][] uv)
  {
    grfArrUtil.fscale(uv[0], -1.f);
    grfArrUtil.fscale(uv[1], -1.f);
  } //flipFlow

  //----------------------------------------------------------------
  //----------------------------------------------------------------

  public static void flo2burkitt(String[] cmdline)
  {
    // avoid the edges
    boolean saveinwin = false;

    try {
      float[][][] f = loadFlow(cmdline[0]);

      int xres = f[0][0].length;
      int yres = f[0].length;
      int wxres = xres;
      int wyres = yres;
      if (saveinwin) {
	wxres = (int)(0.8*xres);
	wyres = (int)(0.8*yres);	// regular
	wyres = (int)(0.4*yres);	// nosky
	System.out.println("subwin = " + wxres +","+ wyres);
      }

      saveBFlow(cmdline[1], f[0], f[1], true,
		wxres, wyres);

      System.exit(0);
    }
    catch(IOException x) { zliberror.die(x); }
  } //flo2burkitt

  //----------------------------------------------------------------

  public static void burkitt2flo(String[] cmdline)
  {
    // avoid the edges
    boolean saveinwin = false;

    try {
      float[][][] f = loadBFlow(cmdline[0]);

      int xres = f[0][0].length;
      int yres = f[0].length;
      int wxres = xres;
      int wyres = yres;
      if (saveinwin) {
	wxres = (int)(0.8*xres);
	wyres = (int)(0.8*yres);	// regular
	wyres = (int)(0.4*yres);	// nosky
	System.out.println("subwin = " + wxres +","+ wyres);
      }

      saveFlow(cmdline[1], f[0], f[1], true, wxres, wyres);

      System.exit(0);
    }
    catch(IOException x) { zliberror.die(x); }
  } //burkitt2flo

  //----------------------------------------------------------------

  static void usage()
  {
    System.err.println("java FlowUtil in.flo out.flow");
    System.err.println("java FlowUtil in.flow out.flo");
    System.err.println("converts .flow(burkitt) to .flo format or v.vs");
    System.exit(1);
  }

  public static void main(String[] cmdline)
  {

    if (cmdline.length != 2) {
      usage();
    }

    String ext = zlib.getExtension(cmdline[0]);
    if (ext.equals(".flo"))
      flo2burkitt(cmdline);
    else if (ext.equals(".flow"))
      burkitt2flo(cmdline);
    else usage();
    
  } //main

} //FlowUtil