// storymorph macros by Karsten Theis, CC-BY-SA 4/28/2021
// Set up, run and save a morph between two conformations of a molecule
//   Intermediate conformations are created by a combination of
//   rigid body transformation and linear interpolation, with anchors placed
//   to maintain connectivity

function validate (models, domains) {
  if (models[1] & model[2]) {return "models not distinct"}
  if (models[1].length == 0) {return "model 1 empty"}
  if (models[2].length == 0) {return "model 2 empty"}
  cumulator = {none}
  for (var j FROM [1,domains.length]) {
    len = domains[j].length
    sel = (domains[j])[1]
    if (len == 1) {domains[j].push(sel)}
    if (len < 3) {domains[j].push(sel)}
    domain = domains[j]
    sel1 = models[1] & domain[1]
    sel2 = models[2] & domain[1]
    if (sel1.length == 0) {return "empty selection model 1 domain " + j}
    if (sel2.length == 0) {return "empty selection model 2 domain " + j}
    if (sel1 & cumulator) {return "not distinct from other domains: domain " + j}
    cumulator |= sel1
    anch1 = models[1] & domain[2]
    if (sel1 & anch1) {
      len2 = (sel1 & anch1).length
      if (len2 != anch1.length) {return "anchor has to be in or out of domain" + j}
    }
    sel1 = models[1] & domain[3]
    sel2 = models[2] & domain[3]
    if (sel1.length < 3) {
       print "oops, check sel1 and sel2"
       pause
       return "need at least 3 atoms for superposition in domain " + j
    }
    if (sel1.length != sel2.length) {return "superposition requires matching atoms, domain " + j}
    if (sel1 & sel2) {return "atoms for superposition have to be distinct, domain " + j}
  }
  if (morph_savefiles) {
    print "morph_savefiles is set, will save frames to individual files"
  }
  if (morph_delay) {
    print "morph_delay is set, overriding 0.05 s delay between frames"
  }
  if (morph_pause) {
    print "morph_pause is set, will pause at specified frame"
  }
  if (morph_skiplinear) {
    print "morph_skiplinear is set, no linear interpolation"
  }
  if (morph_skiprigid) {
    print "morph_skiprigid is set, no rigid body fit of domains"
  }
}

function superimpose(models, domain, flag) {
  sup1 = models[1] & domain[3]
  sup2 = models[2] & domain[3]
  compare @{models[1]} @{models[2]} ATOMS @{sup1} @{sup2} ROTATE TRANSLATE
  if (flag) {
    select @{models[1]}; translate selected X 50%; center visible 
  }
}

function show_domains(models, domains) {
  original = {all}.xyz.all
  problem = validate(models, domains)
  if (problem) {
    print problem
    return
  }
  superimpose(models, domains[1], 1)
  set echo ID 1
  font echo 100
  color echo cyan
  set echo ID 2
  font echo 100
  color echo cyan
  for (var j FROM [1,domains.length]) {
    sel =  (models[1] | models[2]) & (domains[j])[1]
    select @{sel}
    anch1 = models[1] & (domains[j])[2]
    anch2 = models[2] & (domains[j])[2]
    set echo ID 1 @{anch1}
    echo "⚓" // anchor icon
    set echo ID 2 @{anch2}
    echo "⚓"
    dots on
    pause
    dots off
  }
  {all}.xyz = @original
  center visible
}


/**
 * Linear interpolation of coordinates in the atom sets sel1 and sel2
 * @param {atom set} sel1 First set of atoms
 * @param {atom set} sel2 Second related set of atoms
 * @param {float} progress Weight for atoms in selection sel2
**/
function linear(progress, sel1, sel2) {
  // 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, unmatched atoms ***************"
    return
  }
  for (var i FROM [1,len1]) {
    // Weighted average of coordinates for equivalent points
    coord1[i] = (coord2[i] * progress + coord1[i] * (1 - progress))
  }
  // 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) {
  // Calculate lowest RMSD superposition of supsel1 and supsel2 as 4x4 matrix
  fourbyfour = compare({@supsel1}, {@supsel2})

  // 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}
  anchor_pos = {@ansel1}.xyz    // where the anchor is at the moment
  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
}

function morph_step(progress, instructions, timing, morphed_structures, apus1, apus2) {
    flag = timing[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 < timing[1]) {
      p = 0
    } else if (progress > timing[2]) {
      p = 1
    } else {
      p = (progress - timing[1]) / (timing[2] - timing[1])
    }
    if (timing[1] < 0) {
      p = (progress + timing[1]) / (timing[1] + timing[2])
    }
    s = instructions[1] // domain definition
    a = instructions[2] // 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
    sup = instructions[3]
    sup1 = @{morphed_structures[1]} and @{sup}
    sup2 = @{morphed_structures[2]} and @{sup}
    if (!morph_skiprigid) {
      rigid(p, s1, s2, sup1, sup2, apus1, apus2, a1, a2, flag, test)
    }
    if (!morph_skiplinear) {
      if (s1.length == s2.length) {
        linear(p, s1, s2)
      } else {
        linear(p, sup1, sup2)
      }
    }   
}


/** *********************************************************************
 * 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} flag, if 0 nothing, if = domain, superimpose domains, if -1, don't do linear interpolation
********************************************************************* **/
function morph_actual(steps, morphed_structures, recipe, timing) {
 apos1 = []
 apos2 = []
 nr1 = @{morphed_structures[1]}
 nr2 = @{morphed_structures[2]}
 nr = @{nr1} and @{nr2}
 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"
    return
 }
 for (var j FROM [1,recipe.length]) {
   ansel = (recipe[j])[2]
   apos1.push({@ansel and @{morphed_structures[1]}}.xyz)
   apos2.push({@ansel and @{morphed_structures[2]}}.xyz)
 }
 // Save coordinates to use in every step of the interpolation
 original = {all}.xyz.all
 for (var step FROM [0,steps]) {
  // Set coordinates back to original values in preparation of calculating next step in the morph
  {all}.xyz = @original
  progress = step * (1.0 / steps)
  for (var j FROM [1,recipe.length]) {
    //print "ij:" + @step + " " + @j
    morph_step(progress, recipe[j], timing[j], morphed_structures, apos1[j], apos2[j])
  }
  if (morph_savefiles) {
  // To save this step in the morph, un-comment the following three lines while running in local Jmol app
    fname = "morph" + step + ".pdb"
    select @{morphed_structures[1]}
    write @fname
  }
  set refreshing true
  if (morph_delay) {
    delay @morph_delay
  } else {
    delay 0.05
  }
  if (step and step == morph_pause) {
    prompt
  }
  set refreshing false
 }
 set refreshing true
 delay 2.0
 {all}.xyz = @original
}

function morph(steps, structures, domains, timing) {
  if (!domains) {
    print "usage: morph(steps, structures, domains, timing)"
    return
  }
  problem = validate(structures, domains)
  if (problem) {
    print problem
    return
  }
  nr_domains = domains.length
  if not (timing) {
    timing = []
    for (var n FROM [1, domains.length]) {
      timing.push([0,1,0])
    }
  }
  morph_actual(steps, structures, domains, timing)
}

/**
load dna1.pdb
load append dna2.pdb

my_doms = [[{all}]]
my_models = [{1.1}, {2.1}]
**/


print "Functions defined in storymorph.spt:"
print "  validate (conformations, domains)"
print "  morph(steps, conformations, domains, timing)"
print "  linear(progress, sel1, sel2)"
print "  superimpose(conformations, domain, sidebyside)"
print "  show_domains(conformations, domains)"
print "  validate(domains)"
print "Use the following variables for additional control and features:"
print "  morph_delay = 0.5 for 0.5 seconds dalay between frames"
print "  morph_skiplinear = 1 for skipping the linear morph when debugging"
print "  morph_skiprigid = 1 for skipping the rigid body move when debugging"
print "  morph_savefiles = 1 for saving frames to file"
print "  morph_pause = 6 for pausing at frame 6"
print "For examples of the lists describing conformations and domains, start from scratch and type"
print "  macro storymorph"
print "  print conformations"
print "  print domains"

test = {all}
if (test.length == 0) {
  load files "=1prw" "=1cll"
  delete water
  delete protein and not alpha
  model 0
  backbone only
  backbone 0.2
  display 4-147
  delete water
  color group
  structures = [{1.1}, {2.1}]
  domains = [ 
    [{78-147},{78-147}, {(78-138)}], 
    [{4-77}, {78}],
  ]
  superimpose(structures, domains[1])
  center visible
  morph(20,structures, domains)
  print "morph(5,structures, domains) will do another morph"
}