"""
* Copyright (c) 2009 Stanford University, School of Medicine All rights reserved. 
* 
* STANFORD ACADEMIC SOFTWARE LICENSE AGREEMENT 
* FOR AngioLocalize SOURCE CODE: 
* 
* 1.  This is a legal agreement between you, RECIPIENT, and STANFORD UNIVERSITY.  
*       By accepting, receiving, and using AngioLocalize.including any accompanying information, materials or 
*       manuals (Program), you are agreeing to be bound by the terms of this Agreement.  
*       If you do not agree to the terms of this Agreement, promptly return the Program to the STANFORD.

* 2.  STANFORD grants to RECIPIENT a royalty-free, nonexclusive, and nontransferable license to use the Program 
*       and accompanying source code furnished hereunder, upon the terms and conditions set out below. 
*       documentation and/or other materials provided with the distribution. 
* 3.  RECIPIENT acknowledges that the Program is a research tool still in the development stage and that 
*       it is being supplied as is, without any accompanying services, support or improvements from STANFORD.
* 4.  RECIPIENT agrees to use the Program solely for internal non-commercial purposes. 
* 5.  RECIPIENT may create derivatives and further transfer the Program under the following terms and conditions: 
*        A) Recipient agrees to not commercially exploit the Program and agrees all uses will be for non-commercial purposes; and,
*        B) Recipient may not mark any derivatives with the title of Program and may not to use iPad mark for any derivatives; and,
*        C) Recipient agrees to mark and acknowledge all derivatives as follows: "This is a derivative work based on AngioLocalize:
*           Frameless AVM localization with 2D Digital Subtraction Angiography (DSA) and 3D Rotational Angiography (DRA) Images developed
*           by Dimitre Hristov, Lina Liu.
*        D) Recipient agrees to make all improvements, enhancements and derivatives available under 
*           similar terms and conditions to the research community. 
* 6.   RECIPIENT warrants that RECIPIENT will not remove or export any part of the software or Program from the United 
*       States except in full compliance with all United States and other applicable laws and regulations.
* 7.  RECIPIENT will use the Program in compliance with all applicable laws, policies and regulations including, but not 
*       limited to, any approvals, informed consent and patient confidentiality principles.
* 8.   RECIPIENT will indemnify, hold harmless, and defend Stanford against any claim of any kind arising out of or
*       related to the exercise of any rights granted under this Agreement or the breach of this Agreement by RECIPIENT.
* 9.  Title and copyright to the Program and any associated documentation shall at all times remain with STANFORD, 
*       and RECIPIENT agrees to preserve same.
""" 
#----------------------------------------------------------------------------------
#! Macro module Registration2D3D
#/*!
# \file    Registration2D3D.py
# \author  Lina Liu and Dimitre Hristov
# \date    2009-12-10
#
# 
#*/
#----------------------------------------------------------------------------------

from mevis import *
from linear import *
from math import *


## Global varibles
source = None 
dest = None
srcID = 0
targetID = 0
perspectM = ""
affineM = ""
postAffineM = ""

## initialize some varibles when window pops up
def initWindow():
  ctx.field("PatientID").setValue(ctx.parent().field("PatientID").value);
  ctx.field("PatientsName").setValue(ctx.parent().field("PatientsName").value);
  ctx.field("PatientsBirthDate").setValue(ctx.parent().field("PatientsBirthDate").value);
  ctx.field("PatientsSex").setValue(ctx.parent().field("PatientsSex").value);
  ctx.field("ImgLoad2D.autoLoad").touch();
  ctx.field("ImgLoad3D.autoLoad").touch();
  ctx.field("Registration_3D_2D.View2D.ext.annotation.annotationMode").setValue("ANNO_MODE_USER");
  ctx.field("Registration_3D_2D.View2D.ext.annotation.annotationUserBottomLeft").setValue("Timepoint: $(slice) \n$(sizex),$(sizey),$(sizec) \n$(voxelx),$(voxely) \n$(SeriesDescription)");
  ctx.field("Registration_3D_2D.View2D.ext.annotation.showOrientation").setValue("ANNO_ORIENTATION_OFF");
  patientID = getPatientRecordID();
  if (patientID != "0"):
    loadGUILists(patientID);

## cleanUp when windows are closed
def cleanUp():
  ctx.disconnectField("Registration_3D_2D.Image2D");
  ctx.disconnectField("Registration_3D_2D.Image3D");
  ctx.field("ImgLoad2D.close").touch();
  ctx.field("ImgLoad3D.close").touch();
  
## send a query to the Database and return the result set
def sendQuery(query):
  ctx.field("DBConn.Statement").setValue(query);
  dbModule = ctx.module("DBConn");
  result = dbModule.call("sendQuery");
  if (result == -1):
    MLAB.logError("The following query to database failed. \n" + query + "\n.");
  dbModule.call("close");
  return result;

## get the DB record ID for a patient given MRO number 
def getPatientRecordID():
  patientID = 0;
  if (ctx.field("PatientID").value == ""):
    return str(patientID);
  queryString = "SELECT PatientID FROM Patients WHERE MRO = " + "'" + ctx.field("PatientID").value + "'";
  result = sendQuery(queryString);
  if (str(result) != "()"):
    for item in result:
      patientID = item["patientid"];
  return str(patientID);

## load GUI Lists (2D, 3D and registration lists)
def loadGUILists(patientID):
    load3DImageList(patientID);
    load2DImageList(patientID);   
    loadRegItem();
    
## load 3D image List in GUI   
def load3DImageList(patientID):
    imgList = ctx.controlDebug("list3D");
    imgList.clearItems();
    #init DB
    queryString = "select imageID, modality, seriesDescr from images where imagetype = 2 and patientID = " + patientID + ";";
    result = sendQuery(queryString);
    for item in result:
      imgName = item["imageid"];
      imgM = item["modality"];
      series = item["seriesdescr"];
      imgList.insertItem([str(imgName), imgM, series]);

## load 2D image List in GUI
def load2DImageList(patientID):
    imgList = ctx.controlDebug("list2D");
    imgList.clearItems();
    #init DB
    queryString = "select imageID, modality, seriesDescr from images where imagetype = 1 and patientID = " + patientID + ";";
    result = sendQuery(queryString);
    for item in result:
      imgName = item["imageid"];
      imgM = item["modality"];
      series = item["seriesdescr"];
      imgList.insertItem([str(imgName), imgM, series]);

## pop up the registration context menu 
def selectRegImg():
   ctx.control("imgRegOp").popup();
   
## pop up the 3D context menu 
def selectChanged3D():
   ctx.control("popupTgt").popup();
   
## pop up the 2D context menu 
def selectChanged2D():
   ctx.control("popupSrc").popup();
   
## select 2D image as source
def selectAsSource():
  list2D = ctx.control("list2D");
  srcItem = list2D.selectedItem();
  loadImage(srcItem.text(0), 1);
  
## select 3D image as target
def selectAsDest():
  list3D = ctx.control("list3D");
  targetItem = list3D.selectedItem();
  loadImage(targetItem.text(0), 2);
  
  global targetID;
  global srcID;
  redoOp = False;
  
  if (regExist(srcID, targetID)):
    answer = MLAB.showQuestion("The Registration already exists. Do you want to redo or delete?", "", "MeVisLab", ["Redo", "Delete", "Cancel"], 0);
    if (answer== 2):
      return;
    elif (answer == 1):
     deleteReg(srcID, targetID);
     return;
    else:
      redoOp = True;
  if (redoOp == False):
    getInitialRegMatrix();
  global affineM;
  global perspectM;
  global postAffineM;
  register(affineM, perspectM, postAffineM);  

## do registration with given AffineMatrix, perspectiveMatrix and postAffine Matrix
def register(affineM, perspectM, postAffineM):
  ctx.field("ImgLoad3D.load").touch();
  ctx.field("ImgLoad2D.load").touch();
  ctx.connectField("Registration_3D_2D.Image2D", "ImgLoad2D.output0");
  ctx.connectField("Registration_3D_2D.Image3D", "ImgLoad3D.output0");
  wind = ctx.showWindow("Reg2D3DW");
  wind.showMaximized();
  loadMatrices(affineM, perspectM, postAffineM);

## redo the registration
def redoRegFromPopup():
  regList = ctx.control("registeredImg");
  item = regList.selectedItem();
  global targetID;
  global srcID;
  srcID = item.text(0);
  targetID = item.text(3);
  loadImage(item.text(0),1);
  loadImage(item.text(3),2);
  register(item.text(8),item.text(9),item.text(10));
  
## load Image as target or source from a given fileName 
def loadImage(imageID, srcOrTgt):  
  srcF = getImageFileName(imageID);
  src = os.path.normpath(srcF + ".dcm");
  source = os.path.normpath(src);
  if (srcOrTgt == 1):
    global srcID;
    srcID = imageID;
    if source != None:
      ctx.field("ImgLoad2D.filename").setValue(source);
      ctx.field("ImgLoad2D.load").touch();
    ctx.control("srcLabel").setTitle("Source:" + source);
  else:
    global targetID;
    targetID = imageID;
    if source != None:  
      ctx.field("ImgLoad3D.filename").setValue(source);
      ctx.field("ImgLoad2D.load").touch();
    ctx.control("targetLabel").setTitle("Target:" + source);
    
## check if registration exists    
def regExist(srcID, targetID):
  queryString = "SELECT matrixID as Num, affineM, perspectM, postAffineM FROM REGMATRICES WHERE srcImage = " + str(srcID) + " AND type = 1 AND targetImage = " + str(targetID) + ";";
  result = sendQuery(queryString);
  if (str(result) == "()"):
    return 0;
  number = 0;
  number = result[0]["num"];
  global affineM;
  global postAffineM;
  global perspectM;
  affineM = result[0]["affinem"];
  postAffineM = result[0]["postaffinem"];
  perspectM = result[0]["perspectm"];
  return number;    

## generate the initial registration matrix based on Dicom tags
## code tested only for the Siemens Artis System and the dynaCT mode. 
## if different system is used it may have to be rewritten to generate the global matrices 
def getInitialRegMatrix():
  global perspectM;
  global postAffineM;
  global affineM;
  
  #image2D=ctx.field("Registration_3D_2D.Image2D");
  dist_source_detector = float((ctx.field("ImgLoad2D.output0").getDicomTagValueByName("DistanceSourceToDetector")));
  dist_source_patient = float(ctx.field("ImgLoad2D.output0").getDicomTagValueByName("DistanceSourceToPatient"));
  rows = float(ctx.field("ImgLoad2D.output0").getDicomTagValueByName("Rows"));
  columns = float(ctx.field("ImgLoad2D.output0").getDicomTagValueByName("Columns"));
  
  #get the pixel spacing in x and y directions
  pixel_spacing_string = []
  temp = str(ctx.field("ImgLoad2D.output0").getDicomTagValueByName("PixelSpacing"))
  resume = True
  count = 0
  pixel_spacing_string = ''
  while (resume):
    if (temp[count] == '\\'):
      resume = False
      count = count + 1
    else:
      pixel_spacing_string = pixel_spacing_string + temp[count]
      count = count + 1
  pixel_spacing_x = float(pixel_spacing_string);
  
  pixel_spacing_string = ''
  while (count<len(temp)):
    pixel_spacing_string = pixel_spacing_string + temp[count]
    count = count + 1
  pixel_spacing_y = float(pixel_spacing_string);
  
  #get the primary and secondary angles
  pAngle = float(ctx.field("ImgLoad2D.output0").getDicomTagValueByName("PositionerPrimaryAngle"));
  sAngle = float(ctx.field("ImgLoad2D.output0").getDicomTagValueByName("PositionerSecondaryAngle"));
  
  #assign rotation angles: theta_x, theta_y, theta_z
  
  theta_x = -sAngle;
  theta_z =  pAngle;
  
  #assign translations
  t_x = 0;
  t_y = 0;
  t_z = dist_source_patient;
  
  #find the rotation matrices: R_x, R_table_camera, R_z
  #R_table_camera is used to transform from the table/supine patient coordinate system
  #to one linked with the x-ray source (camera)
  R_x = matrix();
  R_table_camera = matrix();
  R_z = matrix();
  R_x = matrix.fromSeq(R_x,[[1,0,0,0],[0,cos(deg2rad(theta_x)),sin(deg2rad(theta_x)),0],[0,-sin(deg2rad(theta_x)),cos(deg2rad(theta_x)),0],[0,0,0,1]]);
  R_table_camera = matrix.fromSeq(R_table_camera,[[-1,0,0,0],[0,0,-1,0],[0,-1,0,0],[0,0,0,1]]);
  R_z = matrix.fromSeq(R_z,[[cos(deg2rad(theta_z)),sin(deg2rad(theta_z)),0,0],[-sin(deg2rad(theta_z)),cos(deg2rad(theta_z)),0,0],[0,0,1,0],[0,0,0,1]]);
  
  #multiply them in the order R_x*R_table_camera*R_z to get the Rotation matrix
  affineM = R_z.multMatrixFromLeft(R_x)
  affineM = affineM.multMatrixFromLeft(R_table_camera)

  #find the K matrix
  perspectM = matrix();
  perspectM.setElement(0,0,dist_source_detector);
  perspectM.setElement(1,1,dist_source_detector);
  perspectM.setElement(2,2,1);
  perspectM.setElement(3,3,1);
  # Not sure whether x corresponds to columns and y to rows
  perspectM.setElement(0,2,columns*pixel_spacing_x/2);
  perspectM.setElement(1,2,rows*pixel_spacing_y/2);
  
  postAffineM=matrix();
  postAffineM.setElement(0,0,-1);
  postAffineM.setElement(0,3,columns*pixel_spacing_x);
  
  #find the A matrix
  affineM.setElement(2,3,dist_source_patient);

## load the Affine Matrix, perspective Matrix and post AffineMatrix into Registration_3D_2D module
def loadMatrices(affineM, perspectM, postAffineM):
  ctx.field("Registration_3D_2D.AffineMatrix").setStringValue(affineM);
  ctx.field("Registration_3D_2D.AffinePostMatrix").setStringValue(postAffineM);
  ctx.field("Registration_3D_2D.PerspectiveMatrix").setStringValue(perspectM);

## delete a registration from registered Image list
def deleteRegFromPopup():
  regList = ctx.control("registeredImg");
  item = regList.selectedItem();
  deleteReg(item.text(0), item.text(3));

## delete a registration from a database record
def deleteReg(srcid, targetid):
  query = "select matrixID from regmatrices where srcImage = " + str(srcid) + " AND targetImage = " + str(targetid) + ";";
  result = sendQuery(query);
  regToDel = result[0]["matrixid"];
  
  query = "DELETE FROM REGMATRICES WHERE srcImage = " + str(srcid) + " AND targetImage = " + str(targetid);
  result = sendQuery(query);
  
  query = "DELETE FROM APPROVALS WHERE objType = 0 and objID = " + str(regToDel) + ";";
  result = sendQuery(query);
  delete2D3DReg(ctx.parent().field("PatientID").value, srcid, targetid);
  loadRegItem();

## retrieve the image file name from a database record ID
def getImageFileName(imgID):
  queryString = "select imageName, imageDir from images where imageID = " + str(imgID);
  result = sendQuery(queryString);
  imgName = None;
  for item in result:
    imgName = item["imagename"];
  return imgName;

## load registration into the GUI registered image list
def loadRegItem():
    regImage = ctx.controlDebug("registeredImg");
    regImage.clearItems();
    patientID = getPatientRecordID();
    if (patientID == "0"):
      return;
    queryString = "select b.affineM as matrixA, b.perspectM as matrixB, b.postAffineM as matrixC, b.srcImage as src, a.modality as srcM, a.seriesDescr as srcDesc, c.modality as trgM, c.seriesDescr as tgtDesc, b.targetImage as target , d.approvalStatus as status, to_char(d.approvaldate,'YYYY-MM-DD HH12:MI:SS') as savedTime, b.status as genre from images a, regmatrices b, images c, approvals d where a.imagetype = 1 and c.imagetype = 2 and b.srcImage = a.imageID and b.targetImage = c.imageID and d.objID = b.matrixID and b.type = 1 and a.patientID = " + patientID + ";"
    result = sendQuery(queryString);
    for item in result:
      matrixReg = item["matrixa"];
      perspectiveM = item["matrixb"];
      postAffineM = item["matrixc"];
      srcImg = item["src"];
      tgtImg = item["target"];
      srcMod = item["srcm"];
      srcDesc = item["srcdesc"];
      targetMod = item["trgm"];
      targetDesc = item["tgtdesc"];
      statusNum = item["status"];
      approvalDate = item["savedtime"];
      genreNum = item["genre"];
      status = None;
      genre = None;
      if (statusNum == 0):
        status = "false";
      else:
        status = "true";
      if (genreNum == 0):
        genre = "direct";
      else:
        genre = "inferred";
      regImage.insertItem([str(srcImg), srcMod, srcDesc, str(tgtImg), targetMod, targetDesc,status, genre, approvalDate, matrixReg, perspectiveM, postAffineM]);

## This function is to check if the registration exists and is direct
def directRegExist(srcID, targetID):
  query = "SELECT matrixID from RegMatrices where srcImage = " + str(srcID) + " and targetImage = " + str(targetID) + " and status = 0;"
  result = sendQuery(query);
  if (str(result) != "()"): 
    return True;
  return False;

## propagate the registration to 3D and 3D registration
##For example if registration for 3D image 1 and 2 exist, 2 and 3 exist, 
## we need to infer registration for 1 and 3 and save it into database 
## the propagation rule is that if the infered pair's registration exists and the status is direct, 
## we will not overwrite it.
def propagate(targetID, srcID, approvalStatus, regMatrix):
  queryString = "SELECT targetImage, affineM, approvalStatus, status FROM REGMATRICES, approvals WHERE objID = matrixID and srcImage = " + str(targetID) + " AND type = 0 " + ";";
  result = sendQuery(queryString);
  newTarget = 0;
  for item in result:
     newTarget = item["targetimage"];
     oldApprovalS = item["approvalstatus"];
     Affine = item["affinem"];
     if (directRegExist(srcID, newTarget)):
         continue;
     insertOp = 0;
     status = 1; # status is inferred
     newAffine = regMultiply(regMatrix, Affine);
     regS = newAffine.replace("\n", " ");
     ret = regExist(srcID, newTarget);
     if (ret == 0):
         insertOp = 1;
     newApprovalS = approvalStatus & oldApprovalS;
     regID = saveReg(insertOp, srcID, newTarget, status, regS);
     saveApprovalStatus(regID, newApprovalS, insertOp);
 
## save or update registration into database   
def saveReg(insertOp, srcID, targetID, status, regS):
 perspectiveMatrix = ctx.field("Registration_3D_2D.PerspectiveMatrix").stringValue();
 perspectS = perspectiveMatrix.replace("\n", " ");
 affineS = regS;
 postAffineMatrix = ctx.field("Registration_3D_2D.AffinePostMatrix").stringValue();
 postAffineS = postAffineMatrix.replace("\n", " ");
 saveProjRegToFile(perspectS,affineS, postAffineS, srcID, targetID, ctx.parent().field("PatientID").value);
 ## affineTransformation regType : 0        
 ## perspectiveTransformation regType: 1
 if (insertOp == 1): 
   queryS = "INSERT INTO REGMATRICES(srcImage, targetImage, affineM, perspectM, postAffineM, status, type)values(" + str(srcID) + "," + str(targetID) + ",'" + affineS + "'" + ",'" + perspectS + "','"+ postAffineS + "'," + str(status) + ", 1" + ");";
 else:
   queryS = "UPDATE REGMATRICES set status = " + str(status) + "," + "perspectM = " + "'" + perspectS + "'" + " , affineM = '" + affineS + "' , postAffineM = '" + postAffineS + "' where srcImage = " + str(srcID) + " AND targetImage = " + str(targetID) + ";"; 
 result = sendQuery(queryS);

 query = "select matrixID from regmatrices where type = 1 and srcImage = " + str(srcID) + " and targetImage = " + str(targetID) + ";";
 result = sendQuery(query);
 imgID = None;
 for item in result:
    imgID = item["matrixid"]; 

 return imgID;

## multiply registration matrices of affine Matrix and registration modification
def regMultiply(regMatrix, Affine):
  ctx.field("Registration_3D_2D.regMultiply.matrixA").setStringValue(regMatrix);
  ctx.field("Registration_3D_2D.regMultiply.matrixB").setStringValue(Affine);
  ctx.field("Registration_3D_2D.regMultiply.operation").setStringValue("Multiply");
  return ctx.field("Registration_3D_2D.regMultiply.outputMatrixC").stringValue();

## save the approval status
def saveApprovalStatus(imgID, status, insertOp):
 if (insertOp == 0):
   queryS = "Update APPROVALS SET approvalStatus = " + str(status) + " where UserID = 1 AND objType = 0 AND objID = " + str(imgID) + ";";
 else:
   queryS = "INSERT INTO APPROVALS(userID, objID, objType, patientID, approvalStatus)VALUES(1," + str(imgID) + ", 0 , " + getPatientRecordID() + " , " + str(status) + ");";
 result = sendQuery(queryS);
 
## the function is called when user clicks "save" button
## it saves the registration and approval status
def saveOnRegManual():
  saveRegAndApprovalStatus(0);
  
## the function is called when user clicks "save and approve" button  
## it saves the registration and approval status  
def approveOnRegManual():
  saveRegAndApprovalStatus(1);
  
## save the registration and approval status    
def saveRegAndApprovalStatus(approvalS):
  global srcID;
  global targetID;
  insertOp = 0;
  status = 0;
  if (regExist(srcID, targetID) == 0):
   insertOp = 1;
  affineMatrix = ctx.field("Registration_3D_2D.AffineMatrix").stringValue();
  affineS = affineMatrix.replace("\n", " ");
  imgID = saveReg(insertOp, srcID, targetID, status, affineS);
  saveApprovalStatus(imgID, str(approvalS), insertOp);
  propagate(targetID, srcID, approvalS, affineS);
  ctx.module("Registration_3D_2D").closeWindow(); 
  loadRegItem();


## show Help
def help():
  ctx.showHtmlHelp();
  
## the function is called when user clicks "cancel" button    
def cancelonRegManual():
  MLAB.log("The user reviewed the image and canceled.");
  ctx.disconnectField("Registration_3D_2D.Image2D");
  ctx.disconnectField("Registration_3D_2D.Image3D");
  ctx.module("Registration_3D_2D").closeWindow(); 
