/** Morph script, combination of rigid body and linear morph, with separate timing for each domain/**

/**
 * Linear interpolation of coordinates in the atom sets sel1 and sel2
 * @param {atom set} sel1 First set of atoms, will be modified by the function
 * @param {atom set} sel2 Second related set of atoms, will not be touched
 * @param {float} progress Weight for atoms in selection sel2
**/
function linear(progress, sel1, sel2) {

  //return
  // Copy coordinates to list of points and determine number of points
  coord1 = {@sel1}.xyz.all
  coord2 = {@sel2}.xyz.all
  len1 = coord1.length
  len2 = coord2.length
  if (len1 != len2) {
    print "************** exiting linear ***************"
    return
  }

  ssergorp = 1 - progress
  for (var i FROM [1,len1]) {

    // Weighted average of coordinates for equivalent points
    coord1[i] = (coord2[i] * progress + coord1[i] * ssergorp)

  }

  // Interpolated coordinates are stored in selection sel1 and sel2
  {@sel1}.xyz = @coord1
  {@sel2}.xyz = @coord1
}



/**
 * Rigid body interpolation of coordinates in the equivalent atom sets sel1 and sel2
 * @param {float} progress Relative weight of sel2 vs sel1 to choose rotation and translation
 * @param {atom set} sel1 First set of atoms, will be modified by the function
 * @param {atom set} sel2 Second related set of atoms, will be modified by the function
 * @param {atom set} supsel1 First set of atoms used for superposition
 * @param {atom set} supsel2 Second set of atoms used for superposition
 * @param {3D positon} anch1 Where the anchor was in original structure 1
 * @param {3D positon} anch2 Where the anchor was in original structure 2
 * @param {atom set} ansel1 Anchor to determine translation of first set 
 * @param {atom set} ansel2 Anchor to determine translation of second set
 * @param {integer} flag, rotate the long way if 1
**/
function rigid(progress, sel1, sel2, supsel1, supsel2, anch1, anch2, ansel1, ansel2, flag, test) {
  // print "function rigid"
  // print ansel1
  // print anch2
  // print ansel2
  // Calculate lowest RMSD superposition of sel1 and sel2 as 4x4 matrix
  // print [{@supsel1}, {@supsel2}]
  fourbyfour = compare({@supsel1}, {@supsel2})
  // print fourbyfour

  // Extract rotation sel1 -> sel2 and convert into quaternion
  total_quat = quaternion(@fourbyfour%1)

  // Rotate in the other direction if flag is set 
  if (flag == 1){
    theta = total_quat %"theta"
    ax = total_quat %"vector"
    total_quat = quaternion(ax, 360 - theta) 
  }
  ssergorp = 1 - progress

  // Partial rotation for atoms in sel1
  quat1 = total_quat * progress

  // Partial rotation for atoms in sel2 in the other direction */
  quat2 = total_quat * (progress - 1)

  // Work on first set
  select sel1

  // Rotate the coordinates into common orientation
  rotateselected @quat1 molecular

  // Calculate and apply translation to match anchor in rotated rigid body
  rot = quat1 %"matrix"
  orig = {0 0 0}
  // print ansel1
  anchor_pos = {@ansel1}.xyz    // where the anchor is at the moment
  // print "anchor_pos " + @anchor_pos
  if (test == 0){
    anchor_is = anch1 * !rot  // where the anchor gets rotated to
    transl = anchor_pos - anchor_is // anchor should move with other domain to maintain connection
  } else {
    transl = anch1 - anchor_pos // anchor should not move
  }
  translateselected @transl

  // Same steps for second (equivalent) set of atoms 
  select sel2

  rotateselected @quat2 molecular

  rot = quat2 %"matrix"
  anchor_pos = {@ansel2}.xyz
  if (test == 0){
    anchor_is = anch2 * !rot
    transl = anchor_pos - anchor_is
  } else {
    transl = anch2 - anchor_pos
  }
  translateselected @transl
}

/**
 * Morph between two structures based on domain and anchor definitions
 * @param {integer} steps number of intermediate conformations
 * @param {list of atom sets} selection of first and second structure
 * @param {list of various} recipe domain, anchor and timing definitions
 * @param {integer} debug flag, if 0 nothing, if 1 skip linear interpolation
**/
function morph(steps, morphed_structures, recipe, debug) {
 // Save coordinates to use in every step of the interpolation
 original = {all}.xyz.all
 apos1 = []
 apos2 = []
 print "Checking input"
 nr1 = @{morphed_structures[1]}
 nr2 = @{morphed_structures[2]}
 nr = @{n1} and @{n2}
 if (nr.length > 0 or nr1.length == 0 or nr2.length == 0) {
    print "morphed structures should be non-overlapping non-empty sets of atoms, please check"
 }

 cumulator = {none}
 for (var j FROM [1,recipe.length]) {
  cumulator = check_domain(j, recipe[j], morphed_structures, cumulator)
  ansel = (recipe[j])[5]
  //print ansel
  apos1.push({@ansel and @{morphed_structures[1]}}.xyz)
  apos2.push({@ansel and @{morphed_structures[2]}}.xyz)
 }
 for (var i FROM [1,steps]) {
  // Set coordinates back to original values in preparation of calculating next step in the morph
  {all}.xyz = @original
  progress = i * (1.0 / steps)
  for (var j FROM [1,recipe.length]) {
    //print "ij:" + @i + " " + @j
    morph_inner(progress, recipe[j], morphed_structures, apos1[j], apos2[j], debug)
  }
  // To save this step in the morph, un-comment the following three lines while running in local Jmol app
  // fname = "morph" + i + ".pdb"
  // select 1.1
  // write @fname
  set refreshing true
  delay 0.05
  set refreshing false

 }
 set refreshing true
 delay 2.0
 {all}.xyz = @original
}

function check_domain(j, instructions, morphed_structures, cumulator) {
    print ""
    print "Report for domain " + j
    s = instructions[4] // domain definition
    s1 = @{morphed_structures[1]} and @{s}
    print "  Number of atoms in domain: " + s.length
    print "  Cumulative number of atoms in previous domains: " + cumulator.length
    nr = @{s} and @{cumulator}
    print "  Overlap with previous domains (should be zero):" + nr.length
    a = instructions[5] // anchor definition
    nr1 = @{a} and @{cumulator}
    nr2 = @{a} and @{s}
    //print [a, cumulator, s, nr1, nr2]
    if (nr1.length > 0) {
        print "  domain is anchored to previous domains; number of atoms = " + nr1.length
        if (nr2.length > 0) {
            print "  anchor should not be part of this domain and other domains at the same time"
        }
    } else {
        print "  domain is not anchored (anchor will proceed on linear path); number of atoms = " + nr2.length
        if (nr2.length == 0) {
            print "    oops, anchor definition is empty" + nr2
        }
    }
    cumulator = @{s} or @{cumulator}
    s2 = @{morphed_structures[2]} and @{s}
    len1 = s1.all.length
    len2 = s2.all.length
    if (len1 != len2) {
         print "  domain " + j + "has unequal number of atoms in the two structures"
    }
    return cumulator
}    

function morph_inner(progress, instructions, morphed_structures, apus1, apus2, debug) {
    //print [progress, instructions, morphed_structures, apus1, apus2, debug]
    //print [progress, instructions, morphed_structures, apus1, apus2, debug]
    flag = instructions[3]
    // progress is overall timer from 0 to 1
    // instructions[1] is when to start moving this domain
    // instructions[2] is when to finish moving this domain
    // p is how far into the morph this domain is
    if (progress < instructions[1]) {
      p = 0
    } else if (progress > instructions[2]) {
      p = 1
    } else {
      p = (progress - instructions[1]) / (instructions[2] - instructions[1])
    }
    s = instructions[4] // domain definition
    a = instructions[5] // anchor definition
    // s1 and s2: selections for the corresponding domains
    s1 = @{morphed_structures[1]} and @{s}
    s2 = @{morphed_structures[2]} and @{s}
    // a1 and a2: anchors for the correspoding domains
    a1 = @{morphed_structures[1]} and @{a}
    a2 = @{morphed_structures[2]} and @{a}
    // Atom sets s1 and s2 will be moved into a common frame on the path from s1 to s2.
    // Position on the path (closer to s1, closer to s2) is determined by parameter p
    test = {@s1 and @a1}.size // whether anchor is within domain or not
    //print "test" + test
    if (instructions.length < 6){
      //print "pre-rigid"
      rigid(p, s1, s2, s1, s2, apus1, apus2, a1, a2, flag, test)
      // Linear interpolation between the superimposed coordinates, weighted by parameter progress
      print "post_rigid"
      if (debug == 0) {
        //print "linear"
        linear(p, s1, s2)
      }
    } else {
      //print "new part"
      sup = instructions[6]
      sup1 = @{morphed_structures[1]} and @{sup}
      sup2 = @{morphed_structures[2]} and @{sup}
      rigid(p, s1, s2, sup1, sup2, apus1, apus2, a1, a2, flag, test)
      if (debug == 0) {
        //print "linear"
        linear(p, s1, s2)
      }
    }   

}


/**
 if (6 == 7) {
   print "morph(steps, structures, domain_table, 0)"
   print "steps is the number of intermediate structures (integer)"
   print "structures is which models to morph, e.g. structures = [{1.1}, {1.2}]"
   print "domain_table defines domains and how to treat them"
   print "for each domain, specify start, end, flag, domain selection, anchor selection"
   print "my_recipe = [\n[0, 0.5, 0, {1-116}, {1-116}],\n[0.5, 1, 0, {117-129}, {116}],\n]"
   print "first domain is 1-116, second domain is 117-129, anchored via 116 of the first domain"
   return
 } 

**/