// storymorph macros by Karsten Theis, CC-BY-SA 5/17/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

/** *********************************************************************
 * Morph between two structures based on domain and anchor definitions
 * @param {integer} steps - Number of intermediate conformations
 * @param {list of atom sets} structures - Selection of first and second structure
 * @param {list of various} domains - Domain, anchor and superposition selections
 * @param {list of various} timing - Start, end, and direction of transition
********************************************************************* **/
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)
}

/** *********************************************************************
 * Superposition of two structures based on domain selection
 * @param {list of atom sets} structures - Selection of first and second structure
 * @param {list of various} domains - Domain, anchor and superposition selections
 * @param {integer} nr - nr of domain to superimpose (rest of structure gets moved too)
********************************************************************* **/
function superimpose(structures, domains, nr) {
  if (!domains.length) {
    print "usage: superimpose(structures, single_domain) or superimpose(structures, domains, nr)"
    return
  }
  if (nr) {
    validate(structures, domains)
    domain = (domains[nr])[3]
  } else {
    domain = domains
  }
  sup1 = structures[1] & domain
  sup2 = structures[2] & domain
  if (sup1.length != sup2.length) {
    print "selections don't match in number of atoms (check altloc, i.e. select %B)"
    return
  }
  if (sup1.length < 3) {
    print "less than 3 atoms in domain"
    return
  }
  compare @{structures[1]} @{structures[2]} ATOMS @{sup1} @{sup2} ROTATE TRANSLATE
}

/** *********************************************************************
 * Visualization of structures, domains and anchors for debugging and documenting
 * @param {list of atom sets} structures - Selection of first and second structure
 * @param {list of various} domains - Domain, anchor and superposition selections
********************************************************************* **/
function show_domains(structures, domains) {
  if (!domains) {
    print "usage: show_domains(structures, domains)"
    return
  }
  original = {all}.xyz.all
  problem = validate(structures, domains)
  if (problem) {
    print problem
    return
  }
  superimpose(structures, domains, 1)
  center visible; select @{structures[1]}; translate selected X 50%; center visible 
  set echo ID 1
  font echo 100
  color echo yellow
  set echo ID 2
  font echo 100
  color echo yellow
  for (var j FROM [1,domains.length]) {
    sel =  (structures[1] | structures[2]) & (domains[j])[1]
    select @{sel}
    anch1 = structures[1] & (domains[j])[2]
    anch2 = structures[2] & (domains[j])[2]
    set echo ID 1 @{anch1}
    echo "⚓" // anchor icon
    set echo ID 2 @{anch2}
    echo "⚓"
    dots on
    pause
    select @{sel}
    dots off
  }
  select all
  dots off
  set echo off
  {all}.xyz = @original
  center visible
}


function atom_order(sel1, sel2, nr) {
  if (nr) {
    validate(sel1, sel2)
    str = sel1
    dom = sel2[nr]
    sel1 = str[1] and dom[3]
    sel2 = str[2] and dom[3]
  }
  if (sel1.length != sel2.length) {
    print "unequal number of atoms, check for altloc"
    print "... pin-pointing offending residues ..."
    matched_residues(sel1, sel2)
    return 0
  }
  for (var i FROM [1,sel1.length]) {
    atom1 = sel1.ident[i]
    atom2 = sel2.ident[i]
    print atom1 + '\tvs.\t' + atom2
  }
}


function matched_residues(sel1, sel2, nr) {
  if (nr) {
    validate(sel1, sel2)
    str = sel1
    dom = sel2[nr]
    sel1 = str[1] and dom[3]
    sel2 = str[2] and dom[3]
  }
  ~residues = (sel1 | sel2).resno.all
  ~first = ~residues.min
  ~last = ~first
  lastresult = 0
  for (var i FROM [~residues.min, ~residues.max]) {
    sresno = {resno=@i}
    s1 = sel1 & sresno
    s2 = sel2 & sresno
    if (s1.length and s2.length) {
      if (s1.length = s2.length) {
        result = 'all'
      } else {
        result = 'some'
      }
    } else {
      result = 0
    }
    if (result == lastresult) {
      ~last = i
    } else {
      if (lastresult) {
        if (~first == ~last) {
          print ' ' + ~first + ':' + lastresult
        } else {
          print ' ' + ~first + '-' + ~last + ': ' + lastresult
        }
      }
      ~first = i
      ~last = i
      lastresult = result
    }
    if (result==25) {
      if (result == 2) {
        print 'good match' + i
      } else {
        print 'partial match' + i
      }
    }
  }
  if (lastresult) {
    print ' ' + ~first + '-' + ~last + ': ' + lastresult
  }
}

/** *********************************************************************
 * Helper functions, not listed, i.e. documented, when loading the macro
********************************************************************* **/

function morph_actual(steps, structures, recipe, timing) {
 apos1 = []
 apos2 = []
 nr1 = structures[1]
 nr2 = structures[2]
 nr = nr1 & 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 @{structures[1]}}.xyz)
   apos2.push({@ansel and @{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], structures, apos1[j], apos2[j])
  }
  if (morph_savefiles) {
    fname = "morph" + step + ".pdb"
    select @{structures[2]} and visible
    write @fname
  }
  set refreshing true
  if (morph_delay) {
    delay @morph_delay
  } else {
    delay 0.05
  }
  if (step == 0) {delay 1.0}
  if (step and step == morph_pause) {
    prompt
  }
  set refreshing false
 }
 if (morph_palindrome) {
  set refreshing true
  delay 2.0
  set refreshing false
  for (var step FROM [steps,0]) {
  // 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], structures, apos1[j], apos2[j])
  }
  set refreshing true
  if (morph_delay) {
    delay @morph_delay
  } else {
    delay 0.05
  }
  set refreshing false
 }}

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

function morph_step(progress, instructions, timing, structures, apos1, apos2) {
    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 = @{structures[1]} and @{s}
    s2 = @{structures[2]} and @{s}
    // a1 and a2: anchors for the correspoding domains
    a1 = @{structures[1]} and @{a}
    a2 = @{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 = @{structures[1]} and @{sup}
    sup2 = @{structures[2]} and @{sup}
    if (!morph_skiprigid) {
      rigid(p, s1, s2, sup1, sup2, apos1, apos2, a1, a2, flag, test)
    }
    if (!morph_skiplinear) {
      if (s1.length == s2.length) {
        linear(p, s1, s2)
      } else {
        linear(p, sup1, sup2)
      }
    }   
}



function validate (structures, domains) {
  if (structures[1] & model[2]) {return "structures not distinct"}
  if (structures[1].length == 0) {return "model 1 empty"}
  if (structures[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 = structures[1] & domain[1]
    sel2 = structures[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 = structures[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 = structures[1] & domain[3]
    sel2 = structures[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"
  }
}



/** ******************************************************************
 * 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 {
    fixed_anchor = anch1 * (1 - progress) + anch2 * progress
    transl = fixed_anchor - 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 = fixed_anchor - anchor_pos
  }
  translateselected @transl
}

/** *********************************************************************
 * Messages upon loading and example morph (plays when no atoms loaded yet)
********************************************************************* **/


print "Functions defined in storymorph.spt:"
print "  morph(steps, structures, domains, [timing])"
print "  show_domains(structures, domains)"
print "  superimpose(structures, domains, domain_nr)"
print "  matched_residues(structures, domains, domain_nr)"
print "  atom_order(structures, domains, domain_nr)"
print "Modify how morph() works through these variables:"
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 (@{structures[2]} and visible)"
print "  morph_pause = 6 for pausing at frame 6"
print "  morph_palindrome = 1 for returning back to the start"
print "Manual, interactive examples and code examples at https://proteopedia.org/w/Jmol/Storymorph#Worked_examples"

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)}], // morph 78-147 as unit, use 78-138 for superposition
    [{4-77}, {78}],                  // morph 4-77 as unit, anchored at 78
  ]
  superimpose(structures, domains, 1)
  center visible
  morph(5,structures, domains)
  print "morph(20,structures, domains) will do a smoother morph"
}